MapMil/centers/parsing/vpn/controller.py

"""Poll CA /internal/vpn and keep mihomo config in sync for subscription mode."""

from __future__ import annotations

import json
import logging
import os
import signal
import subprocess
import sys
import time
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
logger = logging.getLogger("cp-vpn")

CA_API_URL = os.getenv("CA_API_URL", "http://ca-api:8000").rstrip("/")
INTERNAL_TOKEN = os.getenv("INTERNAL_TOKEN", "dev-internal-token")
POLL_SECONDS = int(os.getenv("VPN_POLL_SECONDS", "15"))
CONFIG_DIR = Path(os.getenv("MIHOMO_CONFIG_DIR", "/etc/mihomo"))
CONFIG_PATH = CONFIG_DIR / "config.yaml"
MIHOMO_BIN = os.getenv("MIHOMO_BIN", "/usr/local/bin/mihomo")
MIXED_PORT = int(os.getenv("VPN_SOCKS_PORT", "1080"))
CONTROLLER = "127.0.0.1:9090"

_mihomo_proc: subprocess.Popen | None = None
_last_fingerprint: str | None = None


def fetch_vpn() -> dict:
    req = Request(
        f"{CA_API_URL}/internal/vpn",
        headers={"X-Internal-Token": INTERNAL_TOKEN},
        method="GET",
    )
    with urlopen(req, timeout=15) as resp:
        return json.loads(resp.read().decode("utf-8"))


def fingerprint(cfg: dict) -> str:
    return json.dumps(
        {
            "enabled": cfg.get("enabled"),
            "mode": cfg.get("mode"),
            "subscription_url": cfg.get("subscription_url"),
            "subscription_interval_seconds": cfg.get("subscription_interval_seconds"),
        },
        sort_keys=True,
    )


def write_direct_config() -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    (CONFIG_DIR / "providers").mkdir(parents=True, exist_ok=True)
    CONFIG_PATH.write_text(
        f"""# managed by cp-vpn controller — DIRECT (subscription inactive)
mixed-port: {MIXED_PORT}
allow-lan: true
bind-address: "*"
mode: direct
log-level: warning
external-controller: {CONTROLLER}
""",
        encoding="utf-8",
    )


def write_subscription_config(url: str, interval: int) -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    (CONFIG_DIR / "providers").mkdir(parents=True, exist_ok=True)
    # YAML: quote URL; escape double quotes in URL if any
    safe_url = url.replace("\\", "\\\\").replace('"', '\\"')
    CONFIG_PATH.write_text(
        f"""# managed by cp-vpn controller — subscription
mixed-port: {MIXED_PORT}
allow-lan: true
bind-address: "*"
mode: rule
log-level: info
external-controller: {CONTROLLER}
proxy-providers:
  sub:
    type: http
    url: "{safe_url}"
    interval: {interval}
    path: ./providers/sub.yaml
    health-check:
      enable: true
      url: https://www.gstatic.com/generate_204
      interval: 600
proxy-groups:
  - name: PROXY
    type: select
    use:
      - sub
rules:
  - MATCH,PROXY
""",
        encoding="utf-8",
    )


def start_mihomo() -> None:
    global _mihomo_proc
    if _mihomo_proc is not None and _mihomo_proc.poll() is None:
        return
    logger.info("Starting mihomo (%s)", MIHOMO_BIN)
    _mihomo_proc = subprocess.Popen(
        [MIHOMO_BIN, "-d", str(CONFIG_DIR)],
        stdout=sys.stdout,
        stderr=sys.stderr,
    )


def stop_mihomo() -> None:
    global _mihomo_proc
    if _mihomo_proc is None:
        return
    if _mihomo_proc.poll() is None:
        logger.info("Stopping mihomo")
        _mihomo_proc.send_signal(signal.SIGTERM)
        try:
            _mihomo_proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            _mihomo_proc.kill()
    _mihomo_proc = None


def reload_mihomo() -> None:
    """Force reload via external controller; fall back to process restart."""
    body = json.dumps({"path": str(CONFIG_PATH)}).encode("utf-8")
    req = Request(
        f"http://{CONTROLLER}/configs?force=true",
        data=body,
        headers={"Content-Type": "application/json"},
        method="PUT",
    )
    try:
        with urlopen(req, timeout=10) as resp:
            resp.read()
        logger.info("Mihomo config reloaded")
        return
    except (HTTPError, URLError, OSError) as exc:
        logger.warning("Reload via API failed (%s); restarting mihomo", exc)
    stop_mihomo()
    start_mihomo()


def apply_config(cfg: dict) -> None:
    global _last_fingerprint
    fp = fingerprint(cfg)
    if fp == _last_fingerprint and _mihomo_proc is not None and _mihomo_proc.poll() is None:
        return

    enabled = bool(cfg.get("enabled"))
    mode = (cfg.get("mode") or "").strip()
    url = (cfg.get("subscription_url") or "").strip()
    interval = int(cfg.get("subscription_interval_seconds") or 3600)

    if enabled and mode == "subscription" and url:
        logger.info("Applying subscription config (interval=%ss)", interval)
        write_subscription_config(url, interval)
    else:
        logger.info("Applying DIRECT config (enabled=%s mode=%s)", enabled, mode)
        write_direct_config()

    already_running = _mihomo_proc is not None and _mihomo_proc.poll() is None
    start_mihomo()
    if already_running or _last_fingerprint is not None:
        time.sleep(1)
        reload_mihomo()
    _last_fingerprint = fp


def main() -> None:
    global _mihomo_proc
    logger.info(
        "cp-vpn controller started (poll=%ss api=%s)",
        POLL_SECONDS,
        CA_API_URL,
    )
    write_direct_config()
    start_mihomo()

    while True:
        try:
            cfg = fetch_vpn()
            apply_config(cfg)
        except Exception:
            logger.exception("Failed to sync VPN config")
        if _mihomo_proc is not None and _mihomo_proc.poll() is not None:
            logger.warning("mihomo exited with code %s; restarting", _mihomo_proc.returncode)
            _mihomo_proc = None
            start_mihomo()
        time.sleep(POLL_SECONDS)


if __name__ == "__main__":
    main()