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

Enhance gluetun PIA WireGuard rotator with health check and recovery logic. Introduce new environment variables for health check interval and unhealthy cooldown. Refactor entrypoint.py to manage Gluetun health status and implement two-step recovery for unhealthy states. Update rotation logic to account for health status during scheduled rotations.
Build and Push Docker Images / build-and-push (push) Successful in 51s
Details

Browse Source
This commit is contained in:
Bram
2026-08-16 18:47:29 +02:00
parent f5b9a7fa21
commit 7d6405c776
4 changed files with 310 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockers/gluetun-pia-wireguard-rotator/rotate.py
+150 -66
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import argparse
import json
import os
import random
@@ -82,25 +83,14 @@ def average_tcp_latency_ms(ip: str, port: int, timeout: float, samples: int) ->
return sum(readings) / len(readings)
def pick_fastest(
candidates: list[str],
previous_region: str | None,
previous_server_ip: str | None,
) -> tuple[pia.WgServer, list[dict[str, Any]]]:
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")))
margin = float(os.environ.get("LATENCY_SWITCH_MARGIN_MS", "15"))
serverlist = pia.fetch_serverlist()
servers: list[pia.WgServer] = []
for region_id in candidates:
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]] = {}
with ThreadPoolExecutor(max_workers=min(32, len(servers))) as pool:
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
@@ -108,67 +98,73 @@ def pick_fastest(
for future in as_completed(futures):
server = futures[future]
value = future.result()
current = best_by_region.setdefault(
server.region,
results.append(
{
"region": server.region,
"latency_ms": None,
"server_ip": None,
"server_cn": None,
"servers": 0,
"failures": 0,
},
"server_ip": server.ip,
"server_cn": server.cn,
"latency_ms": None if value is None else round(value, 2),
}
)
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"] = server.ip
current["server_cn"] = server.cn
results = sorted(
best_by_region.values(),
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} "
f"(best={item['server_cn']}/{item['server_ip']}, "
f"servers={item['servers']}, failures={item['failures']})"
)
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 region")
raise SystemExit("All latency probes failed; cannot select fastest server")
winner = reachable[0]
if previous_region:
previous_result = next((item for item in reachable if item["region"] == previous_region), None)
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"]
same_server = (
previous_server_ip
and previous_result["server_ip"] == previous_server_ip
)
if winner["region"] != previous_region and improvement < margin:
if winner["server_ip"] != previous_server_ip and improvement < margin:
log(
f"Keeping current region {previous_region} "
f"Keeping current server {previous_server_ip} "
f"({previous_result['latency_ms']:.2f}ms); "
f"best {winner['region']} only {improvement:.2f}ms faster "
f"best {winner['server_ip']} only {improvement:.2f}ms faster "
f"(margin {margin:g}ms)"
)
winner = previous_result
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")
elif winner["server_ip"] == previous_server_ip:
log(f"Current endpoint {previous_server_ip} is still fastest")
server = pia.WgServer(
region=winner["region"],
@@ -179,7 +175,13 @@ def pick_fastest(
return server, results
def pick_random(candidates: list[str], previous_region: str | None) -> pia.WgServer:
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:
@@ -189,15 +191,24 @@ def pick_random(candidates: list[str], previous_region: str | None) -> pia.WgSer
random.shuffle(pool)
for region_id in pool:
servers = pia.region_wg_servers(serverlist, region_id)
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 found for configured regions")
raise SystemExit("No WireGuard servers left after excludes")
def pick_server(state_path: Path) -> tuple[pia.WgServer, str, list[dict[str, Any]]]:
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
@@ -205,11 +216,17 @@ def pick_server(state_path: Path) -> tuple[pia.WgServer, str, list[dict[str, Any
mode = os.environ.get("REGION_SELECT", "fastest").strip().lower() or "fastest"
if mode == "random":
return pick_random(regions, previous_region), mode, []
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)
server, results = pick_fastest(
regions,
previous_region,
previous_server_ip,
exclude_ips=exclude_ips,
apply_margin=apply_margin,
)
return server, mode, results
@@ -241,14 +258,17 @@ def write_state(
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("region") == server.region), None),
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] = {
@@ -258,11 +278,17 @@ def write_state(
"selection": mode,
"wg_config": os.environ.get("WG_CONFIG_PATH", "/config/wireguard/wg0.conf"),
"restarted_containers": restarted,
"rotated_at": now,
"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
@@ -271,17 +297,70 @@ def write_state(
state_path.chmod(0o644)
def rotate_once() -> None:
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)
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)
@@ -307,7 +386,7 @@ def rotate_once() -> None:
restarted = parse_restart_containers()
restart_containers(restarted)
write_state(state_path, server, mode, restarted, latency_results)
write_state(state_path, server, mode, restarted, latency_results, skipped=False)
log(f"Rotation complete for {server.region}/{server.ip}")
@@ -321,7 +400,12 @@ def list_regions_main() -> None:
if __name__ == "__main__":
if len(sys.argv) > 1 and sys.argv[1] in {"--list-regions", "list-regions"}:
# 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()
rotate_once(args)