import asyncio import importlib from datetime import UTC, datetime import pytest from starlette.websockets import WebSocketDisconnect from app.core.security import get_auth_context, require_authenticated_user from app.models.health import HealthSnapshot from app.models.proxy import ProxyCreateInput, ProxyDocument from app.models.service import ServiceEndpointDocument from app.services.mitmproxy.templates import list_file_templates from app.services.mitmproxy.validator import ProxyCodeValidationError proxy_router_module = importlib.import_module("app.api.proxy.router") def make_proxy_document( *, proxy_port: int, name: str = "svc", code: str = "proxy-code", port: int = 9000, log_path: str | None = None, ) -> ProxyDocument: return ProxyDocument( proxy_port=proxy_port, name=name, code=code, service=ServiceEndpointDocument(port=port), log_path=log_path, ) def make_health_snapshot( proxy_port: int, *, overall: str = "healthy", timestamp: datetime | None = None, ) -> HealthSnapshot: return HealthSnapshot( proxy_port=proxy_port, timestamp=timestamp or datetime.now(UTC), enabled=True, overall=overall, checks=[], ) def test_proxy_requires_auth_without_dependency_overrides(client): client.app.dependency_overrides.pop(require_authenticated_user, None) client.app.dependency_overrides.pop(get_auth_context, None) response = client.get("/proxy") assert response.status_code == 401 assert response.json()["detail"] == "Could not validate credentials" def test_list_proxies_returns_documents(client, monkeypatch: object): docs = [ make_proxy_document(proxy_port=8080, name="api"), make_proxy_document(proxy_port=8081, name="worker"), ] async def fake_list_proxies() -> list[ProxyDocument]: return docs monkeypatch.setattr(proxy_router_module, "list_proxy_documents", fake_list_proxies) response = client.get("/proxy") assert response.status_code == 200 body = response.json() assert [item["proxy_port"] for item in body] == [8080, 8081] def test_list_services_uses_query_scan_root(client, monkeypatch): captured = {} def fake_get_all_services_in_dir(root): captured["root"] = str(root) return [ { "name": "echo", "ports": [9081], "bindings": [], "folder": "/root/play", } ] monkeypatch.setattr( proxy_router_module, "get_all_services_in_dir", fake_get_all_services_in_dir, ) response = client.get("/proxy/services", params={"scan_root": "/root"}) assert response.status_code == 200 assert response.json() == { "backend_port": 3671, "scan_root": "/root", "services": [ {"name": "echo", "ports": [9081], "bindings": [], "folder": "/root/play"} ], } assert captured["root"] == "/root" def test_get_proxy_returns_404_when_not_found(client, monkeypatch: object): async def fake_get_proxy_by_port(proxy_port: int) -> None: return None monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) response = client.get("/proxy/9999") assert response.status_code == 404 assert response.json()["detail"] == "Proxy not found" def test_create_proxy_returns_409_on_port_conflict(client, monkeypatch: object): existing = make_proxy_document(proxy_port=4088) async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4088 return existing monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) response = client.post( "/proxy", json={ "proxy_port": 4088, "name": "new", "code": "abc", "service": {"port": 8080}, }, ) assert response.status_code == 409 assert response.json()["detail"] == "Proxy already exists on port 4088" def test_create_proxy_rejects_backend_api_port(client, monkeypatch: object): create_called = False async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 3671 return None async def fake_create_proxy(payload: ProxyCreateInput) -> ProxyDocument: nonlocal create_called create_called = True return make_proxy_document(proxy_port=payload.proxy_port) monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr(proxy_router_module, "create_proxy", fake_create_proxy) response = client.post( "/proxy", json={ "proxy_port": 3671, "name": "backend-clash", "code": "abc", "service": {"port": 8080}, }, ) assert response.status_code == 422 assert response.json()["detail"] == ( "Proxy port 3671 is reserved for the backend API " "and cannot be used by mitmdump." ) assert create_called is False def test_create_proxy_returns_created_document(client, monkeypatch: object): captured: dict[str, object] = {} async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return None async def fake_create_proxy(payload: ProxyCreateInput) -> ProxyDocument: captured["proxy_port"] = payload.proxy_port captured["name"] = payload.name captured["service"] = payload.service.model_dump(exclude_none=True) captured["code"] = payload.code captured["protocol"] = payload.protocol return make_proxy_document( proxy_port=payload.proxy_port, name=payload.name, code=payload.code, port=payload.service.port, ) async def fake_validate_proxy_code(code: str) -> None: assert code == "def proxy_handler():\n return 'ok'\n" monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr( proxy_router_module, "validate_proxy_code", fake_validate_proxy_code, ) monkeypatch.setattr(proxy_router_module, "create_proxy", fake_create_proxy) response = client.post( "/proxy", json={ "proxy_port": 4040, "name": "edge", "code": "def proxy_handler():\n return 'ok'\n", "service": {"port": 8181}, }, ) assert response.status_code == 201 assert response.json()["proxy_port"] == 4040 assert captured == { "proxy_port": 4040, "name": "edge", "service": {"host": "127.0.0.1", "port": 8181}, "code": "def proxy_handler():\n return 'ok'\n", "protocol": "http", } def test_create_proxy_returns_422_when_code_validation_fails( client, monkeypatch: object, ): create_called = False async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4000 return None async def fake_validate_proxy_code(code: str) -> None: assert code == "invalid-code" raise ProxyCodeValidationError( "Proxy code contains ruff violations", errors=[ { "code": "F821", "message": "Undefined name `missing_symbol`", "location": {"row": 1, "column": 1}, } ], ) async def fake_create_proxy( *, proxy_port: int, name: str, service: dict[str, object], code: str, ) -> ProxyDocument: nonlocal create_called create_called = True return make_proxy_document(proxy_port=proxy_port, name=name, code=code) monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr( proxy_router_module, "validate_proxy_code", fake_validate_proxy_code, ) monkeypatch.setattr(proxy_router_module, "create_proxy", fake_create_proxy) response = client.post( "/proxy", json={ "proxy_port": 4000, "name": "broken", "code": "invalid-code", "service": {"port": 9000}, }, ) assert response.status_code == 422 body = response.json()["detail"] assert body["message"] == "Proxy code contains ruff violations" assert body["errors"][0]["code"] == "F821" assert create_called is False def test_update_proxy_returns_422_when_code_validation_fails( client, monkeypatch: object, ): update_called = False async def fake_validate_proxy_code(code: str) -> None: assert code == "invalid-code" raise ProxyCodeValidationError( "Proxy code has invalid Python syntax", errors=[ { "code": "COMPILE", "message": "invalid syntax", "location": {"row": 1, "column": 1}, } ], ) async def fake_update_proxy( proxy_port: int, payload: object, ) -> ProxyDocument | None: nonlocal update_called update_called = True return make_proxy_document(proxy_port=proxy_port) monkeypatch.setattr( proxy_router_module, "validate_proxy_code", fake_validate_proxy_code, ) monkeypatch.setattr(proxy_router_module, "update_proxy", fake_update_proxy) response = client.put( "/proxy/8123", json={"code": "invalid-code"}, ) assert response.status_code == 422 body = response.json()["detail"] assert body["message"] == "Proxy code has invalid Python syntax" assert body["errors"][0]["code"] == "COMPILE" assert update_called is False def test_update_proxy_persists_code_changes(client, monkeypatch: object): captured: dict[str, object] = {} updated_code = "def response(flow):\n flow.response.text = 'ok'\n" async def fake_validate_proxy_code(code: str) -> None: captured["validated"] = code async def fake_update_proxy( proxy_port: int, payload: object, ) -> ProxyDocument | None: captured["proxy_port"] = proxy_port captured["code"] = payload.code return make_proxy_document( proxy_port=proxy_port, code=payload.code, ) monkeypatch.setattr( proxy_router_module, "validate_proxy_code", fake_validate_proxy_code, ) monkeypatch.setattr(proxy_router_module, "update_proxy", fake_update_proxy) response = client.put( "/proxy/8123", json={"code": updated_code}, ) assert response.status_code == 200 assert response.json()["code"] == updated_code assert captured == { "validated": updated_code, "proxy_port": 8123, "code": updated_code, } def test_start_proxy_updates_runtime_state(client, monkeypatch: object): proxy = make_proxy_document(proxy_port=4040) captured: dict[str, object] = {} async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy class FakeManager: async def start(self, document: ProxyDocument) -> dict[str, object]: assert document.proxy_port == 4040 return { "status": "running", "pid": 1234, "script_path": "/tmp/theplug/addons/4040.py", "log_path": "/tmp/theplug/logs/4040.log", "service.upstream_port": 9000, "last_error": None, } async def fake_update_proxy_runtime( proxy_port: int, values: dict[str, object], ) -> ProxyDocument: captured["proxy_port"] = proxy_port captured["values"] = values return proxy.model_copy(update={"status": "running", "pid": 1234}) monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr(proxy_router_module, "manager", FakeManager()) monkeypatch.setattr( proxy_router_module, "update_proxy_runtime", fake_update_proxy_runtime, ) response = client.post("/proxy/4040/start") assert response.status_code == 200 assert response.json()["status"] == "running" assert captured["proxy_port"] == 4040 assert captured["values"]["pid"] == 1234 def test_start_proxy_records_runtime_error(client, monkeypatch: object): proxy = make_proxy_document(proxy_port=4040) async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy class FakeManager: async def start(self, document: ProxyDocument) -> dict[str, object]: assert document.proxy_port == 4040 raise proxy_router_module.MitmdumpError("mitmdump executable not found") async def fake_update_proxy_runtime( proxy_port: int, values: dict[str, object], ) -> ProxyDocument: assert proxy_port == 4040 assert values == { "status": "error", "last_error": "mitmdump executable not found", } return proxy.model_copy( update={"status": "error", "last_error": values["last_error"]} ) monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr(proxy_router_module, "manager", FakeManager()) monkeypatch.setattr( proxy_router_module, "update_proxy_runtime", fake_update_proxy_runtime, ) response = client.post("/proxy/4040/start") assert response.status_code == 200 assert response.json()["status"] == "error" assert response.json()["last_error"] == "mitmdump executable not found" def test_proxy_logs_socket_rejects_missing_token(client): with pytest.raises(WebSocketDisconnect) as error: with client.websocket_connect("/proxy/4040/logs/ws"): pass assert error.value.code == 4401 def test_proxy_logs_socket_streams_snapshot_and_append(client, monkeypatch, tmp_path): log_path = tmp_path / "proxy.log" log_path.write_text("line one\n", encoding="utf-8") proxy = make_proxy_document(proxy_port=4040, log_path=str(log_path)) async def fake_authenticate(_: object) -> None: return None async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy monkeypatch.setattr( proxy_router_module, "authenticate_proxy_socket", fake_authenticate, ) monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) with client.websocket_connect("/proxy/4040/logs/ws?token=test-token") as socket: snapshot = socket.receive_json() assert snapshot == {"type": "snapshot", "log": "line one\n"} with log_path.open("a", encoding="utf-8") as handle: handle.write("line two\n") handle.flush() appended = socket.receive_json() assert appended == {"type": "append", "log": "line two\n"} def test_proxy_logs_socket_resets_snapshot_after_log_rotation( client, monkeypatch, tmp_path, ): log_path = tmp_path / "proxy.log" log_path.write_text("line one\n", encoding="utf-8") proxy = make_proxy_document(proxy_port=4040, log_path=str(log_path)) async def fake_authenticate(_: object) -> None: return None async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy monkeypatch.setattr( proxy_router_module, "authenticate_proxy_socket", fake_authenticate, ) monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) with client.websocket_connect("/proxy/4040/logs/ws?token=test-token") as socket: snapshot = socket.receive_json() assert snapshot == {"type": "snapshot", "log": "line one\n"} log_path.write_text("fresh start\n", encoding="utf-8") rotated = socket.receive_json() assert rotated == {"type": "snapshot", "log": "fresh start\n"} with log_path.open("a", encoding="utf-8") as handle: handle.write("line two\n") handle.flush() appended = socket.receive_json() assert appended == {"type": "append", "log": "line two\n"} def test_remove_proxy_stops_runtime_and_cleans_artifacts(client, monkeypatch: object): proxy = make_proxy_document( proxy_port=4040, log_path="/tmp/theplug/logs/4040.log", ).model_copy( update={ "script_path": "/tmp/theplug/addons/4040.py", "config_path": "/tmp/theplug/conf/4040/config.yaml", } ) captured: dict[str, object] = {} async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy class FakeManager: async def stop(self, document: ProxyDocument) -> dict[str, object]: captured["stopped"] = document.proxy_port return { "status": "stopped", "pid": None, "service.upstream_port": None, "last_error": None, } def cleanup_runtime_artifacts(self, document: ProxyDocument) -> None: captured["cleaned"] = { "proxy_port": document.proxy_port, "script_path": document.script_path, "config_path": document.config_path, "log_path": document.log_path, } async def fake_delete_proxy(proxy_port: int) -> bool: captured["deleted"] = proxy_port return True async def fake_delete_health_history(proxy_port: int) -> None: captured["history_deleted"] = proxy_port monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr(proxy_router_module, "manager", FakeManager()) monkeypatch.setattr( proxy_router_module, "delete_health_history", fake_delete_health_history, ) monkeypatch.setattr(proxy_router_module, "delete_proxy", fake_delete_proxy) response = client.delete("/proxy/4040") assert response.status_code == 204 assert captured == { "stopped": 4040, "cleaned": { "proxy_port": 4040, "script_path": "/tmp/theplug/addons/4040.py", "config_path": "/tmp/theplug/conf/4040/config.yaml", "log_path": "/tmp/theplug/logs/4040.log", }, "deleted": 4040, "history_deleted": 4040, } def test_proxy_health_runs_immediate_check_and_publishes(client, monkeypatch: object): proxy = make_proxy_document(proxy_port=4040) captured: dict[str, object] = {} snapshot = make_health_snapshot(4040) async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy async def fake_check_proxy(document: ProxyDocument): captured["checked"] = document.proxy_port return snapshot async def fake_record_health_sample(document): captured["recorded"] = document return document class FakeHealthMonitor: async def publish(self, proxy_port: int, document: HealthSnapshot) -> None: captured["published"] = (proxy_port, document) monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr(proxy_router_module, "check_proxy", fake_check_proxy) monkeypatch.setattr( proxy_router_module, "record_health_sample", fake_record_health_sample, ) monkeypatch.setattr(proxy_router_module, "health_monitor", FakeHealthMonitor()) response = client.get("/proxy/4040/health") assert response.status_code == 200 assert response.json()["overall"] == "healthy" assert captured["checked"] == 4040 assert captured["recorded"] == snapshot assert captured["published"] == (4040, snapshot) def test_proxy_health_history_returns_latest_and_history(client, monkeypatch: object): proxy = make_proxy_document(proxy_port=4040) timestamp = datetime.now(UTC) async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy async def fake_get_latest_health_sample(proxy_port: int): assert proxy_port == 4040 return make_health_snapshot(4040, timestamp=timestamp) async def fake_list_health_history(proxy_port: int, limit: int = 40): assert proxy_port == 4040 assert limit == 40 return [make_health_snapshot(4040, timestamp=timestamp)] monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr( proxy_router_module, "get_latest_health_sample", fake_get_latest_health_sample, ) monkeypatch.setattr( proxy_router_module, "list_health_history", fake_list_health_history, ) response = client.get("/proxy/4040/health/history") assert response.status_code == 200 body = response.json() assert body["proxy_port"] == 4040 assert body["latest"]["overall"] == "healthy" assert len(body["history"]) == 1 def test_proxy_health_history_clamps_requested_limit(client, monkeypatch: object): proxy = make_proxy_document(proxy_port=4040) captured: dict[str, int] = {} async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy async def fake_get_latest_health_sample(_proxy_port: int) -> HealthSnapshot | None: return None async def fake_list_health_history(proxy_port: int, limit: int = 40): assert proxy_port == 4040 captured["limit"] = limit return [] monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr( proxy_router_module, "get_latest_health_sample", fake_get_latest_health_sample, ) monkeypatch.setattr( proxy_router_module, "list_health_history", fake_list_health_history, ) response = client.get("/proxy/4040/health/history", params={"limit": 9999}) assert response.status_code == 200 assert captured["limit"] == 200 def test_proxy_health_socket_streams_latest_and_updates(client, monkeypatch: object): proxy = make_proxy_document(proxy_port=4040) timestamp = datetime.now(UTC) follow_up = make_health_snapshot(4040, overall="unhealthy") captured: dict[str, object] = {} async def fake_authenticate(_: object) -> None: return None async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy async def fake_get_latest_health_sample(proxy_port: int): assert proxy_port == 4040 return make_health_snapshot(4040, timestamp=timestamp) class FakeHealthMonitor: def subscribe(self, proxy_port: int): assert proxy_port == 4040 queue = asyncio.Queue() queue.put_nowait(follow_up) return queue def unsubscribe(self, proxy_port: int, _queue) -> None: assert proxy_port == 4040 captured["unsubscribed"] = True monkeypatch.setattr( proxy_router_module, "authenticate_proxy_socket", fake_authenticate, ) monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr( proxy_router_module, "get_latest_health_sample", fake_get_latest_health_sample, ) monkeypatch.setattr(proxy_router_module, "health_monitor", FakeHealthMonitor()) with client.websocket_connect("/proxy/4040/health/ws?token=test-token") as socket: first = socket.receive_json() assert first["overall"] == "healthy" second = socket.receive_json() assert second["overall"] == "unhealthy" assert captured["unsubscribed"] is True def test_stop_proxy_updates_runtime_state(client, monkeypatch: object): proxy = make_proxy_document(proxy_port=4040).model_copy( update={"status": "running", "pid": 1234} ) async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None: assert proxy_port == 4040 return proxy class FakeManager: async def stop(self, document: ProxyDocument) -> dict[str, object]: assert document.pid == 1234 return { "status": "stopped", "pid": None, "service.upstream_port": None, "last_error": None, } async def fake_update_proxy_runtime( proxy_port: int, values: dict[str, object], ) -> ProxyDocument: assert proxy_port == 4040 assert values["status"] == "stopped" return proxy.model_copy(update={"status": "stopped", "pid": None}) monkeypatch.setattr( proxy_router_module, "get_proxy_by_port", fake_get_proxy_by_port, ) monkeypatch.setattr(proxy_router_module, "manager", FakeManager()) monkeypatch.setattr( proxy_router_module, "update_proxy_runtime", fake_update_proxy_runtime, ) response = client.post("/proxy/4040/stop") assert response.status_code == 200 assert response.json()["status"] == "stopped" def test_list_mitmproxy_templates_returns_builtins_and_custom(client, monkeypatch): file_templates = list_file_templates() async def fake_list_custom_templates(): return [] monkeypatch.setattr( proxy_router_module, "list_custom_templates", fake_list_custom_templates, ) monkeypatch.setattr( proxy_router_module, "list_file_templates", lambda: file_templates, ) response = client.get("/proxy/mitmproxy/templates") assert response.status_code == 200 body = response.json() assert len(body) == len(file_templates) assert body[0]["built_in"] is True assert body[0]["origin"] == "file" assert body[0]["protocol"] in {"http", "tcp"} assert isinstance(body[0]["tags"], list) assert body[0]["action"] in {"block", "rewrite"} def test_create_mitmproxy_template_validates_and_persists(client, monkeypatch): captured = {} async def fake_validate_proxy_code(code: str) -> None: captured["validated"] = code async def fake_create_custom_template(payload): captured["payload"] = payload return { "id": payload.id, "name": payload.name, "description": payload.description, "category": payload.category, "protocol": payload.protocol, "family": payload.family, "tags": payload.tags, "action": payload.action, "summary": payload.summary, "behavior": payload.behavior, "code": payload.code, "source_url": payload.source_url, "built_in": False, "origin": "custom", } monkeypatch.setattr( proxy_router_module, "validate_proxy_code", fake_validate_proxy_code, ) monkeypatch.setattr( proxy_router_module, "create_custom_template", fake_create_custom_template, ) response = client.post( "/proxy/mitmproxy/templates", json={ "id": "custom-header", "name": "Custom header", "description": "Adds a custom response header", "category": "http", "protocol": "http", "family": "custom", "tags": ["observe"], "action": "custom", "summary": "Saved custom template.", "behavior": "Applies the current script as-is.", "code": "def response(flow):\n pass\n", }, ) assert response.status_code == 201 assert response.json()["id"] == "custom-header" assert response.json()["protocol"] == "http" assert response.json()["origin"] == "custom" assert captured["validated"] == "def response(flow):\n pass\n" def test_create_mitmproxy_template_rejects_file_backed_id(client, monkeypatch): monkeypatch.setattr( proxy_router_module, "file_template_ids", lambda: {"http-block-path-traversal"}, ) response = client.post( "/proxy/mitmproxy/templates", json={ "id": "http-block-path-traversal", "name": "Custom header", "description": "Adds a custom response header", "category": "http", "protocol": "http", "family": "custom", "tags": [], "action": "custom", "summary": "Saved custom template.", "behavior": "Applies the current script as-is.", "code": "def response(flow):\n pass\n", }, ) assert response.status_code == 409 assert "Template already exists" in response.json()["detail"]