aisbom-cli 1.3.0__tar.gz → 1.3.2__tar.gz

{aisbom_cli-1.3.0 → aisbom_cli-1.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aisbom-cli
-Version: 1.3.0
+Version: 1.3.2
 Summary: An AI Supply Chain security tool that that detects Pickle bombs and generates CycloneDX SBOMs for Machine Learning models.
 License-File: LICENSE
 Author: Ajoy L
@@ -81,7 +81,10 @@ A typical scan against a project with mixed artifacts:
 │                              │             │ __class__, __mro__, __subclasses__)            │                              │
 │ exfil.onnx                   │ ONNX        │ CRITICAL (ONNX External Data Escapes Model     │ UNKNOWN                      │
 │                              │             │ Directory: ../../../etc/passwd)                │                              │
+│ sklearn_pipeline.joblib      │ Joblib      │ CRITICAL (RCE Detected: posix.system)          │ UNKNOWN                      │
+│ features.npy                 │ NumPy       │ CRITICAL (RCE Detected: builtins.eval)         │ UNKNOWN                      │
 │ detector.onnx                │ ONNX        │ LOW                                            │ UNKNOWN                      │
+│ embeddings.npy               │ NumPy       │ LOW                                            │ UNKNOWN                      │
 │ llama-3-quant.gguf           │ GGUF        │ LOW                                            │ LEGAL RISK (cc-by-nc-sa-4.0) │
 │ safe_model.safetensors       │ SafeTensors │ LOW                                            │ PASS                         │
 │ restricted_model.safetensors │ SafeTensors │ LOW                                            │ LEGAL RISK (cc-by-nc-4.0)    │
@@ -94,7 +97,10 @@ A compliant `sbom.json` (CycloneDX v1.6) including SHA256 hashes and license dat
 
 | Format | Extensions | What AIsbom looks for |
 |---|---|---|
-| **PyTorch / Pickle** | `.pt` `.pth` `.bin` | Dangerous globals in the pickle opcodes, including indirect-execution gadgets. Concatenated streams are all scanned — a legacy `torch.save` file hides its object behind three header pickles — and non-standard containers (7z, rar, xz…) are flagged. If a file is so full of streams that the walk hits its work limit, the scan says so rather than reporting clean. |
+| **PyTorch / Pickle** | `.pt` `.pth` `.bin` `.pkl` `.pickle` | Dangerous globals in the pickle opcodes, including indirect-execution gadgets. Concatenated streams are all scanned — a legacy `torch.save` file hides its object behind three header pickles — and non-standard containers (7z, rar, xz…) are flagged. If a file is so full of streams that the walk hits its work limit, the scan says so rather than reporting clean. |
+| **joblib** | `.joblib` | The pickle inside the container, whichever codec joblib chose (zlib, gzip, bz2, lzma/xz, and the legacy `ZF` format), plus uncompressed files. A payload placed *after* an array — where a real model puts its weights — is found. |
+| **dill** | `.dill` | Everything the pickle path finds, plus dill's own code-reconstruction globals: a dill'd function or lambda is a marshalled code object rebuilt on load, and is reported as CRITICAL. |
+| **NumPy** | `.npy` `.npz` | The pickle stream behind an `allow_pickle` object array. Each `.npz` member is opened, including ones whose checksum or header has been tampered with. |
 | **Keras** | `.keras` `.h5` `.hdf5` | `Lambda` layers and embedded marshalled code objects in the model config — an actively exploited RCE vector. |
 | **GGUF** | `.gguf` | License and architecture metadata, plus the embedded Jinja **chat template**, checked for sandbox-escape constructs. |
 | **ONNX** | `.onnx` | Producer/opset/IR metadata, custom operators, and external-data paths that point outside the model directory. |
@@ -134,6 +140,18 @@ xattr -d com.apple.quarantine aisbom-macos-*
 
 ## Common workflows
 
