Add graceful degradation and scan API for robust scanning (#218)

* Add graceful degradation and scan API for robust scanning

- Add ScanResult class with is_safe, severity, results, and errors
- Add scan_file() for graceful single-file scanning
- Add scan_archive() for scanning ZIP archives
- Add RelaxedZipFile class that ignores CRC validation errors
- Add _scan_bytes() helper for in-memory scanning
- Export new functions from package

This API provides picklescan-like graceful degradation, continuing
to scan even when individual files fail to parse.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Fix scan API: broken RelaxedZipFile, unsafe ScanResult, missing tests

- Replace RelaxedZipFile CRC=0 hack (caused BadZipFile on valid files)
  with _expected_crc=None on the returned file handle
- Make ScanResult.is_safe a computed property instead of a stored bool;
  __bool__ now returns False when errors exist (incomplete scan != safe)
- Escalate severity in _scan_bytes when check_safety() or outer parse
  exceptions occur in graceful mode
- Deduplicate scan_file by delegating to _scan_bytes
- Change graceful default to False for scan_file and scan_archive
- Remove dead code: data-is-None branch, is_safe kwarg, read() override
- Add 12 tests covering scan_file, scan_archive, ScanResult, RelaxedZipFile

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Address review: move RelaxedZipFile, fix error handling, add tests

- Move RelaxedZipFile from polyglot.py to loader.py to avoid requiring
  numpy/torch for scan_archive (RelaxedZipFile only needs stdlib zipfile)
- Add RuntimeWarning when _expected_crc attribute is missing on future
  Python versions so CRC bypass degradation is not silent
- Widen scan_archive outer exception handler to catch OSError (covers
  FileNotFoundError, PermissionError) in graceful mode, matching scan_file
- Separate archive.read() try/except from _scan_bytes() call so the
  inner catch only handles I/O errors, not analysis-layer exceptions
- Escalate _scan_bytes outer except Exception to LIKELY_UNSAFE (was
  SUSPICIOUS — for a security scanner, unanalyzable files should be
  treated as dangerous)
- Include exception type in all error messages for debuggability
- Add 8 new tests: archive nonexistent graceful/non-graceful, bad ZIP
  non-graceful raises, mixed good/bad archive members, all pickle
  extensions, is_safe boundary at POSSIBLY_UNSAFE, analysis error
  severity escalation, RelaxedZipFile CRC bypass regression

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Fix Ruff linting failure

* Minor improvements to scan API

- Narrow file-open exception catch from Exception to OSError
- Rename scan_archive to scan_zip_archive to clarify ZIP-only scope
- Use PurePosixPath.suffix for extension parsing
- Remove .pt/.pth from extension filter (these are ZIP containers, not raw pickles)
- Update docstring to reference fickling.polyglot for other archive formats

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Fix test_scans_all_pickle_extensions

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Thomas Chauchefoin <thomas.chauchefoin@trailofbits.com>

Author: 3 people

fickling/__init__.py

@@ -1,5 +1,5 @@
 # fmt: off
-from .loader import load, loads #noqa
+from .loader import load, loads, scan_file, scan_zip_archive, ScanResult  # noqa
 from .context import check_safety #noqa
 from .hook import always_check_safety, activate_safe_ml_environment #noqa
 from .analysis import is_likely_safe # noqa

fickling/loader.py

@@ -1,9 +1,38 @@
 import pickle
+import warnings
+import zipfile
 from io import BytesIO
+from pathlib import PurePosixPath
 
-from fickling.analysis import Severity, check_safety
+from fickling.analysis import AnalysisResults, Severity, check_safety
 from fickling.exception import UnsafeFileError
-from fickling.fickle import Pickled
+from fickling.fickle import Pickled, PickleDecodeError, StackedPickle
+
+
+class RelaxedZipFile(zipfile.ZipFile):
+    """A ZipFile subclass that ignores CRC validation errors.
+
+    Matches PyTorch's lenient ZIP parsing behavior. Uses CPython's
+    internal _expected_crc attribute on ZipExtFile — guarded by hasattr
+    so it degrades to standard CRC-checked behavior if the attribute
+    is renamed in a future Python version.
+    """
+
+    def open(self, name, mode="r", pwd=None, *, force_zip64=False):
+        """Open a member with CRC validation disabled."""
+        f = super().open(name, mode, pwd, force_zip64=force_zip64)
+        if hasattr(f, "_expected_crc"):
+            f._expected_crc = None
+        else:
+            warnings.warn(
+                "RelaxedZipFile: _expected_crc not found on ZipExtFile. "
+                "CRC validation is still enabled. This may cause failures "
+                "scanning PyTorch model files with CRC mismatches.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+        return f
+
 
 # Save the original pickle.loads before any hooks can replace it.
 # loader.load() must use the real pickle.loads for final deserialization,
@@ -75,3 +104,170 @@ def loads(
         json_output_path=json_output_path,
         **kwargs,
     )
+
+
+class ScanResult:
+    """Result of scanning a file or archive for malicious pickle content."""
+
+    def __init__(
+        self,
+        filepath: str,
+        severity: Severity,
+        results: list[AnalysisResults],
+        errors: list[str],
+    ):
+        self.filepath = filepath
+        self.severity = severity
+        self.results = results
+        self.errors = errors
+
+    @property
+    def is_safe(self) -> bool:
+        return self.severity <= Severity.LIKELY_SAFE
+
+    def __bool__(self) -> bool:
+        return self.is_safe and not self.errors
+
+    def __repr__(self) -> str:
+        return (
+            f"ScanResult(filepath={self.filepath!r}, "
+            f"severity={self.severity.name}, "
+            f"results={len(self.results)}, errors={len(self.errors)})"
+        )
+
+
+def scan_file(
+    filepath: str,
+    graceful: bool = False,
+    json_output_path: str | None = None,
+) -> ScanResult:
+    """Scan a file for malicious pickle content.
+
+    Args:
+        filepath: Path to the file to scan
+        graceful: If True, continue on parse errors and report them.
+                  If False, raise exceptions on parse errors.
+        json_output_path: Optional path to write JSON analysis results
+
+    Returns:
+        ScanResult with severity, results list, and errors list
+    """
+    try:
+        with open(filepath, "rb") as f:
+            data = f.read()
+    except OSError as e:
+        if graceful:
+            return ScanResult(
+                filepath=filepath,
+                severity=Severity.SUSPICIOUS,
+                results=[],
+                errors=[f"File error ({type(e).__name__}): {e!s}"],
+            )
+        raise
+    return _scan_bytes(filepath, data, graceful, json_output_path)
+
+
+def scan_zip_archive(
+    filepath: str,
+    graceful: bool = False,
+    json_output_path: str | None = None,
+) -> dict[str, ScanResult]:
+    """Scan a ZIP archive for malicious raw pickle content.
+
+    Scans each file within the archive that has a raw pickle-related extension
+    (.pkl, .pickle, .bin). Only ZIP archives are supported;
+    for TAR or 7z archives, see fickling.polyglot.
+
+    Args:
+        filepath: Path to the ZIP archive to scan
+        graceful: If True, continue on parse errors and report them
+        json_output_path: Optional path to write JSON analysis results
+
+    Returns:
+        Dict mapping archive member names to their ScanResults
+    """
+    results: dict[str, ScanResult] = {}
+    pickle_extensions = {".pkl", ".pickle", ".bin"}
+
+    try:
+        with RelaxedZipFile(filepath, "r") as archive:
+            for info in archive.infolist():
+                if info.is_dir():
+                    continue
+                ext = PurePosixPath(info.filename).suffix.lower()
+                if ext not in pickle_extensions:
+                    continue
+
+                try:
+                    data = archive.read(info)
+                except (zipfile.BadZipFile, OSError) as e:
+                    if graceful:
+                        results[info.filename] = ScanResult(
+                            filepath=info.filename,
+                            severity=Severity.SUSPICIOUS,
+                            results=[],
+                            errors=[f"Read error ({type(e).__name__}): {e!s}"],
+                        )
+                        continue
+                    raise
+
+                results[info.filename] = _scan_bytes(
+                    info.filename, data, graceful, json_output_path
+                )
+    except (zipfile.BadZipFile, OSError) as e:
+        if not graceful:
+            raise
+        results["<archive>"] = ScanResult(
+            filepath=filepath,
+            severity=Severity.SUSPICIOUS,
+            results=[],
+            errors=[f"Archive error ({type(e).__name__}): {e!s}"],
+        )
+
+    return results
+
+
+def _scan_bytes(
+    name: str,
+    data: bytes,
+    graceful: bool,
+    json_output_path: str | None,
+) -> ScanResult:
+    """Scan bytes data for malicious pickle content."""
+    results: list[AnalysisResults] = []
+    errors: list[str] = []
+    overall_severity = Severity.LIKELY_SAFE
+
+    try:
+        stacked = StackedPickle.load(BytesIO(data), fail_on_decode_error=not graceful)
+        for pickled in stacked:
+            try:
+                result = check_safety(pickled, json_output_path=json_output_path)
+                results.append(result)
+                if result.severity > overall_severity:
+                    overall_severity = result.severity
+            except Exception as e:
+                if graceful:
+                    errors.append(f"Analysis error ({type(e).__name__}): {e!s}")
+                    overall_severity = max(overall_severity, Severity.LIKELY_UNSAFE)
+                else:
+                    raise
+    except PickleDecodeError as e:
+        if graceful:
+            errors.append(f"Parse error ({type(e).__name__}): {e!s}")
+            overall_severity = max(overall_severity, Severity.SUSPICIOUS)
+        else:
+            raise
+    except Exception as e:
+        if graceful:
+            errors.append(f"Unexpected error ({type(e).__name__}): {e!s}")
+            overall_severity = max(overall_severity, Severity.LIKELY_UNSAFE)
+        else:
+            raise
+
+    return ScanResult(
+        filepath=name,
+        severity=overall_severity,
+        results=results,
+        errors=errors,
+    )