add share mode

This commit is contained in:
Kenneth Bingham 2025-01-29 04:06:37 -05:00
parent 41ce0809ab
commit 5c64a73aad
No known key found for this signature in database
GPG Key ID: 31709281860130B6
2 changed files with 23 additions and 11 deletions

View File

@ -31,23 +31,30 @@ def main():
parser.add_argument('-f', '--frontends', nargs='+', help='One or more space-separated frontends to use') parser.add_argument('-f', '--frontends', nargs='+', help='One or more space-separated frontends to use')
parser.add_argument('-k', '--insecure', action='store_false', dest='verify_ssl', default=True, parser.add_argument('-k', '--insecure', action='store_false', dest='verify_ssl', default=True,
help='Skip SSL verification') help='Skip SSL verification')
parser.add_argument('-s', '--share-mode', default='public', choices=['public', 'private'],
help='Share mode (default: public)')
args = parser.parse_args() args = parser.parse_args()
logger.info("=== Starting proxy server ===") logger.info("=== Starting proxy server ===")
logger.info(f"Target URL: {args.target_url}") logger.info(f"Target URL: {args.target_url}")
logger.info(f"Share mode: {args.share_mode}")
# Load environment and create proxy share # Load environment and create proxy share
root = zrok.environment.root.Load() root = zrok.environment.root.Load()
proxy_share = ProxyShare.create( proxy_share = ProxyShare.create(
root=root, root=root,
target=args.target_url, target=args.target_url,
share_mode=args.share_mode,
unique_name=args.unique_name, unique_name=args.unique_name,
frontends=args.frontends, frontends=args.frontends,
verify_ssl=args.verify_ssl verify_ssl=args.verify_ssl
) )
# Log access information and start the proxy # Log access information and start the proxy
logger.info(f"Access proxy at: {', '.join(proxy_share.endpoints)}") if args.share_mode == "public":
logger.info(f"Access proxy at: {', '.join(proxy_share.endpoints)}")
elif args.share_mode == "private":
logger.info(f"Run a private access frontend: 'zrok access private {proxy_share.token}'")
proxy_share.run() proxy_share.run()

View File

@ -12,7 +12,7 @@ import requests
from flask import Flask, Response, request from flask import Flask, Response, request
from waitress import serve from waitress import serve
from zrok.environment.root import Root from zrok.environment.root import Root
from zrok.model import (PROXY_BACKEND_MODE, PUBLIC_SHARE_MODE, Share, from zrok.model import (PROXY_BACKEND_MODE, PUBLIC_SHARE_MODE, PRIVATE_SHARE_MODE, Share,
ShareRequest) ShareRequest)
from zrok.overview import Overview from zrok.overview import Overview
from zrok.share import CreateShare, ReleaseReservedShare from zrok.share import CreateShare, ReleaseReservedShare
@ -50,7 +50,7 @@ class ProxyShare:
verify_ssl: bool = True verify_ssl: bool = True
@classmethod @classmethod
def create(cls, root: Root, target: str, unique_name: Optional[str] = None, def create(cls, root: Root, target: str, share_mode: str = PUBLIC_SHARE_MODE, unique_name: Optional[str] = None,
frontends: Optional[List[str]] = None, verify_ssl: bool = True) -> 'ProxyShare': frontends: Optional[List[str]] = None, verify_ssl: bool = True) -> 'ProxyShare':
""" """
Create a new proxy share, handling reservation and cleanup logic based on unique_name. Create a new proxy share, handling reservation and cleanup logic based on unique_name.
@ -79,16 +79,18 @@ class ProxyShare:
) )
# Compose the share request # Compose the share request
if frontends: share_frontends = []
share_frontends = frontends if share_mode == PUBLIC_SHARE_MODE:
elif root.cfg and root.cfg.DefaultFrontend: if frontends:
share_frontends = [root.cfg.DefaultFrontend] share_frontends = frontends
else: elif root.cfg and root.cfg.DefaultFrontend:
share_frontends = ['public'] share_frontends = [root.cfg.DefaultFrontend]
else:
share_frontends = ['public']
share_request = ShareRequest( share_request = ShareRequest(
BackendMode=PROXY_BACKEND_MODE, BackendMode=PROXY_BACKEND_MODE,
ShareMode=PUBLIC_SHARE_MODE, ShareMode=share_mode,
Target=target, Target=target,
Frontends=share_frontends, Frontends=share_frontends,
Reserved=True Reserved=True
@ -98,7 +100,10 @@ class ProxyShare:
# Create the share # Create the share
share = CreateShare(root=root, request=share_request) share = CreateShare(root=root, request=share_request)
logger.debug(f"Created new proxy share with endpoints: {', '.join(share.FrontendEndpoints)}") if share_mode == PUBLIC_SHARE_MODE:
logger.debug(f"Created new proxy share with endpoints: {', '.join(share.FrontendEndpoints)}")
elif share_mode == PRIVATE_SHARE_MODE:
logger.debug(f"Created new private share with token: {share.Token}")
# Create class instance and setup cleanup-at-exit if we reserved a random share token # Create class instance and setup cleanup-at-exit if we reserved a random share token
instance = cls( instance = cls(