import asyncio import os import shutil import signal from pathlib import Path import yaml from app.integrations.docker import change_port from app.integrations.docker.ports import get_random_port from app.models.proxy import ProxyDocument from app.services.mitmproxy.options import MitmOptionError, validate_options from app.services.mitmproxy.runtime_validation import ( ReservedProxyPortError, ensure_proxy_port_not_reserved, ) class MitmdumpError(Exception): pass class MitmdumpManager: def __init__(self, work_dir: str = "/tmp/theplug", command: str = "mitmdump"): self.work_dir = Path(work_dir) self.command = command self.processes: dict[int, asyncio.subprocess.Process] = {} self.attached_pids: dict[int, int] = {} self.locks: dict[int, asyncio.Lock] = {} self.upstream_port_attempts = 8 async def start(self, proxy: ProxyDocument) -> dict[str, object]: async with self.lock_for(proxy.proxy_port): return await self.start_unlocked(proxy) async def stop(self, proxy: ProxyDocument) -> dict[str, object]: async with self.lock_for(proxy.proxy_port): return await self.stop_unlocked(proxy) def cleanup_runtime_artifacts(self, proxy: ProxyDocument) -> None: script_path = ( Path(proxy.script_path) if proxy.script_path else self.script_path_for(proxy) ) log_path = Path(proxy.log_path) if proxy.log_path else self.log_path_for(proxy) config_path = ( Path(proxy.config_path) if proxy.config_path else self.config_path_for(proxy) ) self.remove_file(script_path) self.remove_file(log_path) self.remove_file(config_path) self.remove_dir(self.conf_dir_for(proxy)) async def start_unlocked(self, proxy: ProxyDocument) -> dict[str, object]: try: ensure_proxy_port_not_reserved(proxy.proxy_port) except ReservedProxyPortError as error: raise MitmdumpError(str(error)) from error if await self.status(proxy) == "running": raise MitmdumpError("Proxy is already running") command_path = shutil.which(self.command) if command_path is None: raise MitmdumpError("mitmdump executable not found in PATH") try: await validate_options(proxy.mitm_options) except MitmOptionError as error: raise MitmdumpError(str(error)) from error upstream_port = await self.prepare_upstream(proxy) script_path = await self.write_script(proxy) config_path = await self.write_config(proxy) log_path = await self.log_path(proxy) command = self.build_command( command_path, proxy, upstream_port, script_path, config_path, ) try: process = await self.spawn(command, log_path) except Exception as error: await self.restore_upstream(proxy, upstream_port) raise MitmdumpError(str(error)) from error failed = await self.exited_early(process) if failed: await self.restore_upstream(proxy, upstream_port) reason = ( await self.read_log_tail(log_path) or "mitmdump exited during startup" ) raise MitmdumpError(reason) self.processes[proxy.proxy_port] = process self.attached_pids.pop(proxy.proxy_port, None) return { "status": "running", "pid": process.pid, "script_path": str(script_path), "config_path": str(config_path), "log_path": str(log_path), "service.upstream_port": upstream_port, "last_error": None, } async def stop_unlocked(self, proxy: ProxyDocument) -> dict[str, object]: process = self.processes.pop(proxy.proxy_port, None) self.attached_pids.pop(proxy.proxy_port, None) if process is not None and process.returncode is None: await self.stop_process(process) elif proxy.pid is not None: await self.stop_pid(proxy.pid) await self.restore_upstream(proxy, proxy.service.upstream_port) return { "status": "stopped", "pid": None, "service.upstream_port": None, "last_error": None, } async def status(self, proxy: ProxyDocument) -> str: process = self.processes.get(proxy.proxy_port) if process is not None: if process.returncode is None: return "running" self.processes.pop(proxy.proxy_port, None) self.attached_pids.pop(proxy.proxy_port, None) return "stopped" pid = self.attached_pids.get(proxy.proxy_port, proxy.pid) if pid is None: return "stopped" if proxy.status == "running" else proxy.status if not self.is_process_alive(pid): self.attached_pids.pop(proxy.proxy_port, None) return "stopped" if not self.matches_proxy_process(proxy, pid): self.attached_pids.pop(proxy.proxy_port, None) return "stopped" return "running" async def reconcile(self, proxy: ProxyDocument) -> dict[str, object] | None: if not self.needs_recovery(proxy): return None if proxy.pid is not None and self.is_process_alive(proxy.pid): if self.matches_proxy_process(proxy, proxy.pid): self.reattach(proxy) return { "status": "running", "pid": proxy.pid, "script_path": str(self.script_path_for(proxy)), "config_path": str(self.config_path_for(proxy)), "log_path": str(self.log_path_for(proxy)), "last_error": None, } self.attached_pids.pop(proxy.proxy_port, None) return { "status": "error", "pid": None, "last_error": ( f"Found a live pid {proxy.pid}, but its command does not match " "the expected mitmdump runtime for this proxy." ), } if self.needs_upstream_restore(proxy): restored = await self.restore_upstream(proxy, proxy.service.upstream_port) if not restored: return { "status": "error", "pid": None, "last_error": ( "Proxy process is gone and the upstream service port " "could not be restored automatically." ), } return { "status": "stopped", "pid": None, "service.upstream_port": None, "last_error": None, } if proxy.status == "running" or proxy.pid is not None: return {"status": "stopped", "pid": None, "last_error": None} return None async def prepare_upstream(self, proxy: ProxyDocument) -> int | None: if proxy.mode_spec is not None or proxy.mode != "reverse": return None if proxy.service.upstream_port is not None: return proxy.service.upstream_port if proxy.service.name is None or proxy.service.folder is None: return proxy.service.port attempted_ports: set[int] = {proxy.proxy_port, proxy.service.port} for _attempt in range(self.upstream_port_attempts): upstream_port = self.pick_upstream_port(attempted_ports) attempted_ports.add(upstream_port) changed = await change_port( proxy.service.name, Path(proxy.service.folder), proxy.service.port, upstream_port, host_to="127.0.0.1", rewrite_host=True, ) if changed: return upstream_port raise MitmdumpError("Unable to move Docker service port after retries") async def restore_upstream( self, proxy: ProxyDocument, upstream_port: int | None, ) -> bool: if upstream_port is None: return True if proxy.service.name is None or proxy.service.folder is None: return True if upstream_port == proxy.service.port: return True return await change_port( proxy.service.name, Path(proxy.service.folder), upstream_port, proxy.service.port, host_to=proxy.service.bind_host, rewrite_host=True, ) def build_command( self, command_path: str, proxy: ProxyDocument, upstream_port: int | None, script_path: Path, config_path: Path, ) -> list[str]: command = [ command_path, "--mode", self.mode_for(proxy, upstream_port), "--listen-host", proxy.listen_host, "--set", f"confdir={config_path.parent}", "-s", str(script_path), ] if self.needs_listen_port(proxy): command.extend(["--listen-port", str(proxy.proxy_port)]) return command def mode_for(self, proxy: ProxyDocument, upstream_port: int | None) -> str: if proxy.mode_spec: return proxy.mode_spec match proxy.mode: case "reverse": port = upstream_port or proxy.service.port return ( f"reverse:{proxy.protocol}://{proxy.service.host}:{port}" f"@{proxy.proxy_port}" ) case "upstream": return ( f"upstream:{proxy.protocol}://{proxy.service.host}:" f"{proxy.service.port}@{proxy.proxy_port}" ) case "wireguard": return f"wireguard@{proxy.proxy_port}" case _: return proxy.mode def needs_listen_port(self, proxy: ProxyDocument) -> bool: if proxy.mode_spec is not None: return False return proxy.mode in {"regular", "socks5", "dns", "transparent"} def pick_upstream_port(self, excluded_ports: set[int]) -> int: while True: candidate = get_random_port() if candidate not in excluded_ports: return candidate async def write_script(self, proxy: ProxyDocument) -> Path: path = self.script_path_for(proxy) await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True) await asyncio.to_thread(path.write_text, proxy.code, encoding="utf-8") return path async def write_config(self, proxy: ProxyDocument) -> Path: path = self.config_path_for(proxy) await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True) options = {"block_global": False, **proxy.mitm_options} config_text = yaml.safe_dump(options, sort_keys=True) await asyncio.to_thread(path.write_text, config_text, encoding="utf-8") return path async def log_path(self, proxy: ProxyDocument) -> Path: path = self.log_path_for(proxy) await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True) return path async def spawn(self, command: list[str], log_path: Path): with log_path.open("ab") as log_file: return await asyncio.create_subprocess_exec( *command, stdout=log_file, stderr=log_file, start_new_session=True, ) async def stop_process(self, process) -> None: await self.stop_pid(process.pid) async def exited_early(self, process) -> bool: try: await asyncio.wait_for(process.wait(), timeout=0.5) except TimeoutError: return False return True async def stop_pid(self, pid: int) -> None: try: pgid = os.getpgid(pid) except ProcessLookupError: return except PermissionError: pgid = pid try: os.killpg(pgid, signal.SIGTERM) except ProcessLookupError: return except PermissionError: try: os.kill(pid, signal.SIGTERM) except ProcessLookupError: return for _ in range(10): await asyncio.sleep(0.2) if not self.is_process_alive(pid): return try: os.killpg(pgid, signal.SIGKILL) except ProcessLookupError: pass except PermissionError: try: os.kill(pid, signal.SIGKILL) except ProcessLookupError: return def reattach(self, proxy: ProxyDocument) -> None: if proxy.pid is not None: self.attached_pids[proxy.proxy_port] = proxy.pid def needs_recovery(self, proxy: ProxyDocument) -> bool: return ( proxy.status == "running" or proxy.pid is not None or self.needs_upstream_restore(proxy) ) def needs_upstream_restore(self, proxy: ProxyDocument) -> bool: upstream_port = proxy.service.upstream_port return upstream_port is not None and upstream_port != proxy.service.port def is_process_alive(self, pid: int) -> bool: try: os.kill(pid, 0) except ProcessLookupError: return False except PermissionError: return True return True def matches_proxy_process(self, proxy: ProxyDocument, pid: int) -> bool: cmdline = self.read_cmdline(pid) if not cmdline: return False expected_command = Path(self.command).name if not any(Path(part).name == expected_command for part in cmdline): return False expected_script = str(self.script_path_for(proxy)) expected_confdir = f"confdir={self.conf_dir_for(proxy)}" return expected_script in cmdline and expected_confdir in cmdline def read_cmdline(self, pid: int) -> list[str]: path = Path(f"/proc/{pid}/cmdline") try: raw = path.read_bytes() except OSError: return [] decoded = raw.decode("utf-8", errors="replace") return [part for part in decoded.split("\x00") if part] def script_path_for(self, proxy: ProxyDocument) -> Path: return self.work_dir / "addons" / f"{proxy.proxy_port}.py" def conf_dir_for(self, proxy: ProxyDocument) -> Path: return self.work_dir / "conf" / str(proxy.proxy_port) def config_path_for(self, proxy: ProxyDocument) -> Path: return self.conf_dir_for(proxy) / "config.yaml" def log_path_for(self, proxy: ProxyDocument) -> Path: return self.work_dir / "logs" / f"{proxy.proxy_port}.log" async def read_log_tail(self, path: Path, limit: int = 4000) -> str: if not await asyncio.to_thread(path.exists): return "" content = await asyncio.to_thread( path.read_text, encoding="utf-8", errors="replace", ) return content[-limit:].strip() def remove_file(self, path: Path) -> None: try: path.unlink(missing_ok=True) except OSError: return def remove_dir(self, path: Path) -> None: try: shutil.rmtree(path) except FileNotFoundError: return except OSError: return def lock_for(self, proxy_port: int) -> asyncio.Lock: lock = self.locks.get(proxy_port) if lock is None: lock = asyncio.Lock() self.locks[proxy_port] = lock return lock manager = MitmdumpManager()