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

import asyncio
import ssl
from dataclasses import dataclass

from models.proxy import HealthCheckDocument, ProxyDocument


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


async def check_proxy(proxy: ProxyDocument) -> dict[str, object]:
    if not proxy.healthcheck.enabled:
        return {"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),
            ),
            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 {"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,
) -> dict[str, object]:
    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,
) -> dict[str, object]:
    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})


async def check_tcp(
    target: Target,
    config: HealthCheckDocument,
    *,
    use_tls: bool,
) -> dict[str, object]:
    _reader, writer = await open_connection(target, config, use_tls=use_tls)
    writer.close()
    await writer.wait_closed()
    return result(target, True)


async def check_udp(target: Target, config: HealthCheckDocument) -> dict[str, object]:
    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)})
    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,
) -> dict[str, object]:
    return {
        "target": target.name,
        "host": target.host,
        "port": target.port,
        "protocol": target.protocol,
        "healthy": healthy,
        "error": error,
        "details": details or {},
    }