feat: add proxy template origin field and implement socket utilities for log streaming
This commit is contained in:
gabdevele
+19
-128
@@ -5,15 +5,20 @@ 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.api.proxy.socket_utils import (
|
||||
authenticate_proxy_socket,
|
||||
proxy_log_path,
|
||||
read_log_snapshot,
|
||||
serialize_health_snapshot,
|
||||
snapshot_log_state,
|
||||
stream_proxy_logs,
|
||||
watch_socket_disconnect,
|
||||
)
|
||||
from app.core.config import settings
|
||||
from app.core.security import InvalidTokenError, 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.health import HealthHistoryResponse
|
||||
from app.models.proxy import ProxyCreateInput, ProxyUpdateInput
|
||||
from app.models.template import ProxyTemplateCreateInput
|
||||
from app.repositories.health import (
|
||||
@@ -51,7 +56,7 @@ from app.services.mitmproxy.runtime_validation import (
|
||||
backend_port,
|
||||
ensure_proxy_port_not_reserved,
|
||||
)
|
||||
from app.services.mitmproxy.templates import BUILTIN_TEMPLATES
|
||||
from app.services.mitmproxy.templates import file_template_ids, list_file_templates
|
||||
from app.services.mitmproxy.validator import (
|
||||
ProxyCodeValidationError,
|
||||
validate_proxy_code,
|
||||
@@ -152,125 +157,6 @@ async def get_proxy_or_404(proxy_port: int):
|
||||
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()
|
||||
@@ -295,7 +181,7 @@ async def list_mitmproxy_options():
|
||||
|
||||
@router.get("/mitmproxy/templates", dependencies=auth_dependencies)
|
||||
async def list_mitmproxy_templates():
|
||||
return [*BUILTIN_TEMPLATES, *await list_custom_templates()]
|
||||
return [*list_file_templates(), *await list_custom_templates()]
|
||||
|
||||
|
||||
@router.post(
|
||||
@@ -305,6 +191,11 @@ async def list_mitmproxy_templates():
|
||||
)
|
||||
async def create_mitmproxy_template(payload: ProxyTemplateCreateInput):
|
||||
await _validate_code_or_raise(payload.code)
|
||||
if payload.id in file_template_ids():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=f"Template already exists: {payload.id}",
|
||||
)
|
||||
try:
|
||||
return await create_custom_template(payload)
|
||||
except DuplicateKeyError as error:
|
||||
@@ -320,7 +211,7 @@ async def create_mitmproxy_template(payload: ProxyTemplateCreateInput):
|
||||
dependencies=auth_dependencies,
|
||||
)
|
||||
async def delete_mitmproxy_template(template_id: str):
|
||||
if any(template.id == template_id for template in BUILTIN_TEMPLATES):
|
||||
if template_id in file_template_ids():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="Built-in templates cannot be deleted",
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import WebSocket
|
||||
|
||||
from app.core.security import InvalidTokenError, authenticate_token
|
||||
from app.models.health import HealthSnapshot
|
||||
from app.services.mitmproxy.manager import manager
|
||||
|
||||
|
||||
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)
|
||||
Reference in New Issue
Block a user