weightguard 0.1.0__tar.gz

weightguard-0.1.0/PKG-INFO

@@ -0,0 +1,79 @@
+Metadata-Version: 2.4
+Name: weightguard
+Version: 0.1.0
+Summary: Multi-format security scanner for ML model artifacts (pickle/PyTorch, Keras, ONNX, GGUF, SafeTensors) with provenance and license checks.
+Author: Sanjeev Jaiswal
+Author-email: Sanjeev Jaiswal <jassics@gmail.com>
+License-Expression: MIT
+Requires-Dist: typer>=0.12
+Requires-Dist: rich>=13.7
+Requires-Dist: pydantic>=2.7
+Requires-Dist: fickling>=0.1.12
+Requires-Dist: safetensors>=0.8.0
+Requires-Dist: huggingface-hub>=1.28.0
+Requires-Dist: gguf>=0.19.0
+Requires-Dist: h5py>=3.16.0
+Requires-Dist: onnx>=1.22.0
+Requires-Python: >=3.11
+Project-URL: Homepage, https://github.com/jassics/weightguard
+Project-URL: Repository, https://github.com/jassics/weightguard
+Description-Content-Type: text/markdown
+
+# weightguard
+
+Security scanner for ML model artifacts. Point it at a Hugging Face repo, a git
+URL, or a local path, and it flags known artifact-level risks — malicious
+pickle/PyTorch payloads, Keras `Lambda`-layer code injection, ONNX custom-op
+RCE surface, and anomalous/crafted GGUF headers — with a severity and a
+concrete mitigation for each finding.
+
+Model files are not inert data. Several common serialization formats can
+embed code that executes the moment the file is *loaded*, before any
+inference happens. `weightguard` is a static, pre-load check you run before
+trusting a downloaded model.
+
+## Install
+
+```bash
+pip install weightguard
+```
+
+## Usage
+
+```bash
+# Hugging Face repo
+modelsec scan https://huggingface.co/<org>/<repo>
+
+# git repo
+modelsec scan https://github.com/<org>/<repo>.git
+
+# local path
+modelsec scan /path/to/model
+
+# CI gate — exit non-zero only above a severity threshold (default: HIGH)
+modelsec scan <target> --fail-on CRITICAL
+```
+
+Exit codes: `0` clean, `1` a finding at/above `--fail-on`, `2` target could
+not be resolved.
+
+## What it checks today
+
+| Format | Detector | Technique |
+|---|---|---|
+| Pickle / PyTorch (`.pkl`, `.bin`, `.pt`, `.pth`) | Fickling AST analysis | Detects arbitrary-code-execution opcode chains; resistant to the malformed-opcode-stream evasion that defeats denylist scanners |
+| SafeTensors | Format check | Flags files that fail to parse as valid SafeTensors (renamed/spoofed files) |
+| Keras (`.h5`, `.keras`) | Lambda-layer check | Flags `Lambda` layers, which embed a marshalled Python function executed on load |
+| ONNX | Custom-op check | Flags graphs referencing non-standard operator domains (native-code load surface) |
+| GGUF | Header/KV anomaly check | Flags malformed headers, oversized KV metadata, unknown quantization types, and tensor sizes inconsistent with the file |
+
+## Development
+
+```bash
+uv sync
+uv run pytest -q
+```
+
+## License
+
+MIT

weightguard-0.1.0/README.md

