optimize

2026-08-14 20:03:05 +02:00
parent 3eeb93e225
commit bc9d79de02
5 changed files with 552 additions and 288 deletions
@@ -0,0 +1,360 @@
+#!/usr/bin/env python3
+"""Minimal PIA WireGuard client: serverlist, token, keygen, addKey."""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+import re
+import socket
+import ssl
+import sys
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from dataclasses import dataclass
+from datetime import datetime
+from http.client import HTTPSConnection
+from pathlib import Path
+from typing import Any
+
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
+
+EXIT_RATE_LIMITED = 75
+SERVERLIST_URL = "https://serverlist.piaservers.net/vpninfo/servers/v6"
+TOKEN_URL = "https://www.privateinternetaccess.com/api/client/v2/token"
+PIA_CA_URL = (
+    "https://raw.githubusercontent.com/pia-foss/manual-connections/master/ca.rsa.4096.crt"
+)
+
+_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)
+
+
+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 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 serverlist_cache_path() -> Path:
+    return Path(os.environ.get("SERVERLIST_CACHE_PATH", "/config/cache/pia-serverlist.json"))
+
+
+def token_cache_path() -> Path:
+    return Path(os.environ.get("TOKEN_CACHE_PATH", "/config/cache/pia-token.json"))
+
+
+def ca_cache_path() -> Path:
+    return Path(os.environ.get("PIA_CA_PATH", "/config/cache/ca.rsa.4096.crt"))
+
+
+@dataclass(frozen=True)
+class WgServer:
+    region: str
+    ip: str
+    cn: str
+
+
+@dataclass(frozen=True)
+class WgKeys:
+    private_key: str
+    public_key: str
+
+
+def generate_wg_keys() -> WgKeys:
+    private = X25519PrivateKey.generate()
+    priv_bytes = private.private_bytes(
+        encoding=serialization.Encoding.Raw,
+        format=serialization.PrivateFormat.Raw,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+    pub_bytes = private.public_key().public_bytes(
+        encoding=serialization.Encoding.Raw,
+        format=serialization.PublicFormat.Raw,
+    )
+    return WgKeys(
+        private_key=base64.b64encode(priv_bytes).decode("ascii"),
+        public_key=base64.b64encode(pub_bytes).decode("ascii"),
+    )
+
+
+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 fetch_serverlist() -> dict[str, Any]:
+    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")
+
+    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:
+            text = resp.read().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 list_regions(serverlist: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+    data = serverlist or fetch_serverlist()
+    regions = []
+    for region in data.get("regions", []):
+        regions.append(
+            {
+                "id": region.get("id"),
+                "name": region.get("name"),
+                "country": region.get("country"),
+                "port_forward": bool(region.get("port_forward")),
+                "offline": bool(region.get("offline")),
+                "wg_servers": len((region.get("servers") or {}).get("wg") or []),
+            }
+        )
+    regions.sort(key=lambda item: (item.get("country") or "", item.get("id") or ""))
+    return regions
+
+
+def region_wg_servers(serverlist: dict[str, Any], region_id: str) -> list[WgServer]:
+    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") or {}).get("wg") or []
+        out: list[WgServer] = []
+        for server in servers:
+            if not isinstance(server, dict) or not server.get("ip") or not server.get("cn"):
+                continue
+            out.append(WgServer(region=region_id, ip=server["ip"], cn=server["cn"]))
+        return out
+    log(f"Region {region_id} not found in PIA server list; skipping")
+    return []
+
+
+def ensure_pia_ca() -> Path:
+    path = ca_cache_path()
+    ttl = parse_duration_seconds(os.environ.get("PIA_CA_CACHE_TTL", "30d"), 30 * 86400)
+    if path.is_file() and (time.time() - path.stat().st_mtime) <= ttl:
+        return path
+    try:
+        with urllib.request.urlopen(PIA_CA_URL, timeout=30) as resp:
+            data = resp.read()
+        if b"BEGIN CERTIFICATE" not in data:
+            raise SystemExit("Downloaded PIA CA does not look like a certificate")
+        path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = path.with_suffix(path.suffix + ".tmp")
+        tmp.write_bytes(data)
+        os.replace(tmp, path)
+        log(f"Cached PIA CA certificate -> {path}")
+    except (urllib.error.URLError, OSError) as exc:
+        if path.is_file():
+            log(f"PIA CA download failed ({exc}); using existing {path}")
+            return path
+        raise SystemExit(f"Failed to download PIA CA certificate: {exc}") from exc
+    return path
+
+
+def _rate_limited(body: str, status_code: int) -> bool:
+    return status_code == 429 or "too_many_attempts" in body
+
+
+def get_token(username: str, password: str) -> str:
+    cache_path = token_cache_path()
+    ttl = parse_duration_seconds(os.environ.get("TOKEN_CACHE_TTL", "20h"), 20 * 3600)
+    force = env_bool("FORCE_TOKEN_REFRESH", False)
+
+    if cache_path.is_file() and not force:
+        try:
+            cached = json.loads(cache_path.read_text(encoding="utf-8"))
+            token = cached.get("token")
+            obtained = float(cached.get("obtained_at", 0))
+            if isinstance(token, str) and token and (time.time() - obtained) <= ttl:
+                log(f"Using cached PIA token (age {int(time.time() - obtained)}s <= TTL {ttl}s)")
+                return token
+        except (OSError, json.JSONDecodeError, TypeError, ValueError):
+            pass
+
+    form = urllib.parse.urlencode({"username": username, "password": password}).encode()
+    req = urllib.request.Request(
+        TOKEN_URL,
+        data=form,
+        method="POST",
+        headers={"Content-Type": "application/x-www-form-urlencoded"},
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=30) as resp:
+            body = resp.read().decode("utf-8", errors="replace")
+            status = resp.status
+    except urllib.error.HTTPError as exc:
+        body = exc.read().decode("utf-8", errors="replace")
+        if _rate_limited(body, exc.code):
+            raise SystemExit(EXIT_RATE_LIMITED) from exc
+        raise SystemExit(f"PIA token request failed with status {exc.code}: {body}") from exc
+    except urllib.error.URLError as exc:
+        raise SystemExit(f"PIA token request failed: {exc}") from exc
+
+    if _rate_limited(body, status):
+        raise SystemExit(EXIT_RATE_LIMITED)
+
+    try:
+        payload = json.loads(body)
+    except json.JSONDecodeError as exc:
+        raise SystemExit(f"Invalid PIA token response: {body}") from exc
+
+    token = payload.get("token")
+    if not isinstance(token, str) or not token:
+        raise SystemExit(f"PIA token response missing token: {body}")
+
+    cache_path.parent.mkdir(parents=True, exist_ok=True)
+    cache_path.write_text(
+        json.dumps({"token": token, "obtained_at": time.time()}, indent=2) + "\n",
+        encoding="utf-8",
+    )
+    cache_path.chmod(0o600)
+    log("Fetched and cached new PIA token")
+    return token
+
+
+class _HTTPSConnectionToIP(HTTPSConnection):
+    """HTTPS connection to a fixed IP while presenting server_hostname for SNI/verify."""
+
+    def __init__(self, ip: str, port: int, server_hostname: str, context: ssl.SSLContext):
+        super().__init__(server_hostname, port=port, context=context, timeout=30)
+        self._connect_ip = ip
+        self._server_hostname = server_hostname
+
+    def connect(self) -> None:
+        sock = socket.create_connection((self._connect_ip, self.port), self.timeout)
+        self.sock = self._context.wrap_socket(sock, server_hostname=self._server_hostname)
+
+
+def add_key(server: WgServer, token: str, public_key: str) -> dict[str, Any]:
+    ca_path = ensure_pia_ca()
+    context = ssl.create_default_context(cafile=str(ca_path))
+    query = urllib.parse.urlencode({"pt": token, "pubkey": public_key})
+    path = f"/addKey?{query}"
+
+    conn = _HTTPSConnectionToIP(server.ip, 1337, server.cn, context)
+    try:
+        conn.request("GET", path, headers={"Content-Type": "application/json"})
+        resp = conn.getresponse()
+        body = resp.read().decode("utf-8", errors="replace")
+        status = resp.status
+    finally:
+        conn.close()
+
+    if _rate_limited(body, status):
+        raise SystemExit(EXIT_RATE_LIMITED)
+    if status != 200:
+        raise SystemExit(f"addKey failed for {server.cn}/{server.ip}: status {status}: {body}")
+
+    try:
+        payload = json.loads(body)
+    except json.JSONDecodeError as exc:
+        raise SystemExit(f"Invalid addKey response: {body}") from exc
+
+    if payload.get("status") not in (None, "OK"):
+        raise SystemExit(f"addKey rejected: {body}")
+
+    required = ("server_key", "server_ip", "peer_ip", "dns_servers")
+    missing = [key for key in required if key not in payload]
+    if missing:
+        raise SystemExit(f"addKey response missing {missing}: {body}")
+    if not payload["dns_servers"]:
+        raise SystemExit(f"addKey response missing dns_servers: {body}")
+
+    return payload
+
+
+def render_wg_config(keys: WgKeys, addkey: dict[str, Any], server: WgServer) -> str:
+    endpoint_ip = addkey.get("server_ip") or server.ip
+    endpoint_port = addkey.get("server_port") or 1337
+    dns = addkey["dns_servers"][0]
+    return (
+        "[Interface]\n"
+        f"PrivateKey = {keys.private_key}\n"
+        f"Address = {addkey['peer_ip']}\n"
+        f"DNS = {dns}\n"
+        "[Peer]\n"
+        f"PublicKey = {addkey['server_key']}\n"
+        "AllowedIPs = 0.0.0.0/0\n"
+        f"Endpoint = {endpoint_ip}:{endpoint_port}\n"
+        "PersistentKeepalive = 25\n"
+    )
+
+
+def generate_wg_config(
+    username: str,
+    password: str,
+    server: WgServer,
+    outfile: Path,
+) -> dict[str, Any]:
+    log(f"Requesting PIA token (cached when possible)")
+    token = get_token(username, password)
+    log("Generating local WireGuard keypair")
+    keys = generate_wg_keys()
+    log(f"Registering pubkey via addKey on {server.cn} ({server.ip})")
+    addkey = add_key(server, token, keys.public_key)
+    config = render_wg_config(keys, addkey, server)
+    outfile.write_text(config, encoding="utf-8")
+    outfile.chmod(0o600)
+    return {
+        "region": server.region,
+        "server_ip": server.ip,
+        "server_cn": server.cn,
+        "endpoint_ip": addkey.get("server_ip") or server.ip,
+        "peer_ip": addkey.get("peer_ip"),
+    }