import * as acorn from "npm:acorn";
import { assert, assertEquals } from "https://deno.land/std@0.224.0/testing/asserts.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.
 */

// 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
}
`;

function generateAndAnalyzeCFG(code: string): string[] {
  const ast = acorn.parse(code, { ecmaVersion: 2022 }) as any;
  const vulnerabilities: string[] = [];

  // A very rudimentary data-flow tracker for local variables
  const variableTaints: Record<string, boolean> = {};

  // Walk AST to find variable declarations and function calls
  function walk(node: any) {
    if (!node) return;

    if (node.type === "VariableDeclarator") {
      const varName = node.id.name;
      // Check if it's assigned from req (our entry point)
      let isTainted = false;
      if (node.init && node.init.type === "MemberExpression") {
          // Simplistic check for req.something
          let current = node.init;
          while (current.object) current = current.object;
          if (current.name === "req") isTainted = true;
      }

      // Check if it's assigned from a sanitize call
      if (node.init && node.init.type === "CallExpression") {
          if (node.init.callee.name === "sanitize") {
              isTainted = false; // It's clean
          }
      }

      variableTaints[varName] = isTainted;
    }

    if (node.type === "CallExpression") {
        if (node.callee.name === "db_query") {
            const arg = node.arguments[0];
            if (arg && arg.type === "Identifier") {
                if (variableTaints[arg.name]) {
                    vulnerabilities.push(`Vulnerability: Unsanitized input '${arg.name}' reached sink 'db_query'`);
                }
            }
        }
    }

    // Recurse over common blocks
    for (const key in node) {
        if (node[key] && typeof node[key] === "object") {
            walk(node[key]);
        }
    }
  }

  walk(ast);
  return vulnerabilities;
}

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

  try {
    const vulns = 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);
  }
}