@oculum/scanner 1.0.11 → 1.0.13

package/src/{layer2/ai-fingerprinting.ts → detect/ai-code/fingerprinting.ts}

@@ -3,8 +3,12 @@
  * Detects patterns commonly found in AI-generated code that may indicate security risks
  */
 
-import type { Vulnerability, VulnerabilitySeverity } from '../types'
-import { isExampleFile, isTestOrMockFile, isPlaceholderValue, isScannerOrFixtureFile } from '../utils/context-helpers'
+import type { Vulnerability, VulnerabilitySeverity } from '../../shared/types'
+import type { ParsedFile } from '../../shared/parsed-file'
+import { isExampleFile, isTestOrMockFile, isPlaceholderValue, isScannerOrFixtureFile } from '../../parse/file-classifier'
+import { getEnvironmentContext, isInPlaceholderAttribute, isDefaultParameterValue } from '../../shared/environment-context'
+
+const BASE_CONFIDENCE = 0.30
 
 interface AIFingerprint {
   name: string
@@ -91,35 +95,17 @@ const AI_FINGERPRINTS: AIFingerprint[] = [
   // NOTE: localhost/example URL detection moved to special handling below
   // to allow context-aware skipping for config/example files
   
-  // ==================== AI code smell patterns - heavily downgraded ====================
-  {
-    name: 'AI console.log debugging',
-    pattern: /console\.log\s*\(\s*['"]?(debug|testing|here|check|log|data|result|response)/gi,
-    severity: 'info',  // Downgraded from low - debug logs are common in dev
-    description: 'Debug logging that could be removed before production',
-    suggestedFix: 'Consider removing debug console.log statements or use proper logging',
-    confidence: 'low',  // Downgraded
-  },
-  // NOTE: Removed 'AI magic number' pattern - magic numbers are style, not security
-  {
-    name: 'AI empty function body',
-    pattern: /function\s+\w+\s*\([^)]*\)\s*\{\s*(\/\/.*)?(\n\s*)?\}|=>\s*\{\s*(\/\/.*)?(\n\s*)?\}/gi,
-    severity: 'low',  // Downgraded from medium
-    description: 'Empty function body - may be incomplete implementation',
-    suggestedFix: 'Implement the function or remove if not needed',
-    confidence: 'low',  // Downgraded
-  },
+  // ==================== AI code smell patterns - REMOVED ====================
+  // NOTE: The following patterns have been REMOVED as they are style/code quality issues,
+  // not security vulnerabilities. Reporting these creates excessive noise:
+  // - 'AI console.log debugging' - debug logs are standard development practice
+  // - 'AI empty function body' - empty functions may be intentional stubs/callbacks
+  // - 'AI magic number' - magic numbers are code quality, not security
   
-  // ==================== AI boilerplate patterns - heavily downgraded ====================
-  {
-    name: 'AI boilerplate error message',
-    pattern: /['"]Something went wrong['"]|['"]An error occurred['"]|['"]Error processing request['"]/gi,
-    severity: 'info',  // Downgraded from low - generic messages are acceptable
-    description: 'Generic error message - consider more specific information',
-    suggestedFix: 'Replace with specific, actionable error messages',
-    confidence: 'low',  // Downgraded
-  },
-  // NOTE: Removed 'AI success message' pattern - success messages are style, not security
+  // ==================== AI boilerplate patterns - REMOVED ====================
+  // NOTE: Generic error messages are intentionally vague to avoid information leakage.
+  // Flagging "Something went wrong" as an issue is counterproductive - it's often the
+  // correct security-conscious approach. These patterns have been removed.
   
   // ==================== AI security bypass patterns - moderated ====================
   {
@@ -429,6 +415,8 @@ function detectSmartAnyUsage(
       suggestedFix: 'Replace "any" with explicit types. For request handlers use typed schemas (Zod, Yup). For database queries use typed ORM models.',
       confidence: 'medium',
       layer: 2,
+        source: 'ai_code' as const,
+        baseConfidence: BASE_CONFIDENCE,
     })
   } else {
     // Report individual findings for 1-2 high-priority 'any' usages
@@ -452,6 +440,8 @@ function detectSmartAnyUsage(
         suggestedFix: 'Replace "any" with an explicit type. Use typed request schemas, ORM models, or interface definitions.',
         confidence: 'medium',
         layer: 2,
+        source: 'ai_code' as const,
+        baseConfidence: BASE_CONFIDENCE,
       })
     }
   }
