brainpatch 1.2.0__py3-none-any.whl

brainpatch/patch/loader.py

@@ -0,0 +1,312 @@
+"""Reading and writing ``.brainpatch`` archives.
+
+Security posture
+----------------
+A ``.brainpatch`` file is **untrusted input**. It typically arrives from the
+internet, and the whole reason the format is ZIP-of-JSON-and-safetensors rather
+than a pickle is that applying one must never be able to run code.
+
+So this loader:
+
+* reads members **by exact name**, never by iterating and trusting what it finds
+* rejects absolute paths, ``..`` traversal, symlinks and directory entries
+* enforces per-member and total size ceilings before decompressing (zip bombs)
+* verifies sha256 for every member against ``checksums.json``
+* parses only JSON and safetensors, both of which are inert
+
+Nothing here evaluates, imports, or executes any part of an artifact.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import stat
+import zipfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from brainpatch.patch import tensors as ts
+from brainpatch.patch.format import (
+    CHECKSUMS_NAME,
+    MANIFEST_NAME,
+    README_NAME,
+    SUFFIX,
+    VECTORS_NAME,
+    Manifest,
+    PatchFormatError,
+)
+
+#: Members the loader will read. Anything else in the archive is ignored, and
+#: an *unexpected* member is reported rather than silently accepted.
+KNOWN_MEMBERS = (MANIFEST_NAME, VECTORS_NAME, CHECKSUMS_NAME, README_NAME)
+
+#: A patch is meant to be tiny. These ceilings are generous by two orders of
+#: magnitude and exist purely to bound a hostile archive.
+MAX_MEMBER_BYTES = 256 * 1024 * 1024
+MAX_TOTAL_BYTES = 512 * 1024 * 1024
+MAX_COMPRESSION_RATIO = 200
+
+
+class PatchLoadError(PatchFormatError):
+    """The archive could not be read safely."""
+
+
+@dataclass
+class LoadedPatch:
+    """A parsed, checksum-verified patch artifact."""
+
+    manifest: Manifest
+    vectors: dict[str, ts.Tensor]
+    readme: str | None = None
+    source: str | None = None
+    #: Size of the archive on disk, for honest size reporting.
+    archive_bytes: int = 0
+
+    @property
+    def name(self) -> str:
+        return self.manifest.name
+
+    def vector_for(self, key: str) -> ts.Tensor:
+        if key not in self.vectors:
+            raise PatchLoadError(
+                f"patch {self.name!r} references vector {key!r}, which is not in "
+                f"{VECTORS_NAME} (present: {sorted(self.vectors)})"
+            )
+        return self.vectors[key]
+
+    def describe(self) -> dict[str, Any]:
+        return {
+            "name": self.manifest.name,
+            "format_version": self.manifest.format_version,
+            "description": self.manifest.description,
+            "base_model": self.manifest.base_model.to_dict(),
+            "num_interventions": len(self.manifest.interventions),
+            "layers": self.manifest.layers,
+            "num_vectors": len(self.vectors),
+            "vector_dtype": next(iter(self.vectors.values())).dtype if self.vectors else None,
+            "evidence_level": self.manifest.evidence_level,
+            "compatibility": self.manifest.compatibility,
+            "archive_bytes": self.archive_bytes,
+            "source": self.source,
+        }
+
+
+def _safe_member_names(zf: zipfile.ZipFile) -> list[str]:
+    """Validate every entry's name and size before reading anything."""
+    total = 0
+    names: list[str] = []
+    for info in zf.infolist():
+        name = info.filename
+
+        if info.is_dir():
+            continue
+
+        # Path safety first: these are the checks that actually matter, and
+        # they give the clearest diagnostics.
+        if name.startswith("/") or name.startswith("\\"):
+            raise PatchLoadError(f"archive member {name!r} uses an absolute path")
+        if ".." in Path(name.replace("\\", "/")).parts:
+            raise PatchLoadError(f"archive member {name!r} escapes the archive root")
+        if os.path.isabs(name) or (len(name) > 1 and name[1] == ":"):
+            raise PatchLoadError(f"archive member {name!r} uses an absolute path")
+
+        # Unix mode lives in the high 16 bits of external_attr, but the type
+        # bits are frequently absent -- CPython's own ``writestr`` stores plain
+        # permissions (0o600) with no S_IFREG. So reject only what is positively
+        # identified as a symlink or other special file; treat "no type bits" as
+        # an ordinary member rather than as suspicious.
+        mode = info.external_attr >> 16
+        if mode and stat.S_IFMT(mode):
+            if stat.S_ISLNK(mode):
+                raise PatchLoadError(f"archive member {name!r} is a symlink")
+            if not stat.S_ISREG(mode) and not stat.S_ISDIR(mode):
+                raise PatchLoadError(f"archive member {name!r} is not a regular file")
+
+        if info.file_size > MAX_MEMBER_BYTES:
+            raise PatchLoadError(
+                f"archive member {name!r} is {info.file_size} bytes, over the "
+                f"{MAX_MEMBER_BYTES} limit"
+            )
+        if info.compress_size > 0:
+            ratio = info.file_size / info.compress_size
+            if ratio > MAX_COMPRESSION_RATIO:
+                raise PatchLoadError(
+                    f"archive member {name!r} has a {ratio:.0f}x compression ratio, "
+                    "which looks like a zip bomb"
+                )
+        total += info.file_size
+        if total > MAX_TOTAL_BYTES:
+            raise PatchLoadError(f"archive expands to over {MAX_TOTAL_BYTES} bytes")
+
+        names.append(name)
+    return names
+
+
+def load_patch(path: str | os.PathLike[str], *, verify_checksums: bool = True) -> LoadedPatch:
+    """Load and verify a ``.brainpatch`` archive.
+
+    Parameters
+    ----------
+    verify_checksums:
+        Leave this True. It is a parameter only so the writer can round-trip a
+        freshly-built archive before its checksums exist.
+
+    Raises
+    ------
+    PatchLoadError
+        On a malformed, unsafe, or checksum-mismatched archive.
+    """
+    p = Path(path)
+    if not p.is_file():
+        raise PatchLoadError(f"patch file not found: {p}")
+
+    archive_bytes = p.stat().st_size
+    try:
+        with zipfile.ZipFile(p) as zf:
+            present = set(_safe_member_names(zf))
+
+            unexpected = present - set(KNOWN_MEMBERS)
+            if unexpected:
+                raise PatchLoadError(
+                    f"archive contains unexpected members: {sorted(unexpected)}. "
+                    f"A BrainPatch may only contain {list(KNOWN_MEMBERS)}."
+                )
+            for required in (MANIFEST_NAME, VECTORS_NAME):
+                if required not in present:
+                    raise PatchLoadError(f"archive is missing required member {required!r}")
+
+            raw = {name: zf.read(name) for name in present}
+    except zipfile.BadZipFile as exc:
+        raise PatchLoadError(f"{p} is not a valid archive: {exc}") from exc
+
+    if verify_checksums:
+        if CHECKSUMS_NAME not in raw:
+            raise PatchLoadError(f"archive is missing {CHECKSUMS_NAME!r}")
+        _verify_checksums(raw)
+
+    manifest = Manifest.from_json(raw[MANIFEST_NAME].decode("utf-8"))
+
+    try:
+        vectors, _ = ts.load(raw[VECTORS_NAME])
+    except ts.SafetensorsError as exc:
+        raise PatchLoadError(f"{VECTORS_NAME} is malformed: {exc}") from exc
+
+    # Every referenced vector must exist and be a 1-D tensor of hidden_size.
+    hidden = manifest.base_model.hidden_size
+    for key in manifest.vector_keys:
+        if key not in vectors:
+            raise PatchLoadError(
+                f"manifest references vector {key!r} which is absent from {VECTORS_NAME}"
+            )
+        tensor = vectors[key]
+        if len(tensor.shape) != 1:
+            raise PatchLoadError(
+                f"vector {key!r} must be 1-D, got shape {tensor.shape}"
+            )
+        if tensor.shape[0] != hidden:
+            raise PatchLoadError(
+                f"vector {key!r} has length {tensor.shape[0]} but the patch declares "
+                f"hidden_size {hidden}"
+            )
+
+    readme = raw.get(README_NAME)
+    return LoadedPatch(
+        manifest=manifest,
+        vectors=vectors,
+        readme=readme.decode("utf-8") if readme else None,
+        source=str(p),
+        archive_bytes=archive_bytes,
+    )
+
+
+def _verify_checksums(raw: dict[str, bytes]) -> None:
+    try:
+        recorded = json.loads(raw[CHECKSUMS_NAME].decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise PatchLoadError(f"{CHECKSUMS_NAME} is not valid JSON: {exc}") from exc
+    if not isinstance(recorded, dict):
+        raise PatchLoadError(f"{CHECKSUMS_NAME} must be a JSON object")
+
+    for name, blob in raw.items():
+        if name == CHECKSUMS_NAME:
+            continue
+        expected = recorded.get(name)
+        if expected is None:
+            raise PatchLoadError(f"{CHECKSUMS_NAME} has no entry for member {name!r}")
+        actual = hashlib.sha256(blob).hexdigest()
+        if actual != expected:
+            raise PatchLoadError(
+                f"checksum mismatch for {name!r}: manifest records {expected}, "
+                f"archive contains {actual}. The patch is corrupt or was modified."
+            )
+
+
+def save_patch(
+    manifest: Manifest,
+    vectors: dict[str, ts.Tensor],
+    path: str | os.PathLike[str],
+    *,
+    readme: str | None = None,
+    overwrite: bool = False,
+) -> Path:
+    """Write a ``.brainpatch`` archive, with checksums, deterministically.
+
+    ZIP entries are written with a fixed timestamp and in fixed order, so two
+    builds from identical inputs produce byte-identical archives -- which is
+    what makes a published patch hash verifiable.
+    """
+    manifest.validate()
+
+    p = Path(path)
+    if p.suffix != SUFFIX:
+        p = p.with_suffix(SUFFIX)
+    if p.exists() and not overwrite:
+        raise FileExistsError(f"{p} already exists; pass overwrite=True to replace it")
+
+    missing = [k for k in manifest.vector_keys if k not in vectors]
+    if missing:
+        raise PatchFormatError(f"manifest references vectors not provided: {missing}")
+
+    hidden = manifest.base_model.hidden_size
+    for key, tensor in vectors.items():
+        if len(tensor.shape) != 1 or tensor.shape[0] != hidden:
+            raise PatchFormatError(
+                f"vector {key!r} must be 1-D of length {hidden}, got shape {tensor.shape}"
+            )
+
+    members: dict[str, bytes] = {
+        MANIFEST_NAME: (manifest.to_json() + "\n").encode("utf-8"),
+        VECTORS_NAME: ts.dump(vectors, metadata={"format": "brainpatch-vectors-v1"}),
+    }
+    if readme:
+        members[README_NAME] = readme.encode("utf-8")
+
+    checksums = {name: hashlib.sha256(blob).hexdigest() for name, blob in sorted(members.items())}
+    members[CHECKSUMS_NAME] = (
+        json.dumps(checksums, indent=2, sort_keys=True) + "\n"
+    ).encode("utf-8")
+
+    p.parent.mkdir(parents=True, exist_ok=True)
+    with zipfile.ZipFile(p, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+        for name in sorted(members):
+            info = zipfile.ZipInfo(name, date_time=(1980, 1, 1, 0, 0, 0))
+            info.compress_type = zipfile.ZIP_DEFLATED
+            info.external_attr = 0o100644 << 16  # regular file, rw-r--r--
+            zf.writestr(info, members[name])
+    return p
+
+
+def patch_size_report(loaded: LoadedPatch) -> dict[str, Any]:
+    """Size breakdown, for honest reporting rather than marketing claims."""
+    vector_bytes = sum(t.nbytes for t in loaded.vectors.values())
+    return {
+        "archive_bytes": loaded.archive_bytes,
+        "archive_kb": round(loaded.archive_bytes / 1024, 2),
+        "num_vectors": len(loaded.vectors),
+        "vector_payload_bytes": vector_bytes,
+        "hidden_size": loaded.manifest.base_model.hidden_size,
+        "dtype": next(iter(loaded.vectors.values())).dtype if loaded.vectors else None,
+    }

brainpatch/patch/registry.py

@@ -0,0 +1,300 @@
+"""The local patch registry: ``~/.brainpatch``.
+
+Installing a patch means copying a small verified archive into a local store and
+recording where it came from. It does **not** mean downloading a base model --
+patches are tens of KB and models are gigabytes, and conflating the two is how a
+"quick install" becomes a 3 GB surprise.
+
+Layout::
+
+    ~/.brainpatch/
+        config.json          registry-level settings
+        patches/
+            <name>.brainpatch
+            <name>.source.json    provenance of the install
+        cache/                    downloads, safe to delete
+
+Everything here is plain files. Uninstalling is deleting them; there is no
+database to corrupt and nothing to migrate.
+
+Network access happens only in :func:`install_from_hub`, only via
+``huggingface_hub``, and only for the patch artifact itself.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import shutil
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from brainpatch.patch.format import SUFFIX
+from brainpatch.patch.loader import LoadedPatch, load_patch
+
+#: Overridable for tests and for users with an unusual home directory.
+ENV_HOME = "BRAINPATCH_HOME"
+
+#: ``owner/repo`` or ``owner/repo:file.brainpatch``
+_HUB_REF_RE = re.compile(
+    r"^(?P<repo>[A-Za-z0-9][\w.-]*/[\w.-]+)(?::(?P<file>[\w./-]+))?$"
+)
+
+
+class RegistryError(RuntimeError):
+    """The registry could not satisfy the request."""
+
+
+def registry_home() -> Path:
+    """Root of the local registry, honouring ``BRAINPATCH_HOME``."""
+    override = os.environ.get(ENV_HOME)
+    if override:
+        return Path(override).expanduser()
+    return Path.home() / ".brainpatch"
+
+
+@dataclass
+class InstalledPatch:
+    """A patch present in the local registry."""
+
+    name: str
+    path: Path
+    source: dict[str, Any]
+
+    @property
+    def size_bytes(self) -> int:
+        return self.path.stat().st_size
+
+    def load(self) -> LoadedPatch:
+        return load_patch(self.path)
+
+
+class PatchRegistry:
+    """File-backed store of installed patches."""
+
+    def __init__(self, home: str | os.PathLike[str] | None = None) -> None:
+        self.home = Path(home).expanduser() if home is not None else registry_home()
+
+    # -- paths -----------------------------------------------------------------
+
+    @property
+    def patches_dir(self) -> Path:
+        return self.home / "patches"
+
+    @property
+    def cache_dir(self) -> Path:
+        return self.home / "cache"
+
+    @property
+    def config_path(self) -> Path:
+        return self.home / "config.json"
+
+    def ensure_dirs(self) -> None:
+        self.patches_dir.mkdir(parents=True, exist_ok=True)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def path_for(self, name: str) -> Path:
+        return self.patches_dir / f"{name}{SUFFIX}"
+
+    def _source_path(self, name: str) -> Path:
+        return self.patches_dir / f"{name}.source.json"
+
+    # -- queries ---------------------------------------------------------------
+
+    def list_patches(self) -> list[InstalledPatch]:
+        """Every installed patch, sorted by name. Missing dir -> empty list."""
+        if not self.patches_dir.is_dir():
+            return []
+        out: list[InstalledPatch] = []
+        for path in sorted(self.patches_dir.glob(f"*{SUFFIX}")):
+            name = path.name[: -len(SUFFIX)]
+            out.append(InstalledPatch(name=name, path=path, source=self._read_source(name)))
+        return out
+
+    def is_installed(self, name: str) -> bool:
+        return self.path_for(name).is_file()
+
+    def get(self, name: str) -> InstalledPatch:
+        path = self.path_for(name)
+        if not path.is_file():
+            available = [p.name for p in self.list_patches()]
+            raise RegistryError(
+                f"patch {name!r} is not installed. "
+                + (f"Installed: {', '.join(available)}" if available else "No patches installed.")
+            )
+        return InstalledPatch(name=name, path=path, source=self._read_source(name))
+
+    def resolve(self, ref: str) -> Path:
+        """Resolve a name, a path, or an installed patch to a file on disk.
+
+        Accepts an installed name first, then a filesystem path -- so a bare
+        name never accidentally reads a same-named file from the CWD.
+        """
+        if self.is_installed(ref):
+            return self.path_for(ref)
+        candidate = Path(ref).expanduser()
+        if candidate.is_file():
+            return candidate
+        raise RegistryError(
+            f"{ref!r} is neither an installed patch nor an existing file. "
+            "Use `brainpatch list` to see what is installed."
+        )
+
+    def _read_source(self, name: str) -> dict[str, Any]:
+        path = self._source_path(name)
+        if not path.is_file():
+            return {}
+        try:
+            return json.loads(path.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            return {}
+
+    # -- mutation --------------------------------------------------------------
+
+    def install_file(
+        self,
+        path: str | os.PathLike[str],
+        *,
+        name: str | None = None,
+        overwrite: bool = False,
+        source: dict[str, Any] | None = None,
+    ) -> InstalledPatch:
+        """Verify then install a local ``.brainpatch`` file.
+
+        The archive is fully loaded and checksum-verified *before* anything is
+        written, so a corrupt download cannot land in the registry.
+        """
+        src = Path(path).expanduser()
+        loaded = load_patch(src)  # raises on anything malformed or unsafe
+
+        patch_name = name or loaded.manifest.name
+        target = self.path_for(patch_name)
+        if target.exists() and not overwrite:
+            raise RegistryError(
+                f"patch {patch_name!r} is already installed. "
+                "Pass --force to replace it."
+            )
+
+        self.ensure_dirs()
+        shutil.copyfile(src, target)
+
+        record = {
+            "installed_at": datetime.now(timezone.utc).isoformat(),
+            "origin": str(src),
+            "kind": "file",
+            "format_version": loaded.manifest.format_version,
+            "base_model": loaded.manifest.base_model.model_id,
+            "evidence_level": loaded.manifest.evidence_level,
+            **(source or {}),
+        }
+        self._source_path(patch_name).write_text(
+            json.dumps(record, indent=2, sort_keys=True) + "\n", encoding="utf-8"
+        )
+        return InstalledPatch(name=patch_name, path=target, source=record)
+
+    def install_from_hub(
+        self,
+        ref: str,
+        *,
+        filename: str | None = None,
+        revision: str | None = None,
+        overwrite: bool = False,
+        offline: bool = False,
+    ) -> InstalledPatch:
+        """Install from a Hugging Face repo reference.
+
+        ``ref`` is ``owner/repo`` or ``owner/repo:path/to/file.brainpatch``.
+        Only the patch artifact is downloaded -- never the base model.
+        """
+        match = _HUB_REF_RE.match(ref)
+        if not match:
+            raise RegistryError(
+                f"{ref!r} is not a valid Hugging Face reference. "
+                "Expected 'owner/repo' or 'owner/repo:file.brainpatch'."
+            )
+        repo = match.group("repo")
+        wanted = filename or match.group("file")
+
+        try:
+            from huggingface_hub import hf_hub_download, list_repo_files
+        except ModuleNotFoundError as exc:  # pragma: no cover - depends on env
+            raise RegistryError(
+                "installing from Hugging Face needs the 'huggingface_hub' package.\n"
+                "  pip install 'brainpatch[hub]'"
+            ) from exc
+
+        if offline:
+            raise RegistryError(
+                "cannot install from Hugging Face in offline mode; "
+                "download the .brainpatch file separately and install it by path"
+            )
+
+        self.ensure_dirs()
+
+        if wanted is None:
+            try:
+                files = list_repo_files(repo, revision=revision)
+            except Exception as exc:  # noqa: BLE001 - hub raises many types
+                raise RegistryError(f"could not list files in {repo!r}: {exc}") from exc
+            candidates = [f for f in files if f.endswith(SUFFIX)]
+            if not candidates:
+                raise RegistryError(
+                    f"{repo!r} contains no {SUFFIX} artifact. "
+                    "Specify one explicitly with 'owner/repo:path/to/file.brainpatch'."
+                )
+            if len(candidates) > 1:
+                raise RegistryError(
+                    f"{repo!r} contains several patches: {candidates}. "
+                    "Choose one with 'owner/repo:<file>'."
+                )
+            wanted = candidates[0]
+
+        try:
+            downloaded = hf_hub_download(
+                repo_id=repo,
+                filename=wanted,
+                revision=revision,
+                cache_dir=str(self.cache_dir),
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise RegistryError(f"could not download {wanted!r} from {repo!r}: {exc}") from exc
+
+        return self.install_file(
+            downloaded,
+            overwrite=overwrite,
+            source={"kind": "huggingface", "repo": repo, "file": wanted, "revision": revision},
+        )
+
+    def install(self, ref: str, *, overwrite: bool = False, offline: bool = False) -> InstalledPatch:
+        """Install from a path or a Hugging Face reference, whichever ``ref`` is."""
+        candidate = Path(ref).expanduser()
+        if candidate.is_file():
+            return self.install_file(candidate, overwrite=overwrite)
+        if _HUB_REF_RE.match(ref):
+            return self.install_from_hub(ref, overwrite=overwrite, offline=offline)
+        raise RegistryError(
+            f"{ref!r} is neither an existing file nor an 'owner/repo' Hugging Face reference"
+        )
+
+    def uninstall(self, name: str) -> None:
+        patch = self.get(name)
+        patch.path.unlink()
+        source = self._source_path(name)
+        if source.exists():
+            source.unlink()
+
+    def clear_cache(self) -> int:
+        """Delete the download cache. Returns bytes freed."""
+        if not self.cache_dir.is_dir():
+            return 0
+        freed = sum(p.stat().st_size for p in self.cache_dir.rglob("*") if p.is_file())
+        shutil.rmtree(self.cache_dir, ignore_errors=True)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        return freed
+
+
+def default_registry() -> PatchRegistry:
+    return PatchRegistry()