import asyncio from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, WebSocket, status from pymongo.errors import DuplicateKeyError from starlette.websockets import WebSocketDisconnect from app.api.proxy.socket_utils import ( authenticate_proxy_socket, proxy_log_path, read_log_snapshot, serialize_health_snapshot, snapshot_log_state, stream_proxy_logs, watch_socket_disconnect, ) from app.core.config import settings from app.core.security import InvalidTokenError, require_authenticated_user from app.integrations.docker import get_all_services_in_dir from app.integrations.docker.discovery import service_scan_root from app.models.health import HealthHistoryResponse from app.models.proxy import ProxyCreateInput, ProxyUpdateInput from app.models.template import ProxyTemplateCreateInput from app.repositories.health import ( delete_health_history, get_latest_health_sample, list_health_history, record_health_sample, ) from app.repositories.proxy import ( create_proxy, delete_proxy, get_proxy_by_port, update_proxy, update_proxy_runtime, ) from app.repositories.proxy import ( list_proxies as list_proxy_documents, ) from app.repositories.templates import ( create_custom_template, delete_custom_template, list_custom_templates, ) from app.services.mitmproxy.health import check_proxy from app.services.mitmproxy.health_monitor import health_monitor from app.services.mitmproxy.manager import MitmdumpError, manager from app.services.mitmproxy.options import ( MitmOptionError, load_options, public_options, validate_options, ) from app.services.mitmproxy.runtime_validation import ( ReservedProxyPortError, backend_port, ensure_proxy_port_not_reserved, ) from app.services.mitmproxy.templates import file_template_ids, list_file_templates from app.services.mitmproxy.validator import ( ProxyCodeValidationError, validate_proxy_code, ) SUPPORTED_MODES = { "regular", "local", "wireguard", "reverse", "transparent", "tun", "upstream", "socks5", "dns", } SUPPORTED_PROTOCOLS = { "http", "https", "tcp", "tls", "udp", "dtls", "dns", "http3", "quic", } DEFAULT_HEALTH_HISTORY_LIMIT = 40 MAX_HEALTH_HISTORY_LIMIT = 200 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 def validate_proxy_port(proxy_port: int) -> None: try: ensure_proxy_port_not_reserved(proxy_port) except ReservedProxyPortError 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 @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(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) return { "scan_root": str(root), "services": services, "backend_port": backend_port(), } @router.get("/mitmproxy/options", dependencies=auth_dependencies) async def list_mitmproxy_options(): return public_options(await load_options()) @router.get("/mitmproxy/templates", dependencies=auth_dependencies) async def list_mitmproxy_templates(): return [*list_file_templates(), *await list_custom_templates()] @router.post( "/mitmproxy/templates", status_code=status.HTTP_201_CREATED, dependencies=auth_dependencies, ) async def create_mitmproxy_template(payload: ProxyTemplateCreateInput): await _validate_code_or_raise(payload.code) if payload.id in file_template_ids(): raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=f"Template already exists: {payload.id}", ) try: return await create_custom_template(payload) except DuplicateKeyError as error: raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=f"Template already exists: {payload.id}", ) from error @router.delete( "/mitmproxy/templates/{template_id}", status_code=status.HTTP_204_NO_CONTENT, dependencies=auth_dependencies, ) async def delete_mitmproxy_template(template_id: str): if template_id in file_template_ids(): raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail="Built-in templates cannot be deleted", ) deleted = await delete_custom_template(template_id) if not deleted: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Template 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): validate_proxy_port(payload.proxy_port) 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) snapshot = await check_proxy(proxy) if proxy.healthcheck.enabled: await record_health_sample(snapshot) await health_monitor.publish(proxy.proxy_port, snapshot) return serialize_health_snapshot(snapshot) @router.get("/{proxy_port}/health/history", dependencies=auth_dependencies) async def proxy_health_history( proxy_port: int, limit: int = DEFAULT_HEALTH_HISTORY_LIMIT, ): proxy = await get_proxy_or_404(proxy_port) limit = max(1, min(limit, MAX_HEALTH_HISTORY_LIMIT)) latest = await get_latest_health_sample(proxy.proxy_port) history = await list_health_history(proxy.proxy_port, limit=limit) response = HealthHistoryResponse( proxy_port=proxy.proxy_port, latest=latest, history=history, ) return response.model_dump(mode="json") @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": await 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) initial_snapshot, offset, marker = await snapshot_log_state(path) await websocket.accept() await websocket.send_json({"type": "snapshot", "log": initial_snapshot}) disconnect_task = asyncio.create_task(watch_socket_disconnect(websocket)) stream_task = asyncio.create_task( stream_proxy_logs(websocket, path, offset=offset, marker=marker) ) 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.websocket("/{proxy_port}/health/ws") async def proxy_health_socket(websocket: WebSocket, proxy_port: int): disconnect_task: asyncio.Task[None] | None = None queue_task: asyncio.Task[None] | None = None queue = None async def forward_health() -> None: assert queue is not None while True: snapshot = await queue.get() await websocket.send_json(serialize_health_snapshot(snapshot)) 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 await websocket.accept() latest = await get_latest_health_sample(proxy.proxy_port) if latest is None: latest = await check_proxy(proxy) await websocket.send_json(serialize_health_snapshot(latest)) queue = health_monitor.subscribe(proxy.proxy_port) disconnect_task = asyncio.create_task(watch_socket_disconnect(websocket)) queue_task = asyncio.create_task(forward_health()) done, pending = await asyncio.wait( {disconnect_task, queue_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="health stream failed") return except asyncio.CancelledError: return finally: if queue is not None: health_monitor.unsubscribe(proxy_port, queue) for task in (disconnect_task, queue_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) manager.cleanup_runtime_artifacts(proxy) await delete_health_history(proxy.proxy_port) deleted = await delete_proxy(proxy_port) if not deleted: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Proxy not found", )