@@ -577,6 +567,8 @@ function detectManagedAICostAbuse(
         suggestedFix: 'Consider adding per-user rate limiting (e.g., @upstash/ratelimit) to prevent cost abuse by authenticated users.',
         confidence: 'low',
         layer: 2,
+        source: 'ai_code' as const,
+        baseConfidence: BASE_CONFIDENCE,
       })
     } else {
       // Unauthenticated route - higher risk
@@ -592,6 +584,8 @@ function detectManagedAICostAbuse(
         suggestedFix: 'Add authentication or rate limiting (e.g., @upstash/ratelimit, rate-limiter-flexible) to prevent cost abuse.',
         confidence: 'medium',
         layer: 2,
+        source: 'ai_code' as const,
+        baseConfidence: BASE_CONFIDENCE,
       })
     }
   }
@@ -632,14 +626,15 @@ function isConfigFile(filePath: string): boolean {
 
 export function detectAIFingerprints(
   content: string,
-  filePath: string
+  filePath: string,
+  options?: { parsed?: ParsedFile }
 ): Vulnerability[] {
   const vulnerabilities: Vulnerability[] = []
 
   // Skip scanner/fixture files to avoid self-detection
   if (isScannerOrFixtureFile(filePath)) return vulnerabilities
 
-  const lines = content.split('\n')
+  const lines = options?.parsed?.lines ?? content.split('\n')
 
   // Skip example/demo files entirely - they contain placeholder code by design
   if (isExampleFile(filePath)) {
@@ -696,6 +691,8 @@ export function detectAIFingerprints(
           suggestedFix: fingerprint.suggestedFix,
           confidence,
           layer: 2,
+        source: 'ai_code' as const,
+          baseConfidence: BASE_CONFIDENCE,
         })
         break // Only report once per line
       }
@@ -703,8 +700,11 @@ export function detectAIFingerprints(
   })
 
   // Context-aware localhost/example URL detection
