128 lines
3.7 KiB
Python
128 lines
3.7 KiB
Python
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)
|