284 lines
9.2 KiB
Python
284 lines
9.2 KiB
Python
import asyncio
|
|
import os
|
|
import shutil
|
|
import signal
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
from integrations import docker
|
|
from integrations.docker.ports import get_random_port
|
|
from models.proxy import ProxyDocument
|
|
from services.mitmproxy.options import MitmOptionError, validate_options
|
|
|
|
|
|
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.locks: dict[int, asyncio.Lock] = {}
|
|
|
|
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)
|
|
|
|
async def start_unlocked(self, proxy: ProxyDocument) -> dict[str, object]:
|
|
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 = self.write_script(proxy)
|
|
config_path = self.write_config(proxy)
|
|
log_path = 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 = self.read_log_tail(log_path) or "mitmdump exited during startup"
|
|
raise MitmdumpError(reason)
|
|
|
|
self.processes[proxy.proxy_port] = process
|
|
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)
|
|
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)
|
|
return "stopped"
|
|
|
|
if proxy.pid is None:
|
|
return "stopped" if proxy.status == "running" else proxy.status
|
|
|
|
try:
|
|
os.kill(proxy.pid, 0)
|
|
except ProcessLookupError:
|
|
return "stopped"
|
|
except PermissionError:
|
|
return "running"
|
|
return "running"
|
|
|
|
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
|
|
|
|
upstream_port = get_random_port()
|
|
changed = await docker.change_port(
|
|
proxy.service.name,
|
|
Path(proxy.service.folder),
|
|
proxy.service.port,
|
|
upstream_port,
|
|
)
|
|
if not changed:
|
|
raise MitmdumpError("Unable to move Docker service port")
|
|
return upstream_port
|
|
|
|
async def restore_upstream(
|
|
self,
|
|
proxy: ProxyDocument,
|
|
upstream_port: int | None,
|
|
) -> None:
|
|
if upstream_port is None:
|
|
return
|
|
if proxy.service.name is None or proxy.service.folder is None:
|
|
return
|
|
if upstream_port == proxy.service.port:
|
|
return
|
|
|
|
await docker.change_port(
|
|
proxy.service.name,
|
|
Path(proxy.service.folder),
|
|
upstream_port,
|
|
proxy.service.port,
|
|
)
|
|
|
|
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 write_script(self, proxy: ProxyDocument) -> Path:
|
|
path = self.work_dir / "addons" / f"{proxy.proxy_port}.py"
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(proxy.code, encoding="utf-8")
|
|
return path
|
|
|
|
def write_config(self, proxy: ProxyDocument) -> Path:
|
|
path = self.work_dir / "conf" / str(proxy.proxy_port) / "config.yaml"
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
options = {"block_global": False, **proxy.mitm_options}
|
|
path.write_text(yaml.safe_dump(options, sort_keys=True), encoding="utf-8")
|
|
return path
|
|
|
|
def log_path(self, proxy: ProxyDocument) -> Path:
|
|
path = self.work_dir / "logs" / f"{proxy.proxy_port}.log"
|
|
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:
|
|
try:
|
|
os.killpg(process.pid, signal.SIGTERM)
|
|
except ProcessLookupError:
|
|
return
|
|
|
|
try:
|
|
await asyncio.wait_for(process.wait(), timeout=5)
|
|
except TimeoutError:
|
|
try:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
return
|
|
await process.wait()
|
|
|
|
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:
|
|
os.kill(pid, signal.SIGTERM)
|
|
except ProcessLookupError:
|
|
return
|
|
|
|
for _ in range(10):
|
|
await asyncio.sleep(0.2)
|
|
try:
|
|
os.kill(pid, 0)
|
|
except ProcessLookupError:
|
|
return
|
|
|
|
try:
|
|
os.kill(pid, signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
return
|
|
|
|
def read_log_tail(self, path: Path, limit: int = 4000) -> str:
|
|
if not path.exists():
|
|
return ""
|
|
content = path.read_text(encoding="utf-8", errors="replace")
|
|
return content[-limit:].strip()
|
|
|
|
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()
|