try reducing rate limit chance

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

@@ -23,6 +23,13 @@ from typing import Any
 EXIT_RATE_LIMITED = 75
 SERVERLIST_URL = "https://serverlist.piaservers.net/vpninfo/servers/v6"
 
+_DURATION_UNITS = {
+    "s": 1,
+    "m": 60,
+    "h": 3600,
+    "d": 86400,
+}
+
 
 def log(msg: str) -> None:
     print(f"[{datetime.now().astimezone().isoformat(timespec='seconds')}] {msg}", file=sys.stderr)
@@ -35,6 +42,34 @@ def require_env(name: str) -> str:
     return value
 
 
+def env_bool(name: str, default: bool = False) -> bool:
+    raw = os.environ.get(name)
+    if raw is None or raw.strip() == "":
+        return default
+    return raw.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def parse_duration_seconds(raw: str, default_seconds: int) -> int:
+    value = (raw or "").strip().lower()
+    if not value:
+        return default_seconds
+    if value.isdigit():
+        return int(value)
+    match = re.fullmatch(r"(\d+)([smhd])", value)
+    if not match:
+        raise SystemExit(f"Invalid duration '{raw}' (use seconds or Ns/Nm/Nh/Nd)")
+    return int(match.group(1)) * _DURATION_UNITS[match.group(2)]
+
+
+def duration_go(seconds: int) -> str:
+    """Format seconds as a Go duration string (pia-wg-config uses time.ParseDuration)."""
+    if seconds % 3600 == 0:
+        return f"{seconds // 3600}h"
+    if seconds % 60 == 0:
+        return f"{seconds // 60}m"
+    return f"{seconds}s"
+
+
 def parse_list(raw: str) -> list[str]:
     raw = raw.strip()
     if not raw:
@@ -54,25 +89,72 @@ def parse_regions() -> list[str]:
     return regions
 
 
-def read_previous_region(state_path: Path) -> str | None:
+def read_state(state_path: Path) -> dict[str, Any]:
     if not state_path.is_file():
-        return None
+        return {}
     try:
         data = json.loads(state_path.read_text(encoding="utf-8"))
     except (OSError, json.JSONDecodeError):
-        return None
-    region = data.get("region")
+        return {}
+    return data if isinstance(data, dict) else {}
+
+
+def read_previous_region(state_path: Path) -> str | None:
+    region = read_state(state_path).get("region")
     return region if isinstance(region, str) and region else None
 
 
+def serverlist_cache_path() -> Path:
+    return Path(os.environ.get("SERVERLIST_CACHE_PATH", "/config/cache/pia-serverlist.json"))
+
+
 def fetch_serverlist() -> dict[str, Any]:
+    """Fetch PIA server list, with on-disk cache shared with pia-wg-config."""
+    cache_path = serverlist_cache_path()
+    ttl = parse_duration_seconds(os.environ.get("SERVERLIST_CACHE_TTL", "24h"), 86400)
+    max_age = parse_duration_seconds(os.environ.get("SERVERLIST_CACHE_MAX_AGE", "168h"), 604800)
+    force = env_bool("SERVERLIST_FORCE_REFRESH", False)
+    now = time.time()
+
+    cached_raw: bytes | None = None
+    cache_age: float | None = None
+    if cache_path.is_file() and not force:
+        try:
+            cached_raw = cache_path.read_bytes()
+            cache_age = now - cache_path.stat().st_mtime
+        except OSError as exc:
+            log(f"Server list cache unreadable ({exc}); fetching fresh copy")
+            cached_raw = None
+            cache_age = None
+
+    if cached_raw is not None and cache_age is not None and cache_age <= ttl:
+        log(f"Using cached PIA server list (age {int(cache_age)}s <= TTL {ttl}s)")
+        return parse_serverlist_bytes(cached_raw)
+
     try:
         with urllib.request.urlopen(SERVERLIST_URL, timeout=30) as resp:
