614 lines
19 KiB
Python
614 lines
19 KiB
Python
import asyncio
|
|
from pathlib import Path
|
|
|
|
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 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 app.repositories.proxy import (
|
|
create_proxy,
|
|
delete_proxy,
|
|
get_proxy_by_port,
|
|
update_proxy,
|
|
update_proxy_runtime,
|
|
)
|
|
from app.repositories.proxy import (
|
|
list_proxies as list_proxy_documents,
|
|
)
|
|
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 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",
|
|
"local",
|
|
"wireguard",
|
|
"reverse",
|
|
"transparent",
|
|
"tun",
|
|
"upstream",
|
|
"socks5",
|
|
"dns",
|
|
}
|
|
SUPPORTED_PROTOCOLS = {
|
|
"http",
|
|
"https",
|
|
"tcp",
|
|
"tls",
|
|
"udp",
|
|
"dtls",
|
|
"dns",
|
|
"http3",
|
|
"quic",
|
|
}
|
|
DEFAULT_HEALTH_HISTORY_LIMIT = 40
|
|
MAX_HEALTH_HISTORY_LIMIT = 200
|
|
|
|
router = APIRouter(
|
|
prefix="/proxy",
|
|
tags=["Proxy"],
|
|
)
|
|
auth_dependencies = [Depends(require_authenticated_user)]
|
|
|
|
|
|
async def _validate_code_or_raise(code: str) -> None:
|
|
try:
|
|
await validate_proxy_code(code)
|
|
except ProxyCodeValidationError as error:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
|
detail=error.to_http_detail(),
|
|
) from error
|
|
|
|
|
|
def validate_mode(mode: str) -> None:
|
|
if mode not in SUPPORTED_MODES:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
|
detail=f"Unsupported proxy mode: {mode}",
|
|
)
|
|
|
|
|
|
def validate_mode_spec(mode_spec: str | None) -> None:
|
|
if mode_spec is None:
|
|
return
|
|
mode = mode_spec.split(":", 1)[0].split("@", 1)[0]
|
|
validate_mode(mode)
|
|
|
|
|
|
def validate_protocol(protocol: str) -> None:
|
|
if protocol not in SUPPORTED_PROTOCOLS:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
|
detail=f"Unsupported reverse proxy protocol: {protocol}",
|
|
)
|
|
|
|
|
|
async def validate_mitm_options(options: dict[str, object]) -> None:
|
|
try:
|
|
await validate_options(options)
|
|
except MitmOptionError as error:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
|
detail=str(error),
|
|
) 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:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="Proxy not found",
|
|
)
|
|
return proxy
|
|
|
|
|
|
async def authenticate_proxy_socket(websocket: WebSocket) -> None:
|
|
token = websocket.query_params.get("token")
|
|
if token is None:
|
|
authorization = websocket.headers.get("authorization", "")
|
|
scheme, _, value = authorization.partition(" ")
|
|
if scheme.lower() == "bearer" and value:
|
|
token = value
|
|
|
|
if not token:
|
|
raise InvalidTokenError
|
|
|
|
await authenticate_token(token)
|
|
|
|
|
|
def proxy_log_path(proxy) -> Path:
|
|
if proxy.log_path:
|
|
return Path(proxy.log_path)
|
|
return manager.log_path_for(proxy)
|
|
|
|
|
|
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:
|
|
while True:
|
|
await websocket.receive_text()
|
|
|
|
|
|
async def stream_proxy_logs(
|
|
websocket: WebSocket,
|
|
path: Path,
|
|
*,
|
|
offset: int,
|
|
marker: bytes,
|
|
) -> None:
|
|
|
|
while True:
|
|
if not await asyncio.to_thread(path.exists):
|
|
await asyncio.sleep(0.4)
|
|
continue
|
|
|
|
try:
|
|
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:
|
|
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:
|
|
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})
|
|
|
|
await asyncio.sleep(0.4)
|
|
|
|
|
|
@router.get("", dependencies=auth_dependencies)
|
|
async def list_proxies_handler():
|
|
return await list_proxy_documents()
|
|
|
|
|
|
@router.get("/services", dependencies=auth_dependencies)
|
|
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)
|
|
return {
|
|
"scan_root": str(root),
|
|
"services": services,
|
|
"backend_port": backend_port(),
|
|
}
|
|
|
|
|
|
@router.get("/mitmproxy/options", dependencies=auth_dependencies)
|
|
async def list_mitmproxy_options():
|
|
return public_options(await load_options())
|
|
|
|
|
|
@router.get("/mitmproxy/templates", dependencies=auth_dependencies)
|
|
async def list_mitmproxy_templates():
|
|
return [*BUILTIN_TEMPLATES, *await list_custom_templates()]
|
|
|
|
|
|
@router.post(
|
|
"/mitmproxy/templates",
|
|
status_code=status.HTTP_201_CREATED,
|
|
dependencies=auth_dependencies,
|
|
)
|
|
async def create_mitmproxy_template(payload: ProxyTemplateCreateInput):
|
|
await _validate_code_or_raise(payload.code)
|
|
try:
|
|
return await create_custom_template(payload)
|
|
except DuplicateKeyError as error:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_409_CONFLICT,
|
|
detail=f"Template already exists: {payload.id}",
|
|
) from error
|
|
|
|
|
|
@router.delete(
|
|
"/mitmproxy/templates/{template_id}",
|
|
status_code=status.HTTP_204_NO_CONTENT,
|
|
dependencies=auth_dependencies,
|
|
)
|
|
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 templates cannot be deleted",
|
|
)
|
|
|
|
deleted = await delete_custom_template(template_id)
|
|
if not deleted:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="Template not found",
|
|
)
|
|
|
|
|
|
@router.get("/{proxy_port}", dependencies=auth_dependencies)
|
|
async def get_proxy(proxy_port: int):
|
|
return await get_proxy_or_404(proxy_port)
|
|
|
|
|
|
@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(
|
|
status_code=status.HTTP_409_CONFLICT,
|
|
detail=f"Proxy already exists on port {payload.proxy_port}",
|
|
)
|
|
|
|
await _validate_code_or_raise(payload.code)
|
|
validate_mode(payload.mode)
|
|
validate_mode_spec(payload.mode_spec)
|
|
validate_protocol(payload.protocol)
|
|
await validate_mitm_options(payload.mitm_options)
|
|
|
|
return await create_proxy(payload)
|
|
|
|
|
|
@router.put("/{proxy_port}", dependencies=auth_dependencies)
|
|
async def edit_proxy(proxy_port: int, payload: ProxyUpdateInput):
|
|
if payload.code is not None:
|
|
await _validate_code_or_raise(payload.code)
|
|
if payload.mode is not None:
|
|
validate_mode(payload.mode)
|
|
if "mode_spec" in payload.model_fields_set:
|
|
validate_mode_spec(payload.mode_spec)
|
|
if payload.protocol is not None:
|
|
validate_protocol(payload.protocol)
|
|
if payload.mitm_options is not None:
|
|
await validate_mitm_options(payload.mitm_options)
|
|
|
|
updated_proxy = await update_proxy(proxy_port, payload)
|
|
if updated_proxy is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="Proxy not found",
|
|
)
|
|
return updated_proxy
|
|
|
|
|
|
@router.post("/{proxy_port}/start", dependencies=auth_dependencies)
|
|
async def start_proxy(proxy_port: int):
|
|
proxy = await get_proxy_or_404(proxy_port)
|
|
|
|
try:
|
|
runtime = await manager.start(proxy)
|
|
except MitmdumpError as error:
|
|
updated = await update_proxy_runtime(
|
|
proxy_port,
|
|
{"status": "error", "last_error": str(error)},
|
|
)
|
|
return updated or proxy
|
|
|
|
return await update_proxy_runtime(proxy_port, runtime)
|
|
|
|
|
|
@router.post("/{proxy_port}/stop", dependencies=auth_dependencies)
|
|
async def stop_proxy(proxy_port: int):
|
|
proxy = await get_proxy_or_404(proxy_port)
|
|
runtime = await manager.stop(proxy)
|
|
return await update_proxy_runtime(proxy_port, runtime)
|
|
|
|
|
|
@router.post("/{proxy_port}/restart", dependencies=auth_dependencies)
|
|
async def restart_proxy(proxy_port: int):
|
|
proxy = await get_proxy_or_404(proxy_port)
|
|
stopped = await manager.stop(proxy)
|
|
proxy = await update_proxy_runtime(proxy_port, stopped) or proxy
|
|
try:
|
|
runtime = await manager.start(proxy)
|
|
except MitmdumpError as error:
|
|
updated = await update_proxy_runtime(
|
|
proxy_port,
|
|
{"status": "error", "last_error": str(error)},
|
|
)
|
|
return updated or proxy
|
|
return await update_proxy_runtime(proxy_port, runtime)
|
|
|
|
|
|
@router.get("/{proxy_port}/status", dependencies=auth_dependencies)
|
|
async def proxy_status(proxy_port: int):
|
|
proxy = await get_proxy_or_404(proxy_port)
|
|
runtime_status = await manager.status(proxy)
|
|
if runtime_status != proxy.status:
|
|
values: dict[str, object] = {"status": runtime_status}
|
|
if runtime_status != "running":
|
|
values["pid"] = None
|
|
proxy = await update_proxy_runtime(proxy_port, values) or proxy
|
|
return {"proxy_port": proxy.proxy_port, "status": runtime_status, "pid": proxy.pid}
|
|
|
|
|
|
@router.get("/{proxy_port}/health", dependencies=auth_dependencies)
|
|
async def proxy_health(proxy_port: int):
|
|
proxy = await get_proxy_or_404(proxy_port)
|
|
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)
|
|
async def proxy_logs(proxy_port: int):
|
|
proxy = await get_proxy_or_404(proxy_port)
|
|
path = proxy_log_path(proxy)
|
|
|
|
return {
|
|
"proxy_port": proxy.proxy_port,
|
|
"log": await read_log_snapshot(path),
|
|
}
|
|
|
|
|
|
@router.websocket("/{proxy_port}/logs/ws")
|
|
async def proxy_logs_socket(websocket: WebSocket, proxy_port: int):
|
|
disconnect_task: asyncio.Task[None] | None = None
|
|
stream_task: asyncio.Task[None] | None = None
|
|
|
|
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
|
|
|
|
path = proxy_log_path(proxy)
|
|
initial_snapshot, offset, marker = await snapshot_log_state(path)
|
|
await websocket.accept()
|
|
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, offset=offset, marker=marker)
|
|
)
|
|
|
|
done, pending = await asyncio.wait(
|
|
{disconnect_task, stream_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="log stream failed")
|
|
return
|
|
except asyncio.CancelledError:
|
|
return
|
|
finally:
|
|
for task in (disconnect_task, stream_task):
|
|
if task is None or task.done():
|
|
continue
|
|
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,
|
|
dependencies=auth_dependencies,
|
|
)
|
|
async def remove_proxy(proxy_port: int):
|
|
proxy = await get_proxy_by_port(proxy_port)
|
|
if proxy is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
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(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="Proxy not found",
|
|
)
|