@@ -0,0 +1,58 @@
+# weightguard
+
+Security scanner for ML model artifacts. Point it at a Hugging Face repo, a git
+URL, or a local path, and it flags known artifact-level risks — malicious
+pickle/PyTorch payloads, Keras `Lambda`-layer code injection, ONNX custom-op
+RCE surface, and anomalous/crafted GGUF headers — with a severity and a
+concrete mitigation for each finding.
+
+Model files are not inert data. Several common serialization formats can
+embed code that executes the moment the file is *loaded*, before any
+inference happens. `weightguard` is a static, pre-load check you run before
+trusting a downloaded model.
+
+## Install
+
+```bash
+pip install weightguard
+```
+
+## Usage
+
+```bash
+# Hugging Face repo
+modelsec scan https://huggingface.co/<org>/<repo>
+
+# git repo
+modelsec scan https://github.com/<org>/<repo>.git
+
+# local path
+modelsec scan /path/to/model
+
+# CI gate — exit non-zero only above a severity threshold (default: HIGH)
+modelsec scan <target> --fail-on CRITICAL
+```
+
+Exit codes: `0` clean, `1` a finding at/above `--fail-on`, `2` target could
+not be resolved.
+
+## What it checks today
+
+| Format | Detector | Technique |
+|---|---|---|
+| Pickle / PyTorch (`.pkl`, `.bin`, `.pt`, `.pth`) | Fickling AST analysis | Detects arbitrary-code-execution opcode chains; resistant to the malformed-opcode-stream evasion that defeats denylist scanners |
+| SafeTensors | Format check | Flags files that fail to parse as valid SafeTensors (renamed/spoofed files) |
+| Keras (`.h5`, `.keras`) | Lambda-layer check | Flags `Lambda` layers, which embed a marshalled Python function executed on load |
+| ONNX | Custom-op check | Flags graphs referencing non-standard operator domains (native-code load surface) |
+| GGUF | Header/KV anomaly check | Flags malformed headers, oversized KV metadata, unknown quantization types, and tensor sizes inconsistent with the file |
+
+## Development
+
+```bash
+uv sync
+uv run pytest -q
+```
+
+## License
+
+MIT

weightguard-0.1.0/pyproject.toml

