import asyncio from pathlib import Path from core.config import settings from core.security import ( InvalidTokenError, authenticate_token, require_authenticated_user, ) from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket, status from integrations.docker import get_all_services_in_dir from integrations.docker.discovery import service_scan_root from models.example import AddonExampleCreateInput from models.proxy import ProxyCreateInput, ProxyUpdateInput from pymongo.errors import DuplicateKeyError from repositories.examples import ( create_custom_example, delete_custom_example, list_custom_examples, ) from repositories.proxy import ( create_proxy, delete_proxy, get_proxy_by_port, update_proxy, update_proxy_runtime, ) from repositories.proxy import ( list_proxies as list_proxy_documents, ) from services.mitmproxy.examples import BUILTIN_EXAMPLES from services.mitmproxy.health import check_proxy from services.mitmproxy.manager import MitmdumpError, manager from services.mitmproxy.options import ( MitmOptionError, load_options, public_options, validate_options, ) from services.mitmproxy.validator import ProxyCodeValidationError, validate_proxy_code from starlette.websockets import WebSocketDisconnect SUPPORTED_MODES = { "regular", "local", "wireguard", "reverse", "transparent", "tun", "upstream", "socks5", "dns", } SUPPORTED_PROTOCOLS = { "http", "https", "tcp", "tls", "udp", "dtls", "dns", "http3", "quic", } router = APIRouter( prefix="/proxy", tags=["Proxy"], ) auth_dependencies = [Depends(require_authenticated_user)] async def _validate_code_or_raise(code: str) -> None: try: await validate_proxy_code(code) except ProxyCodeValidationError as error: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=error.to_http_detail(), ) from error def validate_mode(mode: str) -> None: if mode not in SUPPORTED_MODES: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=f"Unsupported proxy mode: {mode}", ) def validate_mode_spec(mode_spec: str | None) -> None: if mode_spec is None: return mode = mode_spec.split(":", 1)[0].split("@", 1)[0] validate_mode(mode) def validate_protocol(protocol: str) -> None: if protocol not in SUPPORTED_PROTOCOLS: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=f"Unsupported reverse proxy protocol: {protocol}", ) async def validate_mitm_options(options: dict[str, object]) -> None: try: await validate_options(options) except MitmOptionError as error: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=str(error), ) from error async def get_proxy_or_404(proxy_port: int): proxy = await get_proxy_by_port(proxy_port) if proxy is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Proxy not found", ) return proxy async def authenticate_proxy_socket(websocket: WebSocket) -> None: token = websocket.query_params.get("token") if token is None: authorization = websocket.headers.get("authorization", "") scheme, _, value = authorization.partition(" ") if scheme.lower() == "bearer" and value: token = value if not token: raise InvalidTokenError await authenticate_token(token) def proxy_log_path(proxy) -> Path: if proxy.log_path: return Path(proxy.log_path) return manager.log_path_for(proxy) def read_log_snapshot(path: Path, limit: int = 8000) -> str: if not path.is_file(): return "" return path.read_bytes()[-limit:].decode("utf-8", errors="replace") async def watch_socket_disconnect(websocket: WebSocket) -> None: while True: await websocket.receive_text() async def stream_proxy_logs(websocket: WebSocket, path: Path) -> None: offset = path.stat().st_size if path.is_file() else 0 while True: if not path.exists(): await asyncio.sleep(0.4) continue try: size = path.stat().st_size except OSError: await asyncio.sleep(0.4) continue if size < offset: offset = size if size > offset: with path.open("rb") as handle: handle.seek(offset) chunk = handle.read(size - offset).decode("utf-8", errors="replace") offset = size if chunk: await websocket.send_json({"type": "append", "log": chunk}) await asyncio.sleep(0.4) @router.get("", dependencies=auth_dependencies) async def list_proxies_handler(): return await list_proxy_documents() @router.get("/services", dependencies=auth_dependencies) async def list_services(request: Request, scan_root: str | None = None): _ = settings root = Path(scan_root).expanduser() if scan_root else service_scan_root() services = get_all_services_in_dir(root) request.app.state.services = services return {"scan_root": str(root), "services": services} @router.get("/mitmproxy/options", dependencies=auth_dependencies) async def list_mitmproxy_options(): return public_options(await load_options()) @router.get("/mitmproxy/examples", dependencies=auth_dependencies) async def list_mitmproxy_examples(): return [*BUILTIN_EXAMPLES, *await list_custom_examples()] @router.post( "/mitmproxy/examples", status_code=status.HTTP_201_CREATED, dependencies=auth_dependencies, ) async def create_mitmproxy_example(payload: AddonExampleCreateInput): await _validate_code_or_raise(payload.code) try: return await create_custom_example(payload) except DuplicateKeyError as error: raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=f"Example already exists: {payload.id}", ) from error @router.delete( "/mitmproxy/examples/{example_id}", status_code=status.HTTP_204_NO_CONTENT, dependencies=auth_dependencies, ) async def delete_mitmproxy_example(example_id: str): if any(example.id == example_id for example in BUILTIN_EXAMPLES): raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail="Built-in examples cannot be deleted", ) deleted = await delete_custom_example(example_id) if not deleted: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Example not found", ) @router.get("/{proxy_port}", dependencies=auth_dependencies) async def get_proxy(proxy_port: int): return await get_proxy_or_404(proxy_port) @router.post("", status_code=status.HTTP_201_CREATED, dependencies=auth_dependencies) async def create_proxy_handler(payload: ProxyCreateInput): existing_proxy = await get_proxy_by_port(payload.proxy_port) if existing_proxy is not None: raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=f"Proxy already exists on port {payload.proxy_port}", ) await _validate_code_or_raise(payload.code) validate_mode(payload.mode) validate_mode_spec(payload.mode_spec) validate_protocol(payload.protocol) await validate_mitm_options(payload.mitm_options) return await create_proxy(payload) @router.put("/{proxy_port}", dependencies=auth_dependencies) async def edit_proxy(proxy_port: int, payload: ProxyUpdateInput): if payload.code is not None: await _validate_code_or_raise(payload.code) if payload.mode is not None: validate_mode(payload.mode) if "mode_spec" in payload.model_fields_set: validate_mode_spec(payload.mode_spec) if payload.protocol is not None: validate_protocol(payload.protocol) if payload.mitm_options is not None: await validate_mitm_options(payload.mitm_options) updated_proxy = await update_proxy(proxy_port, payload) if updated_proxy is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Proxy not found", ) return updated_proxy @router.post("/{proxy_port}/start", dependencies=auth_dependencies) async def start_proxy(proxy_port: int): proxy = await get_proxy_or_404(proxy_port) try: runtime = await manager.start(proxy) except MitmdumpError as error: updated = await update_proxy_runtime( proxy_port, {"status": "error", "last_error": str(error)}, ) return updated or proxy return await update_proxy_runtime(proxy_port, runtime) @router.post("/{proxy_port}/stop", dependencies=auth_dependencies) async def stop_proxy(proxy_port: int): proxy = await get_proxy_or_404(proxy_port) runtime = await manager.stop(proxy) return await update_proxy_runtime(proxy_port, runtime) @router.post("/{proxy_port}/restart", dependencies=auth_dependencies) async def restart_proxy(proxy_port: int): proxy = await get_proxy_or_404(proxy_port) stopped = await manager.stop(proxy) proxy = await update_proxy_runtime(proxy_port, stopped) or proxy try: runtime = await manager.start(proxy) except MitmdumpError as error: updated = await update_proxy_runtime( proxy_port, {"status": "error", "last_error": str(error)}, ) return updated or proxy return await update_proxy_runtime(proxy_port, runtime) @router.get("/{proxy_port}/status", dependencies=auth_dependencies) async def proxy_status(proxy_port: int): proxy = await get_proxy_or_404(proxy_port) runtime_status = await manager.status(proxy) if runtime_status != proxy.status: values: dict[str, object] = {"status": runtime_status} if runtime_status != "running": values["pid"] = None proxy = await update_proxy_runtime(proxy_port, values) or proxy return {"proxy_port": proxy.proxy_port, "status": runtime_status, "pid": proxy.pid} @router.get("/{proxy_port}/health", dependencies=auth_dependencies) async def proxy_health(proxy_port: int): proxy = await get_proxy_or_404(proxy_port) return await check_proxy(proxy) @router.get("/{proxy_port}/logs", dependencies=auth_dependencies) async def proxy_logs(proxy_port: int): proxy = await get_proxy_or_404(proxy_port) path = proxy_log_path(proxy) return { "proxy_port": proxy.proxy_port, "log": read_log_snapshot(path), } @router.websocket("/{proxy_port}/logs/ws") async def proxy_logs_socket(websocket: WebSocket, proxy_port: int): disconnect_task: asyncio.Task[None] | None = None stream_task: asyncio.Task[None] | None = None try: try: await authenticate_proxy_socket(websocket) except InvalidTokenError: await websocket.close(code=4401) return proxy = await get_proxy_by_port(proxy_port) if proxy is None: await websocket.close(code=4404) return path = proxy_log_path(proxy) await websocket.accept() await websocket.send_json({"type": "snapshot", "log": read_log_snapshot(path)}) disconnect_task = asyncio.create_task(watch_socket_disconnect(websocket)) stream_task = asyncio.create_task(stream_proxy_logs(websocket, path)) done, pending = await asyncio.wait( {disconnect_task, stream_task}, return_when=asyncio.FIRST_EXCEPTION, ) for task in pending: task.cancel() if pending: await asyncio.gather(*pending, return_exceptions=True) for task in done: try: task.result() except asyncio.CancelledError: return except WebSocketDisconnect: return except RuntimeError: await websocket.close(code=1011, reason="log stream failed") return except asyncio.CancelledError: return finally: for task in (disconnect_task, stream_task): if task is None or task.done(): continue task.cancel() @router.delete( "/{proxy_port}", status_code=status.HTTP_204_NO_CONTENT, dependencies=auth_dependencies, ) async def remove_proxy(proxy_port: int): proxy = await get_proxy_by_port(proxy_port) if proxy is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Proxy not found", ) await manager.stop(proxy) deleted = await delete_proxy(proxy_port) if not deleted: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Proxy not found", )