planet/scripts/docker_proxy.py

"""Select a reachable build route and manage only Planet-owned daemon proxy settings."""

from __future__ import annotations

import concurrent.futures
import fcntl
import json
import os
from pathlib import Path
import re
import subprocess
import sys
import tempfile
from urllib.parse import urlsplit

CONFIG = Path("/etc/docker/daemon.json")
STATE = Path("/etc/docker/planet-proxy-state.json")
PROXY_KEYS = ("http-proxy", "https-proxy", "no-proxy")
PROXY_ENV = ("HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy")
PROXY_SCHEMES = ("http", "https")
DEFAULT_NO_PROXY = "localhost,127.0.0.1,::1"
CONFIG_ERROR = "PLANET_PROXY_CONFIG_FAILED"
COMMAND_TIMEOUT_SECONDS = 60
CONNECT_TIMEOUT_SECONDS = 3
REQUEST_TIMEOUT_SECONDS = 8
PROBE_PROCESS_TIMEOUT_SECONDS = 10


class ProxyError(Exception):
    pass


def run(args: list[str], **kwargs: object) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        args, text=True, capture_output=True, check=True, timeout=COMMAND_TIMEOUT_SECONDS, **kwargs
    )


def daemon_proxies() -> dict[str, str]:
    info = json.loads(run(["docker", "info", "--format", "{{json .}}"]).stdout)
    values = dict(
        zip(PROXY_KEYS, (info.get("HttpProxy"), info.get("HttpsProxy"), info.get("NoProxy")))
    )
    return {key: value for key, value in values.items() if value}


def registry_urls(environment: dict[str, str]) -> list[str]:
    dockerfile = Path(__file__).resolve().parents[1] / "aiprovider/Dockerfile"
    defaults = dict(re.findall(r"^ARG (PYTHON_IMAGE|UV_IMAGE)=(.+)$", dockerfile.read_text(), re.M))
    registries = set()
    for key in ("PYTHON_IMAGE", "UV_IMAGE"):
        image = environment.get(key) or defaults[key]
        prefix = image.split("/")[0]
        registry = (
            prefix
            if "/" in image and ("." in prefix or ":" in prefix or prefix == "localhost")
            else "docker.io"
        )
        registries.add("registry-1.docker.io" if registry == "docker.io" else registry)
    return [f"https://{registry}/v2/" for registry in sorted(registries)]


def probe_url(url: str, proxy: str, no_proxy: str = "") -> bool:
    environment = {
        key: value
        for key, value in os.environ.items()
        if key.lower() not in ("http_proxy", "https_proxy", "all_proxy", "no_proxy")
    }
    if proxy:
        environment.update(http_proxy=proxy, https_proxy=proxy)
    try:
        result = subprocess.run(
            [
                "curl",
                "-q",
                "--silent",
                "--output",
                "/dev/null",
                "--write-out",
                "%{http_code}",
                "--noproxy",
                no_proxy if proxy else "",
                "--connect-timeout",
                str(CONNECT_TIMEOUT_SECONDS),
                "--max-time",
                str(REQUEST_TIMEOUT_SECONDS),
                url,
            ],
            env=environment,
            text=True,
            capture_output=True,
            timeout=PROBE_PROCESS_TIMEOUT_SECONDS,
        )
        return result.returncode == 0 and (
            result.stdout.startswith("2") or result.stdout in ("401", "403", "429")
        )
    except (OSError, subprocess.TimeoutExpired):
        return False


def reachable(urls: list[str], proxy: str, no_proxy: str = "") -> bool:
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(urls)) as pool:
        return all(pool.map(lambda url: probe_url(url, proxy, no_proxy), urls))


def select_route(environment: dict[str, str], urls: list[str]) -> dict[str, object]:
    candidates = list(dict.fromkeys(environment[key] for key in PROXY_ENV if environment.get(key)))
    no_proxy = environment.get("NO_PROXY") or environment.get("no_proxy") or DEFAULT_NO_PROXY
    for proxy in candidates:
        if urlsplit(proxy).scheme not in PROXY_SCHEMES:
            continue
        if reachable(urls, proxy, no_proxy):
            return {
                "desired": dict(zip(PROXY_KEYS, (proxy, proxy, no_proxy))),
                "mode": "proxy",
                "reachable": True,
            }
    return {"desired": {}, "mode": "direct", "reachable": reachable(urls, "")}


