projects/Dockers/gluetun-pia-wireguard-rotator/entrypoint.py

#!/usr/bin/env python3
"""Schedule PIA WireGuard rotations via cron + Gluetun health polling."""

from __future__ import annotations

import json
import os
import re
import subprocess
import sys
import threading
import time
from datetime import datetime
import http.server
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo

from croniter import croniter

import pia

EXIT_RATE_LIMITED = pia.EXIT_RATE_LIMITED

_ROTATION_LOCK = threading.Lock()
_RESYNC_EVENT = threading.Event()

CRON_MACROS = {
    "@yearly": "0 0 1 1 *",
    "@annually": "0 0 1 1 *",
    "@monthly": "0 0 1 * *",
    "@weekly": "0 0 * * 0",
    "@daily": "0 0 * * *",
    "@midnight": "0 0 * * *",
    "@hourly": "0 * * * *",
}


def log(msg: str) -> None:
    print(f"[{datetime.now().astimezone().isoformat(timespec='seconds')}] {msg}", file=sys.stderr)


def zone() -> ZoneInfo:
    name = os.environ.get("TZ") or "UTC"
    try:
        return ZoneInfo(name)
    except Exception as exc:
        raise SystemExit(f"Invalid TZ '{name}': {exc}") from exc


def hhmm_to_cron(value: str) -> str:
    match = re.fullmatch(r"([0-9]{1,2}):([0-9]{2})", value.strip())
    if not match:
        raise ValueError(f"Invalid ROTATE_AT '{value}' (expected HH:MM)")
    hour = int(match.group(1))
    minute = int(match.group(2))
    if not (0 <= hour <= 23 and 0 <= minute <= 59):
        raise ValueError(f"Invalid ROTATE_AT '{value}' (expected HH:MM)")
    return f"{minute} {hour} * * *"


def resolve_cron_expr() -> str:
    cron = os.environ.get("ROTATE_CRON", "").strip()
    at = os.environ.get("ROTATE_AT", "").strip()
    if cron:
        expr = cron
    elif at:
        expr = hhmm_to_cron(at)
    else:
        expr = "0 3 * * *"

    expr = CRON_MACROS.get(expr.lower(), expr)
    if not croniter.is_valid(expr):
        raise SystemExit(f"Invalid ROTATE_CRON '{expr}' (expected a 5-field cron expression)")
    return expr


def next_run(expr: str, after: datetime) -> datetime:
    return croniter(expr, after).get_next(datetime)


def health_check_interval() -> float:
    return max(1.0, float(os.environ.get("HEALTH_CHECK_INTERVAL", "10")))


def unhealthy_cooldown_seconds() -> float:
    return max(0.0, float(os.environ.get("UNHEALTHY_ROTATE_COOLDOWN", "60")))


def resync_port() -> int | None:
    raw = os.environ.get("RESYNC_PORT", "8080").strip()
    if not raw:
        return None
    try:
        port = int(raw)
    except ValueError as exc:
        raise SystemExit(f"Invalid RESYNC_PORT '{raw}' (expected integer)") from exc
    if not (1 <= port <= 65535):
        raise SystemExit(f"Invalid RESYNC_PORT {port} (expected 1-65535)")
    return port


def resync_bind_address() -> str:
    return os.environ.get("RESYNC_BIND", "0.0.0.0").strip() or "0.0.0.0"


class ResyncHandler(http.server.BaseHTTPRequestHandler):
    def _respond(self, status: int, body: dict[str, Any]) -> None:
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        self.wfile.write(json.dumps(body).encode("utf-8"))

    def _path(self) -> str:
        return self.path.split("?", 1)[0]

    def do_GET(self) -> None:  # noqa: N802
        path = self._path()
        if path == "/resync":
            _RESYNC_EVENT.set()
            self._respond(202, {"status": "accepted", "message": "Resync requested"})
        elif path == "/":
            self._respond(
                200,
                {
                    "service": "gluetun-pia-wireguard-rotator",
                    "endpoints": ["/resync"],
                },
            )
        else:
            self._respond(404, {"error": "not found"})

    def do_POST(self) -> None:  # noqa: N802
        path = self._path()
        if path == "/resync":
            _RESYNC_EVENT.set()
            self._respond(202, {"status": "accepted", "message": "Resync requested"})
        else:
            self._respond(404, {"error": "not found"})

    def log_message(self, fmt: str, *args: Any) -> None:
        log(fmt % args)


def start_resync_server(port: int, bind: str) -> None:
    server = http.server.HTTPServer((bind, port), ResyncHandler)
    thread = threading.Thread(target=server.serve_forever, daemon=True, name="resync-http")
    thread.start()
    log(f"Resync endpoint listening on http://{bind}:{port}/resync")


def gluetun_container_name() -> str:
    return os.environ.get("GLUETUN_CONTAINER", "m3u-filter-vpn").strip() or "m3u-filter-vpn"


def state_path() -> Path:
    return Path(os.environ.get("ROTATOR_STATE_PATH", "/config/rotator-state.json"))


def read_state_server_ip() -> str | None:
    path = state_path()
    if not path.is_file():
        return None
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    ip = data.get("server_ip")
    return ip if isinstance(ip, str) and ip else None


