@oculum/scanner 1.0.14 → 1.0.15

package/src/__tests__/detect/python-helpers.test.ts

@@ -0,0 +1,536 @@
+/**
+ * Python AST Helper Tests
+ *
+ * Tests for all Python-specific AST analysis helpers:
+ * imports, calls, scope, strings, control flow, decorators.
+ */
+
+import { parseFile, clearASTCache, findNodesOfType } from '../../parse/ast'
+import {
+  collectPythonImportBindings,
+  hasPythonImport,
+  getPythonImportedModules,
+  isPythonImportedFrom,
+  resolvePythonCallTarget,
+  isPythonCallTo,
+  getPythonCallArguments,
+  getPythonCallArgument,
+  getPythonKeywordArgument,
+  isPythonMethodCallOn,
+  getPythonAttributeChain,
+  collectPythonVariableDeclarations,
+  collectPythonVariableAliases,
+  extractPythonStringValue,
+  getPythonFirstStringArg,
+  isPythonFString,
+  getPythonFStringParts,
+  isInsidePythonTryExcept,
+  isInsidePythonExceptClause,
+  getEnclosingPythonBlock,
+  getPythonDecorators,
+  hasPythonDecorator,
+} from '../../detect/ast-rules/helpers/python-helpers'
+
+function parse(code: string) {
+  const ast = parseFile(code, 'test.py')
+  if (!ast) throw new Error('Failed to parse Python code')
+  return ast.tree.rootNode
+}
+
+beforeEach(() => {
+  clearASTCache()
+})
+
+// ============================================================================
+// Import Analysis
+// ============================================================================
+
+describe('Python Import Bindings', () => {
+  it('handles "import os"', () => {
+    const root = parse('import os\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('os')).toEqual({ module: 'os', originalName: 'default' })
+  })
+
+  it('handles "import os.path" (local name is first segment)', () => {
+    const root = parse('import os.path\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('os')).toEqual({ module: 'os.path', originalName: 'default' })
+  })
+
+  it('handles "from os import path"', () => {
+    const root = parse('from os import path\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('path')).toEqual({ module: 'os', originalName: 'path' })
+  })
+
+  it('handles "from os.path import join, exists"', () => {
+    const root = parse('from os.path import join, exists\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('join')).toEqual({ module: 'os.path', originalName: 'join' })
+    expect(bindings.get('exists')).toEqual({ module: 'os.path', originalName: 'exists' })
+  })
+
+  it('handles "from os import path as p"', () => {
+    const root = parse('from os import path as p\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('p')).toEqual({ module: 'os', originalName: 'path' })
+    expect(bindings.has('path')).toBe(false)
+  })
+
+  it('handles relative import "from . import utils"', () => {
+    const root = parse('from . import utils\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('utils')).toEqual({ module: '.', originalName: 'utils' })
+  })
+
+  it('handles relative submodule import "from .foo import bar"', () => {
+    const root = parse('from .foo import bar\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('bar')).toEqual({ module: '.foo', originalName: 'bar' })
+  })
+
+  it('handles deep relative import "from ..package.sub import x"', () => {
+    const root = parse('from ..package.sub import x\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('x')).toEqual({ module: '..package.sub', originalName: 'x' })
+  })
+
+  it('handles "import os as o"', () => {
+    const root = parse('import os as o\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('o')).toEqual({ module: 'os', originalName: 'default' })
+    expect(bindings.has('os')).toBe(false)
+  })
+
+  it('handles "from openai import OpenAI"', () => {
+    const root = parse('from openai import OpenAI\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('OpenAI')).toEqual({ module: 'openai', originalName: 'OpenAI' })
+  })
+
+  it('handles "from langchain.chat_models import ChatOpenAI"', () => {
+    const root = parse('from langchain.chat_models import ChatOpenAI\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.get('ChatOpenAI')).toEqual({ module: 'langchain.chat_models', originalName: 'ChatOpenAI' })
+  })
+
+  it('captures multiple imports in one file', () => {
+    const root = parse('import os\nimport sys\nfrom pathlib import Path\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.size).toBe(3)
+    expect(bindings.has('os')).toBe(true)
+    expect(bindings.has('sys')).toBe(true)
+    expect(bindings.has('Path')).toBe(true)
+  })
+
+  it('returns empty map for no imports', () => {
+    const root = parse('x = 1\n')
+    const bindings = collectPythonImportBindings(root)
+    expect(bindings.size).toBe(0)
+  })
+
+  it('hasPythonImport detects module presence', () => {
+    const root = parse('import os\nfrom pathlib import Path\n')
+    expect(hasPythonImport(root, 'os')).toBe(true)
+    expect(hasPythonImport(root, 'pathlib', 'Path')).toBe(true)
+    expect(hasPythonImport(root, 'sys')).toBe(false)
+  })
+
+  it('getPythonImportedModules returns all modules', () => {
+    const root = parse('import os\nfrom sys import argv\n')
+    const modules = getPythonImportedModules(root)
+    expect(modules.has('os')).toBe(true)
+    expect(modules.has('sys')).toBe(true)
+  })
+
+  it('isPythonImportedFrom checks module source', () => {
+    const root = parse('from openai import OpenAI\n')
+    expect(isPythonImportedFrom(root, 'OpenAI', 'openai')).toBe('openai')
+    expect(isPythonImportedFrom(root, 'OpenAI', 'anthropic')).toBeNull()
+    expect(isPythonImportedFrom(root, 'OpenAI', /openai/)).toBe('openai')
+  })
+})
+
+// ============================================================================
+// Call Analysis
+// ============================================================================
+
+describe('Python Call Analysis', () => {
+  it('isPythonCallTo detects direct call', () => {
+    const root = parse('print("hello")\n')
+    const calls = findNodesOfType(root, 'call')
+    expect(calls).toHaveLength(1)
+    expect(isPythonCallTo(calls[0], 'print')).toBe(true)
+    expect(isPythonCallTo(calls[0], 'len')).toBe(false)
+  })
+
+  it('isPythonCallTo detects attribute call (obj.method)', () => {
+    const root = parse('obj.execute(query)\n')
+    const calls = findNodesOfType(root, 'call')
+    expect(isPythonCallTo(calls[0], 'execute')).toBe(true)
+    expect(isPythonCallTo(calls[0], 'obj')).toBe(false)
+  })
+
+  it('resolvePythonCallTarget for member call', () => {
+    const root = parse('os.system(cmd)\n')
+    const calls = findNodesOfType(root, 'call')
+    const target = resolvePythonCallTarget(calls[0])
+    expect(target).toEqual({
+      type: 'member',
+      name: 'system',
+      object: 'os',
+      fullChain: 'os.system',
+      aliased: false,
+    })
+  })
+
+  it('resolvePythonCallTarget for direct call', () => {
+    const root = parse('print("hi")\n')
+    const calls = findNodesOfType(root, 'call')
+    const target = resolvePythonCallTarget(calls[0])
+    expect(target).toEqual({
+      type: 'direct',
+      name: 'print',
+      fullChain: 'print',
+      aliased: false,
+    })
+  })
+
+  it('resolvePythonCallTarget for self.method()', () => {
+    const root = parse('self.process(data)\n')
+    const calls = findNodesOfType(root, 'call')
+    const target = resolvePythonCallTarget(calls[0])
+    expect(target).toEqual({
+      type: 'member',
+      name: 'process',
+      object: 'self',
+      fullChain: 'self.process',
+      aliased: false,
+    })
+  })
+
+  it('resolvePythonCallTarget resolves aliases', () => {
+    const root = parse('e = eval\ne("code")\n')
+    const aliases = collectPythonVariableAliases(root)
+    const calls = findNodesOfType(root, 'call')
+    const target = resolvePythonCallTarget(calls[0], aliases)
+    expect(target!.name).toBe('eval')
+    expect(target!.aliased).toBe(true)
+  })
+
+  it('getPythonCallArguments returns positional args', () => {
+    const root = parse('foo(1, 2, 3)\n')
+    const calls = findNodesOfType(root, 'call')
+    const args = getPythonCallArguments(calls[0])
+    expect(args).toHaveLength(3)
+  })
+
+  it('getPythonCallArguments excludes keyword args', () => {
+    const root = parse('foo(1, name="bar")\n')
+    const calls = findNodesOfType(root, 'call')
+    const args = getPythonCallArguments(calls[0])
+    expect(args).toHaveLength(1)
+  })
+
+  it('getPythonKeywordArgument extracts keyword arg value', () => {
+    const root = parse("client.create(model='gpt-4', temperature=0.5)\n")
+    const calls = findNodesOfType(root, 'call')
+    const modelArg = getPythonKeywordArgument(calls[0], 'model')
+    expect(modelArg).not.toBeNull()
+    expect(modelArg!.type).toBe('string')
+
+    const tempArg = getPythonKeywordArgument(calls[0], 'temperature')
+    expect(tempArg).not.toBeNull()
+    expect(tempArg!.text).toBe('0.5')
+  })
+
+  it('getPythonKeywordArgument returns null for missing keyword', () => {
+    const root = parse("foo(x=1)\n")
+    const calls = findNodesOfType(root, 'call')
+    expect(getPythonKeywordArgument(calls[0], 'y')).toBeNull()
+  })
+
+  it('chained calls: client.chat.completions.create()', () => {
+    const root = parse("client.chat.completions.create(model='gpt-4')\n")
+    const calls = findNodesOfType(root, 'call')
+    const target = resolvePythonCallTarget(calls[0])
+    expect(target!.name).toBe('create')
+    expect(target!.object).toBe('client.chat.completions')
+    expect(target!.fullChain).toBe('client.chat.completions.create')
+  })
+
+  it('isPythonMethodCallOn matches object and method', () => {
+    const root = parse('os.system("ls")\n')
+    const calls = findNodesOfType(root, 'call')
+    expect(isPythonMethodCallOn(calls[0], 'os', 'system')).toBe(true)
+    expect(isPythonMethodCallOn(calls[0], 'os', 'exec')).toBe(false)
+    expect(isPythonMethodCallOn(calls[0], 'sys', 'system')).toBe(false)
+  })
+
+  it('isPythonMethodCallOn supports regex for object', () => {
+    const root = parse('cursor.execute(query)\n')
+    const calls = findNodesOfType(root, 'call')
+    expect(isPythonMethodCallOn(calls[0], /cursor|conn/, 'execute')).toBe(true)
+  })
+
+  it('getPythonAttributeChain resolves nested attributes', () => {
+    const root = parse('a.b.c.d\n')
+    const attrs = findNodesOfType(root, 'attribute')
+    // The outermost attribute node represents the full chain
+    const outermost = attrs.find(a => a.text === 'a.b.c.d')!
+    const chain = getPythonAttributeChain(outermost)
+    expect(chain).toEqual(['a', 'b', 'c', 'd'])
+  })
+
+  it('getPythonCallArgument returns Nth arg', () => {
+    const root = parse('foo("a", "b", "c")\n')
+    const calls = findNodesOfType(root, 'call')
+    const arg0 = getPythonCallArgument(calls[0], 0)
+    const arg1 = getPythonCallArgument(calls[0], 1)
+    const arg3 = getPythonCallArgument(calls[0], 3)
+    expect(arg0).not.toBeNull()
+    expect(arg1).not.toBeNull()
+    expect(arg3).toBeNull()
+  })
+})
+
+// ============================================================================
+// Scope & Variable Analysis
+// ============================================================================
+
+describe('Python Scope Analysis', () => {
+  it('detects simple assignment: x = 1', () => {
+    const root = parse('x = 1\n')
+    const vars = collectPythonVariableDeclarations(root)
+    expect(vars.has('x')).toBe(true)
+    expect(vars.get('x')!.kind).toBe('assignment')
+    expect(vars.get('x')!.initValue).not.toBeNull()
+  })
+
+  it('detects type-annotated assignment: x: int = 1', () => {
+    const root = parse('x: int = 1\n')
+    const vars = collectPythonVariableDeclarations(root)
+    expect(vars.has('x')).toBe(true)
+    expect(vars.get('x')!.kind).toBe('annotated')
+  })
+
+  it('detects function parameters (positional, keyword, *args, **kwargs)', () => {
+    const root = parse('def foo(a, b=1, *args, **kwargs):\n    pass\n')
+    const vars = collectPythonVariableDeclarations(root)
+    expect(vars.get('a')?.kind).toBe('parameter')
+    expect(vars.get('b')?.kind).toBe('parameter')
+    expect(vars.get('args')?.kind).toBe('parameter')
+    expect(vars.get('kwargs')?.kind).toBe('parameter')
+  })
+
+  it('detects for-loop variable', () => {
+    const root = parse('for x in items:\n    pass\n')
+    const vars = collectPythonVariableDeclarations(root)
+    expect(vars.has('x')).toBe(true)
+    expect(vars.get('x')!.kind).toBe('for')
+  })
+
+  it('detects with-statement variable', () => {
+    const root = parse("with open('f') as fh:\n    pass\n")
+    const vars = collectPythonVariableDeclarations(root)
+    expect(vars.has('fh')).toBe(true)
+    expect(vars.get('fh')!.kind).toBe('with')
+  })
+
+  it('detects except clause variable: except Exception as e', () => {
+    const root = parse('try:\n    pass\nexcept Exception as e:\n    pass\n')
+    const vars = collectPythonVariableDeclarations(root)
+    expect(vars.has('e')).toBe(true)
+    expect(vars.get('e')!.kind).toBe('except')
+  })
+
+  it('detects class attribute assignments', () => {
+    const root = parse('class Foo:\n    x = 1\n    y: str = "hello"\n')
+    const vars = collectPythonVariableDeclarations(root)
+    expect(vars.has('x')).toBe(true)
+    expect(vars.has('y')).toBe(true)
+  })
+
+  it('collectPythonVariableAliases tracks aliases', () => {
+    const root = parse('e = eval\nr = os.system\n')
+    const aliases = collectPythonVariableAliases(root)
+    expect(aliases.get('e')).toBe('eval')
+    expect(aliases.get('r')).toBe('os.system')
+  })
+
+  it('collectPythonVariableAliases ignores non-identifier/attribute values', () => {
+    const root = parse('x = 1\ny = [1, 2]\n')
+    const aliases = collectPythonVariableAliases(root)
+    expect(aliases.has('x')).toBe(false)
+    expect(aliases.has('y')).toBe(false)
+  })
+})
+
+// ============================================================================
+// String Analysis
+// ============================================================================
+
+describe('Python String Analysis', () => {
+  it('extracts single-quoted string value', () => {
+    const root = parse("x = 'hello'\n")
+    const strings = findNodesOfType(root, 'string')
+    expect(extractPythonStringValue(strings[0])).toBe('hello')
+  })
+
+  it('extracts double-quoted string value', () => {
+    const root = parse('x = "hello"\n')
+    const strings = findNodesOfType(root, 'string')
+    expect(extractPythonStringValue(strings[0])).toBe('hello')
+  })
+
+  it('extracts triple-quoted string value', () => {
+    const root = parse('x = """hello world"""\n')
+    const strings = findNodesOfType(root, 'string')
+    expect(extractPythonStringValue(strings[0])).toBe('hello world')
+  })
+
+  it('detects f-strings', () => {
+    const root = parse("msg = f'hello {name}'\n")
+    const strings = findNodesOfType(root, 'string')
+    const fstr = strings.find(s => isPythonFString(s))
+    expect(fstr).toBeDefined()
+  })
+
+  it('non-f-strings are not detected as f-strings', () => {
+    const root = parse("msg = 'hello'\n")
+    const strings = findNodesOfType(root, 'string')
+    expect(isPythonFString(strings[0])).toBe(false)
+  })
+
+  it('getPythonFStringParts separates static and dynamic parts', () => {
+    const root = parse("msg = f'hello {name} world'\n")
+    const strings = findNodesOfType(root, 'string')
+    const fstr = strings.find(s => isPythonFString(s))!
+    const parts = getPythonFStringParts(fstr)
+    expect(parts.staticParts).toEqual(['hello ', ' world'])
+    expect(parts.dynamicParts).toHaveLength(1)
+  })
+
+  it('getPythonFirstStringArg extracts first string arg from call', () => {
+    const root = parse("print('hello')\n")
+    const calls = findNodesOfType(root, 'call')
+    expect(getPythonFirstStringArg(calls[0])).toBe('hello')
+  })
+
+  it('getPythonFirstStringArg returns null for non-string first arg', () => {
+    const root = parse('print(123)\n')
+    const calls = findNodesOfType(root, 'call')
+    expect(getPythonFirstStringArg(calls[0])).toBeNull()
+  })
+
+  it('extractPythonStringValue returns null for non-string nodes', () => {
+    const root = parse('x = 123\n')
+    const ints = findNodesOfType(root, 'integer')
+    expect(extractPythonStringValue(ints[0])).toBeNull()
+  })
+})
+
+// ============================================================================
+// Control Flow
+// ============================================================================
+
+describe('Python Control Flow', () => {
+  it('isInsidePythonTryExcept detects try body', () => {
+    const root = parse('try:\n    risky()\nexcept:\n    pass\n')
+    const calls = findNodesOfType(root, 'call')
+    // risky() is inside try body
+    const riskyCall = calls.find(c => c.text.includes('risky'))!
+    expect(isInsidePythonTryExcept(riskyCall)).toBe(true)
+  })
+
+  it('isInsidePythonTryExcept returns false for except body', () => {
+    const root = parse('try:\n    pass\nexcept:\n    fallback()\n')
+    const calls = findNodesOfType(root, 'call')
+    const fallbackCall = calls.find(c => c.text.includes('fallback'))!
+    expect(isInsidePythonTryExcept(fallbackCall)).toBe(false)
+  })
+
+  it('isInsidePythonExceptClause detects except body', () => {
+    const root = parse('try:\n    pass\nexcept Exception as e:\n    handle(e)\n')
+    const calls = findNodesOfType(root, 'call')
+    const handleCall = calls.find(c => c.text.includes('handle'))!
+    expect(isInsidePythonExceptClause(handleCall)).toBe(true)
+  })
+
+  it('isInsidePythonExceptClause returns false for try body', () => {
+    const root = parse('try:\n    risky()\nexcept:\n    pass\n')
+    const calls = findNodesOfType(root, 'call')
+    expect(isInsidePythonExceptClause(calls[0])).toBe(false)
+  })
+
+  it('getEnclosingPythonBlock finds block node', () => {
+    const root = parse('def foo():\n    x = 1\n')
+    const assignments = findNodesOfType(root, 'assignment')
+    const block = getEnclosingPythonBlock(assignments[0])
+    expect(block).not.toBeNull()
+    expect(block!.type).toBe('block')
+  })
+
+  it('getEnclosingPythonBlock returns null at module level', () => {
+    const root = parse('x = 1\n')
+    const assignments = findNodesOfType(root, 'assignment')
+    // Module level: expression_statement > assignment
+    // No block ancestor at module level
+    const block = getEnclosingPythonBlock(assignments[0])
+    expect(block).toBeNull()
+  })
+
+  it('findEnclosingFunction works from inside except clause', () => {
+    const root = parse('def foo():\n    try:\n        pass\n    except:\n        x = 1\n')
+    const assignments = findNodesOfType(root, 'assignment')
+    const fn = findNodesOfType(root, 'function_definition')
+    expect(assignments).toHaveLength(1)
+    // Import from parse/ast
+    const { findEnclosingFunction } = require('../../parse/ast')
+    const enclosing = findEnclosingFunction(assignments[0])
+    expect(enclosing).not.toBeNull()
+    expect(enclosing.childForFieldName('name').text).toBe('foo')
+  })
+})
+
+// ============================================================================
+// Decorator Detection
+// ============================================================================
+
+describe('Python Decorator Detection', () => {
+  it('getPythonDecorators returns decorator names', () => {
+    const root = parse('@app.route("/api")\n@login_required\ndef handler():\n    pass\n')
+    const funcs = findNodesOfType(root, 'function_definition')
+    const decorators = getPythonDecorators(funcs[0])
+    expect(decorators).toContain('app.route')
+    expect(decorators).toContain('login_required')
+  })
+
+  it('getPythonDecorators returns empty for undecorated functions', () => {
+    const root = parse('def plain():\n    pass\n')
+    const funcs = findNodesOfType(root, 'function_definition')
+    expect(getPythonDecorators(funcs[0])).toEqual([])
+  })
+
+  it('hasPythonDecorator checks for specific decorator', () => {
+    const root = parse('@staticmethod\ndef foo():\n    pass\n')
+    const funcs = findNodesOfType(root, 'function_definition')
+    expect(hasPythonDecorator(funcs[0], 'staticmethod')).toBe(true)
+    expect(hasPythonDecorator(funcs[0], 'classmethod')).toBe(false)
+  })
+
+  it('hasPythonDecorator supports regex', () => {
+    const root = parse('@app.route("/api")\ndef handler():\n    pass\n')
+    const funcs = findNodesOfType(root, 'function_definition')
+    expect(hasPythonDecorator(funcs[0], /app\.route/)).toBe(true)
+    expect(hasPythonDecorator(funcs[0], /flask/)).toBe(false)
+  })
+
+  it('handles decorators on class methods', () => {
+    const root = parse('class Foo:\n    @staticmethod\n    def bar():\n        pass\n')
+    const funcs = findNodesOfType(root, 'function_definition')
+    const barFunc = funcs.find(f => f.childForFieldName('name')?.text === 'bar')!
+    expect(getPythonDecorators(barFunc)).toContain('staticmethod')
+  })
+})