auth-yes/forum/poc-g2/cfg_poc.ts

import {
  assert,
  assertEquals,
} from "https://deno.land/std@0.224.0/testing/asserts.ts";
import {
  dirname,
  fromFileUrl,
  join,
} from "https://deno.land/std@0.224.0/path/mod.ts";
import { execTool, requireTool } from "./sys_exec.ts";

/**
 * Proof of Concept: Abstract Syntax Trees & Control Flow Graphs (Gen 2)
 *
 * Demonstrates the Adversary agent consuming a CFG. Instead of a hardcoded JSON,
 * we dynamically generate a basic flow graph by traversing an actual AST of some
 * target code, and then trace if unsanitized user input reaches a sensitive sink.
 *
 * This version uses the native `tree-sitter` CLI to produce an AST representation.
 */

// Simulated malicious or vulnerable code segment
const targetSource = `
function handleRequest(req) {
    let userInput = req.query.id; // entry

    // safe path
    let safeInput = sanitize(userInput);
    db_query(safeInput); // sink

    // vulnerable path
    let rawHeader = req.headers['user-agent']; // entry
    db_query(rawHeader); // sink
}
`;

async function generateAndAnalyzeCFG(code: string): Promise<string[]> {
  const currentDir = dirname(fromFileUrl(import.meta.url));
  const TEMP_FILE = join(currentDir, "dummy_cfg_target.js");
  const vulnerabilities: string[] = [];

  try {
    await Deno.writeTextFile(TEMP_FILE, code);

    // Call native tree-sitter parser to get XML AST
    const { code: exitCode, stdout, stderr } = await execTool("tree-sitter", [
      "parse",
      TEMP_FILE,
      "-x",
    ]);

    if (exitCode !== 0) {
      throw new Error(`Tree-sitter CLI execution failed: ${stderr || stdout}`);
    }

    // A very rudimentary data-flow tracker for local variables based on the tree-sitter XML output
    const variableTaints: Record<string, boolean> = {};

    // 1. Find variable assignments (variable_declarator)
    const varMatches = stdout.matchAll(
      /<variable_declarator.*?<identifier field="name".*?>(.*?)<\/identifier>.*?field="value".*?>(.*?)<\/variable_declarator>/gs,
    );
    for (const match of varMatches) {
      const varName = match[1];
      const valueBlock = match[2];

      let isTainted = false;

      // Simplistic check: is 'req' anywhere inside the value block?
      if (valueBlock.includes(">req<")) {
        isTainted = true;
      }

      // Check if it's assigned from a sanitize call
      if (
        valueBlock.includes("call_expression") &&
        valueBlock.includes(">sanitize<")
      ) {
        isTainted = false; // It's clean
      }

      variableTaints[varName] = isTainted;
    }

    // 2. Find function calls (call_expression)
    const callMatches = stdout.matchAll(
      /<call_expression.*?<identifier field="function".*?>(.*?)<\/identifier>.*?<arguments.*?<identifier.*?>(.*?)<\/identifier>.*?<\/arguments>.*?<\/call_expression>/gs,
    );
    for (const match of callMatches) {
      const funcName = match[1];
      const argName = match[2];

      if (funcName === "db_query") {
        if (variableTaints[argName]) {
          vulnerabilities.push(
            `Vulnerability: Unsanitized input '${argName}' reached sink 'db_query'`,
          );
        }
      }
    }
  } finally {
    try {
      await Deno.remove(TEMP_FILE);
    } catch {
      // ignore
    }
  }

  return vulnerabilities;
}

async function run() {
  const hasTreeSitter = await requireTool(
    "tree-sitter",
    "npm install -g tree-sitter-cli",
  );
  if (!hasTreeSitter) {
    console.warn(
      "⚠️ CFG Security Proving PoC skipped due to missing host dependency.",
    );
    return;
  }

  try {
    const vulns = await generateAndAnalyzeCFG(targetSource);

    console.log("Adversary Agent Dynamic CFG Analysis Results:");
    vulns.forEach((v) => console.log(` - ${v}`));

    assertEquals(vulns.length, 1);
    assert(
      vulns[0].includes("rawHeader"),
      "Expected rawHeader to flag a vulnerability",
    );

    console.log(
      "✅ CFG Security Proving PoC (Gen 2) successful: Real AST traversal traced taint to a sink.",
    );
  } catch (err) {
    console.error("❌ CFG Security Proving PoC (Gen 2) failed:", err);
    Deno.exit(1);
  }
}

if (import.meta.main) {
  console.log("Running CFG Security Proving PoC (Gen 2) tests...");
  run();
}