def atomic_write(path: Path, contents: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    descriptor, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    try:
        with os.fdopen(descriptor, "wb") as output:
            output.write(contents)
        os.replace(temporary, path)
    finally:
        Path(temporary).unlink(missing_ok=True)


def encode(value: dict[str, object]) -> bytes:
    return (json.dumps(value, indent=2) + "\n").encode()


def validate_plan(plan: dict[str, object]) -> None:
    for field in ("current", "desired"):
        values = plan.get(field)
        if not isinstance(values, dict) or set(values) - set(PROXY_KEYS):
            raise ProxyError(CONFIG_ERROR)
        if any(not isinstance(value, str) or "\n" in value for value in values.values()):
            raise ProxyError(CONFIG_ERROR)
    for key in ("http-proxy", "https-proxy"):
        value = plan["desired"].get(key, "")
        if value and urlsplit(value).scheme not in PROXY_SCHEMES:
            raise ProxyError(CONFIG_ERROR)


def restore_containers(containers: list[str]) -> None:
    if containers:
        run(["docker", "start", *containers])


def apply_plan(plan: dict[str, object], config: Path = CONFIG, state: Path = STATE) -> None:
    validate_plan(plan)
    current = daemon_proxies()
    if current != plan["current"]:
        raise ProxyError("PLANET_PROXY_CONFIG_CHANGED")
    if current == plan["desired"]:
        return
    original = config.read_bytes() if config.exists() else None
    data = json.loads(original) if original else {}
    owned = json.loads(state.read_text()) if state.exists() else None
    configured = {key: value for key, value in data.get("proxies", {}).items() if value}
    if (configured or current) and (
        not owned or owned.get("managed") != configured or current != configured
    ):
        raise ProxyError("PLANET_PROXY_EXTERNAL")
    if plan["desired"]:
        data["proxies"] = plan["desired"]
    else:
        data.pop("proxies", None)
    containers = run(["docker", "ps", "-q"]).stdout.split()
    if original is not None:
        atomic_write(config.with_suffix(".json.planet-proxy.bak"), original)
    atomic_write(config, encode(data))
    restart_requested = False
    try:
        run(["dockerd", "--validate", "--config-file", str(config)])
        restart_requested = True
        run(["systemctl", "restart", "docker"])
        restore_containers(containers)
        if daemon_proxies() != plan["desired"]:
            raise ProxyError(CONFIG_ERROR)
        atomic_write(state, encode({"managed": plan["desired"]}))
    except Exception as error:
        rollback(config, original, containers, restart_requested)
        raise ProxyError(CONFIG_ERROR) from error


def rollback(
    config: Path, original: bytes | None, containers: list[str], restart_requested: bool
) -> None:
    if original is None:
        config.unlink(missing_ok=True)
    else:
        atomic_write(config, original)
    if not restart_requested:
        return
    try:
        run(["systemctl", "restart", "docker"])
        restore_containers(containers)
    except Exception as error:
        raise ProxyError("PLANET_PROXY_ROLLBACK_FAILED") from error


def main() -> None:
    action = sys.argv[1]
    if action == "plan":
        plan = select_route(dict(os.environ), registry_urls(dict(os.environ)))
        plan["current"] = daemon_proxies()
        json.dump(plan, sys.stdout)
    elif action == "status":
        plan = json.load(sys.stdin)
        changed = int(plan["current"] != plan["desired"])
        print(plan["mode"], changed, int(plan["reachable"]))
    elif action == "apply":
        if os.geteuid() != 0:
            raise ProxyError(CONFIG_ERROR)
        CONFIG.parent.mkdir(parents=True, exist_ok=True)
        with (CONFIG.parent / "planet-proxy.lock").open("a") as lock:
            fcntl.flock(lock, fcntl.LOCK_EX)
            apply_plan(json.load(sys.stdin))
    else:
        raise ProxyError(CONFIG_ERROR)


if __name__ == "__main__":
    try:
        main()
    except ProxyError as error:
        sys.stderr.write(str(error) + "\n")
        sys.exit(1)
    except Exception:
        # Subprocess output and proxy URLs may contain credentials.
        sys.stderr.write(CONFIG_ERROR + "\n")
        sys.exit(1)