pxe-server/webapp/services/updates_common.py

"""Shared plumbing for the Updates pages (Dell drivers/BIOS and iPXE).

  * settings: the download speed cap, one small JSON file;
  * download(): curl with that cap, resume, and a hash check;
  * background jobs: long downloads run as a transient systemd unit
    (systemd-run), so a webapp restart or a closed browser does not kill a
    40-minute driver pull. Each job writes its progress to a JSON state file
    the page polls.

The server reaches the internet through its corporate NIC (enp0s31f6); GE's
Zscaler re-signs GitHub, so the Zscaler root CA must be in the system trust
store (the playbook installs it). downloads.dell.com is not intercepted.
"""
from __future__ import annotations

import hashlib
import json
import math
import os
import subprocess
import sys
import tempfile
from datetime import datetime, timezone

import config

DEFAULT_SETTINGS = {
    # MB/s; 0 = no cap. 10 MB/s is about 80 Mbps, a fifth of the site link.
    "max_mbps": 10,
}


def state_dir(*parts):
    path = os.path.join(config.UPDATES_DIR, *parts)
    os.makedirs(path, exist_ok=True)
    return path


def now():
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def read_json(path, default=None):
    try:
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    except (OSError, ValueError):
        return default


def write_json(path, data):
    folder = os.path.dirname(path)
    os.makedirs(folder, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=folder, prefix=".tmp-")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=1)
        os.chmod(tmp, 0o644)  # mkstemp makes 0600
        os.replace(tmp, path)
    except BaseException:
        if os.path.exists(tmp):
            os.remove(tmp)
        raise


# --- settings ---------------------------------------------------------------

def settings():
    data = dict(DEFAULT_SETTINGS)
    saved = read_json(os.path.join(state_dir(), "settings.json"), {})
    if isinstance(saved, dict):
        data.update(saved)
    try:
        cap = float(data.get("max_mbps"))
    except (TypeError, ValueError):
        cap = -1
    if not math.isfinite(cap) or cap < 0:
        data["max_mbps"] = DEFAULT_SETTINGS["max_mbps"]  # hand-edited junk
    return data


def save_settings(max_mbps):
    try:
        value = float(str(max_mbps).strip())
    except ValueError:
        raise ValueError("speed cap must be a number of MB/s") from None
    # float() takes "nan", which passes both comparisons below.
    if not math.isfinite(value) or value < 0 or value > 1000:
        raise ValueError("speed cap must be between 0 and 1000 MB/s")
    data = settings()
    # A tiny cap must not round to 0, which means unlimited.
    data["max_mbps"] = max(round(value, 1), 0.1) if value > 0 else 0
    write_json(os.path.join(state_dir(), "settings.json"), data)
    return data


# --- downloads --------------------------------------------------------------

def file_hash(path, algorithm="sha256"):
    h = hashlib.new(algorithm)
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1 << 22), b""):
            h.update(block)
    return h.hexdigest()


def download(url, dest, expect=None, algorithm="sha256", timeout=None):
    """Fetch url to dest at the configured speed cap, resuming a partial file.
    expect: hex digest to verify (case-insensitive). Raises on failure.
    Without expect nothing catches a stale file resumed into, so callers
    that pass no hash should remove dest first."""
    os.makedirs(os.path.dirname(os.path.abspath(dest)), exist_ok=True)
    cmd = ["curl", "-sS", "-L", "--fail", "--retry", "5", "--retry-delay", "5",
           "-C", "-", "-o", dest]
    cap = settings().get("max_mbps") or 0
    if cap:
        cmd += ["--limit-rate", f"{max(int(float(cap) * 1024 * 1024), 1024)}"]
    if timeout:
        cmd += ["-m", str(int(timeout))]
    for attempt in range(2):
        r = subprocess.run(cmd + [url], capture_output=True, text=True)
        if r.returncode == 33:
            # The server will not resume: the partial on disk is not the
            # file, so throw it away and fetch from the start.
            if os.path.exists(dest):
                os.remove(dest)
            continue
        if r.returncode:
            raise RuntimeError(f"download failed ({r.returncode}): {r.stderr.strip()[:300]}")
        if not expect or file_hash(dest, algorithm).lower() == expect.strip().lower():
            return dest
        # A corrupt partial resumes into a corrupt whole: start over once.
        os.remove(dest)
    if expect and os.path.exists(dest):
        os.remove(dest)
    raise RuntimeError(f"could not download a verified copy of {os.path.basename(dest)}"
                       + (f" ({algorithm} mismatch)" if expect else ""))


def fetch_text(url, timeout=30):
    r = subprocess.run(["curl", "-sS", "-L", "--fail", "-m", str(timeout), url],
                       capture_output=True)
    if r.returncode:
        raise RuntimeError(f"cannot fetch {url}: {r.stderr.decode(errors='replace').strip()[:300]}")
    return r.stdout


# --- background jobs --------------------------------------------------------

def job_running(unit):
    r = subprocess.run(["systemctl", "is-active", "--quiet", unit])
    return r.returncode == 0


def _job_env():
    """--setenv flags for the job: a transient unit starts with an empty
    environment, but pxe-webapp.service sets the share paths (SAMBA_SHARE,
    WEB_ROOT, AUDIT_LOG, ...) and TZ, and the job must see the same config."""
    keep = {"TZ", "LANG", "PYTHONPATH"}
    # Never secrets: systemd-run's argv is readable by every local user.
    secret = ("KEY", "SECRET", "PASS", "TOKEN")
    return [f"--setenv={k}={v}" for k, v in sorted(os.environ.items())
            if (k in keep or (k.isupper() and hasattr(config, k)))
            and not any(w in k for w in secret) and "\n" not in v]


def start_job(unit, args):
    """Run `python3 -m <args>` from the webapp directory as a transient
    systemd unit, with the webapp's own interpreter and config environment.
    Refuses when the unit is already running."""
    if job_running(unit):
        raise RuntimeError("a job of this kind is already running")
    subprocess.run(["systemctl", "reset-failed", unit], capture_output=True)
    here = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    r = subprocess.run(
        ["systemd-run", f"--unit={unit}", "--collect", f"--working-directory={here}"]
        + _job_env() + [sys.executable or "/usr/bin/python3", "-m"] + list(args),
        capture_output=True, text=True)
    if r.returncode:
        raise RuntimeError(f"systemd-run failed: {r.stderr.strip()[:300]}")


def stop_job(unit):
    subprocess.run(["systemctl", "stop", unit], capture_output=True)