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/api/proxy/router.py
+9 -2
View File
@@ -1,7 +1,10 @@
from pathlib import Path
from core.config import settings
from core.security import require_authenticated_user
from fastapi import APIRouter, Depends, HTTPException, Request, status
from integrations.docker import get_all_services_in_dir
from integrations.docker.discovery import service_scan_root
from models.example import AddonExampleCreateInput
from models.proxy import ProxyCreateInput, ProxyUpdateInput
from pymongo.errors import DuplicateKeyError
@@ -120,8 +123,12 @@ async def list_proxies_handler():
@router.get("/services")
async def list_services(request: Request):
return getattr(request.app.state, "services", [])
async def list_services(request: Request, scan_root: str | None = None):
_ = settings
root = Path(scan_root).expanduser() if scan_root else service_scan_root()
services = get_all_services_in_dir(root)
request.app.state.services = services
return {"scan_root": str(root), "services": services}
@router.get("/mitmproxy/options")
backend/app/core/config.py
+2 -1
View File
@@ -15,9 +15,10 @@ class Settings(BaseSettings):
REDIS_URL: str
JWT_SECRET_KEY: str
JWT_ALGORITHM: str = "HS256"
JWT_EXPIRE_MINUTES: int = Field(default=30, gt=0)
JWT_EXPIRE_MINUTES: int = Field(default=720, gt=0)
AUTH_USERNAME: str
AUTH_PASSWORD: str
SERVICE_SCAN_ROOT: str | None = None
settings = Settings()
backend/app/integrations/docker/compose.py
+19 -8
View File
@@ -3,7 +3,7 @@ from pathlib import Path
import yaml
from integrations.docker.models import DockerContainer
from integrations.docker.ports import _get_service_port
from integrations.docker.ports import _get_service_binding
COMPOSE_NAMES = (
"compose.yml",
@@ -50,20 +50,31 @@ def get_services(folder: Path) -> list[DockerContainer] | None:
if not isinstance(mappings, list) or not mappings:
continue
ports = [
p
for p in (_get_service_port(mapping) for mapping in mappings)
if p is not None
bindings = [
{"published": published, "host": host}
for published, host in (
_get_service_binding(mapping) for mapping in mappings
)
if published is not None
]
if not ports:
if not bindings:
continue
services.append({"name": service_name, "ports": ports, "folder": folder})
services.append(
{
"name": service_name,
"ports": [binding["published"] for binding in bindings],
"bindings": bindings,
"folder": folder,
}
)
return services
def get_all_services_in_dir(directory: str | Path) -> list[DockerContainer]:
def get_all_services_in_dir(
directory: str | Path,
) -> list[DockerContainer]:
path = Path(directory)
if not path.is_dir():
return []
backend/app/integrations/docker/discovery.py
+10
View File
@@ -0,0 +1,10 @@
from pathlib import Path
from core.config import settings
def service_scan_root() -> Path:
configured = settings.SERVICE_SCAN_ROOT
if configured:
return Path(configured).expanduser()
return Path.home()
backend/app/integrations/docker/editor.py
+9 -1
View File
@@ -16,6 +16,8 @@ async def change_port(
folder: Path,
port_from: int,
port_to: int,
host_to: str | None = None,
rewrite_host: bool = False,
) -> bool:
compose_data = find_docker_compose(folder)
if compose_data is None:
@@ -38,7 +40,13 @@ async def change_port(
if idx is None:
return False
mappings[idx] = _replace_port_in_mapping(mappings[idx], port_from, port_to)
mappings[idx] = _replace_port_in_mapping(
mappings[idx],
port_from,
port_to,
host_to=host_to,
rewrite_host=rewrite_host,
)
backup_path = _backup_path_for(compose_path)
original_text = compose_path.read_text()
backend/app/integrations/docker/models.py
+6
View File
@@ -2,7 +2,13 @@ from pathlib import Path
from typing import TypedDict
class DockerPortBinding(TypedDict):
published: int
host: str | None
class DockerContainer(TypedDict):
name: str
ports: list[int]
bindings: list[DockerPortBinding]
folder: Path
backend/app/integrations/docker/ports.py
+37 -15
View File
@@ -14,48 +14,64 @@ def get_random_port() -> int:
raise RuntimeError(f"Could not find an available port: {e}") from e
def _get_service_port(mapping: str | int | Mapping[str, object]) -> int | None:
def _get_service_binding(
mapping: str | int | Mapping[str, object],
) -> tuple[int | None, str | None]:
if isinstance(mapping, int):
return mapping
return mapping, None
if isinstance(mapping, Mapping):
published = mapping.get("published")
if published is None:
return None
return None, None
try:
return int(published)
except TypeError, ValueError:
return None
port = int(published)
except TypeError:
return None, None
except ValueError:
return None, None
host = mapping.get("host_ip")
return port, str(host) if host not in {None, ""} else None
parts = str(mapping).split(":")
match parts:
case [""]:
return None
return None, None
case [port] | [port, _]:
try:
return int(port)
return int(port), None
except ValueError:
return None
case ["localhost" | "127.0.0.1", _, _]:
return None
case [_, port, _]:
return None, None
case [host, port, _]:
try:
return int(port)
return int(port), host or None
except ValueError:
return None
return None, None
case _:
return None
return None, None
def _get_service_port(mapping: str | int | Mapping[str, object]) -> int | None:
port, _host = _get_service_binding(mapping)
return port
def _replace_port_in_mapping(
mapping: str | int | Mapping[str, object],
port_from: int,
port_to: int,
host_to: str | None = None,
rewrite_host: bool = False,
) -> str | dict[str, object]:
if isinstance(mapping, Mapping):
updated = dict(mapping)
if _get_service_port(updated) == port_from:
updated["published"] = str(port_to)
if rewrite_host:
if host_to is None:
updated.pop("host_ip", None)
else:
updated["host_ip"] = host_to
return updated
parts = str(mapping).split(":")
@@ -63,8 +79,14 @@ def _replace_port_in_mapping(
case [_]:
return str(port_to)
case [_, container_port]:
if rewrite_host and host_to:
return f"{host_to}:{port_to}:{container_port}"
return f"{port_to}:{container_port}"
case [ip, _, container_port]:
if rewrite_host:
if host_to:
return f"{host_to}:{port_to}:{container_port}"
return f"{port_to}:{container_port}"
return f"{ip}:{port_to}:{container_port}"
case _:
return str(port_to)
backend/app/main.py
+19 -5
View File
@@ -1,7 +1,6 @@
import os
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from api.auth.router import router as auth_router
from api.editor.router import router as editor_router
@@ -12,9 +11,11 @@ from core.redis import close_redis, init_redis
from core.security import hash_password
from fastapi import FastAPI
from integrations.docker import get_all_services_in_dir
from integrations.docker.discovery import service_scan_root
from repositories.auth import ensure_auth_indexes, ensure_single_account
from repositories.examples import ensure_example_indexes
from repositories.proxy import ensure_proxy_indexes
from services.mitmproxy.recovery import reconcile_runtime
@asynccontextmanager
@@ -28,14 +29,27 @@ async def lifespan(_: FastAPI):
)
await ensure_proxy_indexes()
await ensure_example_indexes()
recovered = await reconcile_runtime()
if recovered:
recovered_count = len(recovered)
error_count = sum(
1
for _proxy_port, values in recovered
if values.get("status") == "error"
)
print(
f"[thePlug] Reconciled {recovered_count} persisted proxy runtimes"
f" ({error_count} with errors).",
file=sys.stderr,
)
home_dir = str(Path.home())
app.state.services = get_all_services_in_dir(home_dir) or []
scan_root = service_scan_root()
app.state.services = get_all_services_in_dir(scan_root) or []
if not app.state.services:
warning = (
"\n" + "!" * 60 + "\n"
" WARNING: No services found in home directory!\n"
f" Scanned: {home_dir}\n" + "!" * 60 + "\n"
" WARNING: No services found in scan root!\n"
f" Scanned: {scan_root}\n" + "!" * 60 + "\n"
)
print(warning, file=sys.stderr)
backend/app/models/service.py
+4
View File
@@ -6,6 +6,10 @@ class ServiceEndpointDocument(BaseModel):
host: str = Field(default="127.0.0.1", description="Upstream service host")
port: int = Field(..., description="Original upstream service port")
bind_host: str | None = Field(
default=None,
description="Original bind host exposed by Docker Compose",
)
name: str | None = Field(default=None, description="Docker Compose service name")
folder: str | None = Field(
default=None,
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