feat: troppa roba davvero
This commit is contained in:
@@ -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,3 +1,3 @@
|
||||
from api.auth.router import router
|
||||
from app.api.auth.router import router
|
||||
|
||||
__all__ = ["router"]
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
@@ -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(
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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,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:
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,6 +1,6 @@
|
||||
from pathlib import Path
|
||||
|
||||
from core.config import settings
|
||||
from app.core.config import settings
|
||||
|
||||
|
||||
def service_scan_root() -> Path:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
|
||||
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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})
|
||||
@@ -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"
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
]
|
||||
)
|
||||
""",
|
||||
),
|
||||
]
|
||||
@@ -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()
|
||||
@@ -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:
|
||||
|
||||
@@ -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))
|
||||
@@ -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
@@ -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 = [
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
@@ -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
@@ -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():
|
||||
|
||||
@@ -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,4 +1,4 @@
|
||||
from services.editor import python_lsp
|
||||
from app.services.editor import python_lsp
|
||||
|
||||
|
||||
def test_encode_message():
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user