WS14-Docker-Linux-Server/update_build_metadata.py

#!/usr/bin/env python3
import hashlib
import json
import os
import re
import shutil
import sys
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path

try:
    import tomllib  # Python 3.11+
except ModuleNotFoundError:  # pragma: no cover - fallback for older runtimes
    try:
        import tomli as tomllib  # type: ignore
    except ModuleNotFoundError:  # pragma: no cover
        tomllib = None  # type: ignore


def fetch_manifest(url: str) -> dict | None:
    try:
        with urllib.request.urlopen(url, timeout=15) as resp:
            return json.load(resp)
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc:
        print(f"[WARN] Failed to download manifest from {url}: {exc}", file=sys.stderr)
        return None


def select_latest_build(builds: dict) -> dict | None:
    if not builds:
        return None
    # Builds are keyed by commit hash; pick the newest by the ISO8601 timestamp.
    def build_time(item):
        _, data = item
        try:
            return datetime.fromisoformat(data.get("time", "").rstrip("Z"))
        except ValueError:
            return datetime.min

    latest_key, latest_data = max(builds.items(), key=build_time)
    client = latest_data.get("client")
    if not client:
        print(f"[WARN] Manifest entry {latest_key} is missing client info", file=sys.stderr)
        return None
    client.setdefault("version", latest_key)
    return client


def update_config(config_path: str, url: str, sha256: str) -> bool:
    try:
        config_text = Path(config_path).read_text(encoding="utf-8")
    except OSError as exc:
        print(f"[WARN] Cannot read config {config_path}: {exc}", file=sys.stderr)
        return False

    download_pattern = re.compile(r"(?m)^(#\s*)?download_url\s*=.*$")
    build_pattern = re.compile(r"(?m)^(#\s*)?build\s*=.*$")

    new_text = download_pattern.sub(f'download_url = "{url}"', config_text, count=1)
    new_text = build_pattern.sub(f'build = "{sha256.upper()}"', new_text, count=1)

    if new_text == config_text:
        print("[WARN] No download_url/build entries were updated in server_config.toml", file=sys.stderr)
        return False

    try:
        Path(config_path).write_text(new_text, encoding="utf-8")
    except OSError as exc:
        print(f"[WARN] Cannot write config {config_path}: {exc}", file=sys.stderr)
        return False

    print(f"[INFO] Updated build.download_url to {url}")
    return True


def parse_build_section(config_path: str) -> tuple[str | None, str | None] | None:
    if tomllib is None:
        print("[WARN] tomllib/tomli is not available; cannot parse server_config.toml", file=sys.stderr)
        return None

    try:
        contents = Path(config_path).read_text(encoding="utf-8")
    except OSError as exc:
        print(f"[WARN] Cannot read config {config_path}: {exc}", file=sys.stderr)
        return None

    try:
        data = tomllib.loads(contents)
    except (tomllib.TOMLDecodeError, AttributeError) as exc:  # type: ignore[attr-defined]
        print(f"[WARN] Failed to parse TOML {config_path}: {exc}", file=sys.stderr)
        return None

    build_section = data.get("build")
    if not isinstance(build_section, dict):
        print("[WARN] No [build] section found in server_config.toml", file=sys.stderr)
        return None

    download_url = build_section.get("download_url") or build_section.get("download-url")
    build_hash = build_section.get("build")

    download_url = download_url.strip() if isinstance(download_url, str) else None
    build_hash = build_hash.strip() if isinstance(build_hash, str) else None

    return download_url, build_hash


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as file_obj:
        for chunk in iter(lambda: file_obj.read(65536), b""):
            digest.update(chunk)
    return digest.hexdigest().upper()


def download_to_path(url: str, dest: Path) -> bool:
    tmp_path = dest.with_suffix(dest.suffix + ".tmp")
    dest.parent.mkdir(parents=True, exist_ok=True)
    try:
        with urllib.request.urlopen(url, timeout=60) as resp, tmp_path.open("wb") as out_file:
            shutil.copyfileobj(resp, out_file)
        tmp_path.replace(dest)
        return True
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        print(f"[WARN] Failed to download {url}: {exc}", file=sys.stderr)
        try:
            tmp_path.unlink(missing_ok=True)  # type: ignore[call-arg]
        except TypeError:
            if tmp_path.exists():
                tmp_path.unlink()
        return False


def ensure_client_zip(config_path: str) -> bool:
    parsed = parse_build_section(config_path)
    if not parsed:
        return False

    download_url, expected_hash = parsed
    if not download_url:
        print("[WARN] build.download_url is not set; Content.Client.zip cannot be downloaded", file=sys.stderr)
        return False

    client_zip = Path(config_path).with_name("Content.Client.zip")
    expected_hash_upper = expected_hash.upper() if expected_hash else None

    if client_zip.exists() and expected_hash_upper:
        try:
            current_hash = sha256_file(client_zip)
        except OSError as exc:
            print(f"[WARN] Cannot read {client_zip}: {exc}", file=sys.stderr)
        else:
            if current_hash == expected_hash_upper:
                print("[INFO] Content.Client.zip already matches configured build hash")
                return True
            print("[INFO] Content.Client.zip hash mismatch; re-downloading")

    if not download_to_path(download_url, client_zip):
        return False

    if expected_hash_upper:
        try:
            downloaded_hash = sha256_file(client_zip)
        except OSError as exc:
            print(f"[WARN] Cannot read downloaded {client_zip}: {exc}", file=sys.stderr)
            return False
        if downloaded_hash != expected_hash_upper:
            print(
                f"[WARN] Downloaded Content.Client.zip hash {downloaded_hash} does not match expected {expected_hash_upper}",
                file=sys.stderr,
            )
            client_zip.unlink(missing_ok=True)  # type: ignore[call-arg]
            return False

    print(f"[INFO] Downloaded Content.Client.zip to {client_zip}")
    return True


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: update_build_metadata.py <path-to-server_config.toml>", file=sys.stderr)
        sys.exit(1)

    manifest_url = os.environ.get(
        "ROBUST_CDN_MANIFEST", "https://cdn.wylab.me/fork/wylab/manifest"
    )

    manifest = fetch_manifest(manifest_url)
    if not manifest:
        sys.exit(0)

    client_info = select_latest_build(manifest.get("builds", {}))
    if not client_info:
        sys.exit(0)

    client_url = client_info.get("url")
    client_hash = client_info.get("sha256")
    if not (client_url and client_hash):
        print("[WARN] Client entry missing url/sha256", file=sys.stderr)
        sys.exit(0)

    config_path = sys.argv[1]
    if update_config(config_path, client_url, client_hash):
        ensure_client_zip(config_path)
    else:
        print("[WARN] Skipping client download because config update failed", file=sys.stderr)