mirror of
https://github.com/openziti/zrok.git
synced 2025-06-24 19:51:32 +02:00
migrate the proxy example to ProxyShare class so it's a reusable feature of the Py SDK
This commit is contained in:
Kenneth Bingham
parent
9097643617
commit
fe056a2ccc
@ -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()
|
||||
|
||||
47
sdk/python/sdk/zrok/zrok/README.md
Normal file
47
sdk/python/sdk/zrok/zrok/README.md
Normal file
@ -0,0 +1,47 @@
|
||||
# Zrok Python SDK
|
||||
|
||||
## Proxy Facility
|
||||
|
||||
The SDK includes a proxy facility that makes it easy to create and manage proxy shares. This is particularly useful when you need to:
|
||||
|
||||
1. Create an HTTP proxy with zrok
|
||||
2. Optionally reserve the proxy with a unique name for persistence
|
||||
3. Automatically handle cleanup of non-reserved shares
|
||||
|
||||
### Basic Usage
|
||||
|
||||
```python
|
||||
from zrok.proxy import ProxyShare
|
||||
import zrok
|
||||
|
||||
# Load the environment
|
||||
root = zrok.environment.root.Load()
|
||||
|
||||
# Create a temporary proxy share (will be cleaned up on exit)
|
||||
proxy = ProxyShare.create(root=root, target="http://my-target-service")
|
||||
|
||||
# Access the proxy's endpoints and token
|
||||
print(f"Access proxy at: {proxy.endpoints}")
|
||||
print(f"Share token: {proxy.token}")
|
||||
```
|
||||
|
||||
### Creating a Reserved Proxy Share
|
||||
|
||||
To create a proxy share that persists and can be reused:
|
||||
|
||||
```python
|
||||
# Create/retrieve a reserved proxy share with a unique name
|
||||
proxy = ProxyShare.create(
|
||||
root=root,
|
||||
target="http://my-target-service",
|
||||
unique_name="my-persistent-proxy"
|
||||
)
|
||||
```
|
||||
|
||||
When a `unique_name` is provided:
|
||||
|
||||
1. If the zrok environment already has a share with that name, it will be reused
|
||||
2. If no share exists, a new reserved share will be created
|
||||
3. The share will be automatically cleaned up on exit if no `unique_name` is provided
|
||||
|
||||
When a `unique_name` is not provided, the randomly generated share will be cleaned up on exit.
|
||||
178
sdk/python/sdk/zrok/zrok/proxy.py
Normal file
178
sdk/python/sdk/zrok/zrok/proxy.py
Normal file
@ -0,0 +1,178 @@
|
||||
"""
|
||||
Proxy share management functionality for the zrok SDK.
|
||||
"""
|
||||
|
||||
import atexit
|
||||
import logging
|
||||
import urllib.parse
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import requests
|
||||
from flask import Flask, Response, request
|
||||
from waitress import serve
|
||||
from zrok.environment.root import Root
|
||||
from zrok.model import (PROXY_BACKEND_MODE, PUBLIC_SHARE_MODE, Share,
|
||||
ShareRequest)
|
||||
from zrok.overview import Overview
|
||||
from zrok.share import CreateShare, ReleaseReservedShare
|
||||
|
||||
import zrok
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 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'
|
||||
}
|
||||
|
||||
# The proxy only listens on the zrok socket, the port used to initialize the Waitress server is not actually bound or
|
||||
# listening
|
||||
DUMMY_PORT = 18081
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProxyShare:
|
||||
"""Represents a proxy share with its configuration and state."""
|
||||
root: Root
|
||||
share: Share
|
||||
target: str
|
||||
unique_name: Optional[str] = None
|
||||
_cleanup_registered: bool = False
|
||||
_app: Optional[Flask] = None
|
||||
|
||||
@classmethod
|
||||
def create(cls, root: Root, target: str, unique_name: Optional[str] = None) -> 'ProxyShare':
|
||||
"""
|
||||
Create a new proxy share, handling reservation and cleanup logic based on unique_name.
|
||||
|
||||
Args:
|
||||
root: The zrok root environment
|
||||
target: Target URL or service to proxy to
|
||||
unique_name: Optional unique name for a reserved share
|
||||
|
||||
Returns:
|
||||
ProxyShare instance configured with the created share
|
||||
"""
|
||||
# First check if we have an existing reserved share with this name
|
||||
if unique_name:
|
||||
existing_share = cls._find_existing_share(root, unique_name)
|
||||
if existing_share:
|
||||
logger.debug(f"Found existing share with token: {existing_share.Token}")
|
||||
return cls(
|
||||
root=root,
|
||||
share=existing_share,
|
||||
target=target,
|
||||
unique_name=unique_name
|
||||
)
|
||||
|
||||
# Create new share request
|
||||
share_request = ShareRequest(
|
||||
BackendMode=PROXY_BACKEND_MODE,
|
||||
ShareMode=PUBLIC_SHARE_MODE,
|
||||
Target="http-proxy",
|
||||
Frontends=['public'],
|
||||
Reserved=bool(unique_name)
|
||||
)
|
||||
if unique_name:
|
||||
share_request.UniqueName = unique_name
|
||||
|
||||
# Create the share
|
||||
share = CreateShare(root=root, request=share_request)
|
||||
logger.info(f"Created new proxy share with endpoints: {', '.join(share.FrontendEndpoints)}")
|
||||
|
||||
# Create instance and setup cleanup if needed
|
||||
instance = cls(
|
||||
root=root,
|
||||
share=share,
|
||||
target=target,
|
||||
unique_name=unique_name
|
||||
)
|
||||
if not unique_name:
|
||||
instance.register_cleanup()
|
||||
return instance
|
||||
|
||||
@staticmethod
|
||||
def _find_existing_share(root: Root, unique_name: str) -> Optional[Share]:
|
||||
"""Find an existing share with the given unique name."""
|
||||
overview = Overview.create(root=root)
|
||||
for env in overview.environments:
|
||||
if env.environment.z_id == root.env.ZitiIdentity:
|
||||
for share in env.shares:
|
||||
if share.token == unique_name:
|
||||
return Share(Token=share.token, FrontendEndpoints=[share.frontend_endpoint])
|
||||
return None
|
||||
|
||||
def register_cleanup(self):
|
||||
"""Register cleanup handler to release the share on exit."""
|
||||
if not self._cleanup_registered:
|
||||
def cleanup():
|
||||
ReleaseReservedShare(root=self.root, shr=self.share)
|
||||
logger.info(f"Share {self.share.Token} released")
|
||||
atexit.register(cleanup)
|
||||
self._cleanup_registered = True
|
||||
|
||||
def _create_app(self) -> Flask:
|
||||
"""Create and configure the Flask app for proxying."""
|
||||
app = Flask(__name__)
|
||||
|
||||
@app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
|
||||
@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
|
||||
def proxy(path):
|
||||
# Construct the target URL
|
||||
url = urllib.parse.urljoin(self.target, path)
|
||||
|
||||
# Forward the request
|
||||
resp = requests.request(
|
||||
method=request.method,
|
||||
url=url,
|
||||
headers={key: value for (key, value) in request.headers
|
||||
if key.lower() not in HOP_BY_HOP_HEADERS},
|
||||
data=request.get_data(),
|
||||
cookies=request.cookies,
|
||||
allow_redirects=False,
|
||||
stream=True
|
||||
)
|
||||
|
||||
# Create the response
|
||||
excluded_headers = HOP_BY_HOP_HEADERS.union({'host'})
|
||||
headers = [(name, value) for (name, value) in resp.raw.headers.items()
|
||||
if name.lower() not in excluded_headers]
|
||||
|
||||
return Response(
|
||||
resp.iter_content(chunk_size=10*1024),
|
||||
status=resp.status_code,
|
||||
headers=headers
|
||||
)
|
||||
return app
|
||||
|
||||
def run(self):
|
||||
"""Start the proxy server."""
|
||||
if self._app is None:
|
||||
self._app = self._create_app()
|
||||
|
||||
# Create options dictionary for zrok decorator
|
||||
zrok_opts: Dict[str, Any] = {}
|
||||
zrok_opts['cfg'] = zrok.decor.Opts(root=self.root, shrToken=self.token, bindPort=DUMMY_PORT)
|
||||
|
||||
@zrok.decor.zrok(opts=zrok_opts)
|
||||
def run_server():
|
||||
serve(self._app, port=DUMMY_PORT)
|
||||
run_server()
|
||||
|
||||
@property
|
||||
def endpoints(self) -> List[str]:
|
||||
"""Get the frontend endpoints for this share."""
|
||||
return self.share.FrontendEndpoints
|
||||
|
||||
@property
|
||||
def token(self) -> str:
|
||||
"""Get the share token."""
|
||||
return self.share.Token
|
||||
Loading…
x
Reference in New Issue
Block a user