444 lines
15 KiB
Python
444 lines
15 KiB
Python
import asyncio
|
|
import os
|
|
import shutil
|
|
import signal
|
|
from pathlib import Path
|
|
|
|
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_files import (
|
|
conf_dir_for,
|
|
config_path_for,
|
|
ensure_log_path,
|
|
log_path_for,
|
|
read_cmdline,
|
|
read_log_tail,
|
|
remove_dir,
|
|
remove_file,
|
|
script_path_for,
|
|
write_config,
|
|
write_script,
|
|
)
|
|
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:
|
|
return await write_script(self.work_dir, proxy)
|
|
|
|
async def write_config(self, proxy: ProxyDocument) -> Path:
|
|
return await write_config(self.work_dir, proxy)
|
|
|
|
async def log_path(self, proxy: ProxyDocument) -> Path:
|
|
return await ensure_log_path(self.work_dir, proxy)
|
|
|
|
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]:
|
|
return read_cmdline(pid)
|
|
|
|
def script_path_for(self, proxy: ProxyDocument) -> Path:
|
|
return script_path_for(self.work_dir, proxy)
|
|
|
|
def conf_dir_for(self, proxy: ProxyDocument) -> Path:
|
|
return conf_dir_for(self.work_dir, proxy)
|
|
|
|
def config_path_for(self, proxy: ProxyDocument) -> Path:
|
|
return config_path_for(self.work_dir, proxy)
|
|
|
|
def log_path_for(self, proxy: ProxyDocument) -> Path:
|
|
return log_path_for(self.work_dir, proxy)
|
|
|
|
async def read_log_tail(self, path: Path, limit: int = 4000) -> str:
|
|
return await read_log_tail(path, limit=limit)
|
|
|
|
def remove_file(self, path: Path) -> None:
|
|
remove_file(path)
|
|
|
|
def remove_dir(self, path: Path) -> None:
|
|
remove_dir(path)
|
|
|
|
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()
|