412 lines
14 KiB
Python
412 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Pick a PIA region/server and write a WireGuard config for Gluetun."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
import re
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pia
|
|
|
|
|
|
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_state(state_path: Path) -> dict[str, Any]:
|
|
if not state_path.is_file():
|
|
return {}
|
|
try:
|
|
data = json.loads(state_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
return {}
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
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 probe_servers(servers: list[pia.WgServer]) -> list[dict[str, Any]]:
|
|
"""Probe all servers and return per-IP results sorted by latency."""
|
|
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")))
|
|
|
|
results: list[dict[str, Any]] = []
|
|
with ThreadPoolExecutor(max_workers=min(32, max(1, len(servers)))) as pool:
|
|
futures = {
|
|
pool.submit(average_tcp_latency_ms, server.ip, port, timeout, samples): server
|
|
for server in servers
|
|
}
|
|
for future in as_completed(futures):
|
|
server = futures[future]
|
|
value = future.result()
|
|
results.append(
|
|
{
|
|
"region": server.region,
|
|
"server_ip": server.ip,
|
|
"server_cn": server.cn,
|
|
"latency_ms": None if value is None else round(value, 2),
|
|
}
|
|
)
|
|
|
|
results.sort(
|
|
key=lambda item: (
|
|
item["latency_ms"] is None,
|
|
item["latency_ms"] if item["latency_ms"] is not None else float("inf"),
|
|
item["region"],
|
|
item["server_ip"],
|
|
)
|
|
)
|
|
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} ({item['server_cn']}/{item['server_ip']})")
|
|
return results
|
|
|
|
|
|
def pick_fastest(
|
|
candidates: list[str],
|
|
previous_region: str | None,
|
|
previous_server_ip: str | None,
|
|
*,
|
|
exclude_ips: set[str] | None = None,
|
|
apply_margin: bool = True,
|
|
) -> tuple[pia.WgServer, list[dict[str, Any]]]:
|
|
exclude_ips = exclude_ips or set()
|
|
margin = float(os.environ.get("LATENCY_SWITCH_MARGIN_MS", "15"))
|
|
|
|
serverlist = pia.fetch_serverlist()
|
|
servers: list[pia.WgServer] = []
|
|
for region_id in candidates:
|
|
for server in pia.region_wg_servers(serverlist, region_id):
|
|
if server.ip in exclude_ips:
|
|
continue
|
|
servers.append(server)
|
|
if not servers:
|
|
raise SystemExit("No reachable WireGuard servers left after excludes")
|
|
|
|
results = probe_servers(servers)
|
|
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 server")
|
|
|
|
winner = reachable[0]
|
|
if apply_margin and previous_server_ip and previous_server_ip not in exclude_ips:
|
|
previous_result = next(
|
|
(item for item in reachable if item["server_ip"] == previous_server_ip),
|
|
None,
|
|
)
|
|
if previous_result is not None:
|
|
improvement = previous_result["latency_ms"] - winner["latency_ms"]
|
|
if winner["server_ip"] != previous_server_ip and improvement < margin:
|
|
log(
|
|
f"Keeping current server {previous_server_ip} "
|
|
f"({previous_result['latency_ms']:.2f}ms); "
|
|
f"best {winner['server_ip']} only {improvement:.2f}ms faster "
|
|
f"(margin {margin:g}ms)"
|
|
)
|
|
winner = previous_result
|
|
elif winner["server_ip"] == previous_server_ip:
|
|
log(f"Current endpoint {previous_server_ip} is still fastest")
|
|
|
|
server = pia.WgServer(
|
|
region=winner["region"],
|
|
ip=winner["server_ip"],
|
|
cn=winner["server_cn"],
|
|
)
|
|
log(f"Selected {server.region} via {server.cn} ({server.ip}) at {winner['latency_ms']:.2f}ms")
|
|
return server, results
|
|
|
|
|
|
def pick_random(
|
|
candidates: list[str],
|
|
previous_region: str | None,
|
|
*,
|
|
exclude_ips: set[str] | None = None,
|
|
) -> pia.WgServer:
|
|
exclude_ips = exclude_ips or set()
|
|
serverlist = pia.fetch_serverlist()
|
|
pool = [region for region in candidates if region != previous_region] or list(candidates)
|
|
if previous_region and previous_region not in pool and len(candidates) == 1:
|
|
log(f"Only one region configured; reusing previous: {previous_region}")
|
|
elif previous_region and previous_region not in pool:
|
|
log(f"Excluding previous region: {previous_region}")
|
|
|
|
random.shuffle(pool)
|
|
for region_id in pool:
|
|
servers = [
|
|
server
|
|
for server in pia.region_wg_servers(serverlist, region_id)
|
|
if server.ip not in exclude_ips
|
|
]
|
|
if servers:
|
|
server = random.choice(servers)
|
|
log(f"Randomly selected {server.region} via {server.cn} ({server.ip})")
|
|
return server
|
|
raise SystemExit("No WireGuard servers left after excludes")
|
|
|
|
|
|
def pick_server(
|
|
state_path: Path,
|
|
*,
|
|
exclude_ips: set[str] | None = None,
|
|
apply_margin: bool = True,
|
|
) -> tuple[pia.WgServer, str, list[dict[str, Any]]]:
|
|
regions = parse_regions()
|
|
state = read_state(state_path)
|
|
previous_region = state.get("region") if isinstance(state.get("region"), str) else None
|
|
previous_server_ip = state.get("server_ip") if isinstance(state.get("server_ip"), str) else None
|
|
mode = os.environ.get("REGION_SELECT", "fastest").strip().lower() or "fastest"
|
|
|
|
if mode == "random":
|
|
return pick_random(regions, previous_region, exclude_ips=exclude_ips), mode, []
|
|
if mode != "fastest":
|
|
raise SystemExit(f"Invalid REGION_SELECT '{mode}' (expected fastest|random)")
|
|
|
|
server, results = pick_fastest(
|
|
regions,
|
|
previous_region,
|
|
previous_server_ip,
|
|
exclude_ips=exclude_ips,
|
|
apply_margin=apply_margin,
|
|
)
|
|
return server, 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 write_state(
|
|
state_path: Path,
|
|
server: pia.WgServer,
|
|
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") == server.region and item.get("server_ip") == server.ip
|
|
),
|
|
next((item for item in latency_results if item.get("server_ip") == server.ip), None),
|
|
)
|
|
now = datetime.now().astimezone().isoformat(timespec="seconds")
|
|
payload: dict[str, Any] = {
|
|
"region": server.region,
|
|
"server_ip": server.ip,
|
|
"server_cn": server.cn,
|
|
"selection": mode,
|
|
"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"]
|
|
elif skipped and previous.get("latency_ms") is not None:
|
|
payload["latency_ms"] = previous["latency_ms"]
|
|
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 config_looks_valid(wg_path: Path) -> bool:
|
|
if not wg_path.is_file():
|
|
return False
|
|
text = wg_path.read_text(encoding="utf-8", errors="replace")
|
|
return bool(re.search(r"^\[Interface\]", text, flags=re.MULTILINE))
|
|
|
|
|
|
def parse_args(argv: list[str]) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Rotate PIA WireGuard config for Gluetun")
|
|
parser.add_argument("--list-regions", action="store_true")
|
|
parser.add_argument(
|
|
"--skip-same-server",
|
|
action="store_true",
|
|
help="Skip token/addKey/restarts when the selected server is unchanged",
|
|
)
|
|
parser.add_argument(
|
|
"--force",
|
|
action="store_true",
|
|
help="Always generate a new keypair and restart containers",
|
|
)
|
|
parser.add_argument(
|
|
"--exclude-server",
|
|
action="append",
|
|
default=[],
|
|
help="Exclude this server IP from selection (repeatable)",
|
|
)
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def rotate_once(args: argparse.Namespace) -> 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"))
|
|
state = read_state(state_path)
|
|
exclude_ips = {ip for ip in args.exclude_server if ip}
|
|
# Unhealthy/force path wants true #1/#2 ranking without stickiness margin.
|
|
apply_margin = not args.force and not exclude_ips
|
|
|
|
server, mode, latency_results = pick_server(
|
|
state_path,
|
|
exclude_ips=exclude_ips,
|
|
apply_margin=apply_margin,
|
|
)
|
|
log(f"Selected endpoint: {server.region} / {server.cn} / {server.ip} (mode={mode})")
|
|
|
|
previous_ip = state.get("server_ip") if isinstance(state.get("server_ip"), str) else None
|
|
if args.skip_same_server and not args.force and previous_ip == server.ip and config_looks_valid(wg_path):
|
|
log(
|
|
f"Skipping PIA token/addKey; endpoint unchanged ({server.region}/{server.ip})"
|
|
)
|
|
write_state(
|
|
state_path,
|
|
server,
|
|
mode,
|
|
[],
|
|
latency_results,
|
|
skipped=True,
|
|
)
|
|
log(f"Rotation skipped for {server.region}/{server.ip}")
|
|
return
|
|
|
|
# Write temp file on the same filesystem as the destination so os.replace works
|
|
# across Docker bind mounts (/tmp is often a different device than /config).
|
|
wg_path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp_conf = wg_path.with_name(wg_path.name + ".tmp")
|
|
try:
|
|
log("Generating WireGuard config via native PIA client")
|
|
pia.generate_wg_config(
|
|
require_env("PIA_USER"),
|
|
require_env("PIA_PASS"),
|
|
server,
|
|
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")
|
|
|
|
os.replace(tmp_conf, wg_path)
|
|
wg_path.chmod(0o600)
|
|
log(f"Wrote {wg_path}")
|
|
finally:
|
|
tmp_conf.unlink(missing_ok=True)
|
|
|
|
restarted = parse_restart_containers()
|
|
restart_containers(restarted)
|
|
write_state(state_path, server, mode, restarted, latency_results, skipped=False)
|
|
log(f"Rotation complete for {server.region}/{server.ip}")
|
|
|
|
|
|
def list_regions_main() -> None:
|
|
for region in pia.list_regions():
|
|
print(
|
|
f"{region['id']:<35} {region['name']:<35} "
|
|
f"country={region['country']:<4} port-forward={region['port_forward']!s:<5} "
|
|
f"offline={region['offline']} wg={region['wg_servers']}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Compat: `rotate.py list-regions`
|
|
argv = sys.argv[1:]
|
|
if argv and argv[0] == "list-regions":
|
|
argv = ["--list-regions", *argv[1:]]
|
|
args = parse_args(argv)
|
|
if args.list_regions:
|
|
list_regions_main()
|
|
else:
|
|
rotate_once(args)
|