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
+2
View File
@@ -14,3 +14,5 @@ AUTH_PASSWORD=change-me
# Optional: override the directory whose immediate child folders are scanned
# for docker-compose files. Defaults to the current user's home directory.
SERVICE_SCAN_ROOT=
HEALTHCHECK_INTERVAL_SECONDS=15
HEALTH_HISTORY_TTL_SECONDS=86400
+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()
""",
),
]
+10 -1
View File
@@ -29,8 +29,17 @@ indent-style = "space"
[tool.pytest.ini_options]
testpaths = ["tests"]
pythonpath = ["app"]
pythonpath = ["."]
addopts = "-q"
markers = [
"unit: fast isolated tests",
"integration: tests using real app wiring and persistence",
"api: HTTP API behavior tests",
"websocket: websocket behavior tests",
"slow: slower-running tests",
"requires_docker: tests that require Docker or containerized dependencies",
"perf: performance and benchmark scenarios",
]
[dependency-groups]
dev = [
+3 -3
View File
@@ -1,10 +1,10 @@
import importlib
from datetime import UTC, datetime, timedelta
from core.security import AuthContext, get_auth_context
from models.auth import AccessTokenPayload, AuthUserDocument
from app.core.security import AuthContext, get_auth_context
from app.models.auth import AccessTokenPayload, AuthUserDocument
auth_router_module = importlib.import_module("api.auth.router")
auth_router_module = importlib.import_module("app.api.auth.router")
def test_login_success_returns_bearer_token(client, monkeypatch: object):
+13 -6
View File
@@ -1,15 +1,20 @@
import pytest
from services.editor import python_lsp
from starlette.websockets import WebSocketDisconnect
from app.services.editor import python_lsp
def test_python_editor_session(client, monkeypatch):
session = python_lsp.PythonEditorSession(
root_path=python_lsp.WORKSPACE_ROOT,
file_path=python_lsp.WORKSPACE_ROOT / python_lsp.ENTRY_FILE,
)
monkeypatch.setattr("api.editor.router.python_lsp.lsp_binary", lambda: "/bin/lsp")
monkeypatch.setattr("api.editor.router.python_lsp.get_session", lambda: session)
monkeypatch.setattr(
"app.api.editor.router.python_lsp.lsp_binary", lambda: "/bin/lsp"
)
monkeypatch.setattr(
"app.api.editor.router.python_lsp.get_session", lambda: session
)
response = client.get("/proxy/editor/python/session")
@@ -40,11 +45,13 @@ def test_python_editor_socket_accepts_token(client, monkeypatch):
await websocket.close()
monkeypatch.setattr(
"api.editor.router.authenticate_editor_socket",
"app.api.editor.router.authenticate_editor_socket",
fake_authenticate,
)
monkeypatch.setattr("api.editor.router.python_lsp.lsp_binary", lambda: "/bin/lsp")
monkeypatch.setattr("api.editor.router.python_lsp.bridge", fake_bridge)
monkeypatch.setattr(
"app.api.editor.router.python_lsp.lsp_binary", lambda: "/bin/lsp"
)
monkeypatch.setattr("app.api.editor.router.python_lsp.bridge", fake_bridge)
with client.websocket_connect("/proxy/editor/python/ws?token=test-token") as socket:
assert socket.receive_text().endswith("/proxy_addon.py")
+385 -21
View File
@@ -1,14 +1,18 @@
import asyncio
import importlib
from datetime import UTC, datetime
import pytest
from core.security import get_auth_context, require_authenticated_user
from models.proxy import ProxyCreateInput, ProxyDocument
from models.service import ServiceEndpointDocument
from services.mitmproxy.examples import BUILTIN_EXAMPLES
from services.mitmproxy.validator import ProxyCodeValidationError
from starlette.websockets import WebSocketDisconnect
proxy_router_module = importlib.import_module("api.proxy.router")
from app.core.security import get_auth_context, require_authenticated_user
from app.models.health import HealthSnapshot
from app.models.proxy import ProxyCreateInput, ProxyDocument
from app.models.service import ServiceEndpointDocument
from app.services.mitmproxy.templates import BUILTIN_TEMPLATES
from app.services.mitmproxy.validator import ProxyCodeValidationError
proxy_router_module = importlib.import_module("app.api.proxy.router")
def make_proxy_document(
@@ -28,6 +32,21 @@ def make_proxy_document(
)
def make_health_snapshot(
proxy_port: int,
*,
overall: str = "healthy",
timestamp: datetime | None = None,
) -> HealthSnapshot:
return HealthSnapshot(
proxy_port=proxy_port,
timestamp=timestamp or datetime.now(UTC),
enabled=True,
overall=overall,
checks=[],
)
def test_proxy_requires_auth_without_dependency_overrides(client):
client.app.dependency_overrides.pop(require_authenticated_user, None)
client.app.dependency_overrides.pop(get_auth_context, None)
@@ -80,6 +99,7 @@ def test_list_services_uses_query_scan_root(client, monkeypatch):
assert response.status_code == 200
assert response.json() == {
"backend_port": 3671,
"scan_root": "/root",
"services": [
{"name": "echo", "ports": [9081], "bindings": [], "folder": "/root/play"}
@@ -105,10 +125,10 @@ def test_get_proxy_returns_404_when_not_found(client, monkeypatch: object):
def test_create_proxy_returns_409_on_port_conflict(client, monkeypatch: object):
existing = make_proxy_document(proxy_port=3671)
existing = make_proxy_document(proxy_port=4088)
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
assert proxy_port == 3671
assert proxy_port == 4088
return existing
monkeypatch.setattr(
@@ -120,7 +140,7 @@ def test_create_proxy_returns_409_on_port_conflict(client, monkeypatch: object):
response = client.post(
"/proxy",
json={
"proxy_port": 3671,
"proxy_port": 4088,
"name": "new",
"code": "abc",
"service": {"port": 8080},
@@ -128,7 +148,44 @@ def test_create_proxy_returns_409_on_port_conflict(client, monkeypatch: object):
)
assert response.status_code == 409
assert response.json()["detail"] == "Proxy already exists on port 3671"
assert response.json()["detail"] == "Proxy already exists on port 4088"
def test_create_proxy_rejects_backend_api_port(client, monkeypatch: object):
create_called = False
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
assert proxy_port == 3671
return None
async def fake_create_proxy(payload: ProxyCreateInput) -> ProxyDocument:
nonlocal create_called
create_called = True
return make_proxy_document(proxy_port=payload.proxy_port)
monkeypatch.setattr(
proxy_router_module,
"get_proxy_by_port",
fake_get_proxy_by_port,
)
monkeypatch.setattr(proxy_router_module, "create_proxy", fake_create_proxy)
response = client.post(
"/proxy",
json={
"proxy_port": 3671,
"name": "backend-clash",
"code": "abc",
"service": {"port": 8080},
},
)
assert response.status_code == 422
assert response.json()["detail"] == (
"Proxy port 3671 is reserved for the backend API "
"and cannot be used by mitmdump."
)
assert create_called is False
def test_create_proxy_returns_created_document(client, monkeypatch: object):
@@ -470,6 +527,297 @@ def test_proxy_logs_socket_streams_snapshot_and_append(client, monkeypatch, tmp_
assert appended == {"type": "append", "log": "line two\n"}
def test_proxy_logs_socket_resets_snapshot_after_log_rotation(
client,
monkeypatch,
tmp_path,
):
log_path = tmp_path / "proxy.log"
log_path.write_text("line one\n", encoding="utf-8")
proxy = make_proxy_document(proxy_port=4040, log_path=str(log_path))
async def fake_authenticate(_: object) -> None:
return None
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
assert proxy_port == 4040
return proxy
monkeypatch.setattr(
proxy_router_module,
"authenticate_proxy_socket",
fake_authenticate,
)
monkeypatch.setattr(
proxy_router_module,
"get_proxy_by_port",
fake_get_proxy_by_port,
)
with client.websocket_connect("/proxy/4040/logs/ws?token=test-token") as socket:
snapshot = socket.receive_json()
assert snapshot == {"type": "snapshot", "log": "line one\n"}
log_path.write_text("fresh start\n", encoding="utf-8")
rotated = socket.receive_json()
assert rotated == {"type": "snapshot", "log": "fresh start\n"}
with log_path.open("a", encoding="utf-8") as handle:
handle.write("line two\n")
handle.flush()
appended = socket.receive_json()
assert appended == {"type": "append", "log": "line two\n"}
def test_remove_proxy_stops_runtime_and_cleans_artifacts(client, monkeypatch: object):
proxy = make_proxy_document(
proxy_port=4040,
log_path="/tmp/theplug/logs/4040.log",
).model_copy(
update={
"script_path": "/tmp/theplug/addons/4040.py",
"config_path": "/tmp/theplug/conf/4040/config.yaml",
}
)
captured: dict[str, object] = {}
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
assert proxy_port == 4040
return proxy
class FakeManager:
async def stop(self, document: ProxyDocument) -> dict[str, object]:
captured["stopped"] = document.proxy_port
return {
"status": "stopped",
"pid": None,
"service.upstream_port": None,
"last_error": None,
}
def cleanup_runtime_artifacts(self, document: ProxyDocument) -> None:
captured["cleaned"] = {
"proxy_port": document.proxy_port,
"script_path": document.script_path,
"config_path": document.config_path,
"log_path": document.log_path,
}
async def fake_delete_proxy(proxy_port: int) -> bool:
captured["deleted"] = proxy_port
return True
async def fake_delete_health_history(proxy_port: int) -> None:
captured["history_deleted"] = proxy_port
monkeypatch.setattr(
proxy_router_module,
"get_proxy_by_port",
fake_get_proxy_by_port,
)
monkeypatch.setattr(proxy_router_module, "manager", FakeManager())
monkeypatch.setattr(
proxy_router_module,
"delete_health_history",
fake_delete_health_history,
)
monkeypatch.setattr(proxy_router_module, "delete_proxy", fake_delete_proxy)
response = client.delete("/proxy/4040")
assert response.status_code == 204
assert captured == {
"stopped": 4040,
"cleaned": {
"proxy_port": 4040,
"script_path": "/tmp/theplug/addons/4040.py",
"config_path": "/tmp/theplug/conf/4040/config.yaml",
"log_path": "/tmp/theplug/logs/4040.log",
},
"deleted": 4040,
"history_deleted": 4040,
}
def test_proxy_health_runs_immediate_check_and_publishes(client, monkeypatch: object):
proxy = make_proxy_document(proxy_port=4040)
captured: dict[str, object] = {}
snapshot = make_health_snapshot(4040)
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
assert proxy_port == 4040
return proxy
async def fake_check_proxy(document: ProxyDocument):
captured["checked"] = document.proxy_port
return snapshot
async def fake_record_health_sample(document):
captured["recorded"] = document
return document
class FakeHealthMonitor:
async def publish(self, proxy_port: int, document: HealthSnapshot) -> None:
captured["published"] = (proxy_port, document)
monkeypatch.setattr(
proxy_router_module,
"get_proxy_by_port",
fake_get_proxy_by_port,
)
monkeypatch.setattr(proxy_router_module, "check_proxy", fake_check_proxy)
monkeypatch.setattr(
proxy_router_module,
"record_health_sample",
fake_record_health_sample,
)
monkeypatch.setattr(proxy_router_module, "health_monitor", FakeHealthMonitor())
response = client.get("/proxy/4040/health")
assert response.status_code == 200
assert response.json()["overall"] == "healthy"
assert captured["checked"] == 4040
assert captured["recorded"] == snapshot
assert captured["published"] == (4040, snapshot)
def test_proxy_health_history_returns_latest_and_history(client, monkeypatch: object):
proxy = make_proxy_document(proxy_port=4040)
timestamp = datetime.now(UTC)
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
assert proxy_port == 4040
return proxy
async def fake_get_latest_health_sample(proxy_port: int):
assert proxy_port == 4040
return make_health_snapshot(4040, timestamp=timestamp)
async def fake_list_health_history(proxy_port: int, limit: int = 40):
assert proxy_port == 4040
assert limit == 40
return [make_health_snapshot(4040, timestamp=timestamp)]
monkeypatch.setattr(
proxy_router_module,
"get_proxy_by_port",
fake_get_proxy_by_port,
)
monkeypatch.setattr(
proxy_router_module,
"get_latest_health_sample",
fake_get_latest_health_sample,
)
monkeypatch.setattr(
proxy_router_module,
"list_health_history",
fake_list_health_history,
)
response = client.get("/proxy/4040/health/history")
assert response.status_code == 200
body = response.json()
assert body["proxy_port"] == 4040
assert body["latest"]["overall"] == "healthy"
assert len(body["history"]) == 1
def test_proxy_health_history_clamps_requested_limit(client, monkeypatch: object):
proxy = make_proxy_document(proxy_port=4040)
captured: dict[str, int] = {}
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
assert proxy_port == 4040
return proxy
async def fake_get_latest_health_sample(_proxy_port: int) -> HealthSnapshot | None:
return None
async def fake_list_health_history(proxy_port: int, limit: int = 40):
assert proxy_port == 4040
captured["limit"] = limit
return []
monkeypatch.setattr(
proxy_router_module,
"get_proxy_by_port",
fake_get_proxy_by_port,
)
monkeypatch.setattr(
proxy_router_module,
"get_latest_health_sample",
fake_get_latest_health_sample,
)
monkeypatch.setattr(
proxy_router_module,
"list_health_history",
fake_list_health_history,
)
response = client.get("/proxy/4040/health/history", params={"limit": 9999})
assert response.status_code == 200
assert captured["limit"] == 200
def test_proxy_health_socket_streams_latest_and_updates(client, monkeypatch: object):
proxy = make_proxy_document(proxy_port=4040)
timestamp = datetime.now(UTC)
follow_up = make_health_snapshot(4040, overall="unhealthy")
captured: dict[str, object] = {}
async def fake_authenticate(_: object) -> None:
return None
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
assert proxy_port == 4040
return proxy
async def fake_get_latest_health_sample(proxy_port: int):
assert proxy_port == 4040
return make_health_snapshot(4040, timestamp=timestamp)
class FakeHealthMonitor:
def subscribe(self, proxy_port: int):
assert proxy_port == 4040
queue = asyncio.Queue()
queue.put_nowait(follow_up)
return queue
def unsubscribe(self, proxy_port: int, _queue) -> None:
assert proxy_port == 4040
captured["unsubscribed"] = True
monkeypatch.setattr(
proxy_router_module,
"authenticate_proxy_socket",
fake_authenticate,
)
monkeypatch.setattr(
proxy_router_module,
"get_proxy_by_port",
fake_get_proxy_by_port,
)
monkeypatch.setattr(
proxy_router_module,
"get_latest_health_sample",
fake_get_latest_health_sample,
)
monkeypatch.setattr(proxy_router_module, "health_monitor", FakeHealthMonitor())
with client.websocket_connect("/proxy/4040/health/ws?token=test-token") as socket:
first = socket.receive_json()
assert first["overall"] == "healthy"
second = socket.receive_json()
assert second["overall"] == "unhealthy"
assert captured["unsubscribed"] is True
def test_stop_proxy_updates_runtime_state(client, monkeypatch: object):
proxy = make_proxy_document(proxy_port=4040).model_copy(
update={"status": "running", "pid": 1234}
@@ -515,37 +863,46 @@ def test_stop_proxy_updates_runtime_state(client, monkeypatch: object):
assert response.json()["status"] == "stopped"
def test_list_mitmproxy_examples_returns_builtins_and_custom(client, monkeypatch):
async def fake_list_custom_examples():
def test_list_mitmproxy_templates_returns_builtins_and_custom(client, monkeypatch):
async def fake_list_custom_templates():
return []
monkeypatch.setattr(
proxy_router_module,
"list_custom_examples",
fake_list_custom_examples,
"list_custom_templates",
fake_list_custom_templates,
)
response = client.get("/proxy/mitmproxy/examples")
response = client.get("/proxy/mitmproxy/templates")
assert response.status_code == 200
body = response.json()
assert len(body) == len(BUILTIN_EXAMPLES)
assert len(body) == len(BUILTIN_TEMPLATES)
assert body[0]["built_in"] is True
assert body[0]["protocol"] in {"http", "tcp"}
assert isinstance(body[0]["tags"], list)
assert body[0]["action"] in {"block", "rewrite"}
def test_create_mitmproxy_example_validates_and_persists(client, monkeypatch):
def test_create_mitmproxy_template_validates_and_persists(client, monkeypatch):
captured = {}
async def fake_validate_proxy_code(code: str) -> None:
captured["validated"] = code
async def fake_create_custom_example(payload):
async def fake_create_custom_template(payload):
captured["payload"] = payload
return {
"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,
@@ -558,21 +915,28 @@ def test_create_mitmproxy_example_validates_and_persists(client, monkeypatch):
)
monkeypatch.setattr(
proxy_router_module,
"create_custom_example",
fake_create_custom_example,
"create_custom_template",
fake_create_custom_template,
)
response = client.post(
"/proxy/mitmproxy/examples",
"/proxy/mitmproxy/templates",
json={
"id": "custom-header",
"name": "Custom header",
"description": "Adds a custom response header",
"category": "http",
"protocol": "http",
"family": "custom",
"tags": ["observe"],
"action": "custom",
"summary": "Saved custom template.",
"behavior": "Applies the current script as-is.",
"code": "def response(flow):\n pass\n",
},
)
assert response.status_code == 201
assert response.json()["id"] == "custom-header"
assert response.json()["protocol"] == "http"
assert captured["validated"] == "def response(flow):\n pass\n"
+110 -8
View File
@@ -1,9 +1,14 @@
import asyncio
import os
from collections.abc import Iterator
from datetime import UTC, datetime, timedelta
from pathlib import Path
from uuid import uuid4
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from motor.motor_asyncio import AsyncIOMotorClient
os.environ.setdefault(
"MONGODB_URI",
@@ -19,7 +24,7 @@ os.environ.setdefault("AUTH_PASSWORD", "admin")
@pytest.fixture
def fake_user():
from models.auth import AuthUserDocument
from app.models.auth import AuthUserDocument
return AuthUserDocument(
username="admin",
@@ -29,8 +34,8 @@ def fake_user():
@pytest.fixture
def auth_context(fake_user):
from core.security import AuthContext
from models.auth import AccessTokenPayload
from app.core.security import AuthContext
from app.models.auth import AccessTokenPayload
exp = int((datetime.now(UTC) + timedelta(minutes=720)).timestamp())
payload = AccessTokenPayload(sub=fake_user.username, jti="test-jti", exp=exp)
@@ -39,11 +44,15 @@ def auth_context(fake_user):
@pytest.fixture
def app(fake_user, auth_context) -> FastAPI:
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.security import AuthContext, get_auth_context, require_authenticated_user
from models.auth import AuthUserDocument
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.security import (
AuthContext,
get_auth_context,
require_authenticated_user,
)
from app.models.auth import AuthUserDocument
app_instance = FastAPI()
app_instance.include_router(auth_router)
@@ -67,3 +76,96 @@ def app(fake_user, auth_context) -> FastAPI:
def client(app: FastAPI):
with TestClient(app) as test_client:
yield test_client
async def _mongo_ping(uri: str) -> str | None:
client = AsyncIOMotorClient(uri, serverSelectionTimeoutMS=500)
try:
await client.admin.command("ping")
except Exception as error: # pragma: no cover - exercised via skip path
return str(error)
finally:
client.close()
return None
async def _drop_database(uri: str, database_name: str) -> None:
client = AsyncIOMotorClient(uri, serverSelectionTimeoutMS=500)
try:
await client.drop_database(database_name)
finally:
client.close()
@pytest.fixture
def integration_runtime(tmp_path: Path) -> Iterator[dict[str, str]]:
from app.core.config import settings
from app.services.mitmproxy.manager import manager
mongodb_uri = settings.MONGODB_URI
ping_error = asyncio.run(_mongo_ping(mongodb_uri))
if ping_error is not None:
pytest.skip(f"MongoDB is not available for integration tests: {ping_error}")
db_name = f"theplug_test_{uuid4().hex}"
scan_root = tmp_path / "scan-root"
runtime_root = tmp_path / "runtime"
scan_root.mkdir(parents=True, exist_ok=True)
runtime_root.mkdir(parents=True, exist_ok=True)
original_db_name = settings.MONGODB_DB_NAME
original_scan_root = settings.SERVICE_SCAN_ROOT
original_auth_username = settings.AUTH_USERNAME
original_auth_password = settings.AUTH_PASSWORD
original_manager_work_dir = manager.work_dir
settings.MONGODB_DB_NAME = db_name
settings.SERVICE_SCAN_ROOT = str(scan_root)
settings.AUTH_USERNAME = "integration-admin"
settings.AUTH_PASSWORD = "integration-admin"
manager.work_dir = runtime_root
try:
yield {
"db_name": db_name,
"scan_root": str(scan_root),
"runtime_root": str(runtime_root),
}
finally:
manager.work_dir = original_manager_work_dir
manager.processes.clear()
manager.attached_pids.clear()
settings.MONGODB_DB_NAME = original_db_name
settings.SERVICE_SCAN_ROOT = original_scan_root
settings.AUTH_USERNAME = original_auth_username
settings.AUTH_PASSWORD = original_auth_password
asyncio.run(_drop_database(mongodb_uri, db_name))
@pytest.fixture
def real_app(integration_runtime) -> FastAPI:
from app.main import app as real_app_instance
real_app_instance.dependency_overrides.clear()
return real_app_instance
@pytest.fixture
def real_client(real_app: FastAPI):
with TestClient(real_app) as test_client:
yield test_client
@pytest.fixture
def auth_token(real_client: TestClient) -> str:
response = real_client.post(
"/auth/login",
json={"username": "integration-admin", "password": "integration-admin"},
)
assert response.status_code == 200
return response.json()["access_token"]
@pytest.fixture
def auth_headers(auth_token: str) -> dict[str, str]:
return {"Authorization": f"Bearer {auth_token}"}
@@ -0,0 +1,25 @@
import pytest
pytestmark = [pytest.mark.integration, pytest.mark.api]
def test_login_logout_revokes_real_token(real_client):
login_response = real_client.post(
"/auth/login",
json={"username": "integration-admin", "password": "integration-admin"},
)
assert login_response.status_code == 200
token = login_response.json()["access_token"]
headers = {"Authorization": f"Bearer {token}"}
list_response = real_client.get("/proxy", headers=headers)
assert list_response.status_code == 200
assert list_response.json() == []
logout_response = real_client.post("/auth/logout", headers=headers)
assert logout_response.status_code == 204
revoked_response = real_client.get("/proxy", headers=headers)
assert revoked_response.status_code == 401
assert revoked_response.json()["detail"] == "Could not validate credentials"
@@ -0,0 +1,80 @@
import pytest
pytestmark = [pytest.mark.integration, pytest.mark.api]
VALID_PROXY_CODE = """def request(flow):
return None
"""
def test_proxy_crud_round_trip_uses_real_app_and_mongo(real_client, auth_headers):
create_response = real_client.post(
"/proxy",
headers=auth_headers,
json={
"proxy_port": 49081,
"name": "integration-echo",
"code": VALID_PROXY_CODE,
"protocol": "http",
"service": {
"host": "127.0.0.1",
"port": 49082,
},
"healthcheck": {
"enabled": True,
"path": "/healthz",
"benchmark_enabled": True,
},
},
)
assert create_response.status_code == 201
created = create_response.json()
assert created["proxy_port"] == 49081
assert created["name"] == "integration-echo"
assert created["healthcheck"]["benchmark_enabled"] is True
list_response = real_client.get("/proxy", headers=auth_headers)
assert list_response.status_code == 200
assert [item["proxy_port"] for item in list_response.json()] == [49081]
get_response = real_client.get("/proxy/49081", headers=auth_headers)
assert get_response.status_code == 200
assert get_response.json()["service"]["port"] == 49082
update_response = real_client.put(
"/proxy/49081",
headers=auth_headers,
json={
"name": "integration-echo-updated",
"code": VALID_PROXY_CODE.replace("None", '"updated"'),
"healthcheck": {
"enabled": True,
"path": "/ready",
"benchmark_enabled": False,
},
},
)
assert update_response.status_code == 200
updated = update_response.json()
assert updated["name"] == "integration-echo-updated"
assert updated["code"].endswith('return "updated"\n')
assert updated["healthcheck"]["path"] == "/ready"
assert updated["healthcheck"]["benchmark_enabled"] is False
history_response = real_client.get(
"/proxy/49081/health/history",
headers=auth_headers,
)
assert history_response.status_code == 200
history_payload = history_response.json()
assert history_payload["proxy_port"] == 49081
assert history_payload["latest"] is None
assert history_payload["history"] == []
delete_response = real_client.delete("/proxy/49081", headers=auth_headers)
assert delete_response.status_code == 204
missing_response = real_client.get("/proxy/49081", headers=auth_headers)
assert missing_response.status_code == 404
@@ -1,6 +1,8 @@
import asyncio
from pathlib import Path
from integrations.docker.compose import get_all_services_in_dir
from app.integrations.docker.compose import get_all_services_in_dir
from app.integrations.docker.editor import change_port
def test_get_all_services_in_dir_scans_immediate_children(tmp_path: Path):
@@ -48,3 +50,137 @@ services:
services = get_all_services_in_dir(root)
assert services == []
def test_get_all_services_in_dir_returns_empty_when_root_is_inaccessible(
tmp_path: Path,
monkeypatch,
):
root = tmp_path / "workspace"
root.mkdir()
original_iterdir = Path.iterdir
def fake_iterdir(self: Path):
if self == root:
raise PermissionError("denied")
return original_iterdir(self)
monkeypatch.setattr(Path, "iterdir", fake_iterdir)
assert get_all_services_in_dir(root) == []
def test_get_all_services_in_dir_skips_inaccessible_child_folder(
tmp_path: Path,
monkeypatch,
):
root = tmp_path / "workspace"
allowed = root / "allowed"
blocked = root / "blocked"
allowed.mkdir(parents=True)
blocked.mkdir(parents=True)
(allowed / "docker-compose.yml").write_text(
"""
services:
echo:
image: hashicorp/http-echo:1.0
ports:
- "9081:5678"
""".strip(),
encoding="utf-8",
)
original_is_dir = Path.is_dir
def fake_is_dir(self: Path):
if self == blocked:
raise PermissionError("denied")
return original_is_dir(self)
monkeypatch.setattr(Path, "is_dir", fake_is_dir)
services = get_all_services_in_dir(root)
assert services == [
{
"name": "echo",
"ports": [9081],
"bindings": [{"published": 9081, "host": None}],
"folder": allowed,
}
]
def test_change_port_restores_original_compose_when_reload_fails(
tmp_path: Path,
monkeypatch,
):
folder = tmp_path / "workspace"
folder.mkdir()
compose_path = folder / "docker-compose.yml"
original = """
services:
echo:
image: hashicorp/http-echo:1.0
ports:
- "9081:5678"
""".strip()
compose_path.write_text(original, encoding="utf-8")
async def fake_run_docker_command(
_file: Path,
_command: list[str],
) -> dict | None:
return {"status": "error", "reason": "boom"}
monkeypatch.setattr(
"app.integrations.docker.editor.run_docker_command",
fake_run_docker_command,
)
changed = asyncio.run(change_port("echo", folder, 9081, 54001))
assert changed is False
assert compose_path.read_text(encoding="utf-8") == original
assert (folder / "docker-compose.bak.yml").exists() is False
def test_change_port_restores_original_compose_when_reload_raises(
tmp_path: Path,
monkeypatch,
):
folder = tmp_path / "workspace"
folder.mkdir()
compose_path = folder / "docker-compose.yml"
original = """
services:
echo:
image: hashicorp/http-echo:1.0
ports:
- "127.0.0.1:54001:5678"
""".strip()
compose_path.write_text(original, encoding="utf-8")
async def fake_run_docker_command(
_file: Path,
_command: list[str],
) -> dict | None:
raise RuntimeError("docker unavailable")
monkeypatch.setattr(
"app.integrations.docker.editor.run_docker_command",
fake_run_docker_command,
)
changed = asyncio.run(
change_port(
"echo",
folder,
54001,
9081,
host_to=None,
rewrite_host=True,
)
)
assert changed is False
assert compose_path.read_text(encoding="utf-8") == original
@@ -1,4 +1,4 @@
from integrations.docker.ports import _get_service_binding, _replace_port_in_mapping
from app.integrations.docker.ports import _get_service_binding, _replace_port_in_mapping
def test_get_service_binding_includes_loopback_hosts():
+39 -5
View File
@@ -1,8 +1,9 @@
import asyncio
from models.proxy import ProxyDocument
from models.service import ServiceEndpointDocument
from services.mitmproxy.health import check_proxy
from app.models.health import HealthSnapshot
from app.models.proxy import ProxyDocument
from app.models.service import ServiceEndpointDocument
from app.services.mitmproxy.health import check_proxy
def test_check_proxy_reports_upstream_and_proxy_healthy():
@@ -26,11 +27,44 @@ def test_check_proxy_reports_upstream_and_proxy_healthy():
await upstream.wait_closed()
await proxy.wait_closed()
assert health["overall"] == "healthy"
assert [check["target"] for check in health["checks"]] == [
assert isinstance(health, HealthSnapshot)
assert health.overall == "healthy"
assert [check.target for check in health.checks] == [
"upstream",
"proxy",
]
assert all(check.duration_ms is None for check in health.checks)
asyncio.run(run())
def test_check_proxy_reports_benchmark_latency_when_enabled():
async def run():
upstream = await asyncio.start_server(http_handler, "127.0.0.1", 0)
proxy = await asyncio.start_server(http_handler, "127.0.0.1", 0)
try:
upstream_port = upstream.sockets[0].getsockname()[1]
proxy_port = proxy.sockets[0].getsockname()[1]
document = ProxyDocument(
proxy_port=proxy_port,
name="health",
code="",
healthcheck={"benchmark_enabled": True},
service=ServiceEndpointDocument(port=upstream_port),
)
health = await check_proxy(document)
finally:
upstream.close()
proxy.close()
await upstream.wait_closed()
await proxy.wait_closed()
assert isinstance(health, HealthSnapshot)
assert all(
check.duration_ms is not None and check.duration_ms >= 0
for check in health.checks
)
asyncio.run(run())
@@ -0,0 +1,218 @@
import asyncio
from datetime import UTC, datetime
import pytest
from app.models.health import HealthSnapshot
from app.models.proxy import ProxyDocument
from app.models.service import ServiceEndpointDocument
from app.services.mitmproxy.health_monitor import HealthMonitor
def make_proxy(**updates):
data = {
"proxy_port": 8080,
"name": "svc",
"code": "",
"status": "running",
"service": ServiceEndpointDocument(port=9000),
}
data.update(updates)
return ProxyDocument(**data)
def make_snapshot(proxy_port: int, *, overall: str = "healthy") -> HealthSnapshot:
return HealthSnapshot(
proxy_port=proxy_port,
timestamp=datetime.now(UTC),
enabled=True,
overall=overall,
checks=[],
)
def test_run_cycle_records_and_publishes_running_healthchecks(monkeypatch):
monitor = HealthMonitor()
proxy = make_proxy()
captured: dict[str, object] = {}
snapshot = make_snapshot(8080)
async def fake_list_proxies():
return [proxy]
async def fake_check_proxy(document):
captured["checked"] = document.proxy_port
return snapshot
async def fake_record_health_sample(document):
captured["recorded"] = document
return document
async def fake_publish(proxy_port, document):
captured["published"] = (proxy_port, document)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.list_proxies",
fake_list_proxies,
)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.check_proxy",
fake_check_proxy,
)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.record_health_sample",
fake_record_health_sample,
)
monkeypatch.setattr(monitor, "publish", fake_publish)
asyncio.run(monitor.run_cycle())
assert captured["checked"] == 8080
assert captured["recorded"] == snapshot
assert captured["published"] == (8080, snapshot)
def test_run_cycle_skips_stopped_or_disabled_proxies(monkeypatch):
monitor = HealthMonitor()
running_disabled = make_proxy(healthcheck={"enabled": False})
stopped = make_proxy(proxy_port=8081, status="stopped")
called = False
async def fake_list_proxies():
return [running_disabled, stopped]
async def fake_check_proxy(_document):
nonlocal called
called = True
return make_snapshot(8080)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.list_proxies",
fake_list_proxies,
)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.check_proxy",
fake_check_proxy,
)
asyncio.run(monitor.run_cycle())
assert called is False
def test_run_cycle_checks_proxies_in_parallel(monkeypatch):
monitor = HealthMonitor()
proxies = [
make_proxy(proxy_port=8080),
make_proxy(proxy_port=8081),
make_proxy(proxy_port=8082),
]
started: list[int] = []
release = asyncio.Event()
async def fake_list_proxies():
return proxies
async def fake_check_proxy(document):
started.append(document.proxy_port)
if len(started) == len(proxies):
release.set()
await release.wait()
return make_snapshot(document.proxy_port)
async def fake_record_health_sample(document):
return document
async def fake_publish(_proxy_port, _document):
return None
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.list_proxies",
fake_list_proxies,
)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.check_proxy",
fake_check_proxy,
)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.record_health_sample",
fake_record_health_sample,
)
monkeypatch.setattr(monitor, "publish", fake_publish)
asyncio.run(asyncio.wait_for(monitor.run_cycle(), timeout=1))
assert started == [8080, 8081, 8082]
def test_run_cycle_continues_when_one_proxy_check_fails(monkeypatch):
monitor = HealthMonitor()
proxies = [
make_proxy(proxy_port=8080),
make_proxy(proxy_port=8081),
]
captured: list[int] = []
async def fake_list_proxies():
return proxies
async def fake_check_proxy(document):
if document.proxy_port == 8080:
raise RuntimeError("boom")
return make_snapshot(document.proxy_port)
async def fake_record_health_sample(document):
captured.append(document.proxy_port)
return document
async def fake_publish(proxy_port, _document):
captured.append(proxy_port + 10_000)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.list_proxies",
fake_list_proxies,
)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.check_proxy",
fake_check_proxy,
)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.record_health_sample",
fake_record_health_sample,
)
monkeypatch.setattr(monitor, "publish", fake_publish)
asyncio.run(monitor.run_cycle())
assert captured == [8081, 18081]
def test_run_retries_after_cycle_failure(monkeypatch):
monitor = HealthMonitor()
calls = {"run_cycle": 0}
sleep_calls: list[float] = []
async def fake_run_cycle():
calls["run_cycle"] += 1
if calls["run_cycle"] == 1:
raise RuntimeError("temporary failure")
raise asyncio.CancelledError
async def fake_sleep(seconds):
sleep_calls.append(seconds)
return None
monkeypatch.setattr(monitor, "run_cycle", fake_run_cycle)
monkeypatch.setattr(
"app.services.mitmproxy.health_monitor.asyncio.sleep",
fake_sleep,
)
try:
asyncio.run(monitor.run())
except asyncio.CancelledError:
pass
assert calls["run_cycle"] == 2
assert len(sleep_calls) == 1
assert sleep_calls[0] == pytest.approx(1.0, abs=0.01)
@@ -1,9 +1,11 @@
import asyncio
from pathlib import Path
from models.proxy import ProxyDocument
from models.service import ServiceEndpointDocument
from services.mitmproxy.manager import MitmdumpManager
import pytest
from app.models.proxy import ProxyDocument
from app.models.service import ServiceEndpointDocument
from app.services.mitmproxy.manager import MitmdumpError, MitmdumpManager
def make_proxy(**updates):
@@ -65,6 +67,20 @@ def test_regular_mode_uses_proxy_port_as_listen_port():
assert command[-2:] == ["--listen-port", "8080"]
def test_start_rejects_reserved_backend_port():
manager = MitmdumpManager()
proxy = make_proxy(proxy_port=3671)
async def run():
await manager.start(proxy)
with pytest.raises(
MitmdumpError,
match="Proxy port 3671 is reserved for the backend API",
):
asyncio.run(run())
def test_running_document_without_pid_is_treated_as_stopped():
manager = MitmdumpManager()
proxy = make_proxy(status="running", pid=None)
@@ -167,3 +183,99 @@ def test_reconcile_restores_upstream_when_process_is_gone():
"service.upstream_port": None,
"last_error": None,
}
def test_cleanup_runtime_artifacts_removes_generated_files(tmp_path: Path):
manager = MitmdumpManager(work_dir=str(tmp_path / "theplug"))
proxy = make_proxy(
script_path=str(manager.script_path_for(make_proxy())),
config_path=str(manager.config_path_for(make_proxy())),
log_path=str(manager.log_path_for(make_proxy())),
)
script_path = Path(proxy.script_path)
config_path = Path(proxy.config_path)
log_path = Path(proxy.log_path)
script_path.parent.mkdir(parents=True, exist_ok=True)
config_path.parent.mkdir(parents=True, exist_ok=True)
log_path.parent.mkdir(parents=True, exist_ok=True)
script_path.write_text("print('hello')", encoding="utf-8")
config_path.write_text("block_global: false\n", encoding="utf-8")
log_path.write_text("old log\n", encoding="utf-8")
manager.cleanup_runtime_artifacts(proxy)
assert script_path.exists() is False
assert log_path.exists() is False
assert config_path.exists() is False
assert manager.conf_dir_for(proxy).exists() is False
def test_prepare_upstream_retries_until_port_remap_succeeds(monkeypatch):
manager = MitmdumpManager()
proxy = make_proxy(
mode="reverse",
service=ServiceEndpointDocument(
port=9000,
name="echo",
folder="/tmp/playground",
),
)
attempted_ports: list[int] = []
candidate_ports = iter([54001, 54002, 54003])
async def fake_change_port(
_service_name,
_folder,
_port_from,
port_to,
**_kwargs,
):
attempted_ports.append(port_to)
return port_to == 54003
monkeypatch.setattr(
"app.services.mitmproxy.manager.get_random_port",
lambda: next(candidate_ports),
)
monkeypatch.setattr(
"app.services.mitmproxy.manager.change_port",
fake_change_port,
)
upstream_port = asyncio.run(manager.prepare_upstream(proxy))
assert upstream_port == 54003
assert attempted_ports == [54001, 54002, 54003]
def test_prepare_upstream_fails_after_retry_budget(monkeypatch):
manager = MitmdumpManager()
manager.upstream_port_attempts = 3
proxy = make_proxy(
mode="reverse",
service=ServiceEndpointDocument(
port=9000,
name="echo",
folder="/tmp/playground",
),
)
candidate_ports = iter([54001, 54002, 54003])
async def fake_change_port(*_args, **_kwargs):
return False
monkeypatch.setattr(
"app.services.mitmproxy.manager.get_random_port",
lambda: next(candidate_ports),
)
monkeypatch.setattr(
"app.services.mitmproxy.manager.change_port",
fake_change_port,
)
with pytest.raises(
MitmdumpError,
match="Unable to move Docker service port after retries",
):
asyncio.run(manager.prepare_upstream(proxy))
@@ -1,5 +1,6 @@
import pytest
from services.mitmproxy.options import (
from app.services.mitmproxy.options import (
MitmOption,
parse_options,
public_options,
+1 -1
View File
@@ -1,4 +1,4 @@
from services.editor import python_lsp
from app.services.editor import python_lsp
def test_encode_message():
+3 -3
View File
@@ -1,8 +1,8 @@
import asyncio
from models.proxy import ProxyDocument
from models.service import ServiceEndpointDocument
from services.mitmproxy import recovery as recovery_module
from app.models.proxy import ProxyDocument
from app.models.service import ServiceEndpointDocument
from app.services.mitmproxy import recovery as recovery_module
def make_proxy(**updates):
+13
View File
@@ -0,0 +1,13 @@
from app.services.mitmproxy.templates import BUILTIN_TEMPLATES
def test_builtin_templates_use_four_space_indentation():
for template in BUILTIN_TEMPLATES:
for line in template.code.splitlines():
assert "\t" not in line, f"{template.id} contains a tab: {line!r}"
if not line.strip():
continue
indent = len(line) - len(line.lstrip(" "))
assert indent % 4 == 0, (
f"{template.id} has non-4-space indentation: {line!r}"
)
+2 -2
View File
@@ -1,10 +1,10 @@
from datetime import UTC, datetime
from core.security import create_access_token, decode_access_token
from app.core.security import create_access_token, decode_access_token
def test_create_access_token_uses_twelve_hour_default_window(monkeypatch):
monkeypatch.setattr("core.security.settings.JWT_EXPIRE_MINUTES", 720)
monkeypatch.setattr("app.core.security.settings.JWT_EXPIRE_MINUTES", 720)
issued_at = datetime.now(UTC)
token = create_access_token(subject="admin")