-  // Skip for config files and test files (they legitimately contain example URLs)
-  if (!isConfigOrSettings && !isTestFile) {
+  // Use environment context for smarter filtering
+  const envContext = getEnvironmentContext(filePath)
+
+  // Skip for environments where placeholder URLs are expected
+  if (!envContext.allowsPlaceholderUrls && !isConfigOrSettings && !isTestFile) {
     const localhostPattern = /['"]https?:\/\/(localhost|127\.0\.0\.1|example\.com|your-domain|api\.example)[^'"]*['"]/gi
     lines.forEach((line, index) => {
       if (localhostPattern.test(line)) {
@@ -719,6 +719,14 @@ export function detectAIFingerprints(
         if (/process\.env\.\w+\s*\|\|\s*['"]/.test(line)) {
           return
         }
+        // Skip if it's in a placeholder attribute (placeholder="https://example.com")
+        if (isInPlaceholderAttribute(line)) {
+          return
+        }
+        // Skip if it's a default parameter value
+        if (isDefaultParameterValue(line)) {
+          return
+        }
         vulnerabilities.push({
           id: `ai-fingerprint-${filePath}-${index + 1}-localhost-url`,
           filePath,
@@ -731,6 +739,8 @@ export function detectAIFingerprints(
           suggestedFix: 'Replace with actual production URL from environment variable',
           confidence: 'high',
           layer: 2,
+        source: 'ai_code' as const,
+          baseConfidence: BASE_CONFIDENCE,
         })
         aiPatternCount++
       }
@@ -755,6 +765,8 @@ export function detectAIFingerprints(
       suggestedFix: 'Review this file carefully for security issues, incomplete implementations, and placeholder code',
       confidence: 'medium',
       layer: 2,
+        source: 'ai_code' as const,
+        baseConfidence: BASE_CONFIDENCE,
     })
   }
 

package/src/detect/ai-code/index.ts

@@ -0,0 +1,11 @@
+export { detectAIAgentTools } from './agent-tools'
+export { detectBYOKPatterns } from './byok-patterns'
+export { detectAIEndpointProtection } from './endpoint-protection'
+export { detectAIExecutionSinks } from './execution-sinks'
+export { detectAIFingerprints } from './fingerprinting'
+export { detectMCPSecurity } from './mcp-security'
+export { detectModelSupplyChain } from './model-supply-chain'
+export { detectAIPackageHallucination } from './package-hallucination'
+export { detectAIPromptHygiene } from './prompt-hygiene'
+export { detectRAGSafetyIssues } from './rag-safety'
+export { detectAISchemaValidation } from './schema-validation'

package/src/{layer2/ai-mcp-security.ts → detect/ai-code/mcp-security.ts}

@@ -10,7 +10,8 @@
  * Reference: https://modelcontextprotocol.io, 13,000+ MCP servers deployed
  */
 
-import type { Vulnerability, VulnerabilitySeverity, VulnerabilityCategory } from '../types'
+import type { Vulnerability, VulnerabilitySeverity, VulnerabilityCategory } from '../../shared/types'
+import type { ParsedFile } from '../../shared/parsed-file'
 import {
   isComment,
   isTestOrMockFile,
@@ -18,7 +19,9 @@ import {
   isScannerOrFixtureFile,
   isExampleDirectory,
   isLibraryCode,
-} from '../utils/context-helpers'
+} from '../../parse/file-classifier'
+
+const BASE_CONFIDENCE = 0.50
 
 // ============================================================================
 // Context Detection
@@ -154,7 +157,7 @@ function getSurroundingContext(content: string, lineIndex: number, windowSize: n
 interface MCPSecurityPattern {
   name: string
   pattern: RegExp
-  category: 'tool_poisoning' | 'credential_issue' | 'confused_deputy' | 'description_injection' | 'server_shadowing'
+  category: 'tool_poisoning' | 'credential_issue' | 'confused_deputy' | 'description_injection' | 'server_shadowing' | 'schema_bypass' | 'missing_hitl'
   baseSeverity: VulnerabilitySeverity
   description: string
   suggestedFix: string
@@ -528,6 +531,198 @@ const SERVER_SHADOWING_PATTERNS: MCPSecurityPattern[] = [
   },
 ]
 
+/**
+ * Phase 5 Task 5: MCP Schema Validation Patterns
+ * Detect MCP tools that use arguments without schema validation
+ */
+const SCHEMA_VALIDATION_PATTERNS: MCPSecurityPattern[] = [
+  // MCP tool using args directly without validation (JS)
+  {
+    name: 'MCP tool without input validation',
+    pattern: /server\.tool\s*\([^)]+,\s*async\s*\(\s*(?:args|params|input)\s*\)\s*(?:=>|:)[^{]*\{(?![\s\S]{0,100}(?:schema\.parse|safeParse|validate|zod|yup|joi|superstruct|ajv|\.parse\())/gi,
+    category: 'schema_bypass',
+    baseSeverity: 'medium',
+    description: 'MCP tool uses arguments directly without schema validation. Malformed or malicious input could cause unexpected behavior.',
+    suggestedFix: 'Validate inputs with a schema: const validated = schema.parse(args); return runCommand(validated.command)',
+  },
+  // MCP tool accessing args properties without validation
+  {
+    name: 'MCP tool args used without validation',
+    pattern: /server\.tool\s*\([^)]+,\s*async\s*\(\s*(?:args|params)\s*\)[^{]*\{[^}]*(?:args|params)\.(?:command|query|path|url|file|data|input|content|sql|script|code)(?![\s\S]{0,50}(?:validated|parsed|sanitized))/gi,
+    category: 'schema_bypass',
+    baseSeverity: 'high',
+    description: 'MCP tool uses potentially dangerous argument properties directly. Input validation required.',
+    suggestedFix: 'Validate dangerous inputs: const { command } = commandSchema.parse(args)',
+  },
+  // Python MCP tool without type/validation
+  {
+    name: 'Python MCP tool without validation',
+    pattern: /@server\.tool[^)]*\)\s*(?:async\s+)?def\s+\w+\s*\(\s*(?:args|params|kwargs|\*\*)\s*(?::\s*dict)?\s*\)(?![\s\S]{0,50}(?:pydantic|validate|TypedDict|dataclass))/gi,
+    category: 'schema_bypass',
+    baseSeverity: 'medium',
+    description: 'Python MCP tool accepts dict/kwargs without type validation. Use Pydantic or TypedDict.',
+    suggestedFix: 'Use Pydantic model: def tool_name(args: MyInputModel) or validate with TypedDict',
+  },
+  // Args spread into function call
+  {
+    name: 'MCP tool args spread into call',
+    pattern: /(?:runCommand|exec|spawn|query|execute|fetch)\s*\(\s*\.\.\.(?:args|params|input)/gi,
+    category: 'schema_bypass',
+    baseSeverity: 'high',
+    description: 'MCP tool arguments spread directly into function call. All fields pass through unvalidated.',
+    suggestedFix: 'Validate and destructure specific fields: const { field1, field2 } = schema.parse(args); fn(field1, field2)',
+  },
+  // Dynamic property access on args
+  {
+    name: 'Dynamic property access on MCP args',
+    pattern: /(?:args|params|input)\s*\[\s*(?:key|prop|field|name)\s*\]/gi,
+    category: 'schema_bypass',
+    baseSeverity: 'medium',
+    description: 'Dynamic property access on MCP tool arguments. Could access unintended properties.',
+    suggestedFix: 'Use explicit destructuring with validation: const { expectedField } = schema.parse(args)',
+  },
+]
+
+/**
+ * Phase 6 Task 3: MCP Tool Result Injection Patterns
+ * Detect MCP tool results directly interpolated into prompts without sanitization
+ */
+const RESULT_INJECTION_PATTERNS: MCPSecurityPattern[] = [
+  // MCP result interpolated into prompt template literal
+  {
+    name: 'MCP result in prompt template',
+    pattern: /`[^`]*\$\{[^}]*(?:tool|mcp|result|toolResult|mcpResult)[^}]*\}[^`]*`\s*(?:\+\s*)?(?:system|prompt|message|instruction)/gi,
+    category: 'tool_poisoning',
+    baseSeverity: 'high',
+    description: 'MCP tool results interpolated into prompts could contain injection payloads from external sources.',
+    suggestedFix: 'Sanitize MCP tool results before including in prompts. Use structured data extraction: const safeData = extractSafeFields(toolResult)',
+  },
+  // Tool result concatenated with system prompt
+  {
+    name: 'Tool result concatenated with prompt',
+    pattern: /(?:systemPrompt|prompt|message|instruction)\s*(?:\+|\.concat)\s*(?:toolResult|mcpResult|result|tool\.result|mcp\.result)/gi,
+    category: 'tool_poisoning',
+    baseSeverity: 'high',
+    description: 'Tool results concatenated with prompts. External content in results could manipulate model behavior.',
+    suggestedFix: 'Sanitize tool results before concatenation. Consider using delimiters: prompt + "\\n---DATA---\\n" + sanitize(result)',
+  },
+  // Tool result in messages array
+  {
+    name: 'Raw tool result in messages',
+    pattern: /messages\s*(?:\.push|:\s*\[)[^;]*content\s*:\s*(?:toolResult|mcpResult|result|tool\.result)(?!\.sanitized|\.safe)/gi,
+    category: 'tool_poisoning',
+    baseSeverity: 'medium',
+    description: 'Raw tool results added to message content. Results from external tools could contain injection payloads.',
+    suggestedFix: 'Sanitize or structure tool results: messages.push({ content: sanitizeForPrompt(toolResult) })',
+  },
+  // Tool result used as context without processing
+  {
+    name: 'Tool result as unprocessed context',
+    pattern: /context\s*[:=]\s*(?:toolResult|mcpResult|result|tool\.(?:output|result))(?!\s*\.|\.sanitize|\.filter)/gi,
+    category: 'tool_poisoning',
+    baseSeverity: 'medium',
+    description: 'Tool result assigned directly as context. External content should be processed before use.',
+    suggestedFix: 'Process and validate tool results: const context = processToolResult(result)',
+  },
+  // Spread tool result into prompt data
+  {
+    name: 'Tool result spread into prompt',
+    pattern: /\{[^}]*\.\.\.(?:toolResult|mcpResult|result|tool\.result)[^}]*\}\s*(?:as|:|\s+(?:prompt|message|context))/gi,
+    category: 'tool_poisoning',
+    baseSeverity: 'high',
+    description: 'Tool result spread into prompt data. All fields from external tool pass through.',
+    suggestedFix: 'Extract specific fields: const { safeField1, safeField2 } = validateToolResult(result)',
+  },
+  // JSON stringify tool result into prompt
+  {
+    name: 'JSON stringified tool result in prompt',
+    pattern: /JSON\.stringify\s*\(\s*(?:toolResult|mcpResult|result|tool\.result)\s*\)[^;]*(?:prompt|message|context|instruction)/gi,
+    category: 'tool_poisoning',
+    baseSeverity: 'medium',
+    description: 'Tool result JSON-stringified into prompt. Serialized content could contain injection payloads.',
+    suggestedFix: 'Filter tool result before stringification: JSON.stringify(filterSafeFields(result))',
+  },
+  // Format tool result for LLM
+  {
+    name: 'Unvalidated tool result formatting',
+    pattern: /format(?:Tool|Result|Output)?\s*\(\s*(?:toolResult|mcpResult|result|tool\.result)\s*\)(?![\s\S]{0,30}(?:sanitize|validate|filter))/gi,
+    category: 'tool_poisoning',
+    baseSeverity: 'medium',
+    description: 'Tool result formatted without validation. Formatting function should include sanitization.',
+    suggestedFix: 'Include sanitization in formatting: formatToolResult(sanitize(result))',
+  },
+]
+
+/**
+ * Phase 5 Task 6: Human-in-the-Loop for Destructive Operations
+ * Detect destructive operations without confirmation mechanism
+ */
+const DESTRUCTIVE_OPS_PATTERNS: MCPSecurityPattern[] = [
+  // File deletion without confirmation
+  {
+    name: 'MCP file deletion without confirmation',
+    pattern: /server\.tool\s*\([^)]+(?:delete|remove|unlink|rm)[^)]+,\s*async[^{]*\{(?![\s\S]{0,100}(?:confirm|approved|needsConfirmation|requireApproval|humanInLoop))[^}]*(?:fs\.rm|fs\.unlink|unlinkSync|rmSync|remove|rimraf)/gi,
+    category: 'missing_hitl',
+    baseSeverity: 'high',
+    description: 'MCP tool performs file deletion without confirmation mechanism. Destructive operations should require human approval.',
+    suggestedFix: 'Add confirmation: if (!args.confirmed) { return { needsConfirmation: true, action: "delete", path: args.path } }',
+  },
+  // Database deletion without confirmation
+  {
+    name: 'MCP database deletion without confirmation',
+    pattern: /server\.tool\s*\([^)]+(?:delete|drop|truncate|remove)[^)]+,\s*async[^{]*\{(?![\s\S]{0,100}(?:confirm|approved|needsConfirmation))[^}]*(?:\.delete|\.drop|\.truncate|\.destroy|DELETE\s+FROM|DROP\s+TABLE)/gi,
+    category: 'missing_hitl',
+    baseSeverity: 'high',
+    description: 'MCP tool performs database deletion without confirmation. Data loss risk.',
+    suggestedFix: 'Require confirmation for destructive DB operations: if (!args.confirmed) return { needsConfirmation: true }',
+  },
+  // Recursive directory deletion
+  {
+    name: 'MCP recursive deletion without confirmation',
+    pattern: /(?:fs\.rm|rimraf|rmdir)\s*\([^)]*,\s*\{\s*recursive\s*:\s*true/gi,
+    category: 'missing_hitl',
+    baseSeverity: 'critical',
+    description: 'Recursive directory deletion in MCP tool. High risk of unintended data loss.',
+    suggestedFix: 'Add explicit confirmation with path display: if (!args.confirmed) return { needsConfirmation: true, message: `Delete ${path} and all contents?` }',
+  },
+  // Shell command execution without confirmation
+  {
+    name: 'MCP shell execution without confirmation',
+    pattern: /server\.tool\s*\([^)]+(?:exec|run|shell|command)[^)]+,\s*async[^{]*\{(?![\s\S]{0,100}(?:confirm|approved|needsConfirmation))[^}]*(?:exec|spawn|execSync|spawnSync)\s*\(/gi,
+    category: 'missing_hitl',
+    baseSeverity: 'high',
+    description: 'MCP tool executes shell commands without confirmation. Dangerous commands could be executed.',
+    suggestedFix: 'Require confirmation for shell commands: if (!args.confirmed) return { needsConfirmation: true, command: args.command }',
+  },
+  // Send/publish operations without confirmation
+  {
+    name: 'MCP send operation without confirmation',
+    pattern: /server\.tool\s*\([^)]+(?:send|publish|broadcast|notify)[^)]+,\s*async[^{]*\{(?![\s\S]{0,100}(?:confirm|approved|draft))[^}]*(?:\.send|\.publish|sendEmail|sendMessage)/gi,
+    category: 'missing_hitl',
+    baseSeverity: 'medium',
+    description: 'MCP tool sends messages/emails without confirmation. Could send unintended communications.',
+    suggestedFix: 'Add draft/confirmation: if (!args.confirmed) return { needsConfirmation: true, preview: messageContent }',
+  },
+  // Payment/transaction operations
+  {
+    name: 'MCP payment without confirmation',
+    pattern: /server\.tool\s*\([^)]+(?:pay|charge|transfer|transaction)[^)]+,\s*async[^{]*\{(?![\s\S]{0,100}(?:confirm|approved|needsConfirmation))/gi,
+    category: 'missing_hitl',
+    baseSeverity: 'critical',
+    description: 'MCP tool processes payments without confirmation. Financial operations require human approval.',
+    suggestedFix: 'Always require confirmation for financial operations: if (!args.confirmed) return { needsConfirmation: true, amount, recipient }',
+  },
+  // API key/secret deletion
+  {
+    name: 'MCP credential deletion without confirmation',
+    pattern: /server\.tool\s*\([^)]+(?:delete|revoke|remove)[^)]*(?:key|token|secret|credential)[^)]+,\s*async[^{]*\{(?![\s\S]{0,100}(?:confirm|approved))/gi,
+    category: 'missing_hitl',
+    baseSeverity: 'high',
+    description: 'MCP tool deletes credentials without confirmation. Could cause service disruption.',
+    suggestedFix: 'Require explicit confirmation: if (!args.confirmed) return { needsConfirmation: true, warning: "This will revoke access" }',
+  },
+]
+
 // ============================================================================
 // Main Detection Function
 // ============================================================================
@@ -547,6 +742,10 @@ function mapCategory(internal: MCPSecurityPattern['category']): VulnerabilityCat
       return 'ai_mcp_description_injection'
     case 'server_shadowing':
       return 'ai_mcp_server_shadowing'
+    case 'schema_bypass':
+      return 'ai_mcp_tool_poisoning' // Schema bypass leads to tool poisoning risks
+    case 'missing_hitl':
+      return 'ai_excessive_agency' // Missing human-in-the-loop is excessive agency
   }
 }
 
@@ -555,7 +754,8 @@ function mapCategory(internal: MCPSecurityPattern['category']): VulnerabilityCat
  */
 export function detectMCPSecurity(
   content: string,
-  filePath: string
+  filePath: string,
+  options?: { parsed?: ParsedFile }
 ): Vulnerability[] {
   const vulnerabilities: Vulnerability[] = []
 
@@ -568,7 +768,7 @@ export function detectMCPSecurity(
     return vulnerabilities
   }
 
-  const lines = content.split('\n')
+  const lines = options?.parsed?.lines ?? content.split('\n')
   const isTestFile = isTestOrMockFile(filePath)
   const isExample = isExampleDirectory(filePath)
   const isLibrary = isLibraryCode(filePath)
@@ -580,6 +780,11 @@ export function detectMCPSecurity(
     ...CONFUSED_DEPUTY_PATTERNS,
     ...DESCRIPTION_INJECTION_PATTERNS,
     ...SERVER_SHADOWING_PATTERNS,
+    // Phase 5: New detection patterns
+    ...SCHEMA_VALIDATION_PATTERNS,
+    ...DESTRUCTIVE_OPS_PATTERNS,
+    // Phase 6: MCP result injection
+    ...RESULT_INJECTION_PATTERNS,
   ]
 
   // Track findings to avoid duplicates
@@ -721,7 +926,9 @@ export function detectMCPSecurity(
         suggestedFix: pattern.suggestedFix,
         confidence: severity === 'info' ? 'low' : 'medium',
         layer: 2,
+        source: 'ai_code' as const,
         requiresAIValidation: severity !== 'info' && severity !== 'low',
+        baseConfidence: BASE_CONFIDENCE,
       })
     }
   }

package/src/{layer2 → detect/ai-code}/model-supply-chain.ts

@@ -13,14 +13,17 @@
  * - CWE-502: Deserialization of Untrusted Data
  */
 
-import type { Vulnerability, VulnerabilitySeverity, VulnerabilityCategory } from '../types'
+import type { Vulnerability, VulnerabilitySeverity, VulnerabilityCategory } from '../../shared/types'
+import type { ParsedFile } from '../../shared/parsed-file'
 import {
   isComment,
   isTestOrMockFile,
   isScannerOrFixtureFile,
   isExampleDirectory,
   isLibraryCode,
-} from '../utils/context-helpers'
+} from '../../parse/file-classifier'
+
+const BASE_CONFIDENCE = 0.50
 
 // ============================================================================
 // Context Detection
@@ -43,13 +46,20 @@ function isMLContextFile(filePath: string, content: string): boolean {
 
   // Content patterns suggesting ML usage
   const mlContentPatterns = [
-    /import\s+(?:torch|tensorflow|keras|transformers|joblib|pickle)/i,
-    /from\s+(?:torch|tensorflow|keras|transformers|joblib|pickle)\s+import/i,
+    /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,
   ]
@@ -179,6 +189,72 @@ const MODEL_SUPPLY_CHAIN_PATTERNS: ModelSupplyChainPattern[] = [
     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',
@@ -329,7 +405,8 @@ function calculateSeverity(
  */
 export function detectModelSupplyChain(
   content: string,
-  filePath: string
+  filePath: string,
+  options?: { parsed?: ParsedFile }
 ): Vulnerability[] {
   const vulnerabilities: Vulnerability[] = []
 
@@ -341,7 +418,7 @@ export function detectModelSupplyChain(
     return vulnerabilities
   }
 
-  const lines = content.split('\n')
+  const lines = options?.parsed?.lines ?? content.split('\n')
   const isTestFile = isTestOrMockFile(filePath)
   const isExample = isExampleDirectory(filePath)
   const isLibrary = isLibraryCode(filePath)
@@ -447,7 +524,9 @@ export function detectModelSupplyChain(
         suggestedFix: pattern.suggestedFix,
         confidence: severity === 'info' ? 'low' : 'high',
         layer: 2,
+        source: 'ai_code' as const,
         requiresAIValidation: severity !== 'info' && severity !== 'low',
+        baseConfidence: BASE_CONFIDENCE,
       })
     }
   }