-            raw = resp.read().decode("utf-8", errors="replace")
-    except urllib.error.URLError as exc:
+            raw = resp.read()
+        text = raw.decode("utf-8", errors="replace")
+        data, end = json.JSONDecoder().raw_decode(text)
+        stripped = text[:end].encode("utf-8")
+        if not isinstance(data, dict) or "regions" not in data:
+            raise SystemExit("PIA server list missing regions")
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = cache_path.with_suffix(cache_path.suffix + ".tmp")
+        tmp.write_bytes(stripped)
+        os.replace(tmp, cache_path)
+        log(f"Fetched and cached PIA server list -> {cache_path}")
+        return data
+    except (urllib.error.URLError, json.JSONDecodeError, OSError) as exc:
+        if cached_raw is not None and cache_age is not None and cache_age <= max_age:
+            log(f"Server list fetch failed ({exc}); using stale cache (age {int(cache_age)}s)")
+            return parse_serverlist_bytes(cached_raw)
         raise SystemExit(f"Failed to fetch PIA server list: {exc}") from exc
+
+
+def parse_serverlist_bytes(raw: bytes) -> dict[str, Any]:
     try:
-        data, _ = json.JSONDecoder().raw_decode(raw)
+        data, _ = json.JSONDecoder().raw_decode(raw.decode("utf-8", errors="replace"))
     except json.JSONDecodeError as exc:
         raise SystemExit(f"Failed to parse PIA server list: {exc}") from exc
     if not isinstance(data, dict) or "regions" not in data:
@@ -120,10 +202,14 @@ def average_tcp_latency_ms(ip: str, port: int, timeout: float, samples: int) ->
     return sum(readings) / len(readings)
 
 
-def pick_fastest_region(candidates: list[str]) -> tuple[str, list[dict[str, Any]]]:
+def pick_fastest_region(
+    candidates: list[str],
+    previous: str | None,
+) -> tuple[str, list[dict[str, Any]]]:
     port = int(os.environ.get("LATENCY_PORT", "1337"))
     timeout = float(os.environ.get("LATENCY_TIMEOUT_SECONDS", "2"))
     samples = max(1, int(os.environ.get("LATENCY_SAMPLES", "2")))
+    margin = float(os.environ.get("LATENCY_SWITCH_MARGIN_MS", "15"))
 
     serverlist = fetch_serverlist()
     probes: list[tuple[str, str]] = []
