"use strict";
/**
 * Layer 2: Model Supply Chain Security Detection
 * Detects unsafe model loading patterns that can lead to arbitrary code execution
 *
 * Covers AI Detection Roadmap Phase 2:
 * - Pickle/joblib deserialization RCE
 * - torch.load without weights_only=True
 * - Unverified model sources
 * - Unsafe fine-tuning on user data
 *
 * References:
 * - OWASP LLM05: Supply Chain Vulnerabilities
 * - CWE-502: Deserialization of Untrusted Data
 */
Object.defineProperty(exports, "__esModule", { value: true });
exports.detectModelSupplyChain = detectModelSupplyChain;
const file_classifier_1 = require("../../parse/file-classifier");
const BASE_CONFIDENCE = 0.50;
// ============================================================================
// Context Detection
// ============================================================================
/**
 * Check if file is in an ML/model context based on path and content
 */
function isMLContextFile(filePath, content) {
    // File path indicators
    const mlPathPatterns = [
        /\/(models?|ml|training|inference|weights|checkpoints)\//i,
        /\/(transformers|huggingface|pytorch|tensorflow|keras)\//i,
        /(train|model|inference|fine[-_]?tune).*\.(py|ts|js)$/i,
    ];
    if (mlPathPatterns.some(p => p.test(filePath))) {
        return true;
    }
    // Content patterns suggesting ML usage
    const mlContentPatterns = [
        /import\s+(?:torch|tensorflow|keras|transformers|joblib|pickle|dill|cloudpickle|onnx)/i,
        /from\s+(?:torch|tensorflow|keras|transformers|joblib|pickle|dill|cloudpickle|onnx)\s+import/i,
        /\.load_model\s*\(/i,
        /\.from_pretrained\s*\(/i,
        /torch\.load\s*\(/i,
        /torch\.jit\.load\s*\(/i,
        /pickle\.load/i,
        /joblib\.load/i,
        /dill\.load/i,
        /cloudpickle\.load/i,
        /onnx\.load/i,
        /numpy\.load.*allow_pickle/i,
        /np\.load.*allow_pickle/i,
        /yaml\.(?:unsafe_load|full_load|load)/i,
        /Trainer|TrainingArguments/i,
        /model\.save|model\.load/i,
    ];
    return mlContentPatterns.some(p => p.test(content));
}
// ============================================================================
// Safe Pattern Detection
// ============================================================================
/**
 * Check if torch.load has weights_only=True (safe mode)
 */
function hasTorchWeightsOnly(lineContent, surroundingContext) {
    const fullContext = lineContent + '\n' + surroundingContext;
    return /weights_only\s*=\s*True/i.test(fullContext);
}
/**
 * Check if TensorFlow/Keras load has safe_mode=True
 */
function hasKerasSafeMode(lineContent, surroundingContext) {
    const fullContext = lineContent + '\n' + surroundingContext;
    return /safe_mode\s*=\s*True/i.test(fullContext);
}
/**
 * Check if model source is from a trusted provider
 */
function isTrustedModelSource(lineContent) {
    const trustedPatterns = [
        // Official Hugging Face repos
        /huggingface\.co\/(meta-llama|openai|anthropic|google|microsoft|facebook|EleutherAI)/i,
        // Specific trusted model IDs (no http prefix)
        /['"`](meta-llama|openai|anthropic|google|microsoft|facebook)\//i,
        // Local file paths (not URLs)
        /['"`]\.\/|['"`]\.\.\/|['"`]\//,
        // Environment variables for paths
        /os\.environ|process\.env/i,
    ];
    return trustedPatterns.some(p => p.test(lineContent));
}
/**
 * Check if integrity verification is present
 */
function hasIntegrityCheck(content, lineNumber) {
    const lines = content.split('\n');
    const contextStart = Math.max(0, lineNumber - 20);
    const contextEnd = Math.min(lines.length, lineNumber + 10);
    const context = lines.slice(contextStart, contextEnd).join('\n');
    const integrityPatterns = [
        /checksum|sha256|sha512|md5|verify_hash|hash_check/i,
        /verify.*integrity|integrity.*verify/i,
        /revision\s*=\s*['"`][a-f0-9]{40}/i, // Git SHA pinning
        /safetensors/i, // SafeTensors format (safe by design)
        /\.ckpt\.sha256|\.bin\.sha256/i, // Hash files
    ];
    return integrityPatterns.some(p => p.test(context));
}
/**
 * Check if using SafeTensors format (safe by design)
 */
function usesSafeTensors(lineContent, surroundingContext) {
    const fullContext = lineContent + '\n' + surroundingContext;
    return /safetensors|\.safetensors|from_safetensors|load_safetensors/i.test(fullContext);
}
/**
 * Check if data validation is present for training
 */
function hasDataValidation(content, lineNumber) {
    const lines = content.split('\n');
    const contextStart = Math.max(0, lineNumber - 30);
    const contextEnd = Math.min(lines.length, lineNumber + 10);
    const context = lines.slice(contextStart, contextEnd).join('\n');
    const validationPatterns = [
        /validate|sanitize|clean|filter/i,
        /schema\.(parse|validate)/i,
        /content_moderation|moderate_content/i,
        /review.*data|data.*review/i,
        /verified.*dataset|trusted.*data/i,
    ];
    return validationPatterns.some(p => p.test(context));
}
const MODEL_SUPPLY_CHAIN_PATTERNS = [
    // ========== Pickle Deserialization (RCE) ==========
    {
        name: 'Pickle deserialization',
        pattern: /pickle\.load\s*\(|pickle\.loads\s*\(/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'critical',
        description: 'pickle.load() executes arbitrary Python code embedded in pickle files. Attackers can craft malicious pickle files that execute code when loaded.',
        suggestedFix: 'Use SafeTensors format for ML models. For other data, use JSON or a safe serialization format. If pickle is unavoidable, only load from trusted, verified sources.',
    },
    {
        name: 'Joblib deserialization',
        pattern: /joblib\.load\s*\(/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'critical',
        description: 'joblib.load() uses pickle internally and can execute arbitrary code. This is commonly used for sklearn models but is unsafe for untrusted files.',
        suggestedFix: 'Use ONNX format for sklearn models. Alternatively, use skops.io which provides secure model persistence. Only load from verified sources.',
    },
    // ========== Extended Pickle Variants (RCE) ==========
    {
        name: 'Dill deserialization',
        pattern: /\bdill\.(load|loads|load_session)\s*\(/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'critical',
        description: 'dill can execute arbitrary code during deserialization, similar to pickle but with extended capabilities including serializing lambdas and closures.',
        suggestedFix: 'Use SafeTensors format or JSON serialization instead of dill. If dill is unavoidable, only load from trusted, verified sources.',
    },
    {
        name: 'Cloudpickle deserialization',
        pattern: /\bcloudpickle\.(load|loads)\s*\(/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'critical',
        description: 'cloudpickle can serialize and execute arbitrary functions, enabling remote code execution during deserialization.',
        suggestedFix: 'Use SafeTensors format or explicitly define functions instead of deserializing them. Avoid cloudpickle for untrusted data.',
    },
    // ========== YAML Unsafe Loading (RCE) ==========
    {
        name: 'YAML unsafe deserialization',
        pattern: /\byaml\.(unsafe_load|full_load|UnsafeLoader)\s*\(/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'critical',
        description: 'yaml.unsafe_load() and yaml.full_load() can instantiate arbitrary Python objects, enabling remote code execution.',
        suggestedFix: 'Use yaml.safe_load() instead of yaml.unsafe_load() or yaml.full_load(). SafeLoader only loads basic Python types.',
    },
    {
        name: 'YAML load without explicit Loader',
        pattern: /\byaml\.load\s*\(\s*[^,)]+\s*\)(?!\s*,)/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'medium',
        description: 'yaml.load() without explicit Loader argument may use unsafe loading in older PyYAML versions (< 5.1).',
        suggestedFix: 'Use yaml.safe_load() or explicitly specify Loader=yaml.SafeLoader: yaml.load(file, Loader=yaml.SafeLoader)',
    },
    // ========== NumPy Pickle Loading (RCE) ==========
    {
        name: 'NumPy pickle loading',
        pattern: /\b(?:np|numpy)\.load\s*\([^)]*allow_pickle\s*=\s*True/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'high',
        description: 'numpy.load() with allow_pickle=True can execute arbitrary code embedded in .npy/.npz files via pickle.',
        suggestedFix: 'Use allow_pickle=False (default in numpy >= 1.16.3) or use SafeTensors format. Only allow_pickle from verified sources.',
    },
    // ========== TorchScript Loading (RCE) ==========
    {
        name: 'TorchScript model loading',
        pattern: /\btorch\.jit\.load\s*\(/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'high',
        description: 'torch.jit.load() can execute arbitrary code embedded in TorchScript models. TorchScript files can contain custom operators with native code.',
        suggestedFix: 'Use SafeTensors format or verify model source and integrity before loading. Pin to specific model revisions with checksums.',
    },
    // ========== ONNX Model Loading (Code Execution Risk) ==========
    {
        name: 'ONNX model loading',
        pattern: /\bonnx\.load\s*\(/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'medium',
        description: 'ONNX models with custom operators can execute arbitrary code. Custom ops are loaded as shared libraries.',
        suggestedFix: 'Verify ONNX model source and check for custom operators before loading. Use onnx.checker.check_model() and inspect custom ops.',
    },
    // ========== PyTorch Loading ==========
    {
        name: 'torch.load without weights_only',
        pattern: /torch\.load\s*\([^)]*\)/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'high',
        description: 'torch.load() without weights_only=True can execute arbitrary code embedded in model files via pickle.',
        suggestedFix: 'Use torch.load(path, weights_only=True) to load only tensor data safely. Or use SafeTensors format: from safetensors.torch import load_file',
        checkSafeMode: 'torch',
    },
    // ========== TensorFlow/Keras Loading ==========
    {
        name: 'Keras load_model without safe_mode',
        pattern: /(?:keras\.models\.)?load_model\s*\([^)]*\)/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'high',
        description: 'Keras load_model() can execute arbitrary code from Lambda layers or custom objects in model files.',
        suggestedFix: 'Use tf.keras.models.load_model(path, safe_mode=True) in TensorFlow 2.13+. Alternatively, save/load only weights with model.save_weights()/load_weights().',
        checkSafeMode: 'keras',
    },
    // ========== Unverified Model Sources ==========
    {
        name: 'Model from HTTP URL',
        pattern: /from_pretrained\s*\(\s*['"`]https?:\/\/[^'"`)]+['"`]/gi,
        category: 'ai_unverified_model',
        baseSeverity: 'high',
        description: 'Loading models directly from HTTP URLs bypasses the Hugging Face Hub\'s verification. Attackers could serve malicious models via MITM attacks.',
        suggestedFix: 'Load from Hugging Face Hub by model ID instead of direct URL. Pin to specific revision: from_pretrained("org/model", revision="abc123")',
        checkTrustedSource: true,
    },
    {
        name: 'trust_remote_code=True',
        pattern: /trust_remote_code\s*=\s*True/gi,
        category: 'ai_unsafe_model_load',
        baseSeverity: 'high',
        description: 'trust_remote_code=True allows model repos to execute arbitrary Python code during loading. Malicious repos could compromise your system.',
        suggestedFix: 'Avoid trust_remote_code=True. Use models that don\'t require custom code. If necessary, audit the model repo code before enabling.',
    },
    {
        name: 'Model download without verification',
        pattern: /(?:urllib|requests|wget|curl).*(?:\.pth|\.pt|\.bin|\.ckpt|\.h5|\.keras|model)/gi,
        category: 'ai_unverified_model',
        baseSeverity: 'medium',
        description: 'Downloading model files without integrity verification. Attackers could serve modified files via compromised mirrors or MITM.',
        suggestedFix: 'Verify downloaded files against known checksums: sha256sum file.pth. Use official download APIs (transformers, torch.hub) which include verification.',
        checkIntegrity: true,
    },
    // ========== Unsafe Fine-tuning ==========
    {
        name: 'Training on user uploads',
        pattern: /(?:trainer|model)\.(?:train|fit|fine_tune)\s*\([^)]*(?:user_uploads?|user_data|uploaded|untrusted)/gi,
        category: 'ai_unsafe_finetuning',
        baseSeverity: 'high',
        description: 'Training/fine-tuning directly on user-uploaded data without validation enables data poisoning attacks.',
        suggestedFix: 'Validate and sanitize all training data. Implement content moderation. Use separate data pipelines for user content with review steps.',
        checkDataValidation: true,
    },
    {
        name: 'Trainer with user data',
        pattern: /Trainer\s*\([^)]*(?:train_dataset|eval_dataset)\s*=\s*(?:user_uploads?|user_data|uploaded_data|untrusted_data)/gi,
        category: 'ai_unsafe_finetuning',
        baseSeverity: 'high',
        description: 'Trainer initialized with user-uploaded data without validation. Data poisoning attacks can manipulate model behavior.',
        suggestedFix: 'Validate and sanitize all training data before passing to Trainer. Implement content moderation and review pipelines.',
        checkDataValidation: true,
    },
    {
        name: 'Fine-tuning with external dataset',
        pattern: /(?:load_dataset|datasets\.load)\s*\([^)]*(?:path|url|http)/gi,
        category: 'ai_unsafe_finetuning',
        baseSeverity: 'medium',
        description: 'Loading training datasets from external sources without verification. Data could be poisoned.',
        suggestedFix: 'Use verified datasets from trusted sources. Implement data validation pipelines. Consider dataset hashing/signing.',
        checkDataValidation: true,
    },
    {
        name: 'Training config from user input',
        pattern: /TrainingArguments\s*\([^)]*(?:request|user|input|body)\./gi,
        category: 'ai_unsafe_finetuning',
        baseSeverity: 'medium',
        description: 'Training arguments derived from user input. Attackers could manipulate training hyperparameters to compromise model behavior.',
        suggestedFix: 'Validate and sanitize all training arguments. Use allowlists for hyperparameter values. Don\'t allow users to specify output paths.',
    },
    // ========== TensorFlow Specific ==========
    {
        name: 'TensorFlow SavedModel from path',
        pattern: /tf\.saved_model\.load\s*\([^)]*(?:http|ftp|user|upload)/gi,
        category: 'ai_unverified_model',
        baseSeverity: 'high',
        description: 'Loading TensorFlow SavedModel from untrusted path. SavedModels can contain arbitrary ops that execute code.',
        suggestedFix: 'Only load SavedModels from trusted, verified sources. Use TF Serving with model verification for production.',
        checkTrustedSource: true,
    },
];
// ============================================================================
// Helper Functions
// ============================================================================
/**
 * Get surrounding context for analysis
 */
function getSurroundingContext(content, lineIndex, windowSize = 15) {
    const lines = content.split('\n');
    const start = Math.max(0, lineIndex - windowSize);
    const end = Math.min(lines.length, lineIndex + windowSize);
    return lines.slice(start, end).join('\n');
}
/**
 * Calculate severity based on mitigations
 */
function calculateSeverity(baseSeverity, isMitigated, isPartiallyMitigated, isTestFile, isExample, isLibrary) {
    if (isTestFile || isExample || isLibrary) {
        return 'info';
    }
    if (isMitigated) {
        return 'info';
    }
    if (isPartiallyMitigated) {
        if (baseSeverity === 'critical')
            return 'high';
        if (baseSeverity === 'high')
            return 'medium';
        if (baseSeverity === 'medium')
            return 'low';
    }
    return baseSeverity;
}
// ============================================================================
// Main Detection Function
// ============================================================================
/**
 * Main detection function for model supply chain security issues
 */
function detectModelSupplyChain(content, filePath, options) {
    const vulnerabilities = [];
    // Skip non-applicable files
    if ((0, file_classifier_1.isScannerOrFixtureFile)(filePath))
        return vulnerabilities;
    // Only scan files that appear to be in ML context
    if (!isMLContextFile(filePath, content)) {
        return vulnerabilities;
    }
    const lines = options?.parsed?.lines ?? content.split('\n');
    const isTestFile = (0, file_classifier_1.isTestOrMockFile)(filePath);
    const isExample = (0, file_classifier_1.isExampleDirectory)(filePath);
    const isLibrary = (0, file_classifier_1.isLibraryCode)(filePath);
    for (const pattern of MODEL_SUPPLY_CHAIN_PATTERNS) {
        const regex = new RegExp(pattern.pattern.source, pattern.pattern.flags);
        let match;
        while ((match = regex.exec(content)) !== null) {
            const lineNumber = content.substring(0, match.index).split('\n').length;
            const lineContent = lines[lineNumber - 1]?.trim() || '';
            // Skip comments
            if ((0, file_classifier_1.isComment)(lineContent))
                continue;
            const surroundingContext = getSurroundingContext(content, lineNumber - 1);
            // Check for mitigations
            let isMitigated = false;
            let isPartiallyMitigated = false;
            let description = pattern.description;
            // Check SafeTensors usage (mitigates most pickle/torch issues)
            if (usesSafeTensors(lineContent, surroundingContext)) {
                isMitigated = true;
                description += ' (SafeTensors format detected - safe.)';
            }
            // Check torch weights_only
            if (pattern.checkSafeMode === 'torch') {
                if (hasTorchWeightsOnly(lineContent, surroundingContext)) {
                    isMitigated = true;
                    description += ' (weights_only=True detected - safe.)';
                }
            }
            // Check keras safe_mode
            if (pattern.checkSafeMode === 'keras') {
                if (hasKerasSafeMode(lineContent, surroundingContext)) {
                    isMitigated = true;
                    description += ' (safe_mode=True detected - safe.)';
                }
            }
            // Check trusted source
            if (pattern.checkTrustedSource) {
                if (isTrustedModelSource(lineContent)) {
                    isPartiallyMitigated = true;
                    description += ' (Trusted source detected.)';
                }
            }
            // Check integrity verification
            if (pattern.checkIntegrity) {
                if (hasIntegrityCheck(content, lineNumber)) {
                    isMitigated = true;
                    description += ' (Integrity verification detected.)';
                }
            }
            // Check data validation for training
            if (pattern.checkDataValidation) {
                if (hasDataValidation(content, lineNumber)) {
                    isPartiallyMitigated = true;
                    description += ' (Data validation detected nearby.)';
                }
            }
            // Skip if fully mitigated
            if (isMitigated)
                continue;
            // Calculate final severity
            const severity = calculateSeverity(pattern.baseSeverity, isMitigated, isPartiallyMitigated, isTestFile, isExample, isLibrary);
            // Add context notes
            if (isTestFile) {
                description += ' (In test file.)';
            }
            else if (isExample) {
                description += ' (In example/demo directory.)';
            }
            else if (isLibrary) {
                description += ' (Library code.)';
            }
            // Skip info-level in non-ML focused files to reduce noise
            if (severity === 'info' && !isMLContextFile(filePath, content))
                continue;
            vulnerabilities.push({
                id: `model-supply-${filePath}-${lineNumber}-${pattern.name.replace(/\s+/g, '-')}`,
                filePath,
                lineNumber,
                lineContent,
                severity,
                category: pattern.category,
                title: pattern.name,
                description,
                suggestedFix: pattern.suggestedFix,
                confidence: severity === 'info' ? 'low' : 'high',
                layer: 2,
                source: 'ai_code',
                requiresAIValidation: severity !== 'info' && severity !== 'low',
                baseConfidence: BASE_CONFIDENCE,
            });
        }
    }
    return vulnerabilities;
}
//# sourceMappingURL=model-supply-chain.js.map