BIG BIG BIG BIG BIG UPDATES
Co-authored-by: Copilot <[email protected]>
This commit is contained in:
gabdevele
co-authored by
Copilot
Unverified
Copilot
parent
46f7cd6a5f
commit
93ba8e9ac2
@@ -0,0 +1,283 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user