#!/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 time from datetime import datetime from pathlib import Path from zoneinfo import ZoneInfo from croniter import croniter import pia EXIT_RATE_LIMITED = pia.EXIT_RATE_LIMITED 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 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 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( "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( "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() # 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)" ) log(f"Next scheduled rotation at {nxt.isoformat(timespec='seconds')}") run_rotation("Running rotation on startup", ["--force"]) nxt = next_run(expr, datetime.now(tz)) while True: time.sleep(interval) now = datetime.now(tz) status = gluetun_health_status() if status == "unhealthy": handle_unhealthy() nxt = next_run(expr, datetime.now(tz)) log(f"Next scheduled rotation at {nxt.isoformat(timespec='seconds')}") continue if now >= nxt: run_rotation("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()