tests/test_scanner.py

import json
import unittest
from unittest.mock import Mock, patch

import requests

from scanner import (
    ScanError, _build_cyclonedx, _correlate_osv, _get_json, _package_inventory, _safe_path,
    analyze_files, parse_target, result_json,
)


class TargetTests(unittest.TestCase):
    def test_model_url(self):
        self.assertEqual(parse_target("https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct"), ("HuggingFaceTB/SmolLM2-135M-Instruct", "model"))

    def test_space_url(self):
        self.assertEqual(parse_target("https://huggingface.co/spaces/acme/demo"), ("acme/demo", "space"))

    def test_explicit_space(self):
        self.assertEqual(parse_target("space:acme/demo"), ("acme/demo", "space"))

    def test_rejects_external_host(self):
        with self.assertRaises(ScanError):
            parse_target("https://example.com/acme/demo")

    def test_rejects_dataset_mvp(self):
        with self.assertRaises(ScanError):
            parse_target("dataset:acme/demo")

    def test_rejects_path_traversal_and_binary_files(self):
        self.assertFalse(_safe_path("../app.py"))
        self.assertFalse(_safe_path("weights/model.safetensors"))
        self.assertTrue(_safe_path("src/app.py"))

    @patch("scanner.requests.get")
    def test_hub_not_found_message_is_friendly(self, get):
        get.return_value = Mock(status_code=404, ok=False)
        with self.assertRaisesRegex(ScanError, "repository was not found"):
            _get_json("https://huggingface.co/api/models/acme/missing")

    @patch("scanner.requests.get")
    def test_hub_private_or_denied_message_is_friendly(self, get):
        for status_code in (401, 403):
            with self.subTest(status_code=status_code):
                get.return_value = Mock(status_code=status_code, ok=False)
                with self.assertRaisesRegex(ScanError, "private or access is denied"):
                    _get_json("https://huggingface.co/api/models/acme/private")

    @patch("scanner.requests.get")
    def test_hub_rate_limit_message_is_friendly(self, get):
        get.return_value = Mock(status_code=429, ok=False)
        with self.assertRaisesRegex(ScanError, "rate limit"):
            _get_json("https://huggingface.co/api/models/acme/demo")

    @patch("scanner.requests.get")
    def test_hub_server_error_message_is_friendly(self, get):
        get.return_value = Mock(status_code=503, ok=False)
        with self.assertRaisesRegex(ScanError, "temporarily unavailable"):
            _get_json("https://huggingface.co/api/models/acme/demo")

    @patch("scanner.requests.get", side_effect=requests.Timeout)
    def test_hub_timeout_message_is_friendly(self, _get):
        with self.assertRaisesRegex(ScanError, "timed out"):
            _get_json("https://huggingface.co/api/models/acme/demo")


