import asyncio from pathlib import Path from models.proxy import ProxyDocument from models.service import ServiceEndpointDocument from services.mitmproxy.manager import MitmdumpManager def make_proxy(**updates): data = { "proxy_port": 8080, "name": "svc", "code": "", "service": ServiceEndpointDocument(port=9000), } data.update(updates) return ProxyDocument(**data) def test_reverse_mode_is_built_from_service(): manager = MitmdumpManager() proxy = make_proxy(protocol="tcp") command = manager.build_command( "mitmdump", proxy, 9100, Path("/tmp/addon.py"), Path("/tmp/conf/config.yaml"), ) assert "reverse:tcp://127.0.0.1:9100@8080" in command assert "--listen-port" not in command def test_raw_mode_spec_overrides_generated_mode(): manager = MitmdumpManager() proxy = make_proxy(mode="regular", mode_spec="wireguard:/tmp/wg.conf@51820") command = manager.build_command( "mitmdump", proxy, None, Path("/tmp/addon.py"), Path("/tmp/conf/config.yaml"), ) assert "wireguard:/tmp/wg.conf@51820" in command assert "--listen-port" not in command def test_regular_mode_uses_proxy_port_as_listen_port(): manager = MitmdumpManager() proxy = make_proxy(mode="regular") command = manager.build_command( "mitmdump", proxy, None, Path("/tmp/addon.py"), Path("/tmp/conf/config.yaml"), ) assert "regular" in command assert command[-2:] == ["--listen-port", "8080"] def test_running_document_without_pid_is_treated_as_stopped(): manager = MitmdumpManager() proxy = make_proxy(status="running", pid=None) async def run(): return await manager.status(proxy) assert asyncio.run(run()) == "stopped" def test_matches_proxy_process_uses_expected_script_and_confdir(): manager = MitmdumpManager() proxy = make_proxy() manager.read_cmdline = lambda _pid: [ "/home/test/.venv/bin/python", "/home/test/.venv/bin/mitmdump", "--mode", "reverse:http://127.0.0.1:9000@8080", "--listen-host", "0.0.0.0", "--set", "confdir=/tmp/theplug/conf/8080", "-s", "/tmp/theplug/addons/8080.py", ] assert manager.matches_proxy_process(proxy, 1234) is True def test_reconcile_reattaches_matching_process(): manager = MitmdumpManager() proxy = make_proxy( status="running", pid=4321, service=ServiceEndpointDocument(port=9000, upstream_port=54003), ) manager.is_process_alive = lambda _pid: True manager.matches_proxy_process = lambda _proxy, _pid: True async def run(): return await manager.reconcile(proxy) values = asyncio.run(run()) assert values is not None assert values["status"] == "running" assert values["pid"] == 4321 assert manager.attached_pids[8080] == 4321 def test_reconcile_marks_error_for_mismatched_live_process(): manager = MitmdumpManager() proxy = make_proxy( status="running", pid=4321, service=ServiceEndpointDocument(port=9000, upstream_port=54003), ) manager.is_process_alive = lambda _pid: True manager.matches_proxy_process = lambda _proxy, _pid: False async def run(): return await manager.reconcile(proxy) values = asyncio.run(run()) assert values is not None assert values["status"] == "error" assert values["pid"] is None assert "does not match" in str(values["last_error"]) def test_reconcile_restores_upstream_when_process_is_gone(): manager = MitmdumpManager() proxy = make_proxy( status="running", pid=4321, service=ServiceEndpointDocument( port=9000, upstream_port=54003, name="echo", folder="/tmp/playground", ), ) manager.is_process_alive = lambda _pid: False async def fake_restore(_proxy, _upstream_port): return True manager.restore_upstream = fake_restore async def run(): return await manager.reconcile(proxy) values = asyncio.run(run()) assert values == { "status": "stopped", "pid": None, "service.upstream_port": None, "last_error": None, }