981 lines
29 KiB
Python
981 lines
29 KiB
Python
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"]
|