326 lines
11 KiB
Python
326 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate a PIA WireGuard config and restart dependent containers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import random
|
|
import re
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
# EX_TEMPFAIL — entrypoint retries after RATE_LIMIT_WAIT_SECONDS
|
|
EXIT_RATE_LIMITED = 75
|
|
SERVERLIST_URL = "https://serverlist.piaservers.net/vpninfo/servers/v6"
|
|
|
|
|
|
def log(msg: str) -> None:
|
|
print(f"[{datetime.now().astimezone().isoformat(timespec='seconds')}] {msg}", file=sys.stderr)
|
|
|
|
|
|
def require_env(name: str) -> str:
|
|
value = os.environ.get(name, "")
|
|
if not value:
|
|
raise SystemExit(f"Missing required env var: {name}")
|
|
return value
|
|
|
|
|
|
def parse_list(raw: str) -> list[str]:
|
|
raw = raw.strip()
|
|
if not raw:
|
|
return []
|
|
if raw.startswith("["):
|
|
data = json.loads(raw)
|
|
if not isinstance(data, list):
|
|
raise SystemExit("Expected a JSON array")
|
|
return [str(item).strip() for item in data if str(item).strip()]
|
|
return [part.strip() for part in raw.split(",") if part.strip()]
|
|
|
|
|
|
def parse_regions() -> list[str]:
|
|
regions = parse_list(require_env("PIA_REGIONS"))
|
|
if not regions:
|
|
raise SystemExit("PIA_REGIONS is empty after parsing")
|
|
return regions
|
|
|
|
|
|
def read_previous_region(state_path: Path) -> str | None:
|
|
if not state_path.is_file():
|
|
return None
|
|
try:
|
|
data = json.loads(state_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
return None
|
|
region = data.get("region")
|
|
return region if isinstance(region, str) and region else None
|
|
|
|
|
|
def fetch_serverlist() -> dict[str, Any]:
|
|
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:
|
|
raise SystemExit(f"Failed to fetch PIA server list: {exc}") from exc
|
|
try:
|
|
data, _ = json.JSONDecoder().raw_decode(raw)
|
|
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:
|
|
raise SystemExit("PIA server list missing regions")
|
|
return data
|
|
|
|
|
|
def region_wg_servers(serverlist: dict[str, Any], region_id: str) -> list[dict[str, str]]:
|
|
for region in serverlist.get("regions", []):
|
|
if region.get("id") != region_id:
|
|
continue
|
|
if region.get("offline"):
|
|
log(f"Region {region_id} is marked offline; skipping")
|
|
return []
|
|
servers = region.get("servers", {}).get("wg") or []
|
|
return [
|
|
{"ip": s["ip"], "cn": s.get("cn", "")}
|
|
for s in servers
|
|
if isinstance(s, dict) and s.get("ip")
|
|
]
|
|
log(f"Region {region_id} not found in PIA server list; skipping")
|
|
return []
|
|
|
|
|
|
def tcp_latency_ms(ip: str, port: int, timeout: float) -> float | None:
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock.settimeout(timeout)
|
|
started = time.perf_counter()
|
|
try:
|
|
sock.connect((ip, port))
|
|
return (time.perf_counter() - started) * 1000.0
|
|
except OSError:
|
|
return None
|
|
finally:
|
|
sock.close()
|
|
|
|
|
|
def average_tcp_latency_ms(ip: str, port: int, timeout: float, samples: int) -> float | None:
|
|
readings: list[float] = []
|
|
for _ in range(samples):
|
|
value = tcp_latency_ms(ip, port, timeout)
|
|
if value is None:
|
|
return None
|
|
readings.append(value)
|
|
return sum(readings) / len(readings)
|
|
|
|
|
|
def pick_fastest_region(candidates: list[str]) -> 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")))
|
|
|
|
serverlist = fetch_serverlist()
|
|
probes: list[tuple[str, str]] = []
|
|
server_counts: dict[str, int] = {}
|
|
for region_id in candidates:
|
|
servers = region_wg_servers(serverlist, region_id)
|
|
if not servers:
|
|
continue
|
|
server_counts[region_id] = len(servers)
|
|
for server in servers:
|
|
probes.append((region_id, server["ip"]))
|
|
|
|
if not probes:
|
|
raise SystemExit("No reachable WireGuard servers found for configured regions")
|
|
|
|
best_by_region: dict[str, dict[str, Any]] = {
|
|
region_id: {
|
|
"region": region_id,
|
|
"latency_ms": None,
|
|
"server_ip": None,
|
|
"servers": server_counts[region_id],
|
|
"failures": 0,
|
|
}
|
|
for region_id in server_counts
|
|
}
|
|
|
|
with ThreadPoolExecutor(max_workers=min(32, len(probes))) as pool:
|
|
futures = {
|
|
pool.submit(average_tcp_latency_ms, ip, port, timeout, samples): (region_id, ip)
|
|
for region_id, ip in probes
|
|
}
|
|
for future in as_completed(futures):
|
|
region_id, ip = futures[future]
|
|
value = future.result()
|
|
current = best_by_region[region_id]
|
|
if value is None:
|
|
current["failures"] += 1
|
|
continue
|
|
if current["latency_ms"] is None or value < current["latency_ms"]:
|
|
current["latency_ms"] = round(value, 2)
|
|
current["server_ip"] = ip
|
|
|
|
results = sorted(
|
|
best_by_region.values(),
|
|
key=lambda item: (
|
|
item["latency_ms"] is None,
|
|
item["latency_ms"] if item["latency_ms"] is not None else float("inf"),
|
|
item["region"],
|
|
),
|
|
)
|
|
|
|
for item in results:
|
|
latency = "timeout" if item["latency_ms"] is None else f"{item['latency_ms']:.2f}ms"
|
|
log(
|
|
f"Latency {item['region']}: {latency} "
|
|
f"(best_ip={item['server_ip']}, servers={item['servers']}, failures={item['failures']})"
|
|
)
|
|
|
|
reachable = [item for item in results if item["latency_ms"] is not None]
|
|
if not reachable:
|
|
raise SystemExit("All latency probes failed; cannot select fastest region")
|
|
|
|
winner = reachable[0]
|
|
log(f"Fastest region: {winner['region']} ({winner['latency_ms']:.2f}ms via {winner['server_ip']})")
|
|
return winner["region"], results
|
|
|
|
|
|
def pick_random_region(candidates: list[str], previous: str | None) -> str:
|
|
pool = [region for region in candidates if region != previous] or list(candidates)
|
|
if previous and previous not in pool and len(candidates) == 1:
|
|
log(f"Only one region configured; reusing previous: {previous}")
|
|
elif previous and previous not in pool:
|
|
log(f"Excluding previous region: {previous}")
|
|
return random.choice(pool)
|
|
|
|
|
|
def pick_region(state_path: Path) -> tuple[str, str, list[dict[str, Any]]]:
|
|
regions = parse_regions()
|
|
previous = read_previous_region(state_path)
|
|
mode = os.environ.get("REGION_SELECT", "fastest").strip().lower() or "fastest"
|
|
|
|
if mode == "random":
|
|
return pick_random_region(regions, previous), mode, []
|
|
if mode != "fastest":
|
|
raise SystemExit(f"Invalid REGION_SELECT '{mode}' (expected fastest|random)")
|
|
|
|
region, results = pick_fastest_region(regions)
|
|
return region, mode, results
|
|
|
|
|
|
def parse_restart_containers() -> list[str]:
|
|
gluetun = os.environ.get("GLUETUN_CONTAINER", "m3u-filter-vpn").strip() or "m3u-filter-vpn"
|
|
ordered = [gluetun]
|
|
seen = {gluetun}
|
|
for container in parse_list(os.environ.get("RESTART_CONTAINERS", "")):
|
|
if container in seen:
|
|
continue
|
|
ordered.append(container)
|
|
seen.add(container)
|
|
return ordered
|
|
|
|
|
|
def restart_containers(containers: list[str]) -> None:
|
|
if not containers:
|
|
raise SystemExit("No containers configured to restart")
|
|
for container in containers:
|
|
log(f"Restarting container: {container}")
|
|
result = subprocess.run(["docker", "restart", container], check=False)
|
|
if result.returncode != 0:
|
|
raise SystemExit(f"docker restart failed for container={container}")
|
|
|
|
|
|
def generate_wg_config(region: str, outfile: Path) -> None:
|
|
user = require_env("PIA_USER")
|
|
password = require_env("PIA_PASS")
|
|
cmd = [
|
|
"pia-wg-config",
|
|
"-v",
|
|
"-r",
|
|
region,
|
|
"-o",
|
|
str(outfile),
|
|
user,
|
|
password,
|
|
]
|
|
result = subprocess.run(cmd, check=False, capture_output=True, text=True)
|
|
if result.stdout:
|
|
print(result.stdout, end="", file=sys.stderr)
|
|
if result.stderr:
|
|
print(result.stderr, end="", file=sys.stderr)
|
|
|
|
output = f"{result.stdout}{result.stderr}"
|
|
if result.returncode == 0:
|
|
return
|
|
if "too_many_attempts" in output or "status 429" in output:
|
|
print(f"pia-wg-config rate-limited for region={region}", file=sys.stderr)
|
|
raise SystemExit(EXIT_RATE_LIMITED)
|
|
print(f"pia-wg-config failed for region={region}", file=sys.stderr)
|
|
raise SystemExit(1)
|
|
|
|
|
|
def write_state(
|
|
state_path: Path,
|
|
region: str,
|
|
mode: str,
|
|
restarted: list[str],
|
|
latency_results: list[dict[str, Any]],
|
|
) -> None:
|
|
winner = next((item for item in latency_results if item.get("region") == region), None)
|
|
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,
|
|
}
|
|
if winner and winner.get("latency_ms") is not None:
|
|
payload["latency_ms"] = winner["latency_ms"]
|
|
payload["server_ip"] = winner.get("server_ip")
|
|
if latency_results:
|
|
payload["latency_results"] = latency_results
|
|
|
|
state_path.parent.mkdir(parents=True, exist_ok=True)
|
|
state_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
|
state_path.chmod(0o644)
|
|
|
|
|
|
def rotate_once() -> None:
|
|
require_env("PIA_USER")
|
|
require_env("PIA_PASS")
|
|
require_env("PIA_REGIONS")
|
|
|
|
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"))
|
|
|
|
region, mode, latency_results = pick_region(state_path)
|
|
log(f"Selected region: {region} (mode={mode})")
|
|
|
|
with tempfile.TemporaryDirectory(prefix="pia-rotate-") as tmp:
|
|
tmp_conf = Path(tmp) / "wg0.conf"
|
|
log("Generating WireGuard config with pia-wg-config")
|
|
generate_wg_config(region, tmp_conf)
|
|
text = tmp_conf.read_text(encoding="utf-8", errors="replace")
|
|
if not re.search(r"^\[Interface\]", text, flags=re.MULTILINE):
|
|
raise SystemExit("Generated config missing [Interface] section")
|
|
|
|
wg_path.parent.mkdir(parents=True, exist_ok=True)
|
|
os.replace(tmp_conf, wg_path)
|
|
wg_path.chmod(0o600)
|
|
log(f"Wrote {wg_path}")
|
|
|
|
restarted = parse_restart_containers()
|
|
restart_containers(restarted)
|
|
write_state(state_path, region, mode, restarted, latency_results)
|
|
log(f"Rotation complete for region={region}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
rotate_once()
|