@oculum/scanner 1.0.14 → 1.0.15

package/src/detect/ast-rules/model-supply-chain-ast.ts

@@ -0,0 +1,202 @@
+/**
+ * AST-Based Model Supply Chain Security Detection
+ *
+ * Migrates ai-code/model-supply-chain.ts from regex to AST (JS/TS patterns only).
+ * Python-specific patterns (pickle.load, torch.load) remain in regex detector.
+ *
+ * AST advantages:
+ * - Structurally detects trust_remote_code in config objects
+ * - Resolves dynamic import/require of model files
+ * - Checks for integrity verification in surrounding scope
+ */
+
+import type { ParsedAST } from '../../parse/ast'
+import { registerASTRule, type ASTRuleMatch, type NodeIndex, getNodes } from './index'
+import type Parser from 'tree-sitter'
+import { getObjectLiteralProperty, resolveCallTarget, getCallArguments } from './helpers/call-analysis'
+import { hasImport, getImportedModules } from './helpers/import-analysis'
+import { isScannerOrFixtureFile, isTestOrMockFile, isExampleFile } from '../../parse/file-classifier'
+
+// ============================================================================
+// Helpers
+// ============================================================================
+
+/** Quick content pre-check for ML context */
+const ML_CONTENT_QUICK_CHECK = /transformers|huggingface|onnxruntime|tensorflow|ml5|brain\.js|from_pretrained|loadModel|load_model|trust_remote_code/i
+
+/** Check if file has ML/model context */
+function isMLContextFile(root: Parser.SyntaxNode, filePath: string, content: string): boolean {
+  // Content pre-check: skip files without ML-related mentions
+  if (!ML_CONTENT_QUICK_CHECK.test(content) &&
+      !/\/(models?|ml|training|inference|weights)\//i.test(filePath) &&
+      !/(train|model|inference|fine[-_]?tune)/i.test(filePath)) {
+    return false
+  }
+
+  const modules = getImportedModules(root)
+  const mlModules = /^(transformers|@huggingface|onnxruntime|tensorflow|@tensorflow|ml5|brain\.js)/
+  if ([...modules].some(m => mlModules.test(m))) return true
+
+  if (/\/(models?|ml|training|inference|weights)\//i.test(filePath)) return true
+  if (/(train|model|inference|fine[-_]?tune)/i.test(filePath)) return true
+
+  return /from_pretrained|loadModel|\.load_model|trust_remote_code/i.test(content)
+}
+
+// ============================================================================
+// AST Rule: trust_remote_code
+// ============================================================================
+
+registerASTRule({
+  id: 'ast-model-trust-remote-code',
+  title: 'trust_remote_code enabled for model loading',
+  description: 'trust_remote_code=true allows model repos to execute arbitrary Python code during loading',
+  severity: 'high',
+  category: 'ai_unsafe_model_load',
+  suggestedFix: 'Avoid trust_remote_code. Use models that do not require custom code. Audit model repo code if necessary.',
+  languages: ['javascript', 'typescript', 'tsx'],
+  confidence: 'high',
+  baseConfidence: 0.50,
+  layer: 2,
+  source: 'ai_code',
+  detect(ast: ParsedAST, content: string, nodeIndex?: NodeIndex): ASTRuleMatch[] {
+    if (isScannerOrFixtureFile(ast.filePath)) return []
+
+    const root = ast.tree.rootNode
+    const isTestFile = isTestOrMockFile(ast.filePath)
+    const matches: ASTRuleMatch[] = []
+
+    for (const node of getNodes(nodeIndex!, 'pair')) {
+      // trust_remote_code: true in object literals
+      const key = node.childForFieldName('key')
+      const value = node.childForFieldName('value')
+      if (key && value) {
+        const keyText = key.text.replace(/['"`]/g, '')
+        if (/^trust_remote_code$/i.test(keyText) && value.text === 'true') {
+          matches.push({
+            node,
+            severity: isTestFile ? 'info' : 'high',
+            note: 'trust_remote_code=true allows arbitrary code execution during model loading.',
+          })
+        }
+      }
+    }
+
+    return matches
+  },
+})
+
+// ============================================================================
+// AST Rule: Model from HTTP URL
+// ============================================================================
+
+registerASTRule({
+  id: 'ast-model-http-source',
+  title: 'Model loaded from HTTP URL',
+  description: 'Loading models from HTTP URLs bypasses verification. Attackers could serve malicious models.',
+  severity: 'high',
+  category: 'ai_unverified_model',
+  suggestedFix: 'Load models by ID from Hugging Face Hub. Pin to specific revision with checksums.',
+  languages: ['javascript', 'typescript', 'tsx'],
+  confidence: 'medium',
+  baseConfidence: 0.45,
+  layer: 2,
+  source: 'ai_code',
+  detect(ast: ParsedAST, content: string, nodeIndex?: NodeIndex): ASTRuleMatch[] {
+    if (isScannerOrFixtureFile(ast.filePath)) return []
+
+    const root = ast.tree.rootNode
+    if (!isMLContextFile(root, ast.filePath, content)) return []
+    const isTestFile = isTestOrMockFile(ast.filePath)
+    const matches: ASTRuleMatch[] = []
+
+    for (const node of getNodes(nodeIndex!, 'call_expression')) {
+      const target = resolveCallTarget(node)
+      if (!target) continue
+
+      // from_pretrained("https://...")
+      if (target.name === 'from_pretrained' || target.name === 'loadModel' || target.name === 'load_model') {
+        const args = getCallArguments(node)
+        if (args.length > 0 && (args[0].type === 'string' || args[0].type === 'template_string')) {
+          const urlText = args[0].text
+          if (/https?:\/\//i.test(urlText)) {
+            // Check for trusted sources
+            if (/huggingface\.co\/(meta-llama|openai|anthropic|google|microsoft)/i.test(urlText)) {
+              continue // trusted source
+            }
+            matches.push({
+              node,
+              severity: isTestFile ? 'info' : 'high',
+              note: 'Model loaded from HTTP URL without verification.',
+            })
+          }
+        }
+      }
+    }
+
+    return matches
+  },
+})
+
+// ============================================================================
+// AST Rule: Unverified Model Download
+// ============================================================================
+
+registerASTRule({
+  id: 'ast-model-unverified-download',
+  title: 'Model file downloaded without integrity verification',
+  description: 'Downloading model files without checksum verification. Attackers could serve modified files.',
+  severity: 'medium',
+  category: 'ai_unverified_model',
+  suggestedFix: 'Verify downloaded model files against known checksums before loading.',
+  languages: ['javascript', 'typescript', 'tsx'],
+  confidence: 'medium',
+  baseConfidence: 0.40,
+  layer: 2,
+  source: 'ai_code',
+  detect(ast: ParsedAST, content: string, nodeIndex?: NodeIndex): ASTRuleMatch[] {
+    if (isScannerOrFixtureFile(ast.filePath)) return []
+    if (isTestOrMockFile(ast.filePath)) return []
+
+    const root = ast.tree.rootNode
+    if (!isMLContextFile(root, ast.filePath, content)) return []
+
+    const matches: ASTRuleMatch[] = []
+
+    for (const node of getNodes(nodeIndex!, 'call_expression')) {
+      const target = resolveCallTarget(node)
+      if (!target) continue
+
+      // fetch/axios/got downloading model files
+      if (target.name === 'fetch' || (target.object && /^(axios|got|http|https)$/.test(target.object))) {
+        const args = getCallArguments(node)
+        if (args.length > 0) {
+          const urlText = args[0].text
+          if (/\.(?:pth|pt|bin|ckpt|h5|keras|onnx|safetensors|model)/i.test(urlText)) {
+            // Check for integrity verification nearby
+            let fnBody: Parser.SyntaxNode | null = null
+            let current: Parser.SyntaxNode | null = node.parent
+            while (current) {
+              if (/function_declaration|arrow_function|function_expression/.test(current.type)) {
+                fnBody = current.childForFieldName('body') ?? current
+                break
+              }
+              current = current.parent
+            }
+            if (fnBody && /checksum|sha256|sha512|verify_hash|integrity/i.test(fnBody.text)) {
+              continue // has integrity check
+            }
+
+            matches.push({
+              node,
+              severity: 'medium',
+              note: 'Model file downloaded without integrity verification.',
+            })
+          }
+        }
+      }
+    }
+
+    return matches
+  },
+})