migrate the proxy example to ProxyShare class so it's a reusable feature of the Py SDK

sdk/python/examples/proxy/proxy.py

@@ -1,153 +1,50 @@
-import argparse
-import atexit
-import logging
-import sys
-import urllib.parse
+#!/usr/bin/env python3
 
-import requests
-from flask import Flask, Response, request
-from waitress import serve
-from zrok.model import ShareRequest, Share
-from zrok.overview import EnvironmentAndResources, Overview
+"""
+Example of using zrok's proxy facility to create an HTTP proxy server.
+
+This example demonstrates how to:
+1. Create a proxy share (optionally with a unique name for persistence)
+2. Handle HTTP requests/responses through the proxy
+3. Automatically clean up non-reserved shares on exit
+"""
+
+import argparse
+import logging
 
 import zrok
+from zrok.proxy import ProxyShare
 
 # Setup logging
 logging.basicConfig(
     level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s'
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
-app = Flask(__name__)
-target_url = None
-zrok_opts = {}
-bindPort = 18081
 
-# List of hop-by-hop headers that should not be returned to the viewer
-HOP_BY_HOP_HEADERS = {
-    'connection',
-    'keep-alive',
-    'proxy-authenticate',
-    'proxy-authorization',
-    'te',
-    'trailers',
-    'transfer-encoding',
-    'upgrade'
-}
-
-
-@app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'OPTIONS', 'PATCH'])
-@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'OPTIONS', 'PATCH'])
-def proxy(path):
-    global target_url
-    logger.info(f"Incoming {request.method} request to {request.path}")
-    logger.info(f"Headers: {dict(request.headers)}")
-
-    # Forward the request to target URL
-    full_url = urllib.parse.urljoin(target_url, request.path)
-    logger.info(f"Forwarding to: {full_url}")
-
-    # Copy request headers, excluding hop-by-hop headers
-    headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_BY_HOP_HEADERS and k.lower() != 'host'}
-
-    try:
-        response = requests.request(
-            method=request.method,
-            url=full_url,
-            headers=headers,
-            data=request.get_data(),
-            stream=True
-        )
-
-        logger.info(f"Response status: {response.status_code}")
-        logger.info(f"Response headers: {dict(response.headers)}")
-
-        # Filter out hop-by-hop headers from the response
-        filtered_headers = {k: v for k, v in response.headers.items() if k.lower() not in HOP_BY_HOP_HEADERS}
-
-        return Response(
-            response.iter_content(chunk_size=8192),
-            status=response.status_code,
-            headers=filtered_headers
-        )
-
-    except Exception as e:
-        logger.error(f"Proxy error: {str(e)}", exc_info=True)
-        return str(e), 502
-
-
-@zrok.decor.zrok(opts=zrok_opts)
-def run_proxy():
-    # the port is only used to integrate zrok with frameworks that expect a "hostname:port" combo
-    serve(app, port=bindPort)
-
-
-if __name__ == '__main__':
+def main():
+    """Main entry point."""
     parser = argparse.ArgumentParser(description='Start a zrok proxy server')
     parser.add_argument('target_url', help='Target URL to proxy requests to')
     parser.add_argument('-n', '--unique-name', help='Unique name for the proxy instance')
     args = parser.parse_args()
 
-    target_url = args.target_url
     logger.info("=== Starting proxy server ===")
-    logger.info(f"Target URL: {target_url}")
-    logger.info(f"Logging level: {logger.getEffectiveLevel()}")
+    logger.info(f"Target URL: {args.target_url}")
 
+    # Load environment and create proxy share
     root = zrok.environment.root.Load()
-    my_env = EnvironmentAndResources(
-        environment=None,
-        shares=[]
+    proxy_share = ProxyShare.create(
+        root=root,
+        target=args.target_url,
+        unique_name=args.unique_name
     )
-    overview = Overview.create(root=root)
-    for env_stuff in overview.environments:
-        if env_stuff.environment.z_id == root.env.ZitiIdentity:
-            my_env = EnvironmentAndResources(
-                environment=env_stuff.environment,
-                shares=env_stuff.shares
-            )
-            break
+    
+    # Log access information and start the proxy
+    logger.info(f"Access proxy at: {', '.join(proxy_share.endpoints)}")
+    proxy_share.run()
 
-    if my_env:
-        logger.debug(
-            f"Found environment in overview with Ziti identity "
-            f"matching local environment: {my_env.environment.z_id}"
-        )
-    else:
-        logger.error("No matching environment found")
-        sys.exit(1)
 
-    existing_reserved_share = None
-    for share in my_env.shares:
-        if share.token == args.unique_name:
-            existing_reserved_share = share
-            break
-
-    if existing_reserved_share:
-        logger.debug(f"Found existing share with token: {existing_reserved_share.token}")
-        shr = Share(Token=existing_reserved_share.token, FrontendEndpoints=[existing_reserved_share.frontend_endpoint])
-    else:
-        logger.debug(f"No existing share found with token: {args.unique_name}")
-        share_request = ShareRequest(
-            BackendMode=zrok.model.PROXY_BACKEND_MODE,
-            ShareMode=zrok.model.PUBLIC_SHARE_MODE,
-            Frontends=['public'],
-            Target="http-proxy",
-            Reserved=True
-        )
-        if args.unique_name:
-            share_request.UniqueName = args.unique_name
-
-        shr = zrok.share.CreateShare(root=root, request=share_request)
-
-    def cleanup():
-        zrok.share.ReleaseReservedShare(root=root, shr=shr)
-        logger.info(f"Share {shr.Token} released")
-    if not args.unique_name:
-        atexit.register(cleanup)
-
-    zrok_opts['cfg'] = zrok.decor.Opts(root=root, shrToken=shr.Token, bindPort=bindPort)
-
-    logger.info(f"Access proxy at: {', '.join(shr.FrontendEndpoints)}")
-
-    run_proxy()
+if __name__ == '__main__':
+    main()