226 lines
7.1 KiB
Python
226 lines
7.1 KiB
Python
#!/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()
|