Files
nest/uvm/integrity.py
2026-09-04 20:56:18 +00:00

146 lines
4.5 KiB
Python

"""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