fix(scripts): quote whitelist model ids safely

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Bug Fixes
 
+- emit Hugging Face whitelist model IDs and organization summaries as safe generated Python
 - preserve bounded and fail-closed SafeTensors routing when header lengths resemble compression magic
 - allow equal offsets for supported empty SafeTensors tensors while preserving deterministic range checks and rejecting native size overflow
 - attribute core file-type validation failures directly to S901 instead of message-matching unrelated rules

scripts/fetch_hf_org_models.py

@@ -51,7 +51,10 @@ def fetch_organization_models_page(org: str, page: int) -> dict[str, Any]:
     try:
         req = Request(url, headers={"User-Agent": "ModelAudit/1.0"})
         with urlopen(req, timeout=30) as response:
-            return json.loads(response.read().decode("utf-8"))
+            data = json.loads(response.read().decode("utf-8"))
+            if not isinstance(data, dict):
+                raise ValueError("Hugging Face organization response must be a JSON object")
+            return data
     except HTTPError as e:
         if e.code == 404:
             # Organization doesn't exist or has no models
@@ -76,7 +79,7 @@ def fetch_organization_models(org: str, max_models: int | None = None) -> list[s
         List of model IDs
     """
     models_per_page = 30
-    model_ids = []
+    model_ids: list[str] = []
 
     print(f"\nFetching models from organization: {org}")
 
@@ -166,6 +169,11 @@ def generate_whitelist_module(org_models: dict[str, list[str]], output_path: Pat
         org_models: Dictionary mapping org name to model IDs
         output_path: Path where the module should be written
     """
+    if any(not isinstance(org, str) for org in org_models):
+        raise ValueError("Organization names must be strings")
+    if any(not isinstance(model_id, str) for models in org_models.values() for model_id in models):
+        raise ValueError("Model IDs must be strings")
+
     # Ensure parent directory exists
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
@@ -177,7 +185,10 @@ def generate_whitelist_module(org_models: dict[str, list[str]], output_path: Pat
     # Generate statistics
     total_orgs = len(org_models)
     total_models = len(all_model_ids)
-    org_summary = "\n".join(f"  - {org}: {len(models)} models" for org, models in sorted(org_models.items()))
+    org_summary = "\n".join(
+        f"#   - {json.dumps(org, ensure_ascii=False)}: {len(models)} models"
+        for org, models in sorted(org_models.items())
+    )
 
     # Generate the module content
     content = f'''"""
@@ -192,24 +203,22 @@ def generate_whitelist_module(org_models: dict[str, list[str]], output_path: Pat
 Total organizations: {total_orgs}
 Total unique models: {total_models}
 
-Organizations included:
-{org_summary}
-
 This whitelist is used by ModelAudit to reduce false positives when scanning
 models from trusted organizations. Users can disable this behavior via the
 'use_hf_whitelist' configuration option.
 """
 
+# Organizations included:
+{org_summary}
+
 # Set of model IDs from trusted organizations
 ORGANIZATION_MODELS: set[str] = {{
 '''
 
     # Add model IDs (sorted for readability and diff-friendliness)
     sorted_models = sorted(all_model_ids)
     for model_id in sorted_models:
-        # Escape quotes in model IDs if any
-        escaped_id = model_id.replace('"', '\\"')
-        content += f'    "{escaped_id}",\n'
+        content += f"    {json.dumps(model_id, ensure_ascii=False)},\n"
 
     content += '''}
 
@@ -235,7 +244,7 @@ def is_from_trusted_organization(model_id: str | None) -> bool:
 '''
 
     # Write the module
-    with open(output_path, "w") as f:
+    with open(output_path, "w", encoding="utf-8") as f:
         f.write(content)
 
     print(f"\nWhitelist module written to: {output_path}")

scripts/fetch_hf_top_models.py

@@ -27,7 +27,10 @@ def fetch_models_page(page: int, sort: str = "downloads") -> dict[str, Any]:
     try:
         req = Request(url, headers={"User-Agent": "ModelAudit/1.0"})
         with urlopen(req, timeout=30) as response:
-            return json.loads(response.read().decode("utf-8"))
+            data = json.loads(response.read().decode("utf-8"))
+            if not isinstance(data, dict):
+                raise ValueError("Hugging Face models response must be a JSON object")
+            return data
     except (URLError, HTTPError) as e:
         print(f"Error fetching page {page}: {e}", file=sys.stderr)
         raise
@@ -46,7 +49,7 @@ def fetch_top_models(count: int) -> list[str]:
     models_per_page = 30  # HuggingFace API returns 30 models per page
     pages_needed = (count + models_per_page - 1) // models_per_page
 
-    model_ids = []
+    model_ids: list[str] = []
 
     print(f"Fetching top {count} models from HuggingFace...")
     print(f"Will fetch {pages_needed} pages ({models_per_page} models per page)")
@@ -85,6 +88,9 @@ def generate_whitelist_module(model_ids: list[str], output_path: Path) -> None:
         model_ids: List of model IDs to whitelist
         output_path: Path where the module should be written
     """
+    if any(not isinstance(model_id, str) for model_id in model_ids):
+        raise ValueError("Model IDs must be strings")
+
     # Ensure parent directory exists
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
@@ -112,9 +118,7 @@ def generate_whitelist_module(model_ids: list[str], output_path: Path) -> None:
     # Add model IDs (sorted for readability and diff-friendliness)
     sorted_models = sorted(model_ids)
     for model_id in sorted_models:
-        # Escape quotes in model IDs if any
-        escaped_id = model_id.replace('"', '\\"')
-        content += f'    "{escaped_id}",\n'
+        content += f"    {json.dumps(model_id, ensure_ascii=False)},\n"
 
     content += '''}
 
@@ -140,7 +144,7 @@ def is_popular_model(model_id: str | None) -> bool:
 '''
 
     # Write the module
-    with open(output_path, "w") as f:
+    with open(output_path, "w", encoding="utf-8") as f:
         f.write(content)
 
     print(f"\nWhitelist module written to: {output_path}")

tests/conftest.py

@@ -208,6 +208,7 @@ def pytest_runtest_setup(item):
             "test_large_file_handler.py",  # Large file handler regression tests
             "test_file_iterator.py",  # Streaming file iterator memory regression tests
             "test_large_pickle_corpus_qa.py",  # large PickleScan Rust corpus QA harness tests
+            "test_hf_whitelist_generators.py",  # Hugging Face whitelist code-generation security tests
             "test_call_graph_click.py",  # standalone picklescan Click editor call-graph RCE regressions
             "test_call_graph_execnet.py",  # standalone picklescan execnet call-graph RCE regressions
             "test_call_graph_instance_defaults.py",  # standalone picklescan constructor-default alias RCE regressions

tests/scripts/test_hf_whitelist_generators.py

@@ -0,0 +1,149 @@
+"""Security regressions for generated Hugging Face whitelist modules."""
+
+import importlib.util
+import json
+import re
+from collections.abc import Callable
+from io import BytesIO
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+
+FETCH_PAGE_CASES = [
+    (
+        "fetch_hf_top_models.py",
+        "fetch_models_page",
+        (0,),
+        "Hugging Face models response must be a JSON object",
+    ),
+    (
+        "fetch_hf_org_models.py",
+        "fetch_organization_models_page",
+        ("trusted", 0),
+        "Hugging Face organization response must be a JSON object",
+    ),
+]
+
+
+def _load_module(path: Path) -> ModuleType:
+    spec = importlib.util.spec_from_file_location(f"test_{path.stem}", path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.mark.parametrize(
+    ("script_name", "generated_set_name", "generator_args"),
+    [
+        ("fetch_hf_top_models.py", "POPULAR_MODELS", lambda model_ids: (model_ids,)),
+        (
+            "fetch_hf_org_models.py",
+            "ORGANIZATION_MODELS",
+            lambda model_ids: ({"trusted": model_ids},),
+        ),
+    ],
+)
+def test_whitelist_generator_emits_model_ids_as_safe_python_literals(
+    tmp_path: Path,
+    script_name: str,
+    generated_set_name: str,
+    generator_args: Callable[[list[str]], tuple[object, ...]],
+) -> None:
+    repo_root = Path(__file__).resolve().parents[2]
+    script = _load_module(repo_root / "scripts" / script_name)
+    marker = tmp_path / "executed.txt"
+    payload = f"safe\\\",\n    (__import__('pathlib').Path({str(marker)!r}).write_text('executed') or 'x'),\n    #"
+    model_ids = [payload, "z/ordinary", 'a/quote"and\\backslash', "unicode/模型/😀/\u2028", payload]
+    output_path = tmp_path / f"{Path(script_name).stem}_generated.py"
+
+    script.generate_whitelist_module(*generator_args(model_ids), output_path)
+    generated = _load_module(output_path)
+
+    assert marker.exists() is False
+    assert getattr(generated, generated_set_name) == set(model_ids)
+    generated_source = output_path.read_text(encoding="utf-8")
+    for model_id in set(model_ids):
+        assert f"    {json.dumps(model_id, ensure_ascii=False)}," in generated_source
+
+
+def test_org_whitelist_generator_emits_org_names_as_safe_comments(tmp_path: Path) -> None:
+    repo_root = Path(__file__).resolve().parents[2]
+    script = _load_module(repo_root / "scripts" / "fetch_hf_org_models.py")
+    marker = tmp_path / "executed.txt"
+    payload = f'evil"""\n__import__("pathlib").Path({str(marker)!r}).write_text("executed")\n"""'
+    output_path = tmp_path / "organizations_generated.py"
+
+    script.generate_whitelist_module({payload: ["trusted/model"]}, output_path)
+    generated = _load_module(output_path)
+
+    assert marker.exists() is False
+    assert {"trusted/model"} == generated.ORGANIZATION_MODELS
+    assert f"#   - {json.dumps(payload, ensure_ascii=False)}: 1 models" in output_path.read_text(encoding="utf-8")
+
+
+@pytest.mark.parametrize(
+    ("script_name", "generator_args", "error_message"),
+    [
+        ("fetch_hf_top_models.py", ([123],), "Model IDs must be strings"),
+        ("fetch_hf_org_models.py", ({"trusted": [123]},), "Model IDs must be strings"),
+        ("fetch_hf_org_models.py", ({123: ["trusted/model"]},), "Organization names must be strings"),
+    ],
+)
+def test_whitelist_generator_rejects_non_string_codegen_inputs(
+    tmp_path: Path,
+    script_name: str,
+    generator_args: tuple[object, ...],
+    error_message: str,
+) -> None:
+    repo_root = Path(__file__).resolve().parents[2]
+    script = _load_module(repo_root / "scripts" / script_name)
+    output_path = tmp_path / "generated.py"
+
+    with pytest.raises(ValueError, match=error_message):
+        script.generate_whitelist_module(*generator_args, output_path)
+
+    assert output_path.exists() is False
+
+
+@pytest.mark.parametrize(("script_name", "fetch_name", "fetch_args", "error_message"), FETCH_PAGE_CASES)
+@pytest.mark.parametrize("response_body", [b"[]", b"null"])
+def test_fetch_page_rejects_non_object_response(
+    monkeypatch: pytest.MonkeyPatch,
+    script_name: str,
+    fetch_name: str,
+    fetch_args: tuple[object, ...],
+    error_message: str,
+    response_body: bytes,
+) -> None:
+    repo_root = Path(__file__).resolve().parents[2]
+    script = _load_module(repo_root / "scripts" / script_name)
+
+    def fake_urlopen(*_args: object, **_kwargs: object) -> BytesIO:
+        return BytesIO(response_body)
+
+    monkeypatch.setattr(script, "urlopen", fake_urlopen)
+
+    with pytest.raises(ValueError, match=re.escape(error_message)):
+        getattr(script, fetch_name)(*fetch_args)
+
+
+@pytest.mark.parametrize(("script_name", "fetch_name", "fetch_args", "error_message"), FETCH_PAGE_CASES)
+def test_fetch_page_accepts_object_response(
+    monkeypatch: pytest.MonkeyPatch,
+    script_name: str,
+    fetch_name: str,
+    fetch_args: tuple[object, ...],
+    error_message: str,
+) -> None:
+    del error_message
+    repo_root = Path(__file__).resolve().parents[2]
+    script = _load_module(repo_root / "scripts" / script_name)
+
+    def fake_urlopen(*_args: object, **_kwargs: object) -> BytesIO:
+        return BytesIO(b'{"models": []}')
+
+    monkeypatch.setattr(script, "urlopen", fake_urlopen)
+
+    assert getattr(script, fetch_name)(*fetch_args) == {"models": []}