493 lines
14 KiB
Python
493 lines
14 KiB
Python
import importlib
|
|
|
|
from core.security import get_auth_context, require_authenticated_user
|
|
from models.proxy import ProxyCreateInput, ProxyDocument
|
|
from models.service import ServiceEndpointDocument
|
|
from services.mitmproxy.examples import BUILTIN_EXAMPLES
|
|
from services.mitmproxy.validator import ProxyCodeValidationError
|
|
|
|
proxy_router_module = importlib.import_module("api.proxy.router")
|
|
|
|
|
|
def make_proxy_document(
|
|
*,
|
|
proxy_port: int,
|
|
name: str = "svc",
|
|
code: str = "proxy-code",
|
|
port: int = 9000,
|
|
) -> ProxyDocument:
|
|
return ProxyDocument(
|
|
proxy_port=proxy_port,
|
|
name=name,
|
|
code=code,
|
|
service=ServiceEndpointDocument(port=port),
|
|
)
|
|
|
|
|
|
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() == {
|
|
"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=3671)
|
|
|
|
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
|
|
assert proxy_port == 3671
|
|
return existing
|
|
|
|
monkeypatch.setattr(
|
|
proxy_router_module,
|
|
"get_proxy_by_port",
|
|
fake_get_proxy_by_port,
|
|
)
|
|
|
|
response = client.post(
|
|
"/proxy",
|
|
json={
|
|
"proxy_port": 3671,
|
|
"name": "new",
|
|
"code": "abc",
|
|
"service": {"port": 8080},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 409
|
|
assert response.json()["detail"] == "Proxy already exists on port 3671"
|
|
|
|
|
|
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_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_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_examples_returns_builtins_and_custom(client, monkeypatch):
|
|
async def fake_list_custom_examples():
|
|
return []
|
|
|
|
monkeypatch.setattr(
|
|
proxy_router_module,
|
|
"list_custom_examples",
|
|
fake_list_custom_examples,
|
|
)
|
|
|
|
response = client.get("/proxy/mitmproxy/examples")
|
|
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert len(body) == len(BUILTIN_EXAMPLES)
|
|
assert body[0]["built_in"] is True
|
|
|
|
|
|
def test_create_mitmproxy_example_validates_and_persists(client, monkeypatch):
|
|
captured = {}
|
|
|
|
async def fake_validate_proxy_code(code: str) -> None:
|
|
captured["validated"] = code
|
|
|
|
async def fake_create_custom_example(payload):
|
|
captured["payload"] = payload
|
|
return {
|
|
"id": payload.id,
|
|
"name": payload.name,
|
|
"description": payload.description,
|
|
"category": payload.category,
|
|
"code": payload.code,
|
|
"source_url": payload.source_url,
|
|
"built_in": False,
|
|
}
|
|
|
|
monkeypatch.setattr(
|
|
proxy_router_module,
|
|
"validate_proxy_code",
|
|
fake_validate_proxy_code,
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_router_module,
|
|
"create_custom_example",
|
|
fake_create_custom_example,
|
|
)
|
|
|
|
response = client.post(
|
|
"/proxy/mitmproxy/examples",
|
|
json={
|
|
"id": "custom-header",
|
|
"name": "Custom header",
|
|
"description": "Adds a custom response header",
|
|
"category": "http",
|
|
"code": "def response(flow):\n pass\n",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 201
|
|
assert response.json()["id"] == "custom-header"
|
|
assert captured["validated"] == "def response(flow):\n pass\n"
|