to python

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

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""Schedule PIA WireGuard rotations via ROTATE_CRON (croniter)."""
+
+from __future__ import annotations
+
+import os
+import re
+import subprocess
+import sys
+import time
+from datetime import datetime
+from zoneinfo import ZoneInfo
+
+from croniter import croniter
+
+# Matches rotate.sh EXIT_RATE_LIMITED (EX_TEMPFAIL)
+EXIT_RATE_LIMITED = 75
+
+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 sleep_until_next_rotate(expr: str) -> None:
+    tz = zone()
+    now = datetime.now(tz)
+    nxt = next_run(expr, now)
+    wait_s = max(0.0, (nxt - now).total_seconds())
+    log(f"Next rotation at {nxt.isoformat(timespec='seconds')} (cron '{expr}', TZ={tz.key}) in {int(wait_s)}s")
+    time.sleep(wait_s)
+
+
+def run_rotation(reason: str) -> None:
+    wait_s = int(os.environ.get("RATE_LIMIT_WAIT_SECONDS", "3600"))
+    while True:
+        log(reason)
+        result = subprocess.run(["/usr/local/bin/rotate.sh"], check=False)
+        if result.returncode == 0:
+            return
+        if result.returncode == EXIT_RATE_LIMITED:
+            log(f"PIA rate-limited (too many attempts); waiting {wait_s}s before retry")
+            time.sleep(wait_s)
+            reason = "Retrying rotation after rate-limit wait"
+            continue
+        raise SystemExit(f"Rotation failed with exit code {result.returncode}")
+
+
+def main() -> None:
+    expr = resolve_cron_expr()
+    # Fail fast on bad cron / TZ before rotating.
+    next_run(expr, datetime.now(zone()))
+    log(f"Starting gluetun PIA WireGuard rotator (TZ={zone().key}, ROTATE_CRON='{expr}')")
+
+    run_rotation("Running rotation on startup")
+    while True:
+        sleep_until_next_rotate(expr)
+        run_rotation("Running scheduled rotation")
+
+
+if __name__ == "__main__":
+    main()