reloia/thePlug
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: enhance service discovery and proxy management

Browse Source
This commit is contained in:
gabdevele
2026-05-05 13:57:46 +02:00 Unverified
parent 93ba8e9ac2
commit a523c2faa3
53 changed files with 3206 additions and 706 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/services/mitmproxy/examples.py
+36 -18
View File
@@ -4,11 +4,11 @@ DOCS_EXAMPLES_URL = "https://docs.mitmproxy.org/stable/addons/examples/"
BUILTIN_EXAMPLES = [
AddonExampleDocument(
id="http-add-response-header",
name="HTTP add response header",
id="http-tag-traffic",
name="Tag proxied traffic",
description=(
"Injects a fixed response header on every HTTP response so you can "
"verify that traffic is passing through the proxy response path."
"Adds a visible response header on every HTTP reply so you can tell "
"immediately whether the request crossed the proxy."
),
category="http",
source_url=DOCS_EXAMPLES_URL,
@@ -21,11 +21,11 @@ def response(flow: http.HTTPFlow) -> None:
""",
),
AddonExampleDocument(
id="http-modify-query",
name="HTTP modify query",
id="http-force-query-flag",
name="Force request flag",
description=(
"Adds or overwrites a query parameter before each request is "
"forwarded upstream. Useful for forced debug or feature flags."
"Adds or overwrites a request query flag before the flow is sent "
"upstream. Useful when you need a stable marker on every request."
),
category="http",
source_url=DOCS_EXAMPLES_URL,
@@ -38,11 +38,11 @@ def request(flow: http.HTTPFlow) -> None:
""",
),
AddonExampleDocument(
id="http-reply-from-proxy",
name="HTTP reply from proxy",
id="http-local-health-response",
name="Serve local health response",
description=(
"Serves a local HTTP response for a selected path without touching "
"the upstream service. Useful for probes, maintenance, or stubs."
"Answers one health endpoint directly from the proxy without "
"touching the upstream service. Useful for smoke tests and checks."
),
category="http",
source_url=DOCS_EXAMPLES_URL,
@@ -57,14 +57,32 @@ def request(flow: http.HTTPFlow) -> None:
b"ok",
{"content-type": "text/plain"},
)
""",
),
AddonExampleDocument(
id="http-rewrite-health-path",
name="Rewrite one request path",
description=(
"Rewrites a single incoming path before forwarding it upstream. "
"Useful when the public route and upstream route do not match."
),
category="http",
source_url=DOCS_EXAMPLES_URL,
built_in=True,
code="""from mitmproxy import http
def request(flow: http.HTTPFlow) -> None:
if flow.request.path == "/healthz":
flow.request.path = "/status"
""",
),
AddonExampleDocument(
id="websocket-rewrite",
name="WebSocket rewrite",
description=(
"Rewrites matching bytes inside live WebSocket frames to alter "
"runtime flags or payload markers in transit."
"Rewrites matching bytes inside live WebSocket frames to flip a "
"runtime marker while the flow is passing through the proxy."
),
category="websocket",
source_url=DOCS_EXAMPLES_URL,
@@ -83,8 +101,8 @@ def websocket_message(flow: http.HTTPFlow) -> None:
id="tcp-log-messages",
name="TCP log messages",
description=(
"Logs raw TCP message direction and size for each frame so you can "
"observe binary traffic without altering it."
"Logs direction and byte size for every TCP frame so you can see "
"what is crossing the route without modifying the payload."
),
category="tcp",
source_url=DOCS_EXAMPLES_URL,
@@ -104,8 +122,8 @@ def tcp_message(flow: tcp.TCPFlow) -> None:
id="dns-spoof-a-record",
name="DNS spoof A record",
description=(
"Overrides one DNS A record while forwarding every other query "
"unchanged. Useful for selective service redirection."
"Overrides one DNS A record while leaving all other queries alone. "
"Useful for selective redirection during local testing."
),
category="dns",
source_url=DOCS_EXAMPLES_URL,
backend/app/services/mitmproxy/manager.py
+152 -33
View File
@@ -20,6 +20,7 @@ class MitmdumpManager:
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] = {}
async def start(self, proxy: ProxyDocument) -> dict[str, object]:
@@ -67,6 +68,7 @@ class MitmdumpManager:
raise MitmdumpError(reason)
self.processes[proxy.proxy_port] = process
self.attached_pids.pop(proxy.proxy_port, None)
return {
"status": "running",
"pid": process.pid,
@@ -79,6 +81,7 @@ class MitmdumpManager:
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:
@@ -98,19 +101,70 @@ class MitmdumpManager:
if process.returncode is None:
return "running"
self.processes.pop(proxy.proxy_port, None)
self.attached_pids.pop(proxy.proxy_port, None)
return "stopped"
if proxy.pid is None:
pid = self.attached_pids.get(proxy.proxy_port, proxy.pid)
if pid is None:
return "stopped" if proxy.status == "running" else proxy.status
try:
os.kill(proxy.pid, 0)
except ProcessLookupError:
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"
except PermissionError:
return "running"
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
@@ -126,6 +180,8 @@ class MitmdumpManager:
Path(proxy.service.folder),
proxy.service.port,
upstream_port,
host_to="127.0.0.1",
rewrite_host=True,
)
if not changed:
raise MitmdumpError("Unable to move Docker service port")
@@ -135,19 +191,21 @@ class MitmdumpManager:
self,
proxy: ProxyDocument,
upstream_port: int | None,
) -> None:
) -> bool:
if upstream_port is None:
return
return True
if proxy.service.name is None or proxy.service.folder is None:
return
return True
if upstream_port == proxy.service.port:
return
return True
await docker.change_port(
return await docker.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(
@@ -200,20 +258,20 @@ class MitmdumpManager:
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 = self.script_path_for(proxy)
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 = self.config_path_for(proxy)
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 = self.log_path_for(proxy)
path.parent.mkdir(parents=True, exist_ok=True)
return path
@@ -227,19 +285,7 @@ class MitmdumpManager:
)
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()
await self.stop_pid(process.pid)
async def exited_early(self, process) -> bool:
try:
@@ -250,21 +296,94 @@ class MitmdumpManager:
async def stop_pid(self, pid: int) -> None:
try:
os.kill(pid, signal.SIGTERM)
pgid = os.getpgid(pid)
except ProcessLookupError:
return
except PermissionError:
pgid = pid
for _ in range(10):
await asyncio.sleep(0.2)
try:
os.killpg(pgid, signal.SIGTERM)
except ProcessLookupError:
return
except PermissionError:
try:
os.kill(pid, 0)
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.kill(pid, signal.SIGKILL)
os.killpg(pgid, signal.SIGKILL)
except ProcessLookupError:
return
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"
def read_log_tail(self, path: Path, limit: int = 4000) -> str:
if not path.exists():
backend/app/services/mitmproxy/recovery.py
+16
View File
@@ -0,0 +1,16 @@
from repositories.proxy import list_proxies, update_proxy_runtime
from services.mitmproxy.manager import manager
async def reconcile_runtime() -> list[tuple[int, dict[str, object]]]:
changes: list[tuple[int, dict[str, object]]] = []
for proxy in await list_proxies():
updates = await manager.reconcile(proxy)
if updates is None:
continue
await update_proxy_runtime(proxy.proxy_port, updates)
changes.append((proxy.proxy_port, updates))
return changes