+### Scan targets and exit codes
+
+`scan` accepts a directory, a single model file, a `hf://` repo, or an HTTPS URL. A directory is walked recursively; a single file is scanned on its own.
+
+| Exit | Meaning |
+|---|---|
+| `0` | Scan completed, no CRITICAL risk |
+| `1` | The scan could not be completed — target missing or unreadable, a file failed to parse, or a remote fetch failed |
+| `2` | A **CRITICAL** risk was found (suppress with `--no-fail-on-risk`) |
+
+`--no-fail-on-risk` governs risk findings only. An unusable target still exits `1`, so a typo'd path in CI fails loudly instead of passing as a clean scan.
+
 ### Scan a Hugging Face model
 
 ```bash
@@ -390,6 +408,17 @@ Attribute names are matched **exactly, never as substrings**, so ordinary global
 
 Pickle is not the only way a model file gets code to run. Each of the other formats has its own vector, and each is read as inert data:
 
+**joblib, dill and NumPy object arrays** (`.joblib`, `.dill`, `.npy`, `.npz`) are pickle underneath — the everyday serialization of scikit-learn and scientific Python, carrying exactly the arbitrary-code-execution risk of a `.pt`, and opened by almost nothing that calls itself an SBOM tool. AIsbom reads them without importing joblib, dill or numpy: the compression is standard-library, the `.npy` header is parsed by hand, and nothing is ever unpickled.
+
+Two things about these formats are worth stating plainly:
+
+- **A raw array block does not end the scan.** joblib writes its pickle up to an array, dumps the raw buffer inline, then resumes pickling — which stops an opcode disassembly dead a few hundred bytes into the file. Since every real model has weights, "after the first array" is the natural place for a payload. The remaining bytes get a second pass that recovers globals directly, so a dangerous import behind an array block is still reported. That pass reads bytes rather than structure, so it recognises *known* sinks only: in strict mode it does not contribute unrecognized-import findings, and it cannot judge a dual-use constructor like `operator.methodcaller("system")`, because the argument that decides is a stack relationship it has no stack for.
+- **A declared dtype never decides whether to look.** The `.npy` header is attacker-supplied, so a pickle sitting behind a header claiming `'<f8'` would otherwise be a one-line evasion. The data section is disassembled whatever the header says. What the header *does* affect is the reported risk: an array of ordinary numbers with no pickle in it is LOW, not "pickle present".
+
+- **No limit ends in a clean verdict.** Four bounds apply here — the file-read budget, the `.npz` member count, the decompressor's output cap, and the disassembler's stream budget — and each one leaves bytes unexamined. Hitting any of them reports `MEDIUM (Pickle Scan Incomplete)` rather than passing the prefix off as the whole file. A limit that reports "clean" is not a safety measure; it is a hiding place with a length attached. A real finding still outranks the marker.
+
+joblib's `lz4` and `zstd` codecs have no standard-library decompressor. Following the same reasoning as 7z containers, those are **named but not opened** — `MEDIUM (Unscanned Container: lz4)` — rather than pulling a native dependency into every install and every standalone binary. That is an honest "we did not read these bytes", which is a different answer from a clean scan.
+
 **SafeTensors and GGUF** use binary formats with structured headers — AIsbom parses these directly to extract metadata (artifact names, license info, architecture details) without loading tensor weights.
 
 **Keras** (`.keras`, `.h5`, `.hdf5`) carries a different execution vector: a `Lambda` layer stores an arbitrary Python callable in the model config as a base64-encoded marshalled code object, and `load_model` runs it. AIsbom reads the config out of both containers — the `.keras` zip and the legacy HDF5 attribute — and flags `Lambda` layers and embedded code objects as CRITICAL. The payload is identified from its header bytes and **never unmarshalled**. A truncated or corrupted container is still scanned rather than skipped, so damaging a file is not a way to hide a payload.

{aisbom_cli-1.3.0 → aisbom_cli-1.3.2}/README.md

@@ -56,7 +56,10 @@ A typical scan against a project with mixed artifacts:
 │                              │             │ __class__, __mro__, __subclasses__)            │                              │
 │ exfil.onnx                   │ ONNX        │ CRITICAL (ONNX External Data Escapes Model     │ UNKNOWN                      │
 │                              │             │ Directory: ../../../etc/passwd)                │                              │
+│ sklearn_pipeline.joblib      │ Joblib      │ CRITICAL (RCE Detected: posix.system)          │ UNKNOWN                      │
+│ features.npy                 │ NumPy       │ CRITICAL (RCE Detected: builtins.eval)         │ UNKNOWN                      │
 │ detector.onnx                │ ONNX        │ LOW                                            │ UNKNOWN                      │
+│ embeddings.npy               │ NumPy       │ LOW                                            │ UNKNOWN                      │
 │ llama-3-quant.gguf           │ GGUF        │ LOW                                            │ LEGAL RISK (cc-by-nc-sa-4.0) │
 │ safe_model.safetensors       │ SafeTensors │ LOW                                            │ PASS                         │
 │ restricted_model.safetensors │ SafeTensors │ LOW                                            │ LEGAL RISK (cc-by-nc-4.0)    │
@@ -69,7 +72,10 @@ A compliant `sbom.json` (CycloneDX v1.6) including SHA256 hashes and license dat
 
 | Format | Extensions | What AIsbom looks for |
 |---|---|---|
-| **PyTorch / Pickle** | `.pt` `.pth` `.bin` | Dangerous globals in the pickle opcodes, including indirect-execution gadgets. Concatenated streams are all scanned — a legacy `torch.save` file hides its object behind three header pickles — and non-standard containers (7z, rar, xz…) are flagged. If a file is so full of streams that the walk hits its work limit, the scan says so rather than reporting clean. |
+| **PyTorch / Pickle** | `.pt` `.pth` `.bin` `.pkl` `.pickle` | Dangerous globals in the pickle opcodes, including indirect-execution gadgets. Concatenated streams are all scanned — a legacy `torch.save` file hides its object behind three header pickles — and non-standard containers (7z, rar, xz…) are flagged. If a file is so full of streams that the walk hits its work limit, the scan says so rather than reporting clean. |
+| **joblib** | `.joblib` | The pickle inside the container, whichever codec joblib chose (zlib, gzip, bz2, lzma/xz, and the legacy `ZF` format), plus uncompressed files. A payload placed *after* an array — where a real model puts its weights — is found. |
+| **dill** | `.dill` | Everything the pickle path finds, plus dill's own code-reconstruction globals: a dill'd function or lambda is a marshalled code object rebuilt on load, and is reported as CRITICAL. |
+| **NumPy** | `.npy` `.npz` | The pickle stream behind an `allow_pickle` object array. Each `.npz` member is opened, including ones whose checksum or header has been tampered with. |
 | **Keras** | `.keras` `.h5` `.hdf5` | `Lambda` layers and embedded marshalled code objects in the model config — an actively exploited RCE vector. |
 | **GGUF** | `.gguf` | License and architecture metadata, plus the embedded Jinja **chat template**, checked for sandbox-escape constructs. |
 | **ONNX** | `.onnx` | Producer/opset/IR metadata, custom operators, and external-data paths that point outside the model directory. |
@@ -109,6 +115,18 @@ xattr -d com.apple.quarantine aisbom-macos-*
 
 ## Common workflows
 
+### Scan targets and exit codes
+
+`scan` accepts a directory, a single model file, a `hf://` repo, or an HTTPS URL. A directory is walked recursively; a single file is scanned on its own.
+
+| Exit | Meaning |
+|---|---|
+| `0` | Scan completed, no CRITICAL risk |
+| `1` | The scan could not be completed — target missing or unreadable, a file failed to parse, or a remote fetch failed |
+| `2` | A **CRITICAL** risk was found (suppress with `--no-fail-on-risk`) |
+
+`--no-fail-on-risk` governs risk findings only. An unusable target still exits `1`, so a typo'd path in CI fails loudly instead of passing as a clean scan.
+
 ### Scan a Hugging Face model
 
 ```bash
@@ -365,6 +383,17 @@ Attribute names are matched **exactly, never as substrings**, so ordinary global
 
 Pickle is not the only way a model file gets code to run. Each of the other formats has its own vector, and each is read as inert data:
 
+**joblib, dill and NumPy object arrays** (`.joblib`, `.dill`, `.npy`, `.npz`) are pickle underneath — the everyday serialization of scikit-learn and scientific Python, carrying exactly the arbitrary-code-execution risk of a `.pt`, and opened by almost nothing that calls itself an SBOM tool. AIsbom reads them without importing joblib, dill or numpy: the compression is standard-library, the `.npy` header is parsed by hand, and nothing is ever unpickled.
+
+Two things about these formats are worth stating plainly:
+
+- **A raw array block does not end the scan.** joblib writes its pickle up to an array, dumps the raw buffer inline, then resumes pickling — which stops an opcode disassembly dead a few hundred bytes into the file. Since every real model has weights, "after the first array" is the natural place for a payload. The remaining bytes get a second pass that recovers globals directly, so a dangerous import behind an array block is still reported. That pass reads bytes rather than structure, so it recognises *known* sinks only: in strict mode it does not contribute unrecognized-import findings, and it cannot judge a dual-use constructor like `operator.methodcaller("system")`, because the argument that decides is a stack relationship it has no stack for.
+- **A declared dtype never decides whether to look.** The `.npy` header is attacker-supplied, so a pickle sitting behind a header claiming `'<f8'` would otherwise be a one-line evasion. The data section is disassembled whatever the header says. What the header *does* affect is the reported risk: an array of ordinary numbers with no pickle in it is LOW, not "pickle present".
+
+- **No limit ends in a clean verdict.** Four bounds apply here — the file-read budget, the `.npz` member count, the decompressor's output cap, and the disassembler's stream budget — and each one leaves bytes unexamined. Hitting any of them reports `MEDIUM (Pickle Scan Incomplete)` rather than passing the prefix off as the whole file. A limit that reports "clean" is not a safety measure; it is a hiding place with a length attached. A real finding still outranks the marker.
+
+joblib's `lz4` and `zstd` codecs have no standard-library decompressor. Following the same reasoning as 7z containers, those are **named but not opened** — `MEDIUM (Unscanned Container: lz4)` — rather than pulling a native dependency into every install and every standalone binary. That is an honest "we did not read these bytes", which is a different answer from a clean scan.
+
 **SafeTensors and GGUF** use binary formats with structured headers — AIsbom parses these directly to extract metadata (artifact names, license info, architecture details) without loading tensor weights.
 
 **Keras** (`.keras`, `.h5`, `.hdf5`) carries a different execution vector: a `Lambda` layer stores an arbitrary Python callable in the model config as a base64-encoded marshalled code object, and `load_model` runs it. AIsbom reads the config out of both containers — the `.keras` zip and the legacy HDF5 attribute — and flags `Lambda` layers and embedded code objects as CRITICAL. The payload is identified from its header bytes and **never unmarshalled**. A truncated or corrupted container is still scanned rather than skipped, so damaging a file is not a way to hide a payload.

{aisbom_cli-1.3.0 → aisbom_cli-1.3.2}/aisbom/cli.py

@@ -607,7 +607,11 @@ def scan(
             )
         console.print(table)
     else:
-        console.print("[yellow]No AI models found.[/yellow]")
+        # "No AI models found" is a claim about the target's contents, so it
+        # must not be printed when the target was never scanned at all (#125).
+        # The target error itself is rendered below.
+        if not any(e.get('target_error') for e in results['errors']):
+            console.print("[yellow]No AI models found.[/yellow]")
 
     # LINT OUTPUT (Migration Report)
     lint_failures = [a for a in results['artifacts'] if a.get('details', {}).get('lint_report')]
@@ -633,9 +637,20 @@ def scan(
     if results['dependencies']:
         console.print(f"\n📦 Found [bold]{len(results['dependencies'])}[/bold] Python libraries.")
 
-    # Parse errors only — fetch failures already printed their status-aware
-    # message to stderr above and don't fit the "Could not parse" framing.
-    parse_errors = [e for e in results['errors'] if not e.get('fetch_failure')]
+    # Unusable targets (#125): the path was missing or nothing could scan it,
+    # so nothing was examined. Printed to stderr like fetch failures — it is a
+    # failed instruction, not a finding about the model.
+    for err in [e for e in results['errors'] if e.get('target_error')]:
+        err_console.print(
+            f"[bold red]✖[/bold red] Cannot scan [yellow]{err['file']}[/yellow]: {err['error']}"
+        )
+
+    # Parse errors only — fetch failures and target errors already printed
+    # their own message to stderr and don't fit the "Could not parse" framing.
+    parse_errors = [
+        e for e in results['errors']
+        if not e.get('fetch_failure') and not e.get('target_error')
+    ]
     if parse_errors:
         console.print("\n[bold red]⚠️ Errors Encountered:[/bold red]")
         for err in parse_errors:

{aisbom_cli-1.3.0 → aisbom_cli-1.3.2}/aisbom/corpus.py

@@ -216,6 +216,11 @@ class BypassCase:
     builder: Callable[[Path], None]
     expected: str = "detected"
     malicious: bool = True
+    # Why this case is currently not fully caught, for cases that aren't.
+    # Deliberately does *not* change `expected`: every evasion technique here
+    # remains one a correct scanner should catch, so the gate keeps counting it
+    # against us. This field explains a gap; it never excuses one.
+    limitation: str | None = None
 
 
 @dataclass(frozen=True)
@@ -262,6 +267,15 @@ CASES: tuple[BypassCase, ...] = (
             "meant picklescan never opened it, while the model still loaded."
         ),
         builder=_build_nullifai_7z,
+        limitation=(
+            "AIsbom reports `CRITICAL (Non-Standard Container: 7z)` — the right severity, "
+            "but earned from the container rather than the payload. The archive is named, "
+            "never unpacked, so the `os.system` call inside is never disassembled and the "
+            "reported reason is not the real one. Unpacking 7z would mean a native "
+            "dependency in every install to cover one evasion class, which is not a "
+            "trade worth making; a user acting on this verdict is nonetheless correctly "
+            "warned off the file."
+        ),
     ),
     BypassCase(
         id="nullifai-broken-stream",
@@ -386,6 +400,17 @@ CASES: tuple[BypassCase, ...] = (
             "argument, so string-matching on the opcode argument sees nothing."
         ),
         builder=_build_shadowpickle,
+        limitation=(
+            "AIsbom does resolve STACK_GLOBAL and reads the pair off the stack, so it "
+            "sees `collections.OrderedDict` — and that name is legitimately allowlisted, "
+            "because real state_dicts are OrderedDicts. Both modes therefore return only "
+            "`MEDIUM (Pickle Present)`, the baseline every pickle gets, rather than a "
+            "signal specific to this file. This is the ceiling on static allowlist "
+            "analysis: the call is indistinguishable from a legitimate call to an "
+            "allowlisted global, and flagging the shape would flag ordinary checkpoints. "
+            "Closing it needs evidence beyond the resolved name — argument shape, or "
+            "provenance — not a new entry on a blocklist."
+        ),
     ),
 )
 
@@ -546,14 +571,23 @@ def strip_generated_stamp(text: str) -> str:
     ).strip()
 
 
+def is_caught(row: dict) -> bool:
+    """
+    Whether a case counts as caught: named as a threat in at least one mode.
+
+    The headline count and the limitation-note gate must agree on this, or the
+    document can claim a case is uncaught while counting it among the wins.
+    """
+    return row["blocklist"] == "detected" or row["strict"] == "detected"
+
+
 def render_markdown(results: dict) -> str:
     """Render the human-readable scorecard published by the /blog article."""
     cases = results["cases"]
     evasions = [c for c in CASES if not is_control(c)]
     controls = [c for c in CASES if is_control(c)]
 
-    detected = sum(1 for c in evasions if cases[c.id]["blocklist"] == "detected"
-                   or cases[c.id]["strict"] == "detected")
+    detected = sum(1 for c in evasions if is_caught(cases[c.id]))
 
     lines = [
         "# Does AIsbom catch it? — pickle-evasion scorecard",
@@ -605,6 +639,16 @@ def render_markdown(results: dict) -> str:
             case.description,
             "",
         ]
+        # A limitation describes why a case is *not* caught, so it stops being
+        # true the moment detection improves. Gate it on the live verdict rather
+        # than publishing unconditionally: `--write` would otherwise regenerate
+        # the document with a stale claim, and the staleness test cannot catch
+        # that because it compares the committed file against this same
+        # renderer. test_limitation_note_is_dropped_once_a_case_is_caught pins
+        # the gate; test_no_limitation_note_on_a_caught_case fails the build so
+        # the dead text gets removed rather than silently lingering here.
+        if case.limitation and not is_caught(cases[case.id]):
+            lines += [f"**Current limitation:** {case.limitation}", ""]
 
     return "\n".join(lines).rstrip() + "\n"
 

aisbom_cli-1.3.2/aisbom/pickle_containers.py

@@ -0,0 +1,304 @@
+"""Locate the pickle stream inside the wrapper formats that carry one.
+
+joblib, dill and numpy's object arrays are the everyday serialization formats of
+scientific Python, and every one of them is pickle underneath — the same
+arbitrary-code-execution risk as a bare ``.pkl``, wrapped in a container that
+generic tooling does not open.
+
+Nothing in this module imports joblib, numpy or dill. The compression is stdlib,
+the ``.npy`` header is parsed by hand, and the payload is never unpickled — it is
+handed back as bytes for the disassembler to read. That is deliberate twice over:
+a scanner that loads the artifact to inspect it has already lost, and pulling a
+scientific-Python stack into the package would put it into every install and
+every standalone binary for the sake of reading bytes we can read ourselves.
+
+Everything here is bounded. A container is an attacker-controlled input, so a
+decompressor is never handed an unbounded output buffer and a declared header
+length is never trusted far enough to allocate against it.
+"""
+
+from __future__ import annotations
+
+import ast
+import bz2
+import lzma
+import re
+import zlib
+from typing import Any, Dict, Tuple
+
+# Fallback ceiling on any single decompressed payload: a compression bomb must
+# cost the same as a large ordinary file, not more. The decompressors below are
+# all incremental and take this as a `max_length`, so the bomb is never expanded
+# in the first place.
+#
+# Callers pass their own budget — the scanner passes the same figure it used to
+# read the file, so a remote scan that fetched 2MB does not then expand 16.
+DECOMPRESS_MAX_BYTES = 16 * 1024 * 1024
+
+# `.npy` declares its header length in the file. Cap what we will honour: the
+# real headers are a few dozen bytes, and a declared length is exactly the field
+# an attacker would inflate.
+NPY_MAX_HEADER_BYTES = 1 * 1024 * 1024
+
+NPY_MAGIC = b"\x93NUMPY"
+
+# joblib's pre-0.10 container: a two-byte tag, a decimal length, then zlib data.
+# Still readable today, and still a way to carry a pickle past a scanner that
+# only knows the modern layout.
+JOBLIB_ZFILE_MAGIC = b"ZF"
+
+# Compression formats joblib writes that the standard library can open. The
+# label is what gets reported, so it is the name a user would recognise.
+_GZIP = "gzip"
+_BZ2 = "bz2"
+_LZMA = "lzma"
+_XZ = "xz"
+_ZLIB = "zlib"
+
+SUPPORTED_COMPRESSION = (_ZLIB, _GZIP, _BZ2, _LZMA, _XZ)
+
+# Formats joblib also supports whose decompressors are not in the standard
+# library. Naming one is not the same as reading it — see `describe_container`.
+_MAGIC_UNSUPPORTED = (
+    (b"\x04\x22\x4d\x18", "lz4"),
+    (b"\x28\xb5\x2f\xfd", "zstd"),
+)
+
+_MAGIC_SUPPORTED = (
+    (b"\x1f\x8b", _GZIP),
+    (b"BZh", _BZ2),
+    (b"\xfd7zXZ\x00", _XZ),
+    (b"\x5d\x00\x00", _LZMA),
+)
+
+
+def _looks_like_zlib(head: bytes) -> bool:
+    """True for a zlib stream header (RFC 1950).
+
+    zlib has no magic number, only a two-byte header with a checksum property:
+    the low nibble of the first byte is the compression method (8 = deflate) and
+    the 16-bit big-endian pair is a multiple of 31. Testing both is what keeps
+    an ordinary binary file that happens to start with 0x78 from being mistaken
+    for a compressed container.
+    """
+    if len(head) < 2:
+        return False
+    if head[0] & 0x0F != 8:
+        return False
+    return ((head[0] << 8) | head[1]) % 31 == 0
+
+
+def detect_compression(head: bytes) -> Tuple[str | None, bool]:
+    """Return ``(label, readable)`` for the compression ``head`` begins with.
+
+    ``readable`` is False for a format we can name but not open, which is a
+    materially different answer from "not compressed" and is reported as such
+    rather than being quietly treated as a clean scan.
+    """
+    if not head:
+        return None, False
+    for magic, label in _MAGIC_SUPPORTED:
+        if head.startswith(magic):
+            return label, True
+    for magic, label in _MAGIC_UNSUPPORTED:
+        if head.startswith(magic):
+            return label, False
+    if _looks_like_zlib(head):
+        return _ZLIB, True
+    return None, False
+
+
+def decompress(data: bytes, label: str, limit: int = DECOMPRESS_MAX_BYTES) -> Tuple[bytes | None, bool]:
+    """Inflate ``data`` with the named codec; return ``(payload, complete)``.
+
+    ``complete`` is False when the decompressor stopped before the end of the
+    stream — because it hit ``limit``, or because the compressed data itself was
+    cut short. That distinction has to reach the caller. A payload placed after
+    a large compressible value would otherwise sit in the part we never expanded
+    while the prefix scanned clean, and a cap nobody reports is somewhere to hide
+    a payload rather than a safety measure.
+
+    Returns whatever was recovered before an error when a stream is corrupt — a
+    damaged tail must not discard the front, the same reasoning that makes a
+    broken zip member worth reading anyway. ``(None, False)`` means nothing at
+    all could be read.
+    """
+    if not data:
+        return None, False
+    try:
+        if label == _ZLIB:
+            engine = zlib.decompressobj()
+        elif label == _GZIP:
+            # 16 + MAX_WBITS selects the gzip wrapper.
+            engine = zlib.decompressobj(16 + zlib.MAX_WBITS)
+        elif label == _BZ2:
+            engine = bz2.BZ2Decompressor()
+        elif label in (_LZMA, _XZ):
+            # FORMAT_AUTO reads both the `.xz` container and the older
+            # standalone `.lzma` framing joblib still emits.
+            engine = lzma.LZMADecompressor(format=lzma.FORMAT_AUTO)
+        else:
+            return None, False
+
+        out = engine.decompress(data, limit)
+        # Every one of these decompressors exposes `eof`, which is True only
+        # once the stream's own end marker has been consumed.
+        return (out or None), bool(getattr(engine, "eof", False))
+    except Exception:
+        return None, False
+
+
+def unwrap_zfile(data: bytes, limit: int = DECOMPRESS_MAX_BYTES) -> Tuple[bytes | None, bool]:
+    """Inflate joblib's legacy ``ZF`` container; return ``(payload, complete)``.
+
+    Layout is the tag, a space-padded decimal length, then a zlib stream. The
+    declared length is read past rather than trusted — the decompressor stops at
+    the real end of the stream, so a lie in that field buys nothing.
+    """
+    if not data.startswith(JOBLIB_ZFILE_MAGIC):
+        return None, False
+    # The length field is fixed-width ASCII in every version that wrote it.
+    body = data[len(JOBLIB_ZFILE_MAGIC):]
+    match = re.match(rb"\s*([0-9]+)\s*", body[:32])
+    start = match.end() if match else 0
+    return decompress(body[start:], _ZLIB, limit)
+
+
+def describe_container(data: bytes, limit: int = DECOMPRESS_MAX_BYTES) -> Dict[str, Any]:
+    """Unwrap one layer of compression, if any, and say what was found.
+
+    Returns a dict with ``compression`` (label or None), ``readable`` (whether
+    we could open it), ``payload`` (the inner bytes, or None), and ``truncated``
+    (whether the decompressor stopped short of the stream's end). An unreadable
+    container yields a payload of None with the format still named, so the
+    caller can report "we did not read these bytes" instead of "clean" — and a
+    truncated one says so rather than passing off a prefix as the whole file.
+    """
+    result: Dict[str, Any] = {
+        "compression": None, "readable": True, "payload": data, "truncated": False,
+    }
+
+    if data.startswith(JOBLIB_ZFILE_MAGIC):
+        payload, complete = unwrap_zfile(data, limit)
+        result.update(compression="zfile", readable=payload is not None,
+                      payload=payload, truncated=payload is not None and not complete)
+        return result
+
+    label, readable = detect_compression(data[:16])
+    if label is None:
+        return result
+
+    result["compression"] = label
+    if not readable:
+        result.update(readable=False, payload=None)
+        return result
+
+    payload, complete = decompress(data, label, limit)
+    result.update(readable=payload is not None, payload=payload,
+                  truncated=payload is not None and not complete)
+    return result
+
+
+def parse_npy_header(data: bytes) -> Dict[str, Any] | None:
+    """Read a ``.npy`` header, returning where the data section starts.
+
+    The header is a Python dict *literal* — it is parsed with
+    ``ast.literal_eval``, which builds values and cannot call anything, and only
+    after the declared length has been bounds-checked. Nothing here evaluates
+    the file's contents; a malformed header yields None rather than an
+    exception, because a file we cannot parse still gets scanned by the caller.
+    """
+    if not data.startswith(NPY_MAGIC) or len(data) < 10:
+        return None
+
+    major = data[6]
+    if major == 1:
+        length_field, header_start = 2, 10
+    else:
+        # v2 and v3 widened the length field to four bytes.
+        length_field, header_start = 4, 12
+    if len(data) < header_start:
+        return None
+
+    header_len = int.from_bytes(data[8:8 + length_field], "little")
+    if header_len <= 0 or header_len > NPY_MAX_HEADER_BYTES:
+        return None
+
+    raw_header = data[header_start:header_start + header_len]
+    if len(raw_header) < header_len:
+        # Truncated: the data offset is still known, which is what matters.
+        raw_header = data[header_start:]
+
+    descr: Any = None
+    fortran_order = None
+    shape = None
+    try:
+        parsed = ast.literal_eval(raw_header.decode("latin-1").strip())
+        if isinstance(parsed, dict):
+            descr = parsed.get("descr")
+            fortran_order = parsed.get("fortran_order")
+            shape = parsed.get("shape")
+    except Exception:
+        # A header we cannot parse is not a reason to stop: fall back to reading
+        # the dtype string out of it, and failing that, carry on with the offset
+        # alone. The data section is scanned either way.
+        match = re.search(rb"'descr'\s*:\s*'([^']*)'", raw_header)
+        if match:
+            descr = match.group(1).decode("latin-1")
+
+    return {
+        "version": f"{major}.{data[7]}",
+        "descr": descr,
+        "fortran_order": fortran_order,
+        "shape": shape,
+        "data_offset": header_start + header_len,
+        "object_dtype": _is_object_dtype(descr),
+    }
+
+
+def _is_object_dtype(descr: Any) -> bool:
+    """True if a dtype descriptor carries Python objects anywhere inside it.
+
+    ``'|O'`` is the plain object array. A structured dtype is a list of
+    ``(name, format)`` or ``(name, format, shape)`` entries, and only the
+    *format* element is a dtype — searching the whole entry reads field names as
+    type codes, so an ordinary ``[('FOO', '<i4')]`` came back True on the
+    strength of the letter O in its name.
+
+    Note this only chooses the reported risk label, never whether to scan: the
+    data section is disassembled whatever the header claims, because the header
+    is attacker-supplied. Getting this wrong is noise, not a bypass.
+    """
+    if descr is None:
+        return False
+    if isinstance(descr, str):
+        # Strip the byte-order/size prefix, leaving the type character.
+        return "O" in descr.lstrip("|<>=")
+    if isinstance(descr, list):
+        # Structured: inspect each field's format, never its name or shape.
+        return any(
+            _is_object_dtype(field[1])
+            for field in descr
+            if isinstance(field, (list, tuple)) and len(field) >= 2
+        )
+    if isinstance(descr, tuple) and len(descr) == 2:
+        # A subarray format, `(format, shape)` — the dtype is the first element.
+        return _is_object_dtype(descr[0])
+    return False
+
+
+def npy_data_section(data: bytes, limit: int = DECOMPRESS_MAX_BYTES) -> Tuple[bytes | None, Dict[str, Any] | None]:
+    """Return ``(data_section, header)`` for a ``.npy`` buffer.
+
+    The data section is returned whatever the declared dtype says. Deciding by
+    the header first would mean trusting an attacker-supplied field to tell us
+    whether to look — and a pickle behind a header claiming ``'<f8'`` is exactly
+    the shape that trick would take.
+    """
+    header = parse_npy_header(data)
+    if header is None:
+        return None, None
+    offset = header["data_offset"]
+    if offset >= len(data):
+        return b"", header
+    return data[offset:offset + limit], header

{aisbom_cli-1.3.0 → aisbom_cli-1.3.2}/aisbom/properties.py

@@ -23,8 +23,19 @@ _FRAMEWORK_TO_FORMAT = {
     "GGUF": "gguf",
     "Keras": "keras",
     "ONNX": "onnx",
+    # The non-torch pickle carriers. `Pickle` shares PyTorch's token because the
+    # format genuinely is the same one; joblib and numpy get their own, because
+    # the container is a real difference a consumer may want to filter on.
+    "Pickle": "pickle",
+    "Joblib": "joblib",
+    "NumPy": "numpy",
 }
 
+# Formats whose findings are pickle opcodes. They share the `aisbom:pickle:*`
+# vocabulary rather than each inventing a parallel one, so a consumer that knows
+# how to read a threat off a `.pt` reads one off a `.joblib` unchanged.
+_PICKLE_BEARING_FORMATS = {"pickle", "joblib", "numpy"}
+
 
 def _format_for(art: Dict[str, Any]) -> str | None:
     return _FRAMEWORK_TO_FORMAT.get(art.get("framework"))
@@ -62,12 +73,49 @@ def build_component_properties(art: Dict[str, Any]) -> List[Tuple[str, str]]:
     details = art.get("details") or {}
     props.append(("aisbom:format", fmt))
 
-    if fmt == "pickle":
+    if fmt in _PICKLE_BEARING_FORMATS:
         threats = details.get("threats") or []
         for threat in threats:
             props.append(("aisbom:pickle:opcode", str(threat)))
         props.append(("aisbom:pickle:opcode_count", str(len(threats))))
 
+        # The container the stream was found in — `bare`, `npy`, `npz`, or the
+        # compression joblib used. Absent for `.pt`, which has its own shape.
+        container = details.get("container")
+        if container:
+            props.append(("aisbom:pickle:container", str(container)))
+        if details.get("scan_incomplete"):
+            props.append(("aisbom:pickle:scan_incomplete", "true"))
+        if details.get("dill_code_objects"):
+            props.append(("aisbom:pickle:dill_code_objects", "true"))
+
+        if fmt == "joblib":
+            compression = details.get("compression")
+            if compression:
+                props.append(("aisbom:joblib:compression", str(compression)))
+            decompressed = details.get("decompressed_bytes")
+            if decompressed is not None:
+                props.append(("aisbom:joblib:decompressed_bytes", str(decompressed)))
+
+        elif fmt == "numpy":
+            dtype = details.get("dtype")
+            if dtype:
+                props.append(("aisbom:numpy:dtype", str(dtype)))
+            if details.get("object_dtype") is not None:
+                props.append((
+                    "aisbom:numpy:object_dtype",
+                    "true" if details.get("object_dtype") else "false",
+                ))
+            shape = details.get("shape")
+            if shape:
+                props.append(("aisbom:numpy:shape", str(shape)))
+            npy_version = details.get("npy_version")
+            if npy_version:
+                props.append(("aisbom:numpy:npy_version", str(npy_version)))
+            internal_files = details.get("internal_files")
+            if internal_files is not None:
+                props.append(("aisbom:numpy:member_count", str(internal_files)))
+
     elif fmt == "safetensors":
         tensor_count = details.get("tensors")
         if tensor_count is not None: