thePlug/backend/app/services/mitmproxy/health.py

import asyncio
import ssl
import time
from dataclasses import dataclass
from datetime import UTC, datetime

from app.models.health import HealthCheckEntry, HealthSnapshot
from app.models.proxy import HealthCheckDocument, ProxyDocument


@dataclass(slots=True)
class Target:
    name: str
    host: str
    port: int
    protocol: str


async def check_proxy(proxy: ProxyDocument) -> HealthSnapshot:
    timestamp = datetime.now(UTC)
    if not proxy.healthcheck.enabled:
        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 = list(
        await asyncio.gather(
            check_target(
                Target(
                    name="upstream",
                    host=proxy.service.host,
                    port=upstream_port,
                    protocol=health_protocol(proxy),
                ),
                proxy.healthcheck,
            ),
            check_target(
                Target(
                    name="proxy",
                    host=listen_host_for_check(proxy.listen_host),
                    port=proxy.proxy_port,
                    protocol=health_protocol(proxy),
                ),
                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,
    )


def health_protocol(proxy: ProxyDocument) -> str:
    return proxy.healthcheck.protocol or proxy.protocol


def listen_host_for_check(host: str) -> str:
    if host in {"", "0.0.0.0", "::"}:
        return "127.0.0.1"
    return host


async def check_target(
    target: Target,
    config: HealthCheckDocument,
) -> HealthCheckEntry:
    try:
        match target.protocol:
            case "http":
                return await check_http(target, config, use_tls=False)
            case "https":
                return await check_http(target, config, use_tls=True)
            case "tcp":
                return await check_tcp(target, config, use_tls=False)
            case "tls":
                return await check_tcp(target, config, use_tls=True)
            case "udp" | "dtls" | "quic" | "dns" | "http3":
                return await check_udp(target, config)
            case _:
                return result(target, False, f"Unsupported protocol: {target.protocol}")
    except TimeoutError:
        return result(target, False, "Timed out")
    except OSError as error:
        return result(target, False, str(error))


async def check_http(
    target: Target,
    config: HealthCheckDocument,
    *,
    use_tls: bool,
) -> 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}"
        request = (
            f"{config.method} {path} HTTP/1.1\r\n"
            f"Host: {target.host}\r\n"
            "Connection: close\r\n\r\n"
        )
        writer.write(request.encode("ascii"))
        await writer.drain()
        line = await asyncio.wait_for(reader.readline(), config.timeout_seconds)
    finally:
        writer.close()
        await writer.wait_closed()

    parts = line.decode("iso-8859-1", errors="replace").split()
    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},
        benchmark_ms=elapsed_ms(started) if config.benchmark_enabled else None,
    )


async def check_tcp(
    target: Target,
    config: HealthCheckDocument,
    *,
    use_tls: bool,
) -> 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,
        benchmark_ms=elapsed_ms(started) if config.benchmark_enabled else None,
    )


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),
        remote_addr=(target.host, target.port),
    )
    try:
        payload = config.udp_payload.encode()
        transport.sendto(payload)
        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)},
            benchmark_ms=elapsed_ms(started) if config.benchmark_enabled else None,
        )
    finally:
        transport.close()


async def open_connection(
    target: Target,
    config: HealthCheckDocument,
    *,
    use_tls: bool,
):
    if use_tls and config.verify_tls:
        context = ssl.create_default_context()
    elif use_tls:
        context = ssl.create_default_context()
        context.check_hostname = False
        context.verify_mode = ssl.CERT_NONE
    else:
        context = None
    return await asyncio.wait_for(
        asyncio.open_connection(target.host, target.port, ssl=context),
        config.timeout_seconds,
    )


class DatagramProbe(asyncio.DatagramProtocol):
    def __init__(self, expect_response: bool):
        self.response = asyncio.get_running_loop().create_future()
        if not expect_response:
            self.response.set_result(b"")

    def datagram_received(self, data, addr):
        if not self.response.done():
            self.response.set_result(data)

    def error_received(self, exc):
        if not self.response.done():
            self.response.set_exception(exc)


def result(
    target: Target,
    healthy: bool,
    error: str | None = None,
    details: dict[str, object] | None = None,
    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)