Bram/projects
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

optimize
Build and Push Docker Images / build-and-push (push) Successful in 17s
Details

Browse Source
This commit is contained in:
Bram
2026-08-14 20:03:05 +02:00
parent 3eeb93e225
commit bc9d79de02
5 changed files with 552 additions and 288 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockers/gluetun-pia-wireguard-rotator/rotate.py
+154 -237
View File
@@ -1,5 +1,5 @@
#!/usr/bin/env python3
"""Generate a PIA WireGuard config and restart dependent containers."""
"""Pick a PIA region/server and write a WireGuard config for Gluetun."""
from __future__ import annotations
@@ -12,23 +12,12 @@ 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"
_DURATION_UNITS = {
"s": 1,
"m": 60,
"h": 3600,
"d": 86400,
}
import pia
def log(msg: str) -> None:
@@ -42,34 +31,6 @@ 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:
@@ -99,86 +60,6 @@ def read_state(state_path: Path) -> dict[str, Any]:
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()
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.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:
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)
@@ -202,55 +83,51 @@ def average_tcp_latency_ms(ip: str, port: int, timeout: float, samples: int) ->
return sum(readings) / len(readings)
def pick_fastest_region(
def pick_fastest(
candidates: list[str],
previous: str | None,
) -> tuple[str, list[dict[str, Any]]]:
previous_region: str | None,
previous_server_ip: str | None,
) -> tuple[pia.WgServer, 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]] = []
server_counts: dict[str, int] = {}
serverlist = pia.fetch_serverlist()
servers: list[pia.WgServer] = []
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:
servers.extend(pia.region_wg_servers(serverlist, region_id))
if not servers:
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:
best_by_region: dict[str, dict[str, Any]] = {}
with ThreadPoolExecutor(max_workers=min(32, len(servers))) as pool:
futures = {
pool.submit(average_tcp_latency_ms, ip, port, timeout, samples): (region_id, ip)
for region_id, ip in probes
pool.submit(average_tcp_latency_ms, server.ip, port, timeout, samples): server
for server in servers
}
for future in as_completed(futures):
region_id, ip = futures[future]
server = futures[future]
value = future.result()
current = best_by_region[region_id]
current = best_by_region.setdefault(
server.region,
{
"region": server.region,
"latency_ms": None,
"server_ip": None,
"server_cn": None,
"servers": 0,
"failures": 0,
},
)
current["servers"] += 1
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
current["server_ip"] = server.ip
current["server_cn"] = server.cn
results = sorted(
best_by_region.values(),
@@ -260,12 +137,12 @@ def pick_fastest_region(
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']})"
f"(best={item['server_cn']}/{item['server_ip']}, "
f"servers={item['servers']}, failures={item['failures']})"
)
reachable = [item for item in results if item["latency_ms"] is not None]
@@ -273,46 +150,68 @@ def pick_fastest_region(
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_region:
previous_result = next((item for item in reachable if item["region"] == previous_region), None)
if previous_result is not None:
improvement = previous_result["latency_ms"] - winner["latency_ms"]
if winner["region"] != previous and improvement < margin:
same_server = (
previous_server_ip
and previous_result["server_ip"] == previous_server_ip
)
if winner["region"] != previous_region and improvement < margin:
log(
f"Keeping current region {previous} "
f"Keeping current region {previous_region} "
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")
elif winner["region"] == previous_region and same_server:
log(f"Current endpoint {previous_server_ip} is still fastest in {previous_region}")
elif winner["region"] == previous_region:
log(f"Current region {previous_region} is still fastest")
log(f"Fastest region: {winner['region']} ({winner['latency_ms']:.2f}ms via {winner['server_ip']})")
return winner["region"], results
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_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_random(candidates: list[str], previous_region: str | None) -> pia.WgServer:
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 = pia.region_wg_servers(serverlist, region_id)
if servers:
server = random.choice(servers)
log(f"Randomly selected {server.region} via {server.cn} ({server.ip})")
return server
raise SystemExit("No WireGuard servers found for configured regions")
def pick_region(state_path: Path) -> tuple[str, str, list[dict[str, Any]]]:
def pick_server(state_path: Path) -> tuple[pia.WgServer, str, list[dict[str, Any]]]:
regions = parse_regions()
previous = read_previous_region(state_path)
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_region(regions, previous), mode, []
return pick_random(regions, previous_region), mode, []
if mode != "fastest":
raise SystemExit(f"Invalid REGION_SELECT '{mode}' (expected fastest|random)")
region, results = pick_fastest_region(regions, previous)
return region, mode, results
server, results = pick_fastest(regions, previous_region, previous_server_ip)
return server, mode, results
def parse_restart_containers() -> list[str]:
@@ -337,48 +236,6 @@ def restart_containers(containers: list[str]) -> None:
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")
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",
"-r",
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,
]
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 config_age_seconds(wg_path: Path) -> float | None:
if not wg_path.is_file():
return None
@@ -388,19 +245,26 @@ def config_age_seconds(wg_path: Path) -> float | None:
return None
def should_skip_regeneration(region: str, previous: str | None, wg_path: Path) -> bool:
if env_bool("FORCE_ROTATE", False):
def should_skip_regeneration(
server: pia.WgServer,
state: dict[str, Any],
wg_path: Path,
latency_results: list[dict[str, Any]],
) -> bool:
if pia.env_bool("FORCE_ROTATE", False):
log("FORCE_ROTATE=true; regenerating WireGuard config")
return False
if previous != region:
previous_region = state.get("region")
previous_server_ip = state.get("server_ip")
if previous_region != server.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)
max_age = pia.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
@@ -412,16 +276,36 @@ def should_skip_regeneration(region: str, previous: str | None, wg_path: Path) -
):
return False
# Same region, different server: only regenerate if clearly faster.
if previous_server_ip and server.ip != previous_server_ip:
margin = float(os.environ.get("LATENCY_SWITCH_MARGIN_MS", "15"))
previous_latency = state.get("latency_ms")
winner = next((item for item in latency_results if item.get("server_ip") == server.ip), None)
new_latency = winner.get("latency_ms") if winner else None
if (
isinstance(previous_latency, (int, float))
and isinstance(new_latency, (int, float))
and (previous_latency - new_latency) < margin
):
log(
f"Keeping current server {previous_server_ip}; "
f"{server.ip} only {previous_latency - new_latency:.2f}ms faster "
f"(margin {margin:g}ms)"
)
return True
log(f"Switching server within {server.region}: {previous_server_ip} -> {server.ip}")
return False
log(
f"Skipping PIA token/config request; region unchanged ({region}) "
f"and config age {int(age)}s <= {max_age}s"
f"Skipping PIA token/addKey; endpoint unchanged "
f"({server.region}/{server.ip}) and config age {int(age)}s <= {max_age}s"
)
return True
def write_state(
state_path: Path,
region: str,
server: pia.WgServer,
mode: str,
restarted: list[str],
latency_results: list[dict[str, Any]],
@@ -429,10 +313,19 @@ def write_state(
skipped: bool = False,
) -> None:
previous = read_state(state_path)
winner = next((item for item in latency_results if item.get("region") == region), None)
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("region") == server.region), None),
)
now = datetime.now().astimezone().isoformat(timespec="seconds")
payload: dict[str, Any] = {
"region": region,
"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,
@@ -445,7 +338,8 @@ def write_state(
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")
elif skipped and previous.get("latency_ms") is not None:
payload["latency_ms"] = previous["latency_ms"]
if latency_results:
payload["latency_results"] = latency_results
@@ -461,20 +355,31 @@ 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)
state = read_state(state_path)
region, mode, latency_results = pick_region(state_path)
log(f"Selected region: {region} (mode={mode})")
server, mode, latency_results = pick_server(state_path)
log(f"Selected endpoint: {server.region} / {server.cn} / {server.ip} (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}")
if should_skip_regeneration(server, state, wg_path, latency_results):
# If we decided to keep the previous server IP, persist that identity.
keep_ip = state.get("server_ip") if isinstance(state.get("server_ip"), str) else server.ip
keep_cn = state.get("server_cn") if isinstance(state.get("server_cn"), str) else server.cn
if keep_ip != server.ip:
server = pia.WgServer(region=server.region, ip=keep_ip, cn=str(keep_cn or server.cn))
write_state(state_path, server, mode, [], latency_results, skipped=True)
log(f"Rotation skipped for {server.region}/{server.ip}")
return
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)
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")
@@ -486,9 +391,21 @@ def rotate_once() -> None:
restarted = parse_restart_containers()
restart_containers(restarted)
write_state(state_path, region, mode, restarted, latency_results, skipped=False)
log(f"Rotation complete for region={region}")
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__":
rotate_once()
if len(sys.argv) > 1 and sys.argv[1] in {"--list-regions", "list-regions"}:
list_regions_main()
else:
rotate_once()