llm-security/tests/scanners/ast-taint-scanner.test.mjs

// ast-taint-scanner.test.mjs — Tests for the AST Python-taint scanner.
// Fixtures in tests/fixtures/ast-scan/:
//   - creds-net.py : os.environ -> variable -> requests.post (cross-statement taint)
//   - scope.py     : same var name in two functions, only one tainted (scope test)
//   - reassign.py  : tainted name rebound to a non-source value -> taint clears (#29)
//   - sinks.py     : subprocess/os.system command sinks + file-write sink (#28)
//   - sentinel.py  : os.system("touch SENTINEL") — proves the helper PARSES, never runs
//
// python3-absent and malformed-input paths use throwaway temp dirs so they are
// deterministic regardless of the host having python3.

import { describe, it, beforeEach } from 'node:test';
import assert from 'node:assert/strict';
import { resolve, join } from 'node:path';
import { fileURLToPath } from 'node:url';
import { existsSync, mkdtempSync, mkdirSync, writeFileSync, rmSync, readFileSync } from 'node:fs';
import { tmpdir } from 'node:os';
import { spawnSync } from 'node:child_process';
import { resetCounter } from '../../scanners/lib/output.mjs';
import { discoverFiles } from '../../scanners/lib/file-discovery.mjs';
import { scan } from '../../scanners/ast-taint-scanner.mjs';

const __dirname = fileURLToPath(new URL('.', import.meta.url));
const FIXTURE = resolve(__dirname, '../fixtures/ast-scan');

// which-guard (mirrors vsix-sandbox.test.mjs): skip python3-dependent assertions
// when the interpreter is absent.
const HAS_PYTHON3 = !spawnSync('python3', ['--version']).error;

describe('ast-taint-scanner: python3 present', () => {
  let discovery;

  beforeEach(async () => {
    resetCounter();
    discovery = await discoverFiles(FIXTURE);
  });

  it('flags creds -> network through an intermediate variable', { skip: !HAS_PYTHON3 }, async () => {
    const result = await scan(FIXTURE, discovery);
    assert.equal(result.status, 'ok');
    const real = result.findings.filter(f => f.severity !== 'info' && f.file.includes('creds-net.py'));
    assert.ok(real.length >= 1, `expected a taint finding on creds-net.py, got: ${result.findings.map(f => `${f.title}@${f.file}`).join('; ')}`);
    assert.ok(real.some(f => /os\.environ/.test(f.evidence || f.description)), 'should name os.environ as the source');
    assert.ok(real.some(f => /requests\.post/.test(f.evidence || f.description)), 'should name requests.post as the sink');
  });

  it('respects function scope (only the tainted use is flagged)', { skip: !HAS_PYTHON3 }, async () => {
    const result = await scan(FIXTURE, discovery);
    const scopeReal = result.findings.filter(f => f.severity !== 'info' && f.file.includes('scope.py'));
    assert.equal(scopeReal.length, 1, `scope.py should yield exactly 1 real finding, got ${scopeReal.length}: ${scopeReal.map(f => `L${f.line}`).join(', ')}`);
    assert.ok(/input/.test(scopeReal[0].evidence || scopeReal[0].description), 'the flagged finding should trace back to input');
  });

  it('clears taint when a name is reassigned to a non-source value', { skip: !HAS_PYTHON3 }, async () => {
    const result = await scan(FIXTURE, discovery);
    const real = result.findings.filter(f => f.severity !== 'info' && f.file.includes('reassign.py'));
    assert.equal(real.length, 0, `reassign.py must yield no real findings, got ${real.length}: ${real.map(f => `${f.evidence}@L${f.line}`).join('; ')}`);
  });

  it('reports tainted subprocess and os.system command sinks', { skip: !HAS_PYTHON3 }, async () => {
    const result = await scan(FIXTURE, discovery);
    const real = result.findings.filter(f => f.severity !== 'info' && f.file.includes('sinks.py'));
    assert.ok(real.some(f => /AST-CMD-EXEC: input -> subprocess\.run/.test(f.evidence)),
      `expected an input -> subprocess.run finding, got: ${real.map(f => f.evidence).join('; ')}`);
    assert.ok(real.some(f => /AST-CMD-EXEC: os\.getenv -> os\.system/.test(f.evidence)),
      `expected an os.getenv -> os.system finding, got: ${real.map(f => f.evidence).join('; ')}`);
    assert.ok(real.filter(f => /AST-CMD-EXEC/.test(f.evidence)).every(f => f.severity === 'critical'),
      'command-exec findings should be critical');
  });

  it('reports the tainted file-write sink via a write handle', { skip: !HAS_PYTHON3 }, async () => {
    const result = await scan(FIXTURE, discovery);
    const fw = result.findings.filter(f => f.file.includes('sinks.py') && /AST-FILE-WRITE/.test(f.evidence || ''));
    assert.equal(fw.length, 1, `sinks.py should yield exactly 1 AST-FILE-WRITE finding, got ${fw.length}`);
    assert.equal(fw[0].severity, 'high', 'the file-write sink is the high-severity rule');
    assert.match(fw[0].evidence, /os\.environ -> file\.write/, 'should trace os.environ into file.write');
  });

  it('emits DS-AST- ids with scanner AST', { skip: !HAS_PYTHON3 }, async () => {
    const result = await scan(FIXTURE, discovery);
    const wrong = result.findings.filter(f => !f.id.startsWith('DS-AST-') || f.scanner !== 'AST');
    assert.equal(wrong.length, 0, `Wrong id/scanner: ${wrong.map(f => `${f.id}/${f.scanner}`).join(', ')}`);
  });
});