@@ -187,6 +273,21 @@ def pick_fastest_region(candidates: list[str]) -> tuple[str, list[dict[str, Any]
         raise SystemExit("All latency probes failed; cannot select fastest region")
 
     winner = reachable[0]
+    if previous:
+        previous_result = next((item for item in reachable if item["region"] == previous), None)
+        if previous_result is not None:
+            improvement = previous_result["latency_ms"] - winner["latency_ms"]
+            if winner["region"] != previous and improvement < margin:
+                log(
+                    f"Keeping current region {previous} "
+                    f"({previous_result['latency_ms']:.2f}ms); "
+                    f"best {winner['region']} only {improvement:.2f}ms faster "
+                    f"(margin {margin:g}ms)"
+                )
+                winner = previous_result
+            elif winner["region"] == previous:
+                log(f"Current region {previous} is still fastest")
+
     log(f"Fastest region: {winner['region']} ({winner['latency_ms']:.2f}ms via {winner['server_ip']})")
     return winner["region"], results
 
@@ -210,7 +311,7 @@ def pick_region(state_path: Path) -> tuple[str, str, list[dict[str, Any]]]:
     if mode != "fastest":
         raise SystemExit(f"Invalid REGION_SELECT '{mode}' (expected fastest|random)")
 
-    region, results = pick_fastest_region(regions)
+    region, results = pick_fastest_region(regions, previous)
     return region, mode, results
 
 
@@ -239,6 +340,13 @@ def restart_containers(containers: list[str]) -> None:
 def generate_wg_config(region: str, outfile: Path) -> None:
     user = require_env("PIA_USER")
     password = require_env("PIA_PASS")
+    cache_path = serverlist_cache_path()
+    ttl = parse_duration_seconds(os.environ.get("SERVERLIST_CACHE_TTL", "24h"), 86400)
+    max_age = parse_duration_seconds(os.environ.get("SERVERLIST_CACHE_MAX_AGE", "168h"), 604800)
+
+    # Ensure cache exists so pia-wg-config can reuse it.
+    fetch_serverlist()
+
     cmd = [
         "pia-wg-config",
         "-v",
@@ -246,6 +354,12 @@ def generate_wg_config(region: str, outfile: Path) -> None:
         region,
         "-o",
         str(outfile),
+        "--serverlist-cache",
+        str(cache_path),
+        "--serverlist-cache-ttl",
+        duration_go(ttl),
+        "--serverlist-cache-max-age",
+        duration_go(max_age),
         user,
         password,
     ]
@@ -265,21 +379,70 @@ def generate_wg_config(region: str, outfile: Path) -> None:
     raise SystemExit(1)
 
 
+def config_age_seconds(wg_path: Path) -> float | None:
+    if not wg_path.is_file():
+        return None
+    try:
+        return time.time() - wg_path.stat().st_mtime
+    except OSError:
+        return None
+
+
+def should_skip_regeneration(region: str, previous: str | None, wg_path: Path) -> bool:
+    if env_bool("FORCE_ROTATE", False):
+        log("FORCE_ROTATE=true; regenerating WireGuard config")
+        return False
+
+    if previous != region:
+        return False
+
+    age = config_age_seconds(wg_path)
+    if age is None:
+        return False
+
+    max_age = parse_duration_seconds(os.environ.get("WG_CONFIG_MAX_AGE", "7d"), 604800)
+    if age > max_age:
+        log(f"Existing WireGuard config is stale (age {int(age)}s > max {max_age}s); regenerating")
+        return False
+
+    if not re.search(
+        r"^\[Interface\]",
+        wg_path.read_text(encoding="utf-8", errors="replace"),
+        flags=re.MULTILINE,
+    ):
+        return False
+
+    log(
+        f"Skipping PIA token/config request; region unchanged ({region}) "
+        f"and config age {int(age)}s <= {max_age}s"
+    )
+    return True
+
+
 def write_state(
     state_path: Path,
     region: str,
     mode: str,
     restarted: list[str],
     latency_results: list[dict[str, Any]],
+    *,
+    skipped: bool = False,
 ) -> None:
+    previous = read_state(state_path)
     winner = next((item for item in latency_results if item.get("region") == region), None)
+    now = datetime.now().astimezone().isoformat(timespec="seconds")
     payload: dict[str, Any] = {
         "region": region,
         "selection": mode,
-        "rotated_at": datetime.now().astimezone().isoformat(timespec="seconds"),
         "wg_config": os.environ.get("WG_CONFIG_PATH", "/config/wireguard/wg0.conf"),
         "restarted_containers": restarted,
+        "last_checked_at": now,
+        "skipped_regeneration": skipped,
     }
+    if skipped and isinstance(previous.get("rotated_at"), str):
+        payload["rotated_at"] = previous["rotated_at"]
+    else:
+        payload["rotated_at"] = now
     if winner and winner.get("latency_ms") is not None:
         payload["latency_ms"] = winner["latency_ms"]
         payload["server_ip"] = winner.get("server_ip")
@@ -298,10 +461,16 @@ def rotate_once() -> None:
 
     state_path = Path(os.environ.get("ROTATOR_STATE_PATH", "/config/rotator-state.json"))
     wg_path = Path(os.environ.get("WG_CONFIG_PATH", "/config/wireguard/wg0.conf"))
+    previous = read_previous_region(state_path)
 
     region, mode, latency_results = pick_region(state_path)
     log(f"Selected region: {region} (mode={mode})")
 
+    if should_skip_regeneration(region, previous, wg_path):
+        write_state(state_path, region, mode, [], latency_results, skipped=True)
+        log(f"Rotation skipped for region={region}")
+        return
+
     with tempfile.TemporaryDirectory(prefix="pia-rotate-") as tmp:
         tmp_conf = Path(tmp) / "wg0.conf"
         log("Generating WireGuard config with pia-wg-config")
@@ -317,7 +486,7 @@ def rotate_once() -> None:
 
     restarted = parse_restart_containers()
     restart_containers(restarted)
-    write_state(state_path, region, mode, restarted, latency_results)
+    write_state(state_path, region, mode, restarted, latency_results, skipped=False)
     log(f"Rotation complete for region={region}")