__init__

uvm/integrity.py

@@ -0,0 +1,145 @@
+"""Checksum policy for downloaded executable and guest-image artifacts."""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import json
+import os
+import re
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Mapping
+
+from .errors import UvmError
+
+
+_SHA256_PATTERN = re.compile(r"[0-9a-fA-F]{64}")
+
+
+@dataclass(frozen=True, slots=True)
+class IntegrityManifest:
+    checksums: dict[str, str]
+    verified: bool
+
+
+def normalize_sha256(value: str, environment_name: str) -> str:
+    digest = value.strip().lower()
+    if not _SHA256_PATTERN.fullmatch(digest):
+        raise UvmError(f"{environment_name} must contain a 64-character SHA-256 digest")
+    return digest
+
+
+def require_checksum(
+    expected: str | None,
+    label: str,
+    environment_name: str,
+    *,
+    allow_unverified: bool,
+) -> str | None:
+    if expected is not None:
+        return normalize_sha256(expected, environment_name)
+    if allow_unverified:
+        return None
+    raise UvmError(
+        f"{label} checksum is required. Set {environment_name} to a trusted SHA-256 "
+        "or explicitly set UVM_ALLOW_UNVERIFIED_DOWNLOADS=1 for local development."
+    )
+
+
+def verify_file(
+    path: Path,
+    expected: str | None,
+    label: str,
+    environment_name: str,
+    *,
+    allow_unverified: bool,
+) -> None:
+    expected_digest = require_checksum(
+        expected,
+        label,
+        environment_name,
+        allow_unverified=allow_unverified,
+    )
+    if expected_digest is None:
+        return
+
+    actual_digest = sha256_file(path, label)
+    if not hmac.compare_digest(actual_digest, expected_digest):
+        raise UvmError(
+            f"SHA-256 mismatch for {label}: expected {expected_digest}, got {actual_digest}"
+        )
+
+
+def sha256_file(path: Path, label: str) -> str:
+    digest = hashlib.sha256()
+    try:
+        with path.open("rb") as artifact:
+            for chunk in iter(lambda: artifact.read(1024 * 1024), b""):
+                digest.update(chunk)
+    except OSError as error:
+        raise UvmError(f"could not checksum {label} at {path}: {error}") from error
+    return digest.hexdigest()
+
+
+def load_manifest(path: Path) -> IntegrityManifest:
+    if not path.exists():
+        return IntegrityManifest(checksums={}, verified=False)
+    try:
+        with path.open(encoding="utf-8") as manifest_file:
+            raw = json.load(manifest_file)
+    except (OSError, json.JSONDecodeError) as error:
+        raise UvmError(f"cannot read integrity manifest {path}: {error}") from error
+    if not isinstance(raw, dict):
+        raise UvmError(f"integrity manifest {path} is not an object")
+
+    if "checksums" in raw:
+        raw_checksums = raw["checksums"]
+        verified = raw.get("verified")
+        if not isinstance(raw_checksums, dict) or not isinstance(verified, bool):
+            raise UvmError(f"integrity manifest {path} has an invalid format")
+    else:
+        # Flat manifests were written by earlier uvm versions without provenance.
+        raw_checksums = raw
+        verified = False
+
+    manifest: dict[str, str] = {}
+    for key, value in raw_checksums.items():
+        if not isinstance(key, str) or not isinstance(value, str):
+            raise UvmError(f"integrity manifest {path} contains an invalid entry")
+        manifest[key] = normalize_sha256(value, f"integrity manifest entry {key}")
+    return IntegrityManifest(checksums=manifest, verified=verified)
+
+
+def write_manifest(path: Path, values: Mapping[str, str], *, verified: bool) -> None:
+    normalized = {
+        key: normalize_sha256(value, f"integrity manifest entry {key}")
+        for key, value in values.items()
+    }
+    path.parent.mkdir(parents=True, exist_ok=True)
+    descriptor, temporary_name = tempfile.mkstemp(
+        prefix=f".{path.name}.",
+        suffix=".tmp",
+        dir=path.parent,
+    )
+    try:
+        with os.fdopen(descriptor, "w", encoding="utf-8") as temporary_file:
+            json.dump(
+                {"checksums": normalized, "verified": verified},
+                temporary_file,
+                indent=2,
+                sort_keys=True,
+            )
+            temporary_file.write("\n")
+            temporary_file.flush()
+            os.fsync(temporary_file.fileno())
+        os.replace(temporary_name, path)
+        os.chmod(path, 0o644)
+    except OSError as error:
+        raise UvmError(f"cannot write integrity manifest {path}: {error}") from error
+    finally:
+        try:
+            os.unlink(temporary_name)
+        except FileNotFoundError:
+            pass