@@ -0,0 +1,39 @@
+[project]
+name = "weightguard"
+version = "0.1.0"
+description = "Multi-format security scanner for ML model artifacts (pickle/PyTorch, Keras, ONNX, GGUF, SafeTensors) with provenance and license checks."
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.11"
+dependencies = [
+    "typer>=0.12",
+    "rich>=13.7",
+    "pydantic>=2.7",
+    "fickling>=0.1.12",
+    "safetensors>=0.8.0",
+    "huggingface-hub>=1.28.0",
+    "gguf>=0.19.0",
+    "h5py>=3.16.0",
+    "onnx>=1.22.0",
+]
+
+[[project.authors]]
+name = "Sanjeev Jaiswal"
+email = "jassics@gmail.com"
+
+[project.urls]
+Homepage = "https://github.com/jassics/weightguard"
+Repository = "https://github.com/jassics/weightguard"
+
+[project.scripts]
+modelsec = "modelsec.cli:app"
+
+[build-system]
+requires = ["uv_build>=0.9.2,<0.10.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "modelsec"
+
+[dependency-groups]
+dev = ["pytest>=8.0"]

weightguard-0.1.0/pyproject.toml.orig

@@ -0,0 +1,38 @@
+[project]
+name = "weightguard"
+version = "0.1.0"
+description = "Multi-format security scanner for ML model artifacts (pickle/PyTorch, Keras, ONNX, GGUF, SafeTensors) with provenance and license checks."
+readme = "README.md"
+license = "MIT"
+authors = [
+    { name = "Sanjeev Jaiswal", email = "jassics@gmail.com" }
+]
+requires-python = ">=3.11"
+urls.Homepage = "https://github.com/jassics/weightguard"
+urls.Repository = "https://github.com/jassics/weightguard"
+dependencies = [
+    "typer>=0.12",
+    "rich>=13.7",
+    "pydantic>=2.7",
+    "fickling>=0.1.12",
+    "safetensors>=0.8.0",
+    "huggingface-hub>=1.28.0",
+    "gguf>=0.19.0",
+    "h5py>=3.16.0",
+    "onnx>=1.22.0",
+]
+
+[project.scripts]
+modelsec = "modelsec.cli:app"
+
+[build-system]
+requires = ["uv_build>=0.9.2,<0.10.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "modelsec"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.0",
+]

weightguard-0.1.0/src/modelsec/cli.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from modelsec.models import Severity
+from modelsec.resolver import UnresolvableTarget, resolve
+from modelsec.scanner import scan_path
+
+app = typer.Typer(help="modelsec — security scanner for ML model artifacts.", no_args_is_help=True)
+console = Console()
+
+
+@app.command(name="version")
+def version() -> None:
+    """Print the modelsec version."""
+    from importlib.metadata import version as _v
+
+    console.print(_v("weightguard"))
+
+
+@app.command(name="scan")
+def scan(
+    target: str = typer.Argument(..., help="Hugging Face repo URL, git URL, or local path to scan."),
+    fail_on: Severity = typer.Option(Severity.HIGH, help="Exit non-zero if a finding >= this severity is present."),
+) -> None:
+    """Scan a model (HF repo, git repo, or local directory/file) for known artifact-level security risks."""
+    try:
+        path = resolve(target)
+    except UnresolvableTarget as exc:
+        console.print(f"[red]Error:[/red] {exc}")
+        raise typer.Exit(code=2) from exc
+
+    report = scan_path(path)
+
+    if not report.findings:
+        console.print(f"[green]No findings.[/green] Scanned: {target}")
+        raise typer.Exit(code=0)
+
+    table = Table(title=f"modelsec findings — {target}")
+    table.add_column("Severity")
+    table.add_column("Detector")
+    table.add_column("Title")
+    table.add_column("File")
+    table.add_column("Mitigation")
+
+    for f in sorted(report.findings, key=lambda f: -f.severity.rank):
+        table.add_row(f.severity.value, f.detector, f.title, f.file or "-", f.mitigation)
+
+    console.print(table)
+
+    if report.fails(fail_on):
+        raise typer.Exit(code=1)
+
+
+if __name__ == "__main__":
+    app()

weightguard-0.1.0/src/modelsec/detectors/__init__.py

@@ -0,0 +1,13 @@
+from modelsec.detectors.gguf_detector import GgufDetector
+from modelsec.detectors.keras_detector import KerasDetector
+from modelsec.detectors.onnx_detector import OnnxDetector
+from modelsec.detectors.pickle_detector import PickleDetector
+from modelsec.detectors.safetensors_detector import SafeTensorsDetector
+
+ALL_DETECTORS = [
+    PickleDetector(),
+    SafeTensorsDetector(),
+    KerasDetector(),
+    OnnxDetector(),
+    GgufDetector(),
+]

weightguard-0.1.0/src/modelsec/detectors/base.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Protocol
+
+from modelsec.models import Finding
+
+
+class Detector(Protocol):
+    """A detector inspects one file and returns zero or more findings.
+
+    Detectors must never execute or deserialize untrusted content outside
+    a sandbox (Phase 2). Phase-1 detectors are static-analysis only.
+    """
+
+    name: str
+
+    def applies_to(self, path: Path) -> bool: ...
+
+    def scan(self, path: Path) -> list[Finding]: ...

weightguard-0.1.0/src/modelsec/detectors/gguf_detector.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import gguf
+from gguf.constants import GGUFValueType
+
+from modelsec.models import Finding, Severity
+
+GGUF_EXTENSIONS = {".gguf"}
+
+# No dedicated OSS GGUF security scanner exists (see PROPOSAL.md gap analysis).
+# We reuse gguf-py (llama.cpp) for spec-correct header/KV parsing and layer our
+# own anomaly checks on top, rather than hand-rolling a parser.
+MAX_STRING_KV_BYTES = 1_000_000  # 1MB — legitimate KV strings (names, templates) are tiny
+MAX_ARRAY_KV_LEN = 1_000_000
+KNOWN_QUANT_TYPES = {t.value for t in gguf.GGMLQuantizationType}
+
+
+class GgufDetector:
+    """Static checks on GGUF header/KV metadata for anomalies that could
+    indicate a crafted file: oversized KV blobs, unknown quantization types,
+    or tensor metadata inconsistent with the file's actual size."""
+
+    name = "gguf-header-anomaly"
+
+    def applies_to(self, path: Path) -> bool:
+        return path.suffix.lower() in GGUF_EXTENSIONS
+
+    def scan(self, path: Path) -> list[Finding]:
+        try:
+            reader = gguf.GGUFReader(str(path))
+        except Exception as exc:  # noqa: BLE001
+            return [
+                Finding(
+                    detector=self.name,
+                    severity=Severity.HIGH,
+                    title="Malformed or non-standard GGUF header",
+                    description=(
+                        f"gguf-py failed to parse this file's header/KV metadata: {exc}. "
+                        "A file with the .gguf extension that doesn't conform to the GGUF "
+                        "spec may be attempting to evade format-based scanning."
+                    ),
+                    file=str(path),
+                    mitigation=(
+                        "Do not load this file with llama.cpp/Ollama. Verify it is a "
+                        "genuine GGUF export from a trusted conversion tool."
+                    ),
+                )
+            ]
+
+        findings: list[Finding] = []
+        file_size = path.stat().st_size
+
+        for name, field in reader.fields.items():
+            if not field.types:
+                continue
+            value_type = field.types[-1]
+            if value_type == GGUFValueType.STRING:
+                length = sum(len(p) for p in field.parts if hasattr(p, "__len__"))
+                if length > MAX_STRING_KV_BYTES:
+                    findings.append(
+                        _oversized_kv_finding(path, name, f"string KV of {length} bytes")
+                    )
+            elif value_type == GGUFValueType.ARRAY and len(field.parts) > MAX_ARRAY_KV_LEN:
+                findings.append(
+                    _oversized_kv_finding(path, name, f"array KV of {len(field.parts)} elements")
+                )
+
+        for tensor in reader.tensors:
+            if int(tensor.tensor_type) not in KNOWN_QUANT_TYPES:
+                findings.append(
+                    Finding(
+                        detector=self.name,
+                        severity=Severity.HIGH,
+                        title="Unknown tensor quantization type",
+                        description=(
+                            f"Tensor '{tensor.name}' declares quantization type "
+                            f"{tensor.tensor_type!r}, which is not a recognized GGML "
+                            "quantization type. This may indicate a hand-crafted file "
+                            "targeting a parser bug in a specific loader version."
+                        ),
+                        file=str(path),
+                        evidence=f"tensor={tensor.name} type={tensor.tensor_type!r}",
+                        mitigation=(
+                            "Do not load this model with an out-of-date or forked GGUF "
+                            "loader. Re-export from a trusted source using current llama.cpp "
+                            "conversion tooling."
+                        ),
+                    )
+                )
+            # A tensor claiming far more elements than the whole file could hold
+            # is a classic crafted-header trick to trigger an oversized allocation
+            # or out-of-bounds read in a loader.
+            if tensor.n_elements > 0 and tensor.n_elements > file_size:
+                findings.append(
+                    Finding(
+                        detector=self.name,
+                        severity=Severity.CRITICAL,
+                        title="Tensor element count exceeds file size",
+                        description=(
+                            f"Tensor '{tensor.name}' claims {tensor.n_elements} elements, "
+                            f"which exceeds the file size ({file_size} bytes). This is "
+                            "consistent with a crafted header aimed at triggering an "
+                            "out-of-bounds read or oversized allocation in a loader."
+                        ),
+                        file=str(path),
+                        evidence=f"tensor={tensor.name} n_elements={tensor.n_elements} file_size={file_size}",
+                        mitigation=(
+                            "Do not load this file. Report it as a likely malicious/crafted "
+                            "GGUF artifact."
+                        ),
+                    )
+                )
+
+        return findings
+
+
+def _oversized_kv_finding(path: Path, key: str, detail: str) -> Finding:
+    return Finding(
+        detector="gguf-header-anomaly",
+        severity=Severity.MEDIUM,
+        title="Oversized metadata key-value entry",
+        description=(
+            f"KV key '{key}' contains an unusually large {detail}. Legitimate GGUF "
+            "metadata (names, chat templates, tokenizer config) is normally small; "
+            "an oversized blob may be an attempt to smuggle payload data or trigger "
+            "a parser resource-exhaustion bug."
+        ),
+        file=str(path),
+        evidence=f"key={key} {detail}",
+        mitigation="Inspect this KV entry's content manually before loading the model.",
+    )

weightguard-0.1.0/src/modelsec/detectors/keras_detector.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import h5py
+
+from modelsec.models import Finding, Severity
+
+KERAS_EXTENSIONS = {".h5", ".hdf5", ".keras"}
+
+# Layer types that embed and execute arbitrary Python at deserialization time.
+DANGEROUS_LAYER_TYPES = {"Lambda"}
+
+_MITIGATION = (
+    "Do not call keras.models.load_model() on this file. Inspect the model_config "
+    "JSON for the flagged layer(s), rewrite them as named functions registered via "
+    "@keras.saving.register_keras_serializable, or reject the model if you cannot "
+    "verify the embedded code."
+)
+
+
+class KerasDetector:
+    """Keras H5/.keras files can embed a serialized Python function inside a
+    Lambda layer's config, which is deserialized (and effectively executed)
+    on load. This is the same technique ModelScan's Keras checks target."""
+
+    name = "keras-lambda-layer"
+
+    def applies_to(self, path: Path) -> bool:
+        return path.suffix.lower() in KERAS_EXTENSIONS
+
+    def scan(self, path: Path) -> list[Finding]:
+        try:
+            with h5py.File(path, "r") as f:
+                config_raw = f.attrs.get("model_config")
+        except Exception as exc:  # noqa: BLE001
+            return [
+                Finding(
+                    detector=self.name,
+                    severity=Severity.INFO,
+                    title="Could not parse as Keras HDF5",
+                    description=str(exc),
+                    file=str(path),
+                    mitigation="Verify this is a valid Keras H5 file before loading.",
+                )
+            ]
+
+        if not config_raw:
+            return []
+
+        if isinstance(config_raw, bytes):
+            config_raw = config_raw.decode("utf-8", errors="replace")
+
+        try:
+            config = json.loads(config_raw)
+        except json.JSONDecodeError:
+            return []
+
+        findings: list[Finding] = []
+        for layer in _iter_layers(config):
+            layer_type = layer.get("class_name")
+            if layer_type in DANGEROUS_LAYER_TYPES:
+                findings.append(
+                    Finding(
+                        detector=self.name,
+                        severity=Severity.CRITICAL,
+                        title=f"{layer_type} layer embeds arbitrary code",
+                        description=(
+                            f"Layer '{layer.get('config', {}).get('name', '?')}' is a "
+                            f"{layer_type} layer, which stores a marshalled Python function "
+                            "that executes when the model is loaded."
+                        ),
+                        file=str(path),
+                        evidence=json.dumps(layer.get("config", {}))[:300],
+                        mitigation=_MITIGATION,
+                    )
+                )
+        return findings
+
+
+def _iter_layers(config: dict) -> list[dict]:
+    layers = config.get("config", {}).get("layers", [])
+    if isinstance(layers, list):
+        return layers
+    return []

weightguard-0.1.0/src/modelsec/detectors/onnx_detector.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import onnx
+
+from modelsec.models import Finding, Severity
+
+ONNX_EXTENSIONS = {".onnx"}
+
+# Ops in these domains are provided by ONNX Runtime's built-in kernels.
+# Anything else means a custom-op shared library must be loaded at runtime.
+STANDARD_DOMAINS = {"", "ai.onnx", "ai.onnx.ml", "ai.onnx.preview.training", "ai.onnx.training"}
+
+_MITIGATION = (
+    "Custom ops require ai.onnx Runtime to load an external native (.so/.dll) "
+    "library at inference time — this is arbitrary native code execution. Verify "
+    "the referenced custom-op library's provenance before running this model, or "
+    "run inference only inside an isolated sandbox with no filesystem/network access."
+)
+
+
+class OnnxDetector:
+    """Flags ONNX graphs that reference non-standard operator domains, i.e.
+    custom ops capable of loading arbitrary native code at inference time."""
+
+    name = "onnx-custom-op"
+
+    def applies_to(self, path: Path) -> bool:
+        return path.suffix.lower() in ONNX_EXTENSIONS
+
+    def scan(self, path: Path) -> list[Finding]:
+        try:
+            model = onnx.load(str(path), load_external_data=False)
+        except Exception as exc:  # noqa: BLE001
+            return [
+                Finding(
+                    detector=self.name,
+                    severity=Severity.INFO,
+                    title="Could not parse as ONNX",
+                    description=str(exc),
+                    file=str(path),
+                    mitigation="Verify this is a valid ONNX model file before loading.",
+                )
+            ]
+
+        custom_domains = {
+            opset.domain for opset in model.opset_import if opset.domain not in STANDARD_DOMAINS
+        }
+        custom_node_domains = {
+            node.domain for node in model.graph.node if node.domain not in STANDARD_DOMAINS
+        }
+        all_custom = custom_domains | custom_node_domains
+
+        if not all_custom:
+            return []
+
+        return [
+            Finding(
+                detector=self.name,
+                severity=Severity.HIGH,
+                title="Non-standard operator domain(s) referenced",
+                description=(
+                    "This ONNX graph references custom-op domain(s) not provided by "
+                    "ONNX Runtime's built-in kernel set: " + ", ".join(sorted(all_custom))
+                ),
+                file=str(path),
+                evidence=", ".join(sorted(all_custom)),
+                mitigation=_MITIGATION,
+            )
+        ]

weightguard-0.1.0/src/modelsec/detectors/pickle_detector.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from fickling.analysis import Severity as FicklingSeverity
+from fickling.analysis import check_safety
+from fickling.fickle import Pickled
+
+from modelsec.models import Finding, Severity
+
+PICKLE_EXTENSIONS = {".pkl", ".pickle", ".bin", ".pt", ".pth"}
+
+# Fickling's Severity enum instances are not hashable, so map by rank (the
+# int stored in Severity.value[0]) rather than by dict key.
+_SEVERITY_BY_RANK = {
+    FicklingSeverity.LIKELY_SAFE.value[0]: None,
+    FicklingSeverity.POSSIBLY_UNSAFE.value[0]: Severity.LOW,
+    FicklingSeverity.SUSPICIOUS.value[0]: Severity.MEDIUM,
+    FicklingSeverity.LIKELY_UNSAFE.value[0]: Severity.HIGH,
+    FicklingSeverity.LIKELY_OVERTLY_MALICIOUS.value[0]: Severity.CRITICAL,
+    FicklingSeverity.OVERTLY_MALICIOUS.value[0]: Severity.CRITICAL,
+}
+
+_MITIGATION = (
+    "Do not call torch.load()/pickle.load() on this file directly. Use "
+    "`torch.load(..., weights_only=True)`, migrate the model to SafeTensors, "
+    "or inspect it further with `fickling --check-safety` before loading."
+)
+
+
+class PickleDetector:
+    """Wraps Trail of Bits' Fickling for AST-level pickle-bytecode analysis.
+
+    Fickling is used rather than a hand-rolled opcode allowlist because
+    pattern/denylist scanners (e.g. picklescan) are known to be bypassable
+    via malformed-but-valid opcode streams (the "nullifAI" technique).
+    """
+
+    name = "pickle-fickling"
+
+    def applies_to(self, path: Path) -> bool:
+        return path.suffix.lower() in PICKLE_EXTENSIONS
+
+    def scan(self, path: Path) -> list[Finding]:
+        try:
+            data = path.read_bytes()
+        except OSError as exc:
+            return [
+                Finding(
+                    detector=self.name,
+                    severity=Severity.INFO,
+                    title="Could not read file",
+                    description=str(exc),
+                    file=str(path),
+                    mitigation="Verify the file is accessible and not corrupted.",
+                )
+            ]
+
+        try:
+            pickled = Pickled.load(data)
+            result = check_safety(pickled)
+        except Exception as exc:  # noqa: BLE001 - malformed pickle is itself a signal
+            return [
+                Finding(
+                    detector=self.name,
+                    severity=Severity.HIGH,
+                    title="Malformed or non-standard pickle stream",
+                    description=(
+                        "The pickle bytecode could not be fully disassembled by Fickling. "
+                        f"Malformed opcode streams are a known scanner-evasion technique: {exc}"
+                    ),
+                    file=str(path),
+                    mitigation=_MITIGATION,
+                )
+            ]
+
+        mapped = _SEVERITY_BY_RANK[result.severity.value[0]]
+        if mapped is None:
+            return []
+
+        return [
+            Finding(
+                detector=self.name,
+                severity=mapped,
+                title=f"Fickling analysis: {result.severity.value[1]}",
+                description=(
+                    result.to_string()
+                    or "Fickling could not identify overtly unsafe code, but the pickle "
+                    "may still be unsafe if the source is untrusted."
+                ),
+                file=str(path),
+                evidence="; ".join(f"{k}: {v}" for k, v in result.detailed_results().get("AnalysisResult", {}).items()),
+                mitigation=_MITIGATION,
+            )
+        ]

weightguard-0.1.0/src/modelsec/detectors/safetensors_detector.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from safetensors import safe_open
+
+from modelsec.models import Finding, Severity
+
+SAFETENSORS_EXTENSIONS = {".safetensors"}
+
+
+class SafeTensorsDetector:
+    """SafeTensors has no exec surface (header is JSON metadata + raw tensor
+    bytes, no opcode/callable resolution) so a well-formed file is safe by
+    construction. This detector only flags files that fail to parse as valid
+    SafeTensors, since that's either corruption or a format-spoofing attempt
+    (e.g. a renamed pickle)."""
+
+    name = "safetensors-format-check"
+
+    def applies_to(self, path: Path) -> bool:
+        return path.suffix.lower() in SAFETENSORS_EXTENSIONS
+
+    def scan(self, path: Path) -> list[Finding]:
+        try:
+            with safe_open(str(path), framework="numpy") as f:
+                list(f.keys())
+        except Exception as exc:  # noqa: BLE001
+            return [
+                Finding(
+                    detector=self.name,
+                    severity=Severity.HIGH,
+                    title="File claims .safetensors extension but is not valid SafeTensors",
+                    description=(
+                        f"Failed to parse as SafeTensors: {exc}. This may be a "
+                        "misnamed/spoofed file (e.g. a pickle renamed to evade "
+                        "extension-based scanning)."
+                    ),
+                    file=str(path),
+                    mitigation=(
+                        "Do not load this file with a SafeTensors loader assuming it is "
+                        "safe. Re-identify the real format via magic bytes and re-scan."
+                    ),
+                )
+            ]
+        return []

weightguard-0.1.0/src/modelsec/models.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class Severity(str, Enum):
+    CRITICAL = "CRITICAL"
+    HIGH = "HIGH"
+    MEDIUM = "MEDIUM"
+    LOW = "LOW"
+    INFO = "INFO"
+
+    @property
+    def rank(self) -> int:
+        return {
+            Severity.CRITICAL: 4,
+            Severity.HIGH: 3,
+            Severity.MEDIUM: 2,
+            Severity.LOW: 1,
+            Severity.INFO: 0,
+        }[self]
+
+
+class Finding(BaseModel):
+    detector: str
+    severity: Severity
+    title: str
+    description: str
+    file: str | None = None
+    evidence: str | None = None
+    mitigation: str
+
+
+class Report(BaseModel):
+    target: str
+    findings: list[Finding] = Field(default_factory=list)
+
+    @property
+    def max_severity(self) -> Severity:
+        if not self.findings:
+            return Severity.INFO
+        return max((f.severity for f in self.findings), key=lambda s: s.rank)
+
+    def fails(self, threshold: Severity) -> bool:
+        return self.max_severity.rank >= threshold.rank

weightguard-0.1.0/src/modelsec/resolver.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import re
+import subprocess
+import tempfile
+from pathlib import Path
+from urllib.parse import urlparse
+
+HF_HOST = "huggingface.co"
+_HF_REPO_RE = re.compile(r"^/(?P<repo_id>[^/]+/[^/]+?)(?:/tree/[^/]+)?/?$")
+
+
+class UnresolvableTarget(Exception):
+    pass
+
+
+def resolve(target: str) -> Path:
+    """Resolve a CLI target (HF repo URL, generic git URL, or local path) to a
+    local directory. Never executes/imports any downloaded file — this stage
+    only fetches bytes to disk."""
+    parsed = urlparse(target)
+
+    if parsed.scheme in ("http", "https") and parsed.netloc == HF_HOST:
+        return _resolve_huggingface(parsed.path)
+
+    if parsed.scheme in ("http", "https", "git", "ssh") or target.endswith(".git"):
+        return _resolve_git(target)
+
+    path = Path(target)
+    if path.exists():
+        return path
+
+    raise UnresolvableTarget(f"Could not resolve target: {target!r} is not a valid URL or local path")
+
+
+def _resolve_huggingface(url_path: str) -> Path:
+    match = _HF_REPO_RE.match(url_path)
+    if not match:
+        raise UnresolvableTarget(f"Could not parse Hugging Face repo id from path: {url_path!r}")
+    repo_id = match.group("repo_id")
+
+    from huggingface_hub import snapshot_download
+
+    local_dir = Path(tempfile.mkdtemp(prefix="modelsec-hf-"))
+    snapshot_download(repo_id=repo_id, local_dir=str(local_dir))
+    return local_dir
+
+
+def _resolve_git(url: str) -> Path:
+    local_dir = Path(tempfile.mkdtemp(prefix="modelsec-git-"))
+    subprocess.run(
+        ["git", "clone", "--depth", "1", url, str(local_dir)],
+        check=True,
+        capture_output=True,
+    )
+    return local_dir

weightguard-0.1.0/src/modelsec/scanner.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from modelsec.detectors import ALL_DETECTORS
+from modelsec.models import Report
+
+
+def scan_path(target: Path) -> Report:
+    """Static scan of a local directory or file. Never executes/deserializes content."""
+    report = Report(target=str(target))
+    files = [target] if target.is_file() else [p for p in target.rglob("*") if p.is_file()]
+
+    for file in files:
+        for detector in ALL_DETECTORS:
+            if detector.applies_to(file):
+                report.findings.extend(detector.scan(file))
+
+    return report