diff --git a/sdk/python/examples/proxy/proxy.py b/sdk/python/examples/proxy/proxy.py new file mode 100644 index 00000000..97f7d2a4 --- /dev/null +++ b/sdk/python/examples/proxy/proxy.py @@ -0,0 +1,153 @@ +import argparse +import atexit +import logging +import sys +import urllib.parse + +import requests +from flask import Flask, Response, request +from waitress import serve +from zrok.model import ShareRequest, Share +from zrok.overview import EnvironmentAndResources, Overview + +import zrok + +# Setup logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)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('/', 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__': + 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()}") + + root = zrok.environment.root.Load() + my_env = EnvironmentAndResources( + environment=None, + shares=[] + ) + 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 + + 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()