app.py

#!/usr/bin/env python3
"""WAN-Probe-Server: Download-, Upload- und Ping-Endpunkte für WAN-Qualitätsmessung.

Läuft auf dem VPS (Coolify). Lauscht auf mehreren Ports mit identischen
Endpunkten — der Omada-Gateway routet per Policy Routing Traffic zu Port A
über WAN1 (Kabel) und zu Port B über WAN2 (Starlink). So kann der Monitor
hinter dem Gateway beide Leitungen getrennt messen.

Endpunkte:
  GET  /ping                 → 204, minimale Antwort (HTTP-RTT-Messung)
  GET  /down?mb=15           → streamt N MiB Nullbytes (Download-Probe)
  POST /up                   → nimmt Body entgegen, verwirft ihn, misst
                               serverseitig Dauer/Rate (Upload-Probe)
  GET  /healthz              → Status-JSON (für Coolify-Healthcheck/Domain)

Auth: /down und /up verlangen "Authorization: Bearer $PROBE_TOKEN".
Schutz: Größen-Caps, simples Per-IP-Rate-Limit, Volumen-Budget pro Stunde.

Env:
  PROBE_TOKEN   (Pflicht)  Shared Secret für /down und /up
  PORTS         (Default "8441,8442")  Probe-Ports (Policy-Routing-Trennung)
  WEB_PORT      (Default 8080)  Port für /healthz hinter dem Coolify-Proxy
  MAX_DOWN_MB   (Default 50)   Obergrenze pro Download-Probe
  MAX_UP_MB     (Default 50)   Obergrenze pro Upload-Probe
  HOURLY_GB     (Default 5)    Volumen-Budget pro Stunde und IP
"""

import asyncio
import json
import logging
import os
import time
from collections import defaultdict

from aiohttp import web

log = logging.getLogger("probe")

TOKEN = os.environ.get("PROBE_TOKEN", "")
PORTS = [int(p) for p in os.environ.get("PORTS", "8441,8442").split(",") if p.strip()]
WEB_PORT = int(os.environ.get("WEB_PORT", "8080"))
MAX_DOWN_MB = int(os.environ.get("MAX_DOWN_MB", "50"))
MAX_UP_MB = int(os.environ.get("MAX_UP_MB", "50"))
HOURLY_GB = float(os.environ.get("HOURLY_GB", "5"))

CHUNK = b"\0" * 65536
START_TIME = time.time()

# Per-IP-Buchführung: Requests/Minute und Bytes/Stunde
_requests = defaultdict(list)   # ip -> [timestamps]
_volume = defaultdict(list)     # ip -> [(timestamp, bytes)]
_stats = {"down": 0, "up": 0, "ping": 0, "denied": 0}


def client_ip(request: web.Request) -> str:
    fwd = request.headers.get("X-Forwarded-For")
    if fwd:
        return fwd.split(",")[0].strip()
    return request.remote or "?"


def check_limits(ip: str, want_bytes: int) -> str | None:
    """Gibt einen Ablehnungsgrund zurück oder None."""
    now = time.time()
    _requests[ip] = [t for t in _requests[ip] if now - t < 60]
    if len(_requests[ip]) >= 30:
        return "rate limit (30 req/min)"
    _volume[ip] = [(t, b) for t, b in _volume[ip] if now - t < 3600]
    used = sum(b for _, b in _volume[ip])
    if used + want_bytes > HOURLY_GB * 1e9:
        return f"volume budget ({HOURLY_GB} GB/h)"
    _requests[ip].append(now)
    return None


def book_volume(ip: str, nbytes: int):
    _volume[ip].append((time.time(), nbytes))


def authed(request: web.Request) -> bool:
    return TOKEN and request.headers.get("Authorization") == f"Bearer {TOKEN}"


async def handle_ping(request: web.Request) -> web.Response:
    _stats["ping"] += 1
    return web.Response(status=204)


async def handle_down(request: web.Request) -> web.StreamResponse:
    ip = client_ip(request)
    if not authed(request):
        _stats["denied"] += 1
        return web.Response(status=401, text="unauthorized")
    try:
        mb = min(int(request.query.get("mb", "15")), MAX_DOWN_MB)
    except ValueError:
        return web.Response(status=400, text="mb must be an integer")
    nbytes = mb * 1024 * 1024
    if reason := check_limits(ip, nbytes):
        _stats["denied"] += 1
        return web.Response(status=429, text=reason)

    resp = web.StreamResponse(
        status=200,
        headers={"Content-Type": "application/octet-stream",
                 "Content-Length": str(nbytes), "Cache-Control": "no-store"},
    )
    await resp.prepare(request)
    sent = 0
    try:
        while sent < nbytes:
            n = min(len(CHUNK), nbytes - sent)
            await resp.write(CHUNK[:n])
            sent += n
    finally:
        book_volume(ip, sent)
        _stats["down"] += 1
        log.info("down ip=%s port=%s mb=%s sent=%s", ip, request.transport.get_extra_info("sockname"), mb, sent)
    await resp.write_eof()
    return resp


async def handle_up(request: web.Request) -> web.Response:
    ip = client_ip(request)
    if not authed(request):
        _stats["denied"] += 1
        return web.Response(status=401, text="unauthorized")
    cap = MAX_UP_MB * 1024 * 1024
    declared = request.content_length or 0
    if declared > cap:
        return web.Response(status=413, text=f"max {MAX_UP_MB} MB")
    if reason := check_limits(ip, declared or cap):
        _stats["denied"] += 1
        return web.Response(status=429, text=reason)

    t0 = time.monotonic()
    received = 0
    async for chunk in request.content.iter_chunked(65536):
        received += len(chunk)
        if received > cap:
            return web.Response(status=413, text=f"max {MAX_UP_MB} MB")
    dt = max(time.monotonic() - t0, 1e-6)
    book_volume(ip, received)
    _stats["up"] += 1
    mbps = received * 8 / dt / 1e6
    log.info("up ip=%s bytes=%s s=%.2f mbps=%.1f", ip, received, dt, mbps)
    return web.json_response({"bytes": received, "seconds": round(dt, 3), "mbps": round(mbps, 2)})


async def handle_healthz(request: web.Request) -> web.Response:
    return web.json_response({
        "status": "ok",
        "uptime_s": int(time.time() - START_TIME),
        "probe_ports": PORTS,
        "stats": _stats,
    })


def make_app() -> web.Application:
    app = web.Application()
    app.router.add_get("/ping", handle_ping)
    app.router.add_get("/down", handle_down)
    app.router.add_post("/up", handle_up)
    app.router.add_get("/healthz", handle_healthz)
    return app


async def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    if not TOKEN:
        log.error("PROBE_TOKEN ist nicht gesetzt — /down und /up verweigern alles.")
    app = make_app()
    runner = web.AppRunner(app, access_log=None)
    await runner.setup()
    bind_ports = sorted(set(PORTS + [WEB_PORT]))
    for port in bind_ports:
        await web.TCPSite(runner, "0.0.0.0", port).start()
        log.info("lausche auf :%d", port)
    log.info("probe-server bereit (ports=%s, max_down=%dMB, max_up=%dMB, budget=%.1fGB/h)",
             bind_ports, MAX_DOWN_MB, MAX_UP_MB, HOURLY_GB)
    await asyncio.Event().wait()


if __name__ == "__main__":
    asyncio.run(main())