282 lines
7.5 KiB
Python
282 lines
7.5 KiB
Python
import asyncio
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from app.models.proxy import ProxyDocument
|
|
from app.models.service import ServiceEndpointDocument
|
|
from app.services.mitmproxy.manager import MitmdumpError, 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_start_rejects_reserved_backend_port():
|
|
manager = MitmdumpManager()
|
|
proxy = make_proxy(proxy_port=3671)
|
|
|
|
async def run():
|
|
await manager.start(proxy)
|
|
|
|
with pytest.raises(
|
|
MitmdumpError,
|
|
match="Proxy port 3671 is reserved for the backend API",
|
|
):
|
|
asyncio.run(run())
|
|
|
|
|
|
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,
|
|
}
|
|
|
|
|
|
def test_cleanup_runtime_artifacts_removes_generated_files(tmp_path: Path):
|
|
manager = MitmdumpManager(work_dir=str(tmp_path / "theplug"))
|
|
proxy = make_proxy(
|
|
script_path=str(manager.script_path_for(make_proxy())),
|
|
config_path=str(manager.config_path_for(make_proxy())),
|
|
log_path=str(manager.log_path_for(make_proxy())),
|
|
)
|
|
|
|
script_path = Path(proxy.script_path)
|
|
config_path = Path(proxy.config_path)
|
|
log_path = Path(proxy.log_path)
|
|
script_path.parent.mkdir(parents=True, exist_ok=True)
|
|
config_path.parent.mkdir(parents=True, exist_ok=True)
|
|
log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
script_path.write_text("print('hello')", encoding="utf-8")
|
|
config_path.write_text("block_global: false\n", encoding="utf-8")
|
|
log_path.write_text("old log\n", encoding="utf-8")
|
|
|
|
manager.cleanup_runtime_artifacts(proxy)
|
|
|
|
assert script_path.exists() is False
|
|
assert log_path.exists() is False
|
|
assert config_path.exists() is False
|
|
assert manager.conf_dir_for(proxy).exists() is False
|
|
|
|
|
|
def test_prepare_upstream_retries_until_port_remap_succeeds(monkeypatch):
|
|
manager = MitmdumpManager()
|
|
proxy = make_proxy(
|
|
mode="reverse",
|
|
service=ServiceEndpointDocument(
|
|
port=9000,
|
|
name="echo",
|
|
folder="/tmp/playground",
|
|
),
|
|
)
|
|
attempted_ports: list[int] = []
|
|
candidate_ports = iter([54001, 54002, 54003])
|
|
|
|
async def fake_change_port(
|
|
_service_name,
|
|
_folder,
|
|
_port_from,
|
|
port_to,
|
|
**_kwargs,
|
|
):
|
|
attempted_ports.append(port_to)
|
|
return port_to == 54003
|
|
|
|
monkeypatch.setattr(
|
|
"app.services.mitmproxy.manager.get_random_port",
|
|
lambda: next(candidate_ports),
|
|
)
|
|
monkeypatch.setattr(
|
|
"app.services.mitmproxy.manager.change_port",
|
|
fake_change_port,
|
|
)
|
|
|
|
upstream_port = asyncio.run(manager.prepare_upstream(proxy))
|
|
|
|
assert upstream_port == 54003
|
|
assert attempted_ports == [54001, 54002, 54003]
|
|
|
|
|
|
def test_prepare_upstream_fails_after_retry_budget(monkeypatch):
|
|
manager = MitmdumpManager()
|
|
manager.upstream_port_attempts = 3
|
|
proxy = make_proxy(
|
|
mode="reverse",
|
|
service=ServiceEndpointDocument(
|
|
port=9000,
|
|
name="echo",
|
|
folder="/tmp/playground",
|
|
),
|
|
)
|
|
candidate_ports = iter([54001, 54002, 54003])
|
|
|
|
async def fake_change_port(*_args, **_kwargs):
|
|
return False
|
|
|
|
monkeypatch.setattr(
|
|
"app.services.mitmproxy.manager.get_random_port",
|
|
lambda: next(candidate_ports),
|
|
)
|
|
monkeypatch.setattr(
|
|
"app.services.mitmproxy.manager.change_port",
|
|
fake_change_port,
|
|
)
|
|
|
|
with pytest.raises(
|
|
MitmdumpError,
|
|
match="Unable to move Docker service port after retries",
|
|
):
|
|
asyncio.run(manager.prepare_upstream(proxy))
|