/**
 * Model Supply Chain Security Test Fixtures
 * Tests for detecting unsafe model loading and supply chain risks
 * AI Detection Roadmap Phase 2
 */

import type { TestGroup } from '../../types'

export const modelSupplyChainTests: TestGroup = {
  name: 'Model Supply Chain Security',
  tier: 'A',
  layer: 2,
  description: 'Detection of unsafe model loading patterns that can lead to RCE',

  truePositives: [
    {
      name: 'Model Supply Chain - Pickle RCE',
      expectFindings: true,
      expectedCategories: ['ai_unsafe_model_load'],
      description: 'Unsafe pickle/joblib deserialization that MUST be detected',
      file: {
        path: 'src/ml/model_loader.py',
        content: `
import pickle
import joblib
import torch

# Pickle deserialization - CRITICAL RCE risk
def load_user_model(model_path):
    """Load a model file uploaded by user"""
    with open(model_path, 'rb') as f:
        model = pickle.load(f)  # CRITICAL: Arbitrary code execution
    return model

# Joblib deserialization - CRITICAL RCE risk
def load_sklearn_model(model_path):
    """Load sklearn model from user upload"""
    model = joblib.load(model_path)  # CRITICAL: Uses pickle internally
    return model

# Torch.load without weights_only - HIGH
def load_pytorch_model(checkpoint_path):
    """Load PyTorch checkpoint"""
    checkpoint = torch.load(checkpoint_path)  # HIGH: Can execute code
    return checkpoint['model']
`,
        language: 'python',
        size: 600,
      },
    },
    {
      name: 'Model Supply Chain - Unverified Sources',
      expectFindings: true,
      expectedCategories: ['ai_unverified_model', 'ai_unsafe_model_load'],
      description: 'Models loaded from untrusted sources that MUST be detected',
      file: {
        path: 'src/ml/download_model.py',
        content: `
from transformers import AutoModelForCausalLM, AutoTokenizer
import urllib.request

# Model from HTTP URL - HIGH (no integrity verification)
def load_from_url():
    model = AutoModelForCausalLM.from_pretrained("http://malicious-server.com/model")
    return model

# trust_remote_code=True - HIGH (executes arbitrary code)
def load_with_remote_code():
    model = AutoModelForCausalLM.from_pretrained(
        "unknown-org/untrusted-model",
        trust_remote_code=True  # HIGH: Executes arbitrary Python
    )
    return model

# Model download without checksum - MEDIUM
def download_raw_model():
    urllib.request.urlretrieve(
        "https://example.com/model.pth",  # No checksum verification
        "model.pth"
    )
`,
        language: 'python',
        size: 500,
      },
    },
    {
      name: 'Model Supply Chain - Unsafe Fine-tuning',
      expectFindings: true,
      expectedCategories: ['ai_unsafe_finetuning'],
      description: 'Unsafe fine-tuning on user data that MUST be detected',
      file: {
        path: 'src/ml/finetune.py',
        content: `
from transformers import Trainer, TrainingArguments

# Training on user uploads - HIGH (data poisoning risk)
def finetune_on_user_data(user_uploads):
    """Fine-tune model on user-uploaded data"""
    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=3,
    )
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=user_uploads,  # HIGH: Unvalidated user data
    )
    trainer.train()  # Training on untrusted data
`,
        language: 'python',
        size: 400,
      },
    },
    {
      name: 'Model Supply Chain - Unverified Downloads',
      expectFindings: true,
      expectedCategories: ['ai_unverified_model'],
      description: 'Model downloads without integrity verification that MUST be detected',
      file: {
        path: 'src/ml/model_downloader.py',
        content: `
import requests
import wget
from transformers import AutoModelForCausalLM

# Direct model download without checksum - MEDIUM
def download_model_unsafe():
    # Using requests to download model file
    response = requests.get("https://random-server.com/model.pth")
    with open("model.pth", "wb") as f:
        f.write(response.content)  # No integrity check!

# wget without verification - MEDIUM
def download_with_wget():
    wget.download("http://untrusted-host.com/model.bin", "model.bin")

# Model from HTTP (not HTTPS) - HIGH
def load_from_http():
    model = AutoModelForCausalLM.from_pretrained("http://insecure-host.com/models/gpt")
    return model

# Model from unknown Hugging Face org - MEDIUM (no integrity pin)
def load_unknown_model():
    # No revision pin - could be modified by repo owner
    model = AutoModelForCausalLM.from_pretrained("random-unknown-user/suspicious-model")
    return model
`,
        language: 'python',
        size: 700,
      },
    },
  ],

  falseNegatives: [
    {
      name: 'Model Supply Chain - Safe Patterns',
      expectFindings: false,
      description: 'Safe model loading patterns that should NOT be flagged',
      file: {
        path: 'src/ml/safe_model_loader.py',
        content: `
import torch
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM
import hashlib

# SafeTensors - SAFE (no code execution)
def load_safe_model(model_path):
    """Load model using SafeTensors format"""
    state_dict = load_file(model_path)
    model.load_state_dict(state_dict)
    return model

# torch.load with weights_only - SAFE
def load_pytorch_safe(checkpoint_path):
    """Load PyTorch checkpoint safely"""
    checkpoint = torch.load(checkpoint_path, weights_only=True)
    return checkpoint

# Model from trusted source with verification - SAFE
def load_from_huggingface():
    """Load from official Hugging Face repo with pinned revision"""
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-2-7b",
        revision="a1b2c3d4e5f6"  # Pinned to specific commit
    )
    return model

# Model with checksum verification - SAFE
def load_with_verification(model_path, expected_sha256):
    """Load model after verifying checksum"""
    with open(model_path, 'rb') as f:
        file_hash = hashlib.sha256(f.read()).hexdigest()
    if file_hash != expected_sha256:
        raise ValueError("Checksum mismatch!")
    # Safe to load after verification
    return torch.load(model_path, weights_only=True)
`,
        language: 'python',
        size: 800,
      },
    },
  ],
}
