111 lines
3.2 KiB
Python
111 lines
3.2 KiB
Python
#!/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.py"], 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()
|