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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aisbom-cli
-Version: 1.3.2
+Version: 1.3.3
 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

{aisbom_cli-1.3.2 → aisbom_cli-1.3.3}/aisbom/safety.py

@@ -1158,6 +1158,77 @@ class _NullWriter:
         pass
 
 
+# Opcodes whose argument runs to a newline and can plausibly *start* a real
+# pickle while carrying a lot of bytes. Each maps to the characters its
+# argument legitimately begins with, so a file is only re-read when its second
+# byte is consistent with the opcode its first byte claims to be.
+#
+# This guard is why a parquet file is not re-read: it opens `PAR1`, and while
+# `P` is the PERSID opcode, PERSID is not on this list — no real pickle starts
+# with a persistent id, and accepting it would mean re-reading every parquet,
+# ORC and arrow file in a tree up to the sniff cap.
+_NEWLINE_ARG_STARTS = {
+    "S": b"'\"",                  # STRING — always quoted
+    "V": None,                    # UNICODE — raw text, no reliable lead byte
+    "I": b"0123456789+-",         # INT
+    "L": b"0123456789+-",         # LONG
+    "F": b"0123456789+-.",        # FLOAT
+}
+
+
+def _first_argument_overruns(data: bytes) -> bool:
+    """Does `data`'s first opcode declare an argument longer than `data`?
+
+    Discovery reads a bounded head and only re-reads when the head looked like
+    an unfinished pickle. "At least one opcode parsed" is the usual signal, but
+    a pickle that opens with a single huge literal completes *no* opcodes — so
+    without this check a payload hidden behind one 64KB literal is never found,
+    while the documented limit claims 16MB. Measured before it was fixed: a
+    65,000-byte pad was caught and a 70,000-byte pad was not.
+
+    Length-prefixed arguments are read exactly. Newline-terminated ones cannot
+    be measured without the newline, so they are admitted only for the handful
+    of opcodes that can really begin a pickle, and only when the following byte
+    matches what that opcode's argument must start with.
+    """
+    if not data:
+        return False
+    op = pickletools.code2op.get(chr(data[0]))
+    if op is None or op.arg is None:
+        return False
+
+    n = op.arg.n
+    if n >= 0:
+        # A fixed-width argument. These are a handful of bytes; they cannot be
+        # what overran a 64KB read.
+        return False
+
+    if n == pickletools.UP_TO_NEWLINE:
+        if op.code not in _NEWLINE_ARG_STARTS:
+            return False
+        if b"\n" in data:
+            # The terminator is already in view, so the argument is not what
+            # ran off the end — something else failed, and re-reading will not
+            # change that.
+            return False
+        lead = _NEWLINE_ARG_STARTS[op.code]
+        return lead is None or (len(data) > 1 and data[1] in lead)
+
+    # Length-prefixed: read the declared size and believe it only far enough to
+    # decide whether a bigger read would reach the end of the argument.
+    widths = {
+        pickletools.TAKEN_FROM_ARGUMENT1: 1,
+        pickletools.TAKEN_FROM_ARGUMENT4: 4,
+        pickletools.TAKEN_FROM_ARGUMENT4U: 4,
+        pickletools.TAKEN_FROM_ARGUMENT8U: 8,
+    }
+    width = widths.get(n)
+    if width is None or len(data) < 1 + width:
+        return False
+    declared = int.from_bytes(data[1 : 1 + width], "little")
+    return 1 + width + declared > len(data)
+
+
 def head_looks_like_pickle(data: bytes) -> tuple[bool, bool]:
     """Does the head of a file begin with a genuinely valid pickle?
 
@@ -1201,9 +1272,11 @@ def head_looks_like_pickle(data: bytes) -> tuple[bool, bool]:
         pass
 
     if stop_at is None:
-        # No complete pickle in view. Worth a bigger read only if something
-        # actually parsed; a file that yielded nothing is simply not a pickle.
-        return (False, parsed > 0)
+        # No complete pickle in view. Worth a bigger read if something parsed
+        # — or if nothing parsed *because* the very first opcode carries an
+        # argument that runs off the end of the buffer, which is the one way a
+        # genuine pickle yields no opcodes at all.
+        return (False, parsed > 0 or _first_argument_overruns(data))
 
     try:
         pickletools.dis(io.BytesIO(data[:stop_at]), out=_NullWriter())

{aisbom_cli-1.3.2 → aisbom_cli-1.3.3}/aisbom/scanner.py

@@ -354,15 +354,23 @@ class DeepScanner:
     def _sniff_is_pickle(self, full_path: Path) -> bool:
         """Decide by content whether an unclaimed file is a pickle.
 
-        Reads a small head first and gives up on it unless opcodes actually
-        parsed, which is what keeps a walk over a tree full of parquet and PNGs
-        cheap: those yield zero opcodes and are never opened a second time.
-
-        Known limit, stated rather than hidden: a pickle whose *first* opcode
-        carries an argument longer than the whole sniff budget parses zero
-        opcodes here and is not discovered by content. Reaching that requires
-        a single 16MB-plus literal ahead of the payload. Files carrying a
-        recognized model extension are unaffected — they never reach this path.
+        Reads a small head first and re-reads with a larger budget only when
+        that head looked like an unfinished pickle, which is what keeps a walk
+        over a tree full of parquet and PNGs cheap: those settle on the first
+        read and are never opened a second time.
+
+        Two things can leave a head unfinished, and both must count. Opcodes
+        may parse cleanly without reaching STOP — an ordinary large pickle. Or
+        *no* opcode completes because the first one carries an argument running
+        past the buffer; `_first_argument_overruns` recognizes that case. The
+        original rule counted only the first, so one 64KB literal ahead of a
+        payload hid it from discovery entirely while the documented limit
+        claimed 16MB.
+
+        Known limit, stated rather than hidden: past PICKLE_SNIFF_MAX_BYTES we
+        stop looking, so a payload behind a literal larger than that is not
+        discovered by content. Files carrying a recognized model extension are
+        unaffected — they never reach this path.
         """
         try:
             size = full_path.stat().st_size

{aisbom_cli-1.3.2 → aisbom_cli-1.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "aisbom-cli"
-version = "1.3.2"
+version = "1.3.3"
 description = "An AI Supply Chain security tool that that detects Pickle bombs and generates CycloneDX SBOMs for Machine Learning models."
 authors = ["Ajoy L <lab700xdev@gmail.com>"]
 readme = "README.md"