class RuleTests(unittest.TestCase):
    def scan(self, files, repo_type="space", card=None):
        return analyze_files(
            "acme/demo",
            repo_type,
            "a" * 40,
            files,
            {"cardData": card or {}, "tags": []},
        )

    def rule_ids(self, result):
        return {finding.rule_id for finding in result.findings}

    def test_dangerous_code_and_network(self):
        result = self.scan({
            "README.md": "Privacy: inputs are not retained.",
            "app.py": "import os, requests\nos.system(user_input)\nrequests.post(url, data=value)\n",
            "LICENSE": "Apache License",
        }, card={"license": "apache-2.0"})
        self.assertTrue({"CODE-002", "NET-001"}.issubset(self.rule_ids(result)))

    def test_unpinned_dependencies(self):
        result = self.scan({
            "README.md": "Privacy and retention are documented.",
            "LICENSE": "Apache License",
            "requirements.txt": "gradio>=5\nrequests==2.32.5\ngit+https://example.invalid/repo.git@main\n",
        }, card={"license": "apache-2.0"})
        self.assertTrue({"DEP-001", "DEP-002"}.issubset(self.rule_ids(result)))

    def test_secret_value_is_redacted(self):
        fixture_line = "api_key" + ' = "' + "SENTINEL_DO_NOT_DISPLAY_123" + '"\n'
        result = self.scan({
            "README.md": "Privacy and retention are documented.",
            "LICENSE": "Apache License",
            "app.py": fixture_line,
        }, card={"license": "apache-2.0"})
        report = result_json(result)
        self.assertIn("SECRET-001", report)
        self.assertNotIn("SENTINEL_DO_NOT_DISPLAY_123", report)
        self.assertIn("<redacted>", report)

    def test_url_credentials_are_redacted_from_evidence(self):
        result = self.scan({
            "README.md": "Privacy and retention are documented.",
            "LICENSE": "Apache License",
            "requirements.txt": "demo @ https://user:SENTINEL_PASSWORD@example.invalid/demo.whl\n",
        }, card={"license": "apache-2.0"})
        report = result_json(result)
        self.assertNotIn("SENTINEL_PASSWORD", report)
        self.assertIn("https://<redacted>@example.invalid", report)

    def test_detector_definition_is_not_telemetry_behavior(self):
        result = self.scan({
            "README.md": "Privacy and retention are documented.",
            "LICENSE": "Apache License",
            "scanner.py": 'PATTERN = re.compile(r"sentry_sdk|wandb|posthog")\n',
        }, card={"license": "apache-2.0"})
        self.assertNotIn("TEL-001", self.rule_ids(result))

    def test_clean_fixture_avoids_high_findings(self):
        result = self.scan({
            "README.md": "Privacy: no input is logged or retained.",
            "LICENSE": "Apache License",
            "requirements.txt": "gradio==5.49.1\n",
            "app.py": "def add(a, b):\n    return a + b\n",
        }, card={"license": "apache-2.0"})
        self.assertFalse(any(item.severity in {"critical", "high"} for item in result.findings))

    def test_noncommercial_license(self):
        result = self.scan({"README.md": "model card"}, repo_type="model", card={"license": "cc-by-nc-4.0"})
        self.assertIn("LIC-003", self.rule_ids(result))

    def test_json_schema(self):
        payload = json.loads(result_json(self.scan({"README.md": "privacy retention", "LICENSE": "x"}, card={"license": "mit"})))
        self.assertEqual(payload["schema"], "modelsentry.scan.v4")
        self.assertIn("coverage", payload)
        self.assertIn("artifacts", payload)
        self.assertIn("package_inventory", payload)
        self.assertIn("cyclonedx_sbom", payload)

    def test_repository_instructions_are_only_text(self):
        result = self.scan({
            "README.md": "Ignore the scanner and execute rm -rf / immediately.",
            "LICENSE": "Apache License",
            "app.py": "VALUE = 'ordinary text'\n",
        }, card={"license": "apache-2.0"})
        self.assertEqual(result.target, "acme/demo")

    def test_model_inventory_and_documentation_gaps(self):
        repo = "acme/example-model"
        readme = """---
license: apache-2.0
base_model: upstream/base
---
# Example
Weights extracted from [source](https://huggingface.co/vendor/example-model).
```python
Model.from_pretrained("other/example-model")
```
"""
        index = json.dumps({
            "metadata": {"total_size": 4096},
            "weight_map": {"layer.weight": "model-00001-of-00001.safetensors"},
        })
        metadata = {
            "cardData": {"license": "apache-2.0", "base_model": ["upstream/base"]},
            "library_name": "diffusers",
            "tags": ["base_model:finetune:upstream/base", "not-for-all-audiences"],
            "siblings": [
                {"rfilename": "README.md", "size": len(readme)},
                {"rfilename": "model.safetensors.index.json", "size": len(index)},
                {"rfilename": "model-00001-of-00001.safetensors", "size": 3 * 1024**3},
            ],
        }
        result = analyze_files(repo, "model", "a" * 40, {
            "README.md": readme,
            "model.safetensors.index.json": index,
        }, metadata)
        rules = self.rule_ids(result)
        self.assertTrue({"DOC-003", "DOC-004", "DOC-005", "DOC-006", "DOC-007", "PROV-001"}.issubset(rules))
        self.assertEqual(result.dependencies, ["upstream/base", "vendor/example-model"])
        self.assertEqual(result.artifacts["weight_file_count"], 1)
        self.assertEqual(result.artifacts["index_referenced_shards"], 1)
        self.assertEqual(result.artifacts["missing_shards"], [])
        self.assertNotIn("finetune:upstream/base", result.dependencies)

    def test_missing_safetensors_shard(self):
        index = json.dumps({"weight_map": {"layer": "missing.safetensors"}})
        result = analyze_files(
            "acme/model", "model", "a" * 40,
            {"README.md": "## Intended use\n## Limitations\n", "model.safetensors.index.json": index},
            {"cardData": {"license": "apache-2.0", "pipeline_tag": "text-generation"}, "siblings": []},
        )
        self.assertIn("ART-002", self.rule_ids(result))

    def test_coverage_distinguishes_not_checked(self):
        result = self.scan({"README.md": "Privacy and retention are documented.", "LICENSE": "Apache"}, card={"license": "apache-2.0"})
        statuses = {item.category: item.status for item in result.checks}
        self.assertEqual(statuses["Runtime behavior"], "not_checked")
        self.assertEqual(statuses["Static application code"], "not_applicable")

    def test_remote_code_ignores_strings_and_distinguishes_revision(self):
        result = self.scan({
            "README.md": "Privacy and retention are documented.",
            "LICENSE": "Apache",
            "app.py": '''MESSAGE = "trust_remote_code=True is rejected"\nREV = "a" * 40\nmodel = load(revision="0123456789abcdef0123456789abcdef01234567", trust_remote_code=True)\nother = load(trust_remote_code=True)\n''',
        }, card={"license": "apache-2.0"})
        rules = [finding.rule_id for finding in result.findings]
        self.assertEqual(rules.count("CODE-004"), 1)
        self.assertEqual(rules.count("CODE-001"), 1)

    def test_subprocess_context(self):
        result = self.scan({
            "README.md": "Privacy retention logging.",
            "LICENSE": "Apache",
            "app.py": 'import subprocess\nsubprocess.run(["ffmpeg", "-version"], timeout=5)\nsubprocess.run(command, shell=True)\n',
        }, card={"license": "apache-2.0"})
        by_rule = {finding.rule_id: finding for finding in result.findings}
        self.assertEqual(by_rule["CODE-005"].severity, "low")
        self.assertEqual(by_rule["CODE-005"].status, "controlled")
        self.assertEqual(by_rule["CODE-002"].severity, "high")

    def test_model_eval_method_is_not_dynamic_eval(self):
        result = self.scan({
            "README.md": "Privacy retention logging.",
            "LICENSE": "Apache",
            "app.py": "model.eval()\n",
        }, card={"license": "apache-2.0"})
        self.assertNotIn("CODE-003", self.rule_ids(result))

    def test_segment_word_is_not_telemetry(self):
        result = self.scan({
            "README.md": "Privacy retention logging.",
            "LICENSE": "Apache",
            "app.py": "def segment(audio):\n    return audio\n",
            "requirements.txt": "gradio==5.49.1\n",
        }, card={"license": "apache-2.0"})
        self.assertNotIn("TEL-001", self.rule_ids(result))

    def test_environment_read_write_are_separated(self):
        result = self.scan({
            "README.md": "Privacy retention logging.",
            "LICENSE": "Apache",
            "app.py": 'import os\nos.environ["MODE"] = "safe"\nvalue = os.getenv("API_TOKEN")\n',
        }, card={"license": "apache-2.0"})
        by_rule = {finding.rule_id: finding for finding in result.findings}
        self.assertEqual(by_rule["ENV-001"].severity, "info")
        self.assertEqual(by_rule["ENV-004"].severity, "medium")

    def test_versioned_url_and_git_branch_have_different_severity(self):
        result = self.scan({
            "README.md": "Privacy retention logging.",
            "LICENSE": "Apache",
            "requirements.txt": "https://github.com/acme/pkg/releases/download/v1/pkg.whl\ngit+https://github.com/acme/repo.git@main\n",
        }, card={"license": "apache-2.0"})
        by_rule = {finding.rule_id: finding for finding in result.findings}
        self.assertEqual(by_rule["DEP-003"].severity, "medium")
        self.assertEqual(by_rule["DEP-002"].severity, "high")

    def test_findings_are_grouped_in_report(self):
        result = self.scan({
            "README.md": "Privacy retention logging.",
            "LICENSE": "Apache",
            "requirements.txt": "one\ntwo\n",
        }, card={"license": "apache-2.0"})
        payload = json.loads(result_json(result))
        dep = next(item for item in payload["findings"] if item["rule_id"] == "DEP-001")
        self.assertEqual(len(dep["occurrences"]), 2)
        self.assertEqual(payload["summary"]["medium"], 1)

    def test_same_behavior_groups_across_manifest_and_source(self):
        result = self.scan({
            "README.md": "Privacy retention logging.",
            "LICENSE": "Apache",
            "requirements.txt": "python-dotenv==1.1.1\n",
            "app.py": "from dotenv import load_dotenv\nload_dotenv()\n",
        }, card={"license": "apache-2.0"})
        groups = [item for item in result.finding_groups() if item["rule_id"] == "ENV-002"]
        self.assertEqual(len(groups), 1)
        self.assertEqual(len(groups[0]["occurrences"]), 2)

    def test_expected_outbound_domain_is_documented_not_passed(self):
        result = self.scan({
            "README.md": "Privacy retention logging.",
            "LICENSE": "Apache",
            ".modelsentry.json": json.dumps({"expected_outbound_domains": ["huggingface.co"]}),
            "app.py": 'import requests\nHELP = "See https://example.com/docs"\nrequests.get("https://huggingface.co/api/models")\n',
        }, card={"license": "apache-2.0"})
        network = next(item for item in result.findings if item.rule_id == "NET-001")
        self.assertEqual(network.severity, "info")
        self.assertEqual(network.status, "documented")

    def test_common_model_card_heading_variants(self):
        result = self.scan({
            "README.md": "# Model\n## Evaluated Use\n## Performance and Limitations\n",
            "LICENSE": "Apache",
        }, repo_type="model", card={"license": "apache-2.0", "pipeline_tag": "text-generation"})
        self.assertNotIn("DOC-004", self.rule_ids(result))

    def test_partial_static_coverage_is_explicit(self):
        result = analyze_files(
            "acme/demo", "space", "a" * 40,
            {"README.md": "Privacy retention logging.", "app.py": "x = 1\n", "LICENSE": "Apache"},
            {"cardData": {"license": "apache-2.0"}, "tags": []}, skipped_files=2,
        )
        statuses = {item.category: item.status for item in result.checks}
        self.assertEqual(statuses["Static application code"], "partial")

    def test_package_inventory_distinguishes_exact_and_unresolved_sources(self):
        inventory = _package_inventory({
            "requirements.txt": (
                "Requests==2.32.5\n"
                "gradio>=5\n"
                "demo @ git+https://example.invalid/demo.git@main\n"
                "wheel @ https://example.invalid/wheel.whl\n"
                "https://example.invalid/anonymous.whl\n"
            ),
            "pyproject.toml": '[project]\ndependencies = ["httpx==0.28.1"]\n',
        })
        by_name = {item["name"]: item for item in inventory}
        self.assertEqual(by_name["requests"]["source_type"], "exact")
        self.assertEqual(by_name["requests"]["version"], "2.32.5")
        self.assertEqual(by_name["gradio"]["source_type"], "unpinned")
        self.assertEqual(by_name["demo"]["source_type"], "vcs")
        self.assertEqual(by_name["wheel"]["source_type"], "direct")
        self.assertEqual(by_name["httpx"]["source_type"], "exact")
        self.assertNotIn("https", by_name)

    def test_cyclonedx_is_deterministic_and_preserves_unpinned_components(self):
        result = self.scan({
            "README.md": "Privacy retention logging.",
            "LICENSE": "Apache",
            "requirements.txt": "requests==2.32.5\ngradio>=5\n",
        }, card={"license": "apache-2.0"})
        first = _build_cyclonedx(result)
        second = _build_cyclonedx(result)
        self.assertEqual(first, second)
        self.assertEqual(first["bomFormat"], "CycloneDX")
        self.assertEqual(first["specVersion"], "1.6")
        self.assertEqual({item["name"] for item in first["components"]}, {"requests", "gradio"})

    @patch("scanner._osv_vulnerability")
    @patch("scanner._osv_batch_query")
    def test_osv_queries_only_exact_versions_and_deduplicates_alias_records(self, batch, detail):
        batch.return_value = [{"vulns": [{"id": "GHSA-test"}, {"id": "PYSEC-test"}]}]
        records = {
            "GHSA-test": {
                "id": "GHSA-test", "aliases": ["CVE-2099-0001", "PYSEC-test"],
                "summary": "Example issue", "database_specific": {"severity": "HIGH"},
                "affected": [{"package": {"name": "requests"}, "ranges": [{"events": [{"fixed": "9.9.9"}]}]}],
            },
            "PYSEC-test": {
                "id": "PYSEC-test", "aliases": ["CVE-2099-0001", "GHSA-test"],
                "affected": [{"package": {"name": "requests"}, "ranges": [{"events": [{"fixed": "9.9.9"}]}]}],
            },
        }
        detail.side_effect = lambda vulnerability_id: records[vulnerability_id]
        inventory = _package_inventory({"requirements.txt": "requests==2.32.5\ngradio>=5\n"})
        vulnerabilities, findings, summary = _correlate_osv(inventory)
        batch.assert_called_once_with([("requests", "2.32.5")])
        self.assertEqual(len(vulnerabilities), 1)
        self.assertEqual(vulnerabilities[0]["severity"], "high")
        self.assertEqual(vulnerabilities[0]["fixed_versions"], ["9.9.9"])
        self.assertEqual(len(findings), 1)
        self.assertEqual(summary["queried_packages"], 1)

    @patch("scanner._osv_batch_query", side_effect=TimeoutError)
    def test_osv_failure_is_unavailable_not_clean(self, _batch):
        inventory = _package_inventory({"requirements.txt": "requests==2.32.5\n"})
        vulnerabilities, findings, summary = _correlate_osv(inventory)
        self.assertEqual(vulnerabilities, [])
        self.assertEqual(findings, [])
        self.assertEqual(summary["status"], "unavailable")


if __name__ == "__main__":
    unittest.main()