describe('ast-taint-scanner: parse-only safety', () => {
  it('never executes the target (no SENTINEL file appears)', async () => {
    resetCounter();
    const discovery = await discoverFiles(FIXTURE);
    await scan(FIXTURE, discovery);
    // The canary would land in the process cwd or the fixture dir if executed.
    assert.equal(existsSync(resolve(process.cwd(), 'SENTINEL')), false, 'SENTINEL must not be created in cwd');
    assert.equal(existsSync(join(FIXTURE, 'SENTINEL')), false, 'SENTINEL must not be created in the fixture dir');
  });
});

describe('ast-taint-scanner: python3 absent', () => {
  it('returns status skipped when the interpreter is unavailable', async () => {
    const dir = mkdtempSync(join(tmpdir(), 'ast-nopy-'));
    try {
      mkdirSync(join(dir, '.llm-security'), { recursive: true });
      writeFileSync(
        join(dir, '.llm-security', 'policy.json'),
        JSON.stringify({ ast: { python_path: 'definitely-not-a-real-python-xyz' } }),
      );
      writeFileSync(join(dir, 'a.py'), 'import os\nk = os.environ["X"]\neval(k)\n');
      resetCounter();
      const discovery = await discoverFiles(dir);
      const result = await scan(dir, discovery);
      assert.equal(result.status, 'skipped');
      assert.equal(result.findings.length, 0);
    } finally {
      rmSync(dir, { recursive: true, force: true });
    }
  });
});

describe('ast-taint-scanner: malformed input resilience', () => {
  it('completes without throwing on a syntactically invalid .py', async () => {
    const dir = mkdtempSync(join(tmpdir(), 'ast-bad-'));
    try {
      writeFileSync(join(dir, 'broken.py'), 'def (:\n    this is not python @@@\n');
      resetCounter();
      const discovery = await discoverFiles(dir);
      const result = await scan(dir, discovery);
      assert.ok(['ok', 'skipped'].includes(result.status), `unexpected status ${result.status}`);
      assert.ok(Array.isArray(result.findings));
    } finally {
      rmSync(dir, { recursive: true, force: true });
    }
  });
});

// ---------------------------------------------------------------------------
// Wiring — OWASP map + orchestrator registration (Step 6)
// ---------------------------------------------------------------------------

describe('ast-taint-scanner: OWASP map registration', () => {
  it('AST is present in all four OWASP maps with the right codes', async () => {
    const { OWASP_MAP, OWASP_AGENTIC_MAP, OWASP_SKILLS_MAP, OWASP_MCP_MAP } =
      await import('../../scanners/lib/severity.mjs');
    assert.deepEqual(OWASP_MAP.AST, ['LLM01', 'LLM02']);
    assert.deepEqual(OWASP_AGENTIC_MAP.AST, []);
    assert.deepEqual(OWASP_SKILLS_MAP.AST, ['AST02']);
    assert.deepEqual(OWASP_MCP_MAP.AST, []);
  });
});

describe('ast-taint-scanner: orchestrator registration', () => {
  it('scan-orchestrator imports and lists the AST scanner', () => {
    const orchPath = resolve(__dirname, '../../scanners/scan-orchestrator.mjs');
    const text = readFileSync(orchPath, 'utf8');
    assert.match(text, /import\s+\{\s*scan as astScan\s*\}\s+from\s+['"]\.\/ast-taint-scanner\.mjs['"]/);
    assert.match(text, /\{\s*name:\s*'ast',\s*fn:\s*astScan\s*\}/);
  });
});

// ---------------------------------------------------------------------------
// Policy — enabled:false short-circuits to skipped (Step 6)
// ---------------------------------------------------------------------------

describe('ast-taint-scanner: policy disable', () => {
  it('enabled:false in policy.json short-circuits to skipped', async () => {
    const dir = mkdtempSync(join(tmpdir(), 'ast-off-'));
    try {
      mkdirSync(join(dir, '.llm-security'), { recursive: true });
      writeFileSync(
        join(dir, '.llm-security', 'policy.json'),
        JSON.stringify({ ast: { enabled: false } }),
      );
      writeFileSync(join(dir, 'a.py'), 'import os\nk = os.environ["X"]\neval(k)\n');
      resetCounter();
      const discovery = await discoverFiles(dir);
      const result = await scan(dir, discovery);
      assert.equal(result.status, 'skipped');
      assert.equal(result.findings.length, 0);
    } finally {
      rmSync(dir, { recursive: true, force: true });
    }
  });
});