def gluetun_health_status() -> str | None:
    """Return Docker health status, or None if unavailable / no healthcheck."""
    name = gluetun_container_name()
    result = subprocess.run(
        ["docker", "inspect", "-f", "{{if .State.Health}}{{.State.Health.Status}}{{end}}", name],
        check=False,
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        err = (result.stderr or result.stdout or "").strip()
        log(f"Could not inspect health of {name}: {err or f'exit {result.returncode}'}")
        return None
    status = (result.stdout or "").strip()
    return status or None


def wait_for_rate_limit_cooldown() -> None:
    remaining = pia.rate_limit_remaining_seconds()
    if remaining <= 0:
        return
    log(f"PIA rate-limit cooldown active; waiting {remaining}s before retry")
    time.sleep(remaining)


def run_rotation(reason: str, extra_args: list[str] | None = None) -> None:
    args = ["/usr/local/bin/rotate.py", *(extra_args or [])]
    while True:
        wait_for_rate_limit_cooldown()
        log(reason)
        result = subprocess.run(args, check=False)
        if result.returncode == 0:
            return
        if result.returncode == EXIT_RATE_LIMITED:
            if pia.rate_limit_remaining_seconds() <= 0:
                pia.mark_rate_limited(reason="rotation exit code 75")
            reason = "Retrying rotation after rate-limit wait"
            continue
        raise SystemExit(f"Rotation failed with exit code {result.returncode}")


def run_rotation_locked(reason: str, extra_args: list[str] | None = None) -> None:
    """Run a rotation while holding the global lock to avoid concurrent runs."""
    with _ROTATION_LOCK:
        run_rotation(reason, extra_args)


def wait_until_healthy_or_timeout(timeout_s: float) -> str | None:
    """Sleep up to timeout_s, returning early if Gluetun becomes non-unhealthy."""
    deadline = time.time() + timeout_s
    interval = min(health_check_interval(), max(1.0, timeout_s))
    while True:
        remaining = deadline - time.time()
        if remaining <= 0:
            break
        time.sleep(min(interval, remaining))
        status = gluetun_health_status()
        if status != "unhealthy":
            return status
    return gluetun_health_status()


def handle_unhealthy() -> None:
    """Two-step recovery: re-auth best server, then runner-up if still unhealthy."""
    log(f"Gluetun container {gluetun_container_name()} is unhealthy; starting recovery")

    run_rotation_locked(
        "Unhealthy recovery step 1: force rotate to best server (new keypair)",
        ["--force"],
    )
    step1_ip = read_state_server_ip()
    cooldown = unhealthy_cooldown_seconds()
    log(f"Waiting up to {int(cooldown)}s for Gluetun to become healthy after step 1")
    status = wait_until_healthy_or_timeout(cooldown)
    if status != "unhealthy":
        log(f"Gluetun health after step 1: {status or 'unknown/no-healthcheck'}")
        return

    exclude_args: list[str] = ["--force"]
    if step1_ip:
        exclude_args.extend(["--exclude-server", step1_ip])
        log(f"Still unhealthy; step 2 excluding failed server {step1_ip}")
    else:
        log("Still unhealthy; step 2 without exclude (no server_ip in state)")

    run_rotation_locked(
        "Unhealthy recovery step 2: force rotate to runner-up",
        exclude_args,
    )
    log(f"Waiting up to {int(cooldown)}s for Gluetun to become healthy after step 2")
    status = wait_until_healthy_or_timeout(cooldown)
    log(f"Gluetun health after step 2: {status or 'unknown/no-healthcheck'}")


def main() -> None:
    expr = resolve_cron_expr()
    tz = zone()
    interval = health_check_interval()
    port = resync_port()
    bind = resync_bind_address()
    # Fail fast on bad cron / TZ before rotating.
    nxt = next_run(expr, datetime.now(tz))
    log(
        f"Starting gluetun PIA WireGuard rotator "
        f"(TZ={tz.key}, ROTATE_CRON='{expr}', HEALTH_CHECK_INTERVAL={interval:g}s, "
        f"RESYNC_PORT={port or 'off'})"
    )
    log(f"Next scheduled rotation at {nxt.isoformat(timespec='seconds')}")

    if port:
        start_resync_server(port, bind)

    run_rotation_locked("Running rotation on startup", ["--force"])
    nxt = next_run(expr, datetime.now(tz))

    step = min(interval, 1.0)
    while True:
        # Sleep in short chunks so a resync request wakes us early.
        deadline = time.time() + interval
        while time.time() < deadline and not _RESYNC_EVENT.is_set():
            remaining = deadline - time.time()
            time.sleep(min(step, remaining, 1.0))

        now = datetime.now(tz)

        if _RESYNC_EVENT.is_set():
            _RESYNC_EVENT.clear()
            run_rotation_locked("Resync requested via HTTP", ["--force"])
            nxt = next_run(expr, datetime.now(tz))
            log(f"Next scheduled rotation at {nxt.isoformat(timespec='seconds')}")
            continue

        status = gluetun_health_status()
        if status == "unhealthy":
            handle_unhealthy()
            _RESYNC_EVENT.clear()
            nxt = next_run(expr, datetime.now(tz))
            log(f"Next scheduled rotation at {nxt.isoformat(timespec='seconds')}")
            continue

        if now >= nxt:
            run_rotation_locked("Running scheduled rotation", ["--skip-same-server"])
            nxt = next_run(expr, datetime.now(tz))
            log(f"Next scheduled rotation at {nxt.isoformat(timespec='seconds')}")


if __name__ == "__main__":
    main()