feat: troppa roba davvero

This commit is contained in:
gabdevele
2026-05-07 09:21:45 +02:00 Unverified
parent d977a953cc
commit 53dcdafdb1
73 changed files with 4772 additions and 960 deletions
+1 -1
View File
@@ -1,3 +1,3 @@
from api.auth.router import router
from app.api.auth.router import router
__all__ = ["router"]
+5 -4
View File
@@ -1,15 +1,16 @@
from datetime import UTC, datetime
from typing import Annotated
from core.security import (
from fastapi import APIRouter, Depends, HTTPException, Response, status
from app.core.security import (
AuthContext,
create_access_token,
get_auth_context,
verify_password,
)
from fastapi import APIRouter, Depends, HTTPException, Response, status
from models.auth import LoginInput, TokenResponse
from repositories.auth import get_user_by_username, revoke_token
from app.models.auth import LoginInput, TokenResponse
from app.repositories.auth import get_user_by_username, revoke_token
router = APIRouter(
prefix="/auth",
+5 -4
View File
@@ -1,11 +1,12 @@
from core.security import (
from fastapi import APIRouter, Depends, HTTPException, WebSocket, status
from starlette.websockets import WebSocketDisconnect
from app.core.security import (
InvalidTokenError,
authenticate_token,
require_authenticated_user,
)
from fastapi import APIRouter, Depends, HTTPException, WebSocket, status
from services.editor import python_lsp
from starlette.websockets import WebSocketDisconnect
from app.services.editor import python_lsp
router = APIRouter(prefix="/proxy/editor/python", tags=["Editor"])
+236 -52
View File
@@ -1,44 +1,61 @@
import asyncio
from pathlib import Path
from core.config import settings
from core.security import (
from fastapi import APIRouter, Depends, HTTPException, WebSocket, status
from pymongo.errors import DuplicateKeyError
from starlette.websockets import WebSocketDisconnect
from app.core.config import settings
from app.core.security import (
InvalidTokenError,
authenticate_token,
require_authenticated_user,
)
from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket, 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
from repositories.examples import (
create_custom_example,
delete_custom_example,
list_custom_examples,
from app.integrations.docker import get_all_services_in_dir
from app.integrations.docker.discovery import service_scan_root
from app.models.health import HealthHistoryResponse, HealthSnapshot
from app.models.proxy import ProxyCreateInput, ProxyUpdateInput
from app.models.template import ProxyTemplateCreateInput
from app.repositories.health import (
delete_health_history,
get_latest_health_sample,
list_health_history,
record_health_sample,
)
from repositories.proxy import (
from app.repositories.proxy import (
create_proxy,
delete_proxy,
get_proxy_by_port,
update_proxy,
update_proxy_runtime,
)
from repositories.proxy import (
from app.repositories.proxy import (
list_proxies as list_proxy_documents,
)
from services.mitmproxy.examples import BUILTIN_EXAMPLES
from services.mitmproxy.health import check_proxy
from services.mitmproxy.manager import MitmdumpError, manager
from services.mitmproxy.options import (
from app.repositories.templates import (
create_custom_template,
delete_custom_template,
list_custom_templates,
)
from app.services.mitmproxy.health import check_proxy
from app.services.mitmproxy.health_monitor import health_monitor
from app.services.mitmproxy.manager import MitmdumpError, manager
from app.services.mitmproxy.options import (
MitmOptionError,
load_options,
public_options,
validate_options,
)
from services.mitmproxy.validator import ProxyCodeValidationError, validate_proxy_code
from starlette.websockets import WebSocketDisconnect
from app.services.mitmproxy.runtime_validation import (
ReservedProxyPortError,
backend_port,
ensure_proxy_port_not_reserved,
)
from app.services.mitmproxy.templates import BUILTIN_TEMPLATES
from app.services.mitmproxy.validator import (
ProxyCodeValidationError,
validate_proxy_code,
)
SUPPORTED_MODES = {
"regular",
@@ -62,6 +79,8 @@ SUPPORTED_PROTOCOLS = {
"http3",
"quic",
}
DEFAULT_HEALTH_HISTORY_LIMIT = 40
MAX_HEALTH_HISTORY_LIMIT = 200
router = APIRouter(
prefix="/proxy",
@@ -113,6 +132,16 @@ async def validate_mitm_options(options: dict[str, object]) -> None:
) from error
def validate_proxy_port(proxy_port: int) -> None:
try:
ensure_proxy_port_not_reserved(proxy_port)
except ReservedProxyPortError as error:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
detail=str(error),
) from error
async def get_proxy_or_404(proxy_port: int):
proxy = await get_proxy_by_port(proxy_port)
if proxy is None:
@@ -143,10 +172,46 @@ def proxy_log_path(proxy) -> Path:
return manager.log_path_for(proxy)
def read_log_snapshot(path: Path, limit: int = 8000) -> str:
if not path.is_file():
return ""
return path.read_bytes()[-limit:].decode("utf-8", errors="replace")
def serialize_health_snapshot(
snapshot: HealthSnapshot | None,
) -> dict[str, object] | None:
if snapshot is None:
return None
return snapshot.model_dump(mode="json")
async def read_log_snapshot(path: Path, limit: int = 8000) -> str:
snapshot, _offset, _marker = await snapshot_log_state(path, limit=limit)
return snapshot
async def snapshot_log_state(
path: Path,
limit: int = 8000,
marker_width: int = 64,
) -> tuple[str, int, bytes]:
if not await asyncio.to_thread(path.is_file):
return "", 0, b""
data = await asyncio.to_thread(path.read_bytes)
size = len(data)
marker = data[max(0, size - marker_width) : size]
snapshot = data[-limit:].decode("utf-8", errors="replace")
return snapshot, size, marker
def _read_log_marker_bytes(path: Path, offset: int, width: int) -> bytes:
if offset <= 0 or not path.is_file():
return b""
start = max(0, offset - width)
with path.open("rb") as handle:
handle.seek(start)
return handle.read(offset - start)
async def read_log_marker(path: Path, offset: int, width: int = 64) -> bytes:
return await asyncio.to_thread(_read_log_marker_bytes, path, offset, width)
async def watch_socket_disconnect(websocket: WebSocket) -> None:
@@ -154,28 +219,52 @@ async def watch_socket_disconnect(websocket: WebSocket) -> None:
await websocket.receive_text()
async def stream_proxy_logs(websocket: WebSocket, path: Path) -> None:
offset = path.stat().st_size if path.is_file() else 0
async def stream_proxy_logs(
websocket: WebSocket,
path: Path,
*,
offset: int,
marker: bytes,
) -> None:
while True:
if not path.exists():
if not await asyncio.to_thread(path.exists):
await asyncio.sleep(0.4)
continue
try:
size = path.stat().st_size
size = await asyncio.to_thread(lambda: path.stat().st_size)
except OSError:
await asyncio.sleep(0.4)
continue
if marker and size >= offset and await read_log_marker(path, offset) != marker:
snapshot, offset, marker = await snapshot_log_state(path)
await websocket.send_json({"type": "snapshot", "log": snapshot})
await asyncio.sleep(0.4)
continue
if size < offset:
offset = size
snapshot, offset, marker = await snapshot_log_state(path)
await websocket.send_json({"type": "snapshot", "log": snapshot})
await asyncio.sleep(0.4)
continue
if size > offset:
with path.open("rb") as handle:
handle.seek(offset)
chunk = handle.read(size - offset).decode("utf-8", errors="replace")
def read_log_chunk(
chunk_offset: int = offset,
chunk_size: int = size,
) -> str:
with path.open("rb") as handle:
handle.seek(chunk_offset)
return handle.read(chunk_size - chunk_offset).decode(
"utf-8",
errors="replace",
)
chunk = await asyncio.to_thread(read_log_chunk)
offset = size
marker = await read_log_marker(path, offset)
if chunk:
await websocket.send_json({"type": "append", "log": chunk})
@@ -188,12 +277,15 @@ async def list_proxies_handler():
@router.get("/services", dependencies=auth_dependencies)
async def list_services(request: Request, scan_root: str | None = None):
async def list_services(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}
return {
"scan_root": str(root),
"services": services,
"backend_port": backend_port(),
}
@router.get("/mitmproxy/options", dependencies=auth_dependencies)
@@ -201,44 +293,44 @@ async def list_mitmproxy_options():
return public_options(await load_options())
@router.get("/mitmproxy/examples", dependencies=auth_dependencies)
async def list_mitmproxy_examples():
return [*BUILTIN_EXAMPLES, *await list_custom_examples()]
@router.get("/mitmproxy/templates", dependencies=auth_dependencies)
async def list_mitmproxy_templates():
return [*BUILTIN_TEMPLATES, *await list_custom_templates()]
@router.post(
"/mitmproxy/examples",
"/mitmproxy/templates",
status_code=status.HTTP_201_CREATED,
dependencies=auth_dependencies,
)
async def create_mitmproxy_example(payload: AddonExampleCreateInput):
async def create_mitmproxy_template(payload: ProxyTemplateCreateInput):
await _validate_code_or_raise(payload.code)
try:
return await create_custom_example(payload)
return await create_custom_template(payload)
except DuplicateKeyError as error:
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail=f"Example already exists: {payload.id}",
detail=f"Template already exists: {payload.id}",
) from error
@router.delete(
"/mitmproxy/examples/{example_id}",
"/mitmproxy/templates/{template_id}",
status_code=status.HTTP_204_NO_CONTENT,
dependencies=auth_dependencies,
)
async def delete_mitmproxy_example(example_id: str):
if any(example.id == example_id for example in BUILTIN_EXAMPLES):
async def delete_mitmproxy_template(template_id: str):
if any(template.id == template_id for template in BUILTIN_TEMPLATES):
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail="Built-in examples cannot be deleted",
detail="Built-in templates cannot be deleted",
)
deleted = await delete_custom_example(example_id)
deleted = await delete_custom_template(template_id)
if not deleted:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Example not found",
detail="Template not found",
)
@@ -249,6 +341,7 @@ async def get_proxy(proxy_port: int):
@router.post("", status_code=status.HTTP_201_CREATED, dependencies=auth_dependencies)
async def create_proxy_handler(payload: ProxyCreateInput):
validate_proxy_port(payload.proxy_port)
existing_proxy = await get_proxy_by_port(payload.proxy_port)
if existing_proxy is not None:
raise HTTPException(
@@ -341,7 +434,28 @@ async def proxy_status(proxy_port: int):
@router.get("/{proxy_port}/health", dependencies=auth_dependencies)
async def proxy_health(proxy_port: int):
proxy = await get_proxy_or_404(proxy_port)
return await check_proxy(proxy)
snapshot = await check_proxy(proxy)
if proxy.healthcheck.enabled:
await record_health_sample(snapshot)
await health_monitor.publish(proxy.proxy_port, snapshot)
return serialize_health_snapshot(snapshot)
@router.get("/{proxy_port}/health/history", dependencies=auth_dependencies)
async def proxy_health_history(
proxy_port: int,
limit: int = DEFAULT_HEALTH_HISTORY_LIMIT,
):
proxy = await get_proxy_or_404(proxy_port)
limit = max(1, min(limit, MAX_HEALTH_HISTORY_LIMIT))
latest = await get_latest_health_sample(proxy.proxy_port)
history = await list_health_history(proxy.proxy_port, limit=limit)
response = HealthHistoryResponse(
proxy_port=proxy.proxy_port,
latest=latest,
history=history,
)
return response.model_dump(mode="json")
@router.get("/{proxy_port}/logs", dependencies=auth_dependencies)
@@ -351,7 +465,7 @@ async def proxy_logs(proxy_port: int):
return {
"proxy_port": proxy.proxy_port,
"log": read_log_snapshot(path),
"log": await read_log_snapshot(path),
}
@@ -373,11 +487,14 @@ async def proxy_logs_socket(websocket: WebSocket, proxy_port: int):
return
path = proxy_log_path(proxy)
initial_snapshot, offset, marker = await snapshot_log_state(path)
await websocket.accept()
await websocket.send_json({"type": "snapshot", "log": read_log_snapshot(path)})
await websocket.send_json({"type": "snapshot", "log": initial_snapshot})
disconnect_task = asyncio.create_task(watch_socket_disconnect(websocket))
stream_task = asyncio.create_task(stream_proxy_logs(websocket, path))
stream_task = asyncio.create_task(
stream_proxy_logs(websocket, path, offset=offset, marker=marker)
)
done, pending = await asyncio.wait(
{disconnect_task, stream_task},
@@ -408,6 +525,71 @@ async def proxy_logs_socket(websocket: WebSocket, proxy_port: int):
task.cancel()
@router.websocket("/{proxy_port}/health/ws")
async def proxy_health_socket(websocket: WebSocket, proxy_port: int):
disconnect_task: asyncio.Task[None] | None = None
queue_task: asyncio.Task[None] | None = None
queue = None
async def forward_health() -> None:
assert queue is not None
while True:
snapshot = await queue.get()
await websocket.send_json(serialize_health_snapshot(snapshot))
try:
try:
await authenticate_proxy_socket(websocket)
except InvalidTokenError:
await websocket.close(code=4401)
return
proxy = await get_proxy_by_port(proxy_port)
if proxy is None:
await websocket.close(code=4404)
return
await websocket.accept()
latest = await get_latest_health_sample(proxy.proxy_port)
if latest is None:
latest = await check_proxy(proxy)
await websocket.send_json(serialize_health_snapshot(latest))
queue = health_monitor.subscribe(proxy.proxy_port)
disconnect_task = asyncio.create_task(watch_socket_disconnect(websocket))
queue_task = asyncio.create_task(forward_health())
done, pending = await asyncio.wait(
{disconnect_task, queue_task},
return_when=asyncio.FIRST_EXCEPTION,
)
for task in pending:
task.cancel()
if pending:
await asyncio.gather(*pending, return_exceptions=True)
for task in done:
try:
task.result()
except asyncio.CancelledError:
return
except WebSocketDisconnect:
return
except RuntimeError:
await websocket.close(code=1011, reason="health stream failed")
return
except asyncio.CancelledError:
return
finally:
if queue is not None:
health_monitor.unsubscribe(proxy_port, queue)
for task in (disconnect_task, queue_task):
if task is None or task.done():
continue
task.cancel()
@router.delete(
"/{proxy_port}",
status_code=status.HTTP_204_NO_CONTENT,
@@ -421,6 +603,8 @@ async def remove_proxy(proxy_port: int):
detail="Proxy not found",
)
await manager.stop(proxy)
manager.cleanup_runtime_artifacts(proxy)
await delete_health_history(proxy.proxy_port)
deleted = await delete_proxy(proxy_port)
if not deleted:
raise HTTPException(
+4 -1
View File
@@ -17,7 +17,10 @@ class Settings(BaseSettings):
JWT_EXPIRE_MINUTES: int = Field(default=720, gt=0)
AUTH_USERNAME: str
AUTH_PASSWORD: str
API_PORT: int = Field(default=3671, ge=1, le=65535)
SERVICE_SCAN_ROOT: str | None = None
HEALTHCHECK_INTERVAL_SECONDS: int = Field(default=15, ge=5, le=300)
HEALTH_HISTORY_TTL_SECONDS: int = Field(default=86400, ge=300)
settings = Settings()
settings = Settings() # pyright: ignore[reportCallIssue]
+7 -4
View File
@@ -1,11 +1,14 @@
from core.config import settings
from typing import Any
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
from app.core.config import settings
class _MongoDB:
client: AsyncIOMotorClient | None = None
client: AsyncIOMotorClient[dict[str, Any]] | None = None
def get_database(self) -> AsyncIOMotorDatabase:
def get_database(self) -> AsyncIOMotorDatabase[dict[str, Any]]:
if self.client is None:
raise RuntimeError("MongoDB client has not been initialized")
return self.client[settings.MONGODB_DB_NAME]
@@ -26,7 +29,7 @@ async def close_mongodb() -> None:
client.close()
async def get_database() -> AsyncIOMotorDatabase:
async def get_database() -> AsyncIOMotorDatabase[dict[str, Any]]:
return mongodb.get_database()
+5 -6
View File
@@ -5,13 +5,14 @@ from uuid import uuid4
from authlib.jose import jwt # type: ignore
from authlib.jose.errors import JoseError # type: ignore
from core.config import settings
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from models.auth import AccessTokenPayload, AuthUserDocument
from pwdlib import PasswordHash
from pwdlib.hashers.argon2 import Argon2Hasher
from repositories.auth import get_user_by_username, is_token_revoked
from app.core.config import settings
from app.models.auth import AccessTokenPayload, AuthUserDocument
from app.repositories.auth import get_user_by_username, is_token_revoked
pwd_context = PasswordHash((Argon2Hasher(),))
bearer_scheme = HTTPBearer(auto_error=False)
@@ -41,9 +42,7 @@ def hash_password(password: str) -> str:
def verify_password(plain_password: str, password_hash: str) -> bool:
if hasattr(pwd_context, "verify"):
return pwd_context.verify(plain_password, password_hash)
return pwd_context.check(plain_password, password_hash)
return pwd_context.verify(plain_password, password_hash)
def create_access_token(*, subject: str) -> str:
+3 -3
View File
@@ -1,10 +1,10 @@
from integrations.docker.compose import (
from app.integrations.docker.compose import (
find_docker_compose,
get_all_services_in_dir,
get_services,
)
from integrations.docker.editor import change_port
from integrations.docker.models import DockerContainer
from app.integrations.docker.editor import change_port
from app.integrations.docker.models import DockerContainer
__all__ = [
"DockerContainer",
+42 -16
View File
@@ -2,8 +2,9 @@ import asyncio
from pathlib import Path
import yaml
from integrations.docker.models import DockerContainer
from integrations.docker.ports import _get_service_binding
from app.integrations.docker.models import DockerContainer, DockerPortBinding
from app.integrations.docker.ports import _get_service_binding
COMPOSE_NAMES = (
"compose.yml",
@@ -12,24 +13,43 @@ COMPOSE_NAMES = (
"docker-compose.yaml",
)
type ComposeDocument = dict[str, object]
def _get_visible_folders(directory: Path) -> list[Path]:
return [
p
for p in Path(directory).iterdir()
if p.is_dir() and not p.name.startswith(".")
]
try:
entries = list(Path(directory).iterdir())
except OSError:
return []
folders: list[Path] = []
for entry in entries:
if entry.name.startswith("."):
continue
try:
if entry.is_dir():
folders.append(entry)
except OSError:
continue
return folders
def find_docker_compose(path: Path) -> tuple[dict, Path] | None:
if not path or not path.is_dir():
def find_docker_compose(path: Path) -> tuple[ComposeDocument, Path] | None:
try:
if not path or not path.is_dir():
return None
except OSError:
return None
for name in COMPOSE_NAMES:
file = path / name
if file.is_file():
content = yaml.safe_load(file.read_text()) or {}
return content, file
try:
if file.is_file():
content = yaml.safe_load(file.read_text()) or {}
return content, file
except OSError:
continue
return None
@@ -42,7 +62,7 @@ def get_services(folder: Path) -> list[DockerContainer] | None:
docker_compose, _ = compose_data
compose_services = docker_compose.get("services") or {}
if not compose_services:
if not isinstance(compose_services, dict) or not compose_services:
return None
for service_name, config in compose_services.items():
@@ -50,7 +70,7 @@ def get_services(folder: Path) -> list[DockerContainer] | None:
if not isinstance(mappings, list) or not mappings:
continue
bindings = [
bindings: list[DockerPortBinding] = [
{"published": published, "host": host}
for published, host in (
_get_service_binding(mapping) for mapping in mappings
@@ -76,7 +96,10 @@ def get_all_services_in_dir(
directory: str | Path,
) -> list[DockerContainer]:
path = Path(directory)
if not path.is_dir():
try:
if not path.is_dir():
return []
except OSError:
return []
services: list[DockerContainer] = []
@@ -88,7 +111,10 @@ def get_all_services_in_dir(
return services
async def run_docker_command(file: Path, command: list[str]) -> dict | None:
async def run_docker_command(
file: Path,
command: list[str],
) -> dict[str, str] | None:
if not command:
return None
+1 -1
View File
@@ -1,6 +1,6 @@
from pathlib import Path
from core.config import settings
from app.core.config import settings
def service_scan_root() -> Path:
+36 -19
View File
@@ -1,14 +1,26 @@
import asyncio
from pathlib import Path
import yaml
from integrations.docker.compose import find_docker_compose, run_docker_command
from integrations.docker.ports import _get_port_index, _replace_port_in_mapping
from app.integrations.docker.compose import (
ComposeDocument,
find_docker_compose,
run_docker_command,
)
from app.integrations.docker.ports import _get_port_index, _replace_port_in_mapping
def _backup_path_for(compose_path: Path) -> Path:
if compose_path.suffix in {".yml", ".yaml"}:
return compose_path.with_name(f"{compose_path.stem}.bak{compose_path.suffix}")
return compose_path.with_name(f"{compose_path.name}.bak")
def _render_compose(document: ComposeDocument) -> str:
return yaml.dump(document, default_flow_style=False, sort_keys=False)
def _write_compose(path: Path, content: str) -> None:
path.write_text(content, encoding="utf-8")
def _read_compose_text(path: Path) -> str:
return path.read_text(encoding="utf-8")
async def change_port(
@@ -25,7 +37,7 @@ async def change_port(
docker_compose, compose_path = compose_data
compose_services = docker_compose.get("services") or {}
if not compose_services:
if not isinstance(compose_services, dict) or not compose_services:
return False
service_config = compose_services.get(service_name)
@@ -48,21 +60,26 @@ async def change_port(
rewrite_host=rewrite_host,
)
backup_path = _backup_path_for(compose_path)
original_text = compose_path.read_text()
original_text = await asyncio.to_thread(_read_compose_text, compose_path)
updated_text = _render_compose(docker_compose)
backup_path.write_text(compose_path.read_text())
compose_path.write_text(
yaml.dump(docker_compose, default_flow_style=False, sort_keys=False)
)
result = await run_docker_command(
compose_path,
["up", "-d", service_name],
)
try:
await asyncio.to_thread(_write_compose, compose_path, updated_text)
result = await run_docker_command(
compose_path,
["up", "-d", service_name],
)
except Exception:
try:
await asyncio.to_thread(_write_compose, compose_path, original_text)
except OSError as restore_error:
raise RuntimeError(
f"Failed to restore original compose file: {compose_path}"
) from restore_error
return False
if result is None or result.get("status") != "success":
compose_path.write_text(original_text)
await asyncio.to_thread(_write_compose, compose_path, original_text)
return False
return True
+2 -2
View File
@@ -24,10 +24,10 @@ def _get_service_binding(
published = mapping.get("published")
if published is None:
return None, None
if not isinstance(published, (str, int)):
return None, None
try:
port = int(published)
except TypeError:
return None, None
except ValueError:
return None, None
host = mapping.get("host_ip")
+21 -15
View File
@@ -2,19 +2,22 @@ import os
import sys
from contextlib import asynccontextmanager
from api.auth.router import router as auth_router
from api.editor.router import router as editor_router
from api.proxy.router import router as proxy_router
from core.config import settings
from core.mongodb import close_mongodb, init_mongodb
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
from app.api.auth.router import router as auth_router
from app.api.editor.router import router as editor_router
from app.api.proxy.router import router as proxy_router
from app.core.config import settings
from app.core.mongodb import close_mongodb, init_mongodb
from app.core.security import hash_password
from app.integrations.docker import get_all_services_in_dir
from app.integrations.docker.discovery import service_scan_root
from app.repositories.auth import ensure_auth_indexes, ensure_single_account
from app.repositories.health import ensure_health_indexes
from app.repositories.proxy import ensure_proxy_indexes
from app.repositories.templates import ensure_template_indexes
from app.services.mitmproxy.health_monitor import health_monitor
from app.services.mitmproxy.recovery import reconcile_runtime
@asynccontextmanager
@@ -26,7 +29,8 @@ async def lifespan(_: FastAPI):
password_hash=hash_password(settings.AUTH_PASSWORD),
)
await ensure_proxy_indexes()
await ensure_example_indexes()
await ensure_template_indexes()
await ensure_health_indexes()
recovered = await reconcile_runtime()
if recovered:
recovered_count = len(recovered)
@@ -42,8 +46,8 @@ async def lifespan(_: FastAPI):
)
scan_root = service_scan_root()
app.state.services = get_all_services_in_dir(scan_root) or []
if not app.state.services:
services = get_all_services_in_dir(scan_root) or []
if not services:
warning = (
"\n" + "!" * 60 + "\n"
" WARNING: No services found in scan root!\n"
@@ -52,8 +56,10 @@ async def lifespan(_: FastAPI):
print(warning, file=sys.stderr)
try:
await health_monitor.start()
yield
finally:
await health_monitor.stop()
await close_mongodb()
-28
View File
@@ -1,28 +0,0 @@
from datetime import datetime
from pydantic import BaseModel, ConfigDict, Field
class AddonExampleDocument(BaseModel):
model_config = ConfigDict(extra="forbid")
id: str = Field(..., description="Stable example identifier")
name: str = Field(..., description="Display name")
description: str = Field(default="", description="Short explanation")
category: str = Field(default="custom", description="Example category")
code: str = Field(..., description="Addon source code")
source_url: str | None = Field(default=None, description="Reference URL")
built_in: bool = Field(default=False, description="Provided by thePlug")
created_at: datetime | None = None
updated_at: datetime | None = None
class AddonExampleCreateInput(BaseModel):
model_config = ConfigDict(extra="forbid")
id: str = Field(..., description="Stable example identifier")
name: str = Field(..., description="Display name")
description: str = Field(default="", description="Short explanation")
category: str = Field(default="custom", description="Example category")
code: str = Field(..., description="Addon source code")
source_url: str | None = Field(default=None, description="Reference URL")
+41
View File
@@ -0,0 +1,41 @@
from datetime import datetime
from typing import Any
from pydantic import BaseModel, ConfigDict, Field
class HealthCheckEntry(BaseModel):
model_config = ConfigDict(extra="forbid")
target: str = Field(..., description="Logical target name")
host: str = Field(..., description="Resolved host checked at runtime")
port: int = Field(..., description="Resolved port checked at runtime")
protocol: str = Field(..., description="Protocol used for the check")
healthy: bool = Field(..., description="Whether the check succeeded")
error: str | None = Field(default=None, description="Optional failure reason")
duration_ms: float | None = Field(
default=None,
description=(
"Measured response latency in milliseconds when benchmarking "
"is enabled"
),
)
details: dict[str, Any] = Field(default_factory=dict)
class HealthSnapshot(BaseModel):
model_config = ConfigDict(extra="forbid")
proxy_port: int = Field(..., description="Proxy port this snapshot belongs to")
timestamp: datetime = Field(..., description="When the snapshot was produced")
enabled: bool = Field(..., description="Whether healthchecks are enabled")
overall: str = Field(..., description="Overall health status")
checks: list[HealthCheckEntry] = Field(default_factory=list)
class HealthHistoryResponse(BaseModel):
model_config = ConfigDict(extra="forbid")
proxy_port: int
latest: HealthSnapshot | None = None
history: list[HealthSnapshot] = Field(default_factory=list)
+6 -1
View File
@@ -1,9 +1,10 @@
from datetime import datetime
from typing import Any
from models.service import ServiceEndpointDocument
from pydantic import BaseModel, ConfigDict, Field
from app.models.service import ServiceEndpointDocument
class HealthCheckDocument(BaseModel):
model_config = ConfigDict(extra="forbid")
@@ -28,6 +29,10 @@ class HealthCheckDocument(BaseModel):
default=False,
description="Verify TLS certificates during health checks",
)
benchmark_enabled: bool = Field(
default=False,
description="Measure response latency during health checks when supported",
)
timeout_seconds: float = Field(
default=2.0,
gt=0,
+40
View File
@@ -0,0 +1,40 @@
from datetime import datetime
from pydantic import BaseModel, ConfigDict, Field
class ProxyTemplateDocument(BaseModel):
model_config = ConfigDict(extra="forbid")
id: str = Field(..., description="Stable template identifier")
name: str = Field(..., description="Display name")
description: str = Field(default="", description="Short explanation")
category: str = Field(default="custom", description="Template category")
protocol: str = Field(default="custom", description="Primary proxy protocol")
family: str = Field(default="custom", description="Template family")
tags: list[str] = Field(default_factory=list, description="Filter tags")
action: str = Field(default="custom", description="Primary template action")
summary: str = Field(default="", description="Short operational summary")
behavior: str = Field(default="", description="Runtime behavior description")
code: str = Field(..., description="Addon source code")
source_url: str | None = Field(default=None, description="Reference URL")
built_in: bool = Field(default=False, description="Provided by thePlug")
created_at: datetime | None = None
updated_at: datetime | None = None
class ProxyTemplateCreateInput(BaseModel):
model_config = ConfigDict(extra="forbid")
id: str = Field(..., description="Stable template identifier")
name: str = Field(..., description="Display name")
description: str = Field(default="", description="Short explanation")
category: str = Field(default="custom", description="Template category")
protocol: str = Field(default="custom", description="Primary proxy protocol")
family: str = Field(default="custom", description="Template family")
tags: list[str] = Field(default_factory=list, description="Filter tags")
action: str = Field(default="custom", description="Primary template action")
summary: str = Field(default="", description="Short operational summary")
behavior: str = Field(default="", description="Runtime behavior description")
code: str = Field(..., description="Addon source code")
source_url: str | None = Field(default=None, description="Reference URL")
+3 -2
View File
@@ -2,10 +2,11 @@ from collections.abc import Mapping
from datetime import UTC, datetime
from typing import Any
from core.mongodb import get_database
from models.auth import AuthUserDocument
from motor.motor_asyncio import AsyncIOMotorCollection
from app.core.mongodb import get_database
from app.models.auth import AuthUserDocument
USERS_COLLECTION_NAME = "auth_users"
REVOKED_COLLECTION_NAME = "revoked_tokens"
-70
View File
@@ -1,70 +0,0 @@
from collections.abc import Mapping
from datetime import UTC, datetime
from typing import Any
from core.mongodb import get_database
from models.example import AddonExampleCreateInput, AddonExampleDocument
from motor.motor_asyncio import AsyncIOMotorCollection
COLLECTION_NAME = "addon_examples"
def utcnow() -> datetime:
return datetime.now(UTC)
async def collection() -> AsyncIOMotorCollection[dict[str, Any]]:
database = await get_database()
return database[COLLECTION_NAME]
def deserialize(raw: Mapping[str, Any] | None) -> AddonExampleDocument | None:
if raw is None:
return None
document = dict(raw)
document.pop("_id", None)
return AddonExampleDocument.model_validate(document)
async def ensure_example_indexes() -> None:
examples = await collection()
await examples.create_index("id", unique=True, name="uq_addon_example_id")
async def list_custom_examples() -> list[AddonExampleDocument]:
examples = await collection()
cursor = examples.find({}).sort("name", 1)
result: list[AddonExampleDocument] = []
async for raw in cursor:
example = deserialize(raw)
if example is not None:
result.append(example)
return result
async def create_custom_example(
payload: AddonExampleCreateInput,
) -> AddonExampleDocument:
now = utcnow()
example = AddonExampleDocument(
id=payload.id,
name=payload.name,
description=payload.description,
category=payload.category,
code=payload.code,
source_url=payload.source_url,
built_in=False,
created_at=now,
updated_at=now,
)
examples = await collection()
await examples.insert_one(example.model_dump(exclude_none=True))
return example
async def delete_custom_example(example_id: str) -> bool:
examples = await collection()
result = await examples.delete_one({"id": example_id})
return result.deleted_count > 0
+80
View File
@@ -0,0 +1,80 @@
from datetime import UTC, datetime, timedelta
from typing import Any
from motor.motor_asyncio import AsyncIOMotorCollection
from app.core.config import settings
from app.core.mongodb import get_database
from app.models.health import HealthSnapshot
COLLECTION_NAME = "health_history"
def utcnow() -> datetime:
return datetime.now(UTC)
async def collection() -> AsyncIOMotorCollection[dict[str, Any]]:
database = await get_database()
return database[COLLECTION_NAME]
async def ensure_health_indexes() -> None:
samples = await collection()
await samples.create_index(
[("proxy_port", 1), ("timestamp", -1)],
name="ix_health_proxy_port_timestamp",
)
await samples.create_index(
"expires_at",
expireAfterSeconds=0,
name="ttl_health_history",
)
def deserialize_health_sample(raw: dict[str, Any] | None) -> HealthSnapshot | None:
if raw is None:
return None
document = dict(raw)
document.pop("_id", None)
document.pop("expires_at", None)
return HealthSnapshot.model_validate(document)
async def record_health_sample(snapshot: HealthSnapshot) -> HealthSnapshot:
samples = await collection()
document = snapshot.model_dump(mode="python")
document["expires_at"] = document["timestamp"] + timedelta(
seconds=settings.HEALTH_HISTORY_TTL_SECONDS,
)
await samples.insert_one(document)
return snapshot
async def get_latest_health_sample(proxy_port: int) -> HealthSnapshot | None:
samples = await collection()
raw = await samples.find_one(
{"proxy_port": proxy_port},
sort=[("timestamp", -1)],
)
return deserialize_health_sample(raw)
async def list_health_history(
proxy_port: int,
limit: int = 40,
) -> list[HealthSnapshot]:
samples = await collection()
cursor = samples.find({"proxy_port": proxy_port}).sort("timestamp", -1).limit(limit)
history: list[HealthSnapshot] = []
async for raw in cursor:
snapshot = deserialize_health_sample(raw)
if snapshot is not None:
history.append(snapshot)
history.reverse()
return history
async def delete_health_history(proxy_port: int) -> None:
samples = await collection()
await samples.delete_many({"proxy_port": proxy_port})
+3 -2
View File
@@ -2,11 +2,12 @@ from collections.abc import Mapping
from datetime import UTC, datetime
from typing import Any
from core.mongodb import get_database
from models.proxy import ProxyCreateInput, ProxyDocument, ProxyUpdateInput
from motor.motor_asyncio import AsyncIOMotorCollection
from pymongo import ReturnDocument
from app.core.mongodb import get_database
from app.models.proxy import ProxyCreateInput, ProxyDocument, ProxyUpdateInput
COLLECTION_NAME = "proxies"
+77
View File
@@ -0,0 +1,77 @@
from collections.abc import Mapping
from datetime import UTC, datetime
from typing import Any
from motor.motor_asyncio import AsyncIOMotorCollection
from app.core.mongodb import get_database
from app.models.template import ProxyTemplateCreateInput, ProxyTemplateDocument
COLLECTION_NAME = "addon_templates"
def utcnow() -> datetime:
return datetime.now(UTC)
async def collection() -> AsyncIOMotorCollection[dict[str, Any]]:
database = await get_database()
return database[COLLECTION_NAME]
def deserialize(raw: Mapping[str, Any] | None) -> ProxyTemplateDocument | None:
if raw is None:
return None
document = dict(raw)
document.pop("_id", None)
return ProxyTemplateDocument.model_validate(document)
async def ensure_template_indexes() -> None:
templates = await collection()
await templates.create_index("id", unique=True, name="uq_addon_template_id")
async def list_custom_templates() -> list[ProxyTemplateDocument]:
templates = await collection()
cursor = templates.find({}).sort("name", 1)
result: list[ProxyTemplateDocument] = []
async for raw in cursor:
template = deserialize(raw)
if template is not None:
result.append(template)
return result
async def create_custom_template(
payload: ProxyTemplateCreateInput,
) -> ProxyTemplateDocument:
now = utcnow()
template = ProxyTemplateDocument(
id=payload.id,
name=payload.name,
description=payload.description,
category=payload.category,
protocol=payload.protocol,
family=payload.family,
tags=payload.tags,
action=payload.action,
summary=payload.summary,
behavior=payload.behavior,
code=payload.code,
source_url=payload.source_url,
built_in=False,
created_at=now,
updated_at=now,
)
templates = await collection()
await templates.insert_one(template.model_dump(exclude_none=True))
return template
async def delete_custom_template(template_id: str) -> bool:
templates = await collection()
result = await templates.delete_one({"id": template_id})
return result.deleted_count > 0
-152
View File
@@ -1,152 +0,0 @@
from models.example import AddonExampleDocument
DOCS_EXAMPLES_URL = "https://docs.mitmproxy.org/stable/addons/examples/"
BUILTIN_EXAMPLES = [
AddonExampleDocument(
id="http-tag-traffic",
name="Tag proxied traffic",
description=(
"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,
built_in=True,
code="""from mitmproxy import http
def response(flow: http.HTTPFlow) -> None:
flow.response.headers["x-theplug"] = "enabled"
""",
),
AddonExampleDocument(
id="http-force-query-flag",
name="Force request flag",
description=(
"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,
built_in=True,
code="""from mitmproxy import http
def request(flow: http.HTTPFlow) -> None:
flow.request.query["debug"] = "1"
""",
),
AddonExampleDocument(
id="http-local-health-response",
name="Serve local health response",
description=(
"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,
built_in=True,
code="""from mitmproxy import http
def request(flow: http.HTTPFlow) -> None:
if flow.request.path == "/healthz":
flow.response = http.Response.make(
200,
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 flip a "
"runtime marker while the flow is passing through the proxy."
),
category="websocket",
source_url=DOCS_EXAMPLES_URL,
built_in=True,
code="""from mitmproxy import http
def websocket_message(flow: http.HTTPFlow) -> None:
if flow.websocket is None:
return
message = flow.websocket.messages[-1]
message.content = message.content.replace(b"debug=false", b"debug=true")
""",
),
AddonExampleDocument(
id="tcp-log-messages",
name="TCP log messages",
description=(
"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,
built_in=True,
code="""import logging
from mitmproxy import tcp
def tcp_message(flow: tcp.TCPFlow) -> None:
message = flow.messages[-1]
direction = "client" if message.from_client else "server"
logging.info("%s sent %d bytes", direction, len(message.content))
""",
),
AddonExampleDocument(
id="dns-spoof-a-record",
name="DNS spoof A record",
description=(
"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,
built_in=True,
code="""import ipaddress
from mitmproxy import dns
def dns_request(flow: dns.DNSFlow) -> None:
question = flow.request.question
if question and question.name == "target.local" and question.type == dns.types.A:
flow.response = flow.request.succeed(
[
dns.ResourceRecord(
name="target.local",
type=dns.types.A,
class_=dns.classes.IN,
ttl=60,
data=ipaddress.ip_address("127.0.0.1").packed,
)
]
)
""",
),
]
+83 -40
View File
@@ -1,8 +1,11 @@
import asyncio
import ssl
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from models.proxy import HealthCheckDocument, ProxyDocument
from app.models.health import HealthCheckEntry, HealthSnapshot
from app.models.proxy import HealthCheckDocument, ProxyDocument
@dataclass(slots=True)
@@ -13,33 +16,48 @@ class Target:
protocol: str
async def check_proxy(proxy: ProxyDocument) -> dict[str, object]:
async def check_proxy(proxy: ProxyDocument) -> HealthSnapshot:
timestamp = datetime.now(UTC)
if not proxy.healthcheck.enabled:
return {"enabled": False, "overall": "disabled", "checks": []}
return HealthSnapshot(
proxy_port=proxy.proxy_port,
timestamp=timestamp,
enabled=False,
overall="disabled",
checks=[],
)
upstream_port = proxy.service.upstream_port or proxy.service.port
checks = await asyncio.gather(
check_target(
Target(
name="upstream",
host=proxy.service.host,
port=upstream_port,
protocol=health_protocol(proxy),
checks = list(
await asyncio.gather(
check_target(
Target(
name="upstream",
host=proxy.service.host,
port=upstream_port,
protocol=health_protocol(proxy),
),
proxy.healthcheck,
),
proxy.healthcheck,
),
check_target(
Target(
name="proxy",
host=listen_host_for_check(proxy.listen_host),
port=proxy.proxy_port,
protocol=health_protocol(proxy),
check_target(
Target(
name="proxy",
host=listen_host_for_check(proxy.listen_host),
port=proxy.proxy_port,
protocol=health_protocol(proxy),
),
proxy.healthcheck,
),
proxy.healthcheck,
),
)
)
overall = "healthy" if all(check.healthy for check in checks) else "unhealthy"
return HealthSnapshot(
proxy_port=proxy.proxy_port,
timestamp=timestamp,
enabled=True,
overall=overall,
checks=checks,
)
overall = "healthy" if all(check["healthy"] for check in checks) else "unhealthy"
return {"enabled": True, "overall": overall, "checks": checks}
def health_protocol(proxy: ProxyDocument) -> str:
@@ -55,7 +73,7 @@ def listen_host_for_check(host: str) -> str:
async def check_target(
target: Target,
config: HealthCheckDocument,
) -> dict[str, object]:
) -> HealthCheckEntry:
try:
match target.protocol:
case "http":
@@ -81,7 +99,8 @@ async def check_http(
config: HealthCheckDocument,
*,
use_tls: bool,
) -> dict[str, object]:
) -> HealthCheckEntry:
started = time.perf_counter()
reader, writer = await open_connection(target, config, use_tls=use_tls)
try:
path = config.path if config.path.startswith("/") else f"/{config.path}"
@@ -101,7 +120,13 @@ async def check_http(
status = int(parts[1]) if len(parts) >= 2 and parts[1].isdigit() else None
expected = config.expected_status
healthy = status == expected if expected is not None else status is not None
return result(target, healthy, None, {"status": status})
return result(
target,
healthy,
None,
{"status": status},
benchmark_ms=elapsed_ms(started) if config.benchmark_enabled else None,
)
async def check_tcp(
@@ -109,14 +134,20 @@ async def check_tcp(
config: HealthCheckDocument,
*,
use_tls: bool,
) -> dict[str, object]:
) -> HealthCheckEntry:
started = time.perf_counter()
_reader, writer = await open_connection(target, config, use_tls=use_tls)
writer.close()
await writer.wait_closed()
return result(target, True)
return result(
target,
True,
benchmark_ms=elapsed_ms(started) if config.benchmark_enabled else None,
)
async def check_udp(target: Target, config: HealthCheckDocument) -> dict[str, object]:
async def check_udp(target: Target, config: HealthCheckDocument) -> HealthCheckEntry:
started = time.perf_counter()
loop = asyncio.get_running_loop()
transport, protocol = await loop.create_datagram_endpoint(
lambda: DatagramProbe(config.expect_response),
@@ -128,7 +159,13 @@ async def check_udp(target: Target, config: HealthCheckDocument) -> dict[str, ob
if not config.expect_response:
return result(target, True, None, {"mode": "send-only"})
data = await asyncio.wait_for(protocol.response, config.timeout_seconds)
return result(target, True, None, {"bytes": len(data)})
return result(
target,
True,
None,
{"bytes": len(data)},
benchmark_ms=elapsed_ms(started) if config.benchmark_enabled else None,
)
finally:
transport.close()
@@ -159,7 +196,7 @@ class DatagramProbe(asyncio.DatagramProtocol):
if not expect_response:
self.response.set_result(b"")
def datagram_received(self, data, _addr):
def datagram_received(self, data, addr):
if not self.response.done():
self.response.set_result(data)
@@ -173,13 +210,19 @@ def result(
healthy: bool,
error: str | None = None,
details: dict[str, object] | None = None,
) -> dict[str, object]:
return {
"target": target.name,
"host": target.host,
"port": target.port,
"protocol": target.protocol,
"healthy": healthy,
"error": error,
"details": details or {},
}
benchmark_ms: float | None = None,
) -> HealthCheckEntry:
return HealthCheckEntry(
target=target.name,
host=target.host,
port=target.port,
protocol=target.protocol,
healthy=healthy,
error=error,
duration_ms=benchmark_ms,
details=details or {},
)
def elapsed_ms(started: float) -> float:
return round((time.perf_counter() - started) * 1000, 2)
@@ -0,0 +1,117 @@
import asyncio
import logging
from app.core.config import settings
from app.models.health import HealthSnapshot
from app.models.proxy import ProxyDocument
from app.repositories.health import record_health_sample
from app.repositories.proxy import list_proxies
from app.services.mitmproxy.health import check_proxy
logger = logging.getLogger(__name__)
class HealthMonitor:
def __init__(self):
self.task: asyncio.Task[None] | None = None
self.subscribers: dict[int, set[asyncio.Queue[HealthSnapshot]]] = {}
async def start(self) -> None:
if self.task is None or self.task.done():
self.task = asyncio.create_task(self.run())
async def stop(self) -> None:
if self.task is None:
return
self.task.cancel()
await asyncio.gather(self.task, return_exceptions=True)
self.task = None
async def run(self) -> None:
loop = asyncio.get_running_loop()
next_cycle = loop.time()
while True:
delay = next_cycle - loop.time()
if delay > 0:
await asyncio.sleep(delay)
cycle_started = loop.time()
try:
await self.run_cycle()
except asyncio.CancelledError:
raise
except Exception:
logger.exception("Health monitor cycle failed")
next_cycle = loop.time() + min(
1.0,
float(settings.HEALTHCHECK_INTERVAL_SECONDS),
)
continue
next_cycle = cycle_started + settings.HEALTHCHECK_INTERVAL_SECONDS
async def run_cycle(self) -> None:
proxies = [
proxy
for proxy in await list_proxies()
if proxy.status == "running" and proxy.healthcheck.enabled
]
async with asyncio.TaskGroup() as group:
for proxy in proxies:
group.create_task(self.run_check(proxy))
async def run_check(self, proxy: ProxyDocument) -> None:
try:
snapshot = await check_proxy(proxy)
except asyncio.CancelledError:
raise
except Exception:
logger.exception(
"Health check failed for proxy %s",
proxy.proxy_port,
)
return
try:
await record_health_sample(snapshot)
except asyncio.CancelledError:
raise
except Exception:
logger.exception(
"Health snapshot persistence failed for proxy %s",
proxy.proxy_port,
)
await self.publish(proxy.proxy_port, snapshot)
async def publish(self, proxy_port: int, snapshot: HealthSnapshot) -> None:
for queue in list(self.subscribers.get(proxy_port, set())):
try:
queue.put_nowait(snapshot)
except asyncio.QueueFull:
logger.debug(
"Dropped health update for slow subscriber on proxy %s",
proxy_port,
)
continue
def subscribe(self, proxy_port: int) -> asyncio.Queue[HealthSnapshot]:
queue: asyncio.Queue[HealthSnapshot] = asyncio.Queue(maxsize=8)
listeners = self.subscribers.setdefault(proxy_port, set())
listeners.add(queue)
return queue
def unsubscribe(
self,
proxy_port: int,
queue: asyncio.Queue[HealthSnapshot],
) -> None:
listeners = self.subscribers.get(proxy_port)
if listeners is None:
return
listeners.discard(queue)
if not listeners:
self.subscribers.pop(proxy_port, None)
health_monitor = HealthMonitor()
+93 -32
View File
@@ -5,10 +5,15 @@ 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
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_validation import (
ReservedProxyPortError,
ensure_proxy_port_not_reserved,
)
class MitmdumpError(Exception):
@@ -22,6 +27,7 @@ class MitmdumpManager:
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):
@@ -31,7 +37,30 @@ class MitmdumpManager:
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")
@@ -44,9 +73,9 @@ class MitmdumpManager:
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)
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,
@@ -64,7 +93,10 @@ class MitmdumpManager:
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"
reason = (
await self.read_log_tail(log_path)
or "mitmdump exited during startup"
)
raise MitmdumpError(reason)
self.processes[proxy.proxy_port] = process
@@ -174,18 +206,22 @@ class MitmdumpManager:
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,
host_to="127.0.0.1",
rewrite_host=True,
)
if not changed:
raise MitmdumpError("Unable to move Docker service port")
return upstream_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,
@@ -199,7 +235,7 @@ class MitmdumpManager:
if upstream_port == proxy.service.port:
return True
return await docker.change_port(
return await change_port(
proxy.service.name,
Path(proxy.service.folder),
upstream_port,
@@ -257,22 +293,29 @@ class MitmdumpManager:
return False
return proxy.mode in {"regular", "socks5", "dns", "transparent"}
def write_script(self, proxy: ProxyDocument) -> Path:
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:
path = self.script_path_for(proxy)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(proxy.code, encoding="utf-8")
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
await asyncio.to_thread(path.write_text, proxy.code, encoding="utf-8")
return path
def write_config(self, proxy: ProxyDocument) -> Path:
async def write_config(self, proxy: ProxyDocument) -> Path:
path = self.config_path_for(proxy)
path.parent.mkdir(parents=True, exist_ok=True)
await asyncio.to_thread(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")
config_text = yaml.safe_dump(options, sort_keys=True)
await asyncio.to_thread(path.write_text, config_text, encoding="utf-8")
return path
def log_path(self, proxy: ProxyDocument) -> Path:
async def log_path(self, proxy: ProxyDocument) -> Path:
path = self.log_path_for(proxy)
path.parent.mkdir(parents=True, exist_ok=True)
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
return path
async def spawn(self, command: list[str], log_path: Path):
@@ -385,12 +428,30 @@ class MitmdumpManager:
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():
async def read_log_tail(self, path: Path, limit: int = 4000) -> str:
if not await asyncio.to_thread(path.exists):
return ""
content = path.read_text(encoding="utf-8", errors="replace")
content = await asyncio.to_thread(
path.read_text,
encoding="utf-8",
errors="replace",
)
return content[-limit:].strip()
def remove_file(self, path: Path) -> None:
try:
path.unlink(missing_ok=True)
except OSError:
return
def remove_dir(self, path: Path) -> None:
try:
shutil.rmtree(path)
except FileNotFoundError:
return
except OSError:
return
def lock_for(self, proxy_port: int) -> asyncio.Lock:
lock = self.locks.get(proxy_port)
if lock is None:
+2 -2
View File
@@ -1,5 +1,5 @@
from repositories.proxy import list_proxies, update_proxy_runtime
from services.mitmproxy.manager import manager
from app.repositories.proxy import list_proxies, update_proxy_runtime
from app.services.mitmproxy.manager import manager
async def reconcile_runtime() -> list[tuple[int, dict[str, object]]]:
@@ -0,0 +1,21 @@
from app.core.config import settings
class ReservedProxyPortError(ValueError):
pass
def backend_port() -> int:
return settings.API_PORT
def reserved_proxy_port_message(proxy_port: int) -> str:
return (
f"Proxy port {proxy_port} is reserved for the backend API "
"and cannot be used by mitmdump."
)
def ensure_proxy_port_not_reserved(proxy_port: int) -> None:
if proxy_port == backend_port():
raise ReservedProxyPortError(reserved_proxy_port_message(proxy_port))
+524
View File
@@ -0,0 +1,524 @@
from app.models.template import ProxyTemplateDocument
DOCS_TEMPLATES_URL = "https://docs.mitmproxy.org/stable/addons/examples/"
def builtin_template(
*,
id: str,
name: str,
protocol: str,
family: str,
tags: list[str],
action: str,
summary: str,
behavior: str,
code: str,
) -> ProxyTemplateDocument:
category = protocol if protocol in {"http", "tcp"} else family
return ProxyTemplateDocument(
id=id,
name=name,
description=summary,
category=category,
protocol=protocol,
family=family,
tags=tags,
action=action,
summary=summary,
behavior=behavior,
source_url=DOCS_TEMPLATES_URL,
built_in=True,
code=code,
)
BUILTIN_TEMPLATES = [
builtin_template(
id="http-block-path-traversal",
name="Block path traversal",
protocol="http",
family="web-hardening",
tags=["path-traversal", "web", "block"],
action="block",
summary="Stops common path traversal probes in path and query.",
behavior=(
"Logs the matched traversal token and immediately drops the client "
"connection before the request reaches the upstream service."
),
code="""import logging
from urllib.parse import unquote
from mitmproxy import http
TRAVERSAL_MARKERS = (
"../",
"..\\\\",
"%2e%2e",
"%252e%252e",
"..%2f",
"..%5c",
)
def request(flow: http.HTTPFlow) -> None:
raw_target = (
f"{flow.request.path}?{flow.request.query}"
if flow.request.query
else flow.request.path
)
normalized = unquote(unquote(raw_target)).lower()
for marker in TRAVERSAL_MARKERS:
if marker in normalized:
logging.warning(
"[http-block-path-traversal] blocked %s %s from %s",
flow.request.method,
flow.request.pretty_url,
flow.client_conn.peername,
)
flow.kill()
return
""",
),
builtin_template(
id="http-block-blacklisted-substrings",
name="Block blacklisted substrings",
protocol="http",
family="web-hardening",
tags=["blacklist-string", "web", "block"],
action="block",
summary="Blocks requests containing forbidden strings in path, query, or body.",
behavior=(
"Scans URL and request body for blacklisted markers, then serves a "
"blank HTML page with 200 OK and logs the first hit."
),
code="""import logging
from mitmproxy import http
BLOCKED_STRINGS = (
"flag{",
"give_me_shell",
"debug=1",
"/etc/passwd",
)
def request(flow: http.HTTPFlow) -> None:
body = flow.request.get_text(strict=False)
haystack = "\\n".join(
[flow.request.path, str(flow.request.query), body]
).lower()
for marker in BLOCKED_STRINGS:
if marker.lower() in haystack:
logging.warning(
"[http-block-blacklisted-substrings] blocked marker=%r url=%s",
marker,
flow.request.pretty_url,
)
flow.response = http.Response.make(
200,
b"",
{"content-type": "text/html; charset=utf-8"},
)
return
""",
),
builtin_template(
id="http-block-large-body",
name="Block large request bodies",
protocol="http",
family="web-hardening",
tags=["length-limit", "body-size", "web", "block"],
action="block",
summary="Rejects oversized HTTP request bodies before they hit the service.",
behavior=(
"Checks the buffered request body length, returns 413 Payload Too "
"Large, and logs the observed size."
),
code="""import logging
from mitmproxy import http
MAX_BODY_BYTES = 8192
def request(flow: http.HTTPFlow) -> None:
size = len(flow.request.raw_content or b"")
if size > MAX_BODY_BYTES:
logging.warning(
"[http-block-large-body] blocked %s bytes on %s",
size,
flow.request.pretty_url,
)
flow.response = http.Response.make(
413,
b"",
{"content-type": "text/plain"},
)
""",
),
builtin_template(
id="http-block-suspicious-methods",
name="Block suspicious HTTP methods",
protocol="http",
family="web-hardening",
tags=["method-filter", "web", "block"],
action="block",
summary="Denies uncommon or dangerous HTTP methods.",
behavior=(
"Returns 405 Method Not Allowed for methods outside the allowlist "
"and logs the blocked method and target URL."
),
code="""import logging
from mitmproxy import http
ALLOWED_METHODS = {"GET", "POST", "HEAD"}
def request(flow: http.HTTPFlow) -> None:
if flow.request.method not in ALLOWED_METHODS:
logging.warning(
"[http-block-suspicious-methods] blocked method=%s url=%s",
flow.request.method,
flow.request.pretty_url,
)
flow.response = http.Response.make(
405,
b"",
{"allow": "GET, POST, HEAD"},
)
""",
),
builtin_template(
id="http-block-admin-paths",
name="Block admin and debug paths",
protocol="http",
family="web-hardening",
tags=["admin-path", "debug-path", "web", "block"],
action="block",
summary="Protects common admin, console, and debug endpoints.",
behavior=(
"Matches the request path against a small denylist and returns a "
"blank 404 page before the request is forwarded."
),
code="""import logging
from mitmproxy import http
BLOCKED_PREFIXES = (
"/admin",
"/debug",
"/console",
"/actuator",
)
def request(flow: http.HTTPFlow) -> None:
if flow.request.path.startswith(BLOCKED_PREFIXES):
logging.warning(
"[http-block-admin-paths] blocked path=%s",
flow.request.path,
)
flow.response = http.Response.make(
404,
b"",
{"content-type": "text/html; charset=utf-8"},
)
""",
),
builtin_template(
id="http-sanitize-headers-basic",
name="Sanitize risky forwarding headers",
protocol="http",
family="web-hardening",
tags=["rewrite", "headers", "web"],
action="rewrite",
summary=(
"Removes a few risky forwarding headers before the request is "
"sent upstream."
),
behavior=(
"Deletes selected override or forwarding headers and logs what was "
"rewritten without blocking the request."
),
code="""import logging
from mitmproxy import http
STRIP_HEADERS = (
"x-forwarded-host",
"x-http-method-override",
"x-original-url",
"x-rewrite-url",
)
def request(flow: http.HTTPFlow) -> None:
removed = []
for header in STRIP_HEADERS:
if header in flow.request.headers:
removed.append(header)
del flow.request.headers[header]
if removed:
logging.info(
"[http-sanitize-headers-basic] stripped headers=%s url=%s",
",".join(removed),
flow.request.pretty_url,
)
""",
),
builtin_template(
id="http-block-command-injection-markers",
name="Block command injection markers",
protocol="http",
family="web-hardening",
tags=["command-injection", "web", "block"],
action="block",
summary=(
"Rejects obvious shell metacharacters and command-substitution "
"markers."
),
behavior=(
"Scans URL and body for shell-style injection markers, logs the "
"first hit, and drops the client connection."
),
code="""import logging
from mitmproxy import http
BLOCKED_MARKERS = (";", "&&", "||", "$(", "`", "|sh", "|bash")
def request(flow: http.HTTPFlow) -> None:
body = flow.request.get_text(strict=False)
haystack = "\\n".join(
[flow.request.pretty_url, body]
).lower()
for marker in BLOCKED_MARKERS:
if marker in haystack:
logging.warning(
"[http-block-command-injection-markers] blocked marker=%r url=%s",
marker,
flow.request.pretty_url,
)
flow.kill()
return
""",
),
builtin_template(
id="http-block-sqli-markers",
name="Block obvious SQLi markers",
protocol="http",
family="web-hardening",
tags=["sqli", "web", "block"],
action="block",
summary="Stops common SQL injection probes before they reach the app.",
behavior=(
"Matches a short list of high-signal SQLi markers in path, query, "
"and body, then serves a blank HTML page and logs the hit."
),
code="""import logging
from mitmproxy import http
BLOCKED_MARKERS = (
"' or 1=1",
'" or 1=1',
"union select",
"sleep(",
"information_schema",
"--",
)
def request(flow: http.HTTPFlow) -> None:
body = flow.request.get_text(strict=False)
haystack = "\\n".join(
[flow.request.path, str(flow.request.query), body]
).lower()
for marker in BLOCKED_MARKERS:
if marker in haystack:
logging.warning(
"[http-block-sqli-markers] blocked marker=%r url=%s",
marker,
flow.request.pretty_url,
)
flow.response = http.Response.make(
200,
b"",
{"content-type": "text/html; charset=utf-8"},
)
return
""",
),
builtin_template(
id="tcp-block-repeated-byte-threshold",
name="Block repeated-byte flood",
protocol="tcp",
family="tcp-hardening",
tags=["bof", "repeated-byte", "tcp", "block"],
action="block",
summary="Closes connections that send suspiciously long runs of the same byte.",
behavior=(
"Measures the longest repeated-byte run in each client message, logs "
"the byte and run length, and closes the TCP flow on threshold hit."
),
code="""import logging
from mitmproxy import tcp
MAX_REPEAT = 128
def longest_run(data: bytes) -> tuple[int, int]:
best_byte = -1
best_run = 0
current_byte = -1
current_run = 0
for byte in data:
if byte == current_byte:
current_run += 1
else:
current_byte = byte
current_run = 1
if current_run > best_run:
best_byte = byte
best_run = current_run
return best_byte, best_run
def tcp_message(flow: tcp.TCPFlow) -> None:
message = flow.messages[-1]
if not message.from_client:
return
repeated_byte, run_length = longest_run(message.content)
if run_length > MAX_REPEAT:
logging.warning(
"[tcp-block-repeated-byte-threshold] blocked byte=0x%02x run=%d",
repeated_byte,
run_length,
)
message.drop()
flow.kill()
""",
),
builtin_template(
id="tcp-block-message-length-threshold",
name="Block oversized TCP messages",
protocol="tcp",
family="tcp-hardening",
tags=["length-limit", "tcp", "block"],
action="block",
summary="Cuts off client messages that exceed a fixed size threshold.",
behavior=(
"Checks every client TCP message size, logs the offending length, "
"drops the message, and kills the flow when the threshold is hit."
),
code="""import logging
from mitmproxy import tcp
MAX_MESSAGE_BYTES = 4096
def tcp_message(flow: tcp.TCPFlow) -> None:
message = flow.messages[-1]
if not message.from_client:
return
size = len(message.content)
if size > MAX_MESSAGE_BYTES:
logging.warning(
"[tcp-block-message-length-threshold] blocked message size=%d",
size,
)
message.drop()
flow.kill()
""",
),
builtin_template(
id="tcp-block-blacklisted-byte-sequences",
name="Block blacklisted TCP byte sequences",
protocol="tcp",
family="tcp-hardening",
tags=["blacklist-bytes", "tcp", "block"],
action="block",
summary=(
"Blocks payloads containing forbidden byte sequences or marker "
"strings."
),
behavior=(
"Looks for dangerous byte patterns in client messages, then drops "
"the message and closes the flow."
),
code="""import logging
from mitmproxy import tcp
BLOCKED_SEQUENCES = (
b"/bin/sh",
b"cat flag",
b"AAAA%p%p",
b"\\x90\\x90\\x90\\x90",
)
def tcp_message(flow: tcp.TCPFlow) -> None:
message = flow.messages[-1]
if not message.from_client:
return
for sequence in BLOCKED_SEQUENCES:
if sequence in message.content:
logging.warning(
"[tcp-block-blacklisted-byte-sequences] blocked sequence=%r",
sequence,
)
message.drop()
flow.kill()
return
""",
),
builtin_template(
id="tcp-block-non-printable-threshold",
name="Block binary-noise threshold",
protocol="tcp",
family="tcp-hardening",
tags=["non-printable", "binary-noise", "tcp", "block"],
action="block",
summary=(
"Blocks client messages with an excessive amount of non-printable "
"bytes."
),
behavior=(
"Treats unusually binary-heavy client payloads as suspicious, logs "
"the ratio, and closes the flow when the threshold is exceeded."
),
code="""import logging
from mitmproxy import tcp
NON_PRINTABLE_RATIO = 0.35
def tcp_message(flow: tcp.TCPFlow) -> None:
message = flow.messages[-1]
if not message.from_client or not message.content:
return
non_printable = sum(
1
for byte in message.content
if byte < 0x20 and byte not in {0x09, 0x0a, 0x0d}
)
ratio = non_printable / len(message.content)
if ratio > NON_PRINTABLE_RATIO:
logging.warning(
"[tcp-block-non-printable-threshold] blocked ratio=%.2f size=%d",
ratio,
len(message.content),
)
message.drop()
flow.kill()
""",
),
]