fix: recurse into NumPy object pickle payloads#699
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (5)
WalkthroughAdds recursive scanning of embedded pickle payloads in object-dtype Changes
Sequence DiagramsequenceDiagram
participant Client
participant ZipScanner
participant NumPyScanner
participant PickleScanner
participant FileIO as File I/O
participant Core
Client->>ZipScanner: scan(archive)
ZipScanner->>ZipScanner: extract member (e.g., payload.npy)
ZipScanner->>NumPyScanner: scan(extracted_file, context=archive/member)
NumPyScanner->>FileIO: read headers / dtype info
alt object-dtype found
NumPyScanner->>FileIO: seek(data_offset)
NumPyScanner->>PickleScanner: scan(payload_stream, context_path)
PickleScanner->>NumPyScanner: return ScanResult (issues/checks)
NumPyScanner->>NumPyScanner: merge embedded results
end
NumPyScanner->>ZipScanner: return member ScanResult
ZipScanner->>ZipScanner: _rewrite_nested_result_context(nested_result, tmp, archive, entry)
ZipScanner->>Core: submit(scan_result)
Core->>Client: consolidated ScanResult (grouped by archive/member)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
📝 Coding Plan
Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/scanners/test_numpy_scanner.py`:
- Around line 107-120: Add Python type hints to the new helpers: annotate
_ExecPayload.__reduce__ and _SSLPayload.__reduce__ to return a
Tuple[Callable[..., Any], Tuple[Any, ...]] and annotate their self parameter as
usual; import the needed typing names (e.g., Callable, Tuple, Any). Also
annotate _failed_checks to accept result: Any (or the specific result type if
available) and return List[Any] (or List[CheckType] if you have a Check type);
import List/Any as needed. Ensure all new function/method signatures use these
type annotations to satisfy the repository typing rule.
- Around line 123-174: Add a regression test that ensures embedding a single
comment token does not bypass detection: create a new test (e.g.,
test_object_with_comment_token_still_detected) following the pattern of
test_object_dtype_numpy_recurses_into_pickle_exec and
test_object_npz_member_recurses_into_pickle_exec_with_member_context, but
construct the malicious payload (using _ExecPayload or _SSLPayload) with a
single comment token embedded inside its payload representation; save it to a
.npy and to a .npz member as appropriate, run NumPyScanner and ZipScanner
(NumPyScanner.scan and ZipScanner().scan), then assert the same detections as
the other tests (presence of "CVE-2019-6446" and "exec" or
"ssl.get_server_certificate" in check messages, and for zip case check the
zip_entry/member context like result.issues[*].details.get("zip_entry") ==
"payload.npy"). Ensure the new test mirrors existing assertion styles and uses
_failed_checks/result.issues checks so the single-comment-token case is covered.
- Around line 149-159: Update test_numeric_npz_has_no_pickle_recursion_findings
to also assert that result.issues contains no pickle/exec findings: mirror the
existing negative assertions on result.checks by adding the same two checks
against result.issues (use the same string checks for "CVE-2019-6446" in
(issue.name + issue.message) and "exec" in issue.message.lower()). This ensures
the test fails if the scanner reports findings as issues rather than checks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b3dbca4b-8cc3-4ee5-8d0f-800d35dbd8c2
📒 Files selected for processing (1)
tests/scanners/test_numpy_scanner.py
| class _ExecPayload: | ||
| def __reduce__(self): | ||
| return (exec, ("print('owned')",)) | ||
|
|
||
|
|
||
| class _SSLPayload: | ||
| def __reduce__(self): | ||
| import ssl | ||
|
|
||
| return (ssl.get_server_certificate, (("example.com", 443),)) | ||
|
|
||
|
|
||
| def _failed_checks(result): | ||
| return [c for c in result.checks if c.status.value == "failed"] |
There was a problem hiding this comment.
Add type hints to newly introduced helpers.
The new helper methods/functions are untyped, which violates the repository Python typing rule.
Proposed fix
+from typing import Any, Callable
+
class _ExecPayload:
- def __reduce__(self):
+ def __reduce__(self) -> tuple[Callable[..., object], tuple[str]]:
return (exec, ("print('owned')",))
@@
class _SSLPayload:
- def __reduce__(self):
+ def __reduce__(self) -> tuple[Callable[..., object], tuple[tuple[str, int]]]:
import ssl
return (ssl.get_server_certificate, (("example.com", 443),))
@@
-def _failed_checks(result):
+def _failed_checks(result: Any) -> list[Any]:
return [c for c in result.checks if c.status.value == "failed"]As per coding guidelines: "Always include type hints in Python code".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/scanners/test_numpy_scanner.py` around lines 107 - 120, Add Python type
hints to the new helpers: annotate _ExecPayload.__reduce__ and
_SSLPayload.__reduce__ to return a Tuple[Callable[..., Any], Tuple[Any, ...]]
and annotate their self parameter as usual; import the needed typing names
(e.g., Callable, Tuple, Any). Also annotate _failed_checks to accept result: Any
(or the specific result type if available) and return List[Any] (or
List[CheckType] if you have a Check type); import List/Any as needed. Ensure all
new function/method signatures use these type annotations to satisfy the
repository typing rule.
…ickle-recursion' into review-pr-699-head
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (3)
tests/scanners/test_numpy_scanner.py (3)
157-159:⚠️ Potential issue | 🟡 MinorStrengthen the negative NPZ test to assert against
result.issuestoo.Right now the regression can pass even if findings move from
checkstoissues. Add parallel negative assertions onresult.issuesfor CVE/exec indicators.Proposed fix
assert not any("CVE-2019-6446" in (c.name + c.message) for c in result.checks) assert not any("exec" in c.message.lower() for c in result.checks) + assert not any(i.details.get("cve_id") == "CVE-2019-6446" for i in result.issues) + assert not any("exec" in i.message.lower() for i in result.issues)Based on learnings: "Preserve or strengthen security detections; test both benign and malicious samples when adding scanner/feature changes".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/scanners/test_numpy_scanner.py` around lines 157 - 159, The test currently only asserts negatives against result.checks; extend it to also assert the same negatives against result.issues to prevent regressions that move findings from checks to issues. Update the assertions after the scanner run so you have parallel negative checks: assert not any("CVE-2019-6446" in (i.name + i.message) for i in result.issues) and assert not any("exec" in i.message.lower() for i in result.issues), using the existing result variable to mirror the checks assertions (result.checks) with result.issues.
123-174:⚠️ Potential issue | 🟠 MajorAdd an object-dtype comment-token bypass regression case.
The new recursion tests cover malicious payloads, but there is still no explicit single-comment-token bypass case for this object-dtype path.
As per coding guidelines: "Test bypass prevention by verifying that embedding a single comment token in a malicious payload does NOT suppress detection".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/scanners/test_numpy_scanner.py` around lines 123 - 174, Add a regression test that ensures a single comment-token inside an object-dtype pickle payload does not bypass detection: create a new test (e.g., test_object_dtype_numpy_comment_token_bypass) that mirrors test_object_dtype_numpy_recurses_into_pickle_exec but uses an _ExecPayload variant containing a single comment token embedded in the payload, save with np.save(..., allow_pickle=True), run NumPyScanner().scan on the file, and assert via _failed_checks(result) that the CVE-2019-6446 finding and the "exec" indication are still present (same style of assertions as in test_object_dtype_numpy_recurses_into_pickle_exec and test_object_npz_member_recurses_into_pickle_exec_with_member_context); reference _ExecPayload, NumPyScanner, _failed_checks and ensure the test explicitly checks that the comment token does not suppress detection.
108-116: 🛠️ Refactor suggestion | 🟠 MajorType hints are still missing on new payload helpers.
Line 108 and Line 113 add untyped
__reduce__methods; these should be annotated.As per coding guidelines: "Always include type hints in Python code".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/scanners/test_numpy_scanner.py` around lines 108 - 116, Add explicit return type annotations to the new __reduce__ methods: import typing names (e.g., from typing import Any, Callable, Tuple) and annotate both helpers' signatures to return Tuple[Callable[..., Any], Tuple[Any, ...]] (e.g., def __reduce__(self) -> Tuple[Callable[..., Any], Tuple[Any, ...]]:). Keep existing bodies (including the inline ssl import in _SSLPayload.__reduce__) and ensure the typing imports are added at top of the test module if not already present.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/scanners/test_zip_scanner.py`:
- Around line 320-322: The test class _ExecPayload defines __reduce__ without a
type annotation; add a return type hint to __reduce__ (and annotate self if
desired) to satisfy typing rules—e.g., import typing and annotate __reduce__ to
return a Tuple[typing.Callable[..., typing.Any], tuple] or typing.Any so the
signature becomes def __reduce__(self) -> Tuple[Callable[..., Any], tuple]: (or
an appropriate typing.Any alternative) while keeping the existing return value
unchanged.
In `@tests/test_core_asset_extraction.py`:
- Around line 79-81: The test helper ExecPayload.__reduce__ is missing a return
type annotation; update the method signature for ExecPayload.__reduce__ to
include a precise typing return such as -> Tuple[Callable[..., Any], Tuple[Any,
...]] and add the necessary typing imports (Tuple, Callable, Any) at the top of
the file so the repository typing rule is satisfied while preserving the
existing returned value (exec, ("print('owned')",)).
---
Duplicate comments:
In `@tests/scanners/test_numpy_scanner.py`:
- Around line 157-159: The test currently only asserts negatives against
result.checks; extend it to also assert the same negatives against result.issues
to prevent regressions that move findings from checks to issues. Update the
assertions after the scanner run so you have parallel negative checks: assert
not any("CVE-2019-6446" in (i.name + i.message) for i in result.issues) and
assert not any("exec" in i.message.lower() for i in result.issues), using the
existing result variable to mirror the checks assertions (result.checks) with
result.issues.
- Around line 123-174: Add a regression test that ensures a single comment-token
inside an object-dtype pickle payload does not bypass detection: create a new
test (e.g., test_object_dtype_numpy_comment_token_bypass) that mirrors
test_object_dtype_numpy_recurses_into_pickle_exec but uses an _ExecPayload
variant containing a single comment token embedded in the payload, save with
np.save(..., allow_pickle=True), run NumPyScanner().scan on the file, and assert
via _failed_checks(result) that the CVE-2019-6446 finding and the "exec"
indication are still present (same style of assertions as in
test_object_dtype_numpy_recurses_into_pickle_exec and
test_object_npz_member_recurses_into_pickle_exec_with_member_context); reference
_ExecPayload, NumPyScanner, _failed_checks and ensure the test explicitly checks
that the comment token does not suppress detection.
- Around line 108-116: Add explicit return type annotations to the new
__reduce__ methods: import typing names (e.g., from typing import Any, Callable,
Tuple) and annotate both helpers' signatures to return Tuple[Callable[..., Any],
Tuple[Any, ...]] (e.g., def __reduce__(self) -> Tuple[Callable[..., Any],
Tuple[Any, ...]]:). Keep existing bodies (including the inline ssl import in
_SSLPayload.__reduce__) and ensure the typing imports are added at top of the
test module if not already present.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 3f2cd2db-d2f7-44ca-998d-ea2bdebde626
📒 Files selected for processing (5)
modelaudit/core.pymodelaudit/scanners/zip_scanner.pytests/scanners/test_numpy_scanner.pytests/scanners/test_zip_scanner.pytests/test_core_asset_extraction.py
Summary\n- reuse the pickle scanner on object-dtype .npy payloads by scanning the remaining file stream instead of copying the payload into memory\n- preserve CVE-2019-6446 warnings while surfacing nested pickle findings for malicious object arrays and object members inside .npz archives\n- add NumPy and ZIP regression coverage for malicious, benign, truncated, and corrupted samples\n\n## Validation\n- uv run pytest tests/scanners/test_numpy_scanner.py -q\n- uv run pytest tests/scanners/test_zip_scanner.py -q -k "npz or pickle"\n- uv run ruff format modelaudit/ tests/\n- uv run ruff check --fix modelaudit/ tests/\n- uv run ruff check modelaudit/ tests/\n- uv run ruff format --check modelaudit/ tests/\n- uv run mypy modelaudit/\n- uv run pytest -n auto -m "not slow and not integration" --maxfail=1
Summary by CodeRabbit
Bug Fixes
Documentation
Tests