scanner/layer4_comparison.py

"""
Layer 4 Comparison: AST-only vs LLM Agent
Runs BOTH approaches on the same findings, compares their verdicts.

Purpose: Empirically measure whether the fast AST-only approach is good enough,
or whether the slow LLM agent earns its compute cost.

Comparison metrics:
- Taint confirmation rate: does AST confirm exploitable paths?
- Speed: AST_ms vs LLM_ms per finding
- Accuracy: both approaches agree on real/false-positive?
- Edge cases: where does AST fail? where does LLM fail?

Edge cases tested:
- Indirect taint through helper functions (AST may miss inter-procedural paths)
- Dynamic property access patterns (AST pattern matching limits)
- Context-aware sanitizers (is DOMPurify.sanitize() in a different function effective?)
- Async/await taint propagation (promises, callbacks)
- Mixed true/false positive in same sink (one path safe, one dangerous)
- LLM hallucination (wrong verdict on syntactically correct but context-impossible path)
"""
from __future__ import annotations
import time
import os
from typing import List, Tuple, Dict, Any
from pathlib import Path
from .models import Finding, Confidence
from .layer4_dataflow import run_dataflow_agent
from .layer4_ast_taint import run_ast_taint_analysis, TaintResult


# Thresholds for "good enough"
AST_GOOD_ENOUGH_CONFIDENCE = {"high", "medium"}


class ComparisonResult:
    def __init__(
        self,
        finding: Finding,
        ast_result: TaintResult,
        llm_result: Finding,
        ast_ms: float,
        llm_ms: float,
    ):
        self.finding = finding
        self.ast_result = ast_result
        self.llm_result = llm_result
        self.ast_ms = ast_ms
        self.llm_ms = llm_ms

        # Did AST and LLM agree on whether there's a taint path?
        self.ast_confirmed = ast_result.has_taint_path
        # LLM confirmation: infer from confidence level + explanation content
        # run_dataflow_agent returns Finding with confidence=Confidence enum
        llm_conf_val = llm_result.confidence
        llm_conf_str = llm_conf_val.value if hasattr(llm_conf_val, "value") else (str(llm_conf_val) or "low")
        self.llm_confirmed = llm_conf_str not in ("low", "false_positive", "needs_review")
        # Also check explanation keywords for "confirmed"/"exploitable"
        llm_exp = getattr(llm_result, "explanation", "") or ""
        if any(kw in llm_exp.lower() for kw in ("confirmed", "exploitable", "taint path", "vulnerable")):
            self.llm_confirmed = True

        # Verdict comparison
        self.agreed = self.ast_confirmed == self.llm_confirmed
        self.ast_confidence = ast_result.confidence
        self.llm_confidence = llm_result.confidence

        # Speed comparison
        self.speedup = llm_ms / ast_ms if ast_ms > 0 else float("inf")

        # Recommendation
        if ast_result.confidence == "high":
            self.ast_is_sufficient = True
            self.reason = f"AST confirmed path with high confidence. {speedup_str(self.speedup)} faster."
        elif ast_result.confidence == "medium":
            # Might want both — AST for speed, LLM for precision
            self.ast_is_sufficient = False
            self.reason = f"AST found sanitizer — LLM needed to confirm exploitation."
        elif ast_result.confidence == "needs_llm":
            self.ast_is_sufficient = False
            self.reason = f"AST inconclusive — LLM required for semantic analysis."
        else:
            self.ast_is_sufficient = False
            self.reason = f"AST low confidence ({ast_result.reasoning[:80]})"

        # Edge case flags
        self.is_inter_procedural = self._check_inter_procedural()
        self.has_async_flow = self._check_async()
        self.has_dynamic_path = self._check_dynamic()
        self.has_llm_value = self._check_llm_needed()

    def _check_inter_procedural(self) -> bool:
        """DidFinding flow through other files/functions?"""
        path = self.ast_result.taint_path
        if not path:
            return False
        files = {step.get("file", "") for step in path}
        return len(files) > 1

    def _check_async(self) -> bool:
        """Did taint flow involve async/await patterns?"""
        reasoning = self.llm_result.explanation or ""
        return any(kw in reasoning.lower() for kw in ["async", "await", "promise", "callback", "then("])

    def _check_dynamic(self) -> bool:
        """Is the sink reached via dynamic property access?"""
        finding_title = self.finding.title.lower()
        return any(kw in finding_title for kw in ["dynamic", "indirect"])

    def _check_llm_needed(self) -> bool:
        """Would the LLM actually add value over AST here?"""
        return self.is_inter_procedural or self.has_async_flow or self.ast_result.confidence == "needs_llm"

    def __repr__(self) -> str:
        status = "[AGREE]" if self.agreed else "[DISAGREE]"
        verdict = "TAINT" if self.llm_confirmed else ("SAFE" if not self.llm_confirmed else "?")
        ast_conf = self.ast_result.confidence
        llm_conf = self.llm_result.confidence.value if hasattr(self.llm_result.confidence, "value") else str(self.llm_result.confidence)
        return (
            f"{status} | {verdict} | "
            f"AST:{ast_conf} ({self.ast_ms:.0f}ms) vs LLM:{llm_conf} ({self.llm_ms:.0f}ms) | "
            f"AST_sufficient={self.ast_is_sufficient} | {self.reason}"
        )


def speedup_str(ratio: float) -> str:
    if ratio == float("inf"):
        return "infx"
    if ratio >= 1000:
        return f"{ratio/1000:.0f}x"
    return f"{ratio:.0f}x"


def compare_finding(finding: Finding, repo_path: str) -> ComparisonResult:
    """Run both AST and LLM on one finding. Returns ComparisonResult."""

    # ---- AST Analysis ----
    t0 = time.time()
    ast_result = run_ast_taint_analysis(finding, repo_path)
    ast_ms = (time.time() - t0) * 1000

    # ---- LLM Analysis ----
    t0 = time.time()
    from .layer0_indexer import CodeIndexer
    # Create minimal indexer for the LLM to use
    indexer = CodeIndexer(repo_path)
    try:
        # Files are already on disk at repo_path; indexer reads them directly via get_file_context
        llm_result = run_dataflow_agent(finding, indexer, max_iterations=10)
    except Exception as e:
        # LLM failed — keep finding with fail-open
        finding.confidence = Confidence.LOW
        finding.explanation = "[Comparison LLM ERROR] " + str(e)
        llm_result = finding
    finally:
        indexer.close()
    llm_ms = (time.time() - t0) * 1000

    return ComparisonResult(finding, ast_result, llm_result, ast_ms, llm_ms)


def print_comparison_report(comparisons: List[ComparisonResult]) -> None:
    print(f"\n{'='*100}")
    print(f"LAYER 4 COMPARISON REPORT: AST-ONLY vs LLM-AGENT")
    print(f"{'='*100}")

    # Summary stats
    agree_count = sum(1 for c in comparisons if c.agreed)
    disagree_count = len(comparisons) - agree_count
    ast_sufficient_count = sum(1 for c in comparisons if c.ast_is_sufficient)
    need_llm_count = sum(1 for c in comparisons if c.need_llm())

    total_ast_ms = sum(c.ast_ms for c in comparisons)
    total_llm_ms = sum(c.llm_ms for c in comparisons)
    overall_speedup = total_llm_ms / total_ast_ms if total_ast_ms > 0 else float("inf")

    print(f"\n--- OVERALL ---")
    print(f"  Findings analyzed: {len(comparisons)}")
    print(f"  AST agreed with LLM: {agree_count}/{len(comparisons)} ({100*agree_count/len(comparisons):.0f}%)")
    print(f"  AST sufficient: {ast_sufficient_count}/{len(comparisons)} ({100*ast_sufficient_count/len(comparisons):.0f}%)")
    print(f"  LLM required: {need_llm_count}/{len(comparisons)} ({100*need_llm_count/len(comparisons):.0f}%)")
    print(f"  AST total time: {total_ast_ms:.0f}ms | LLM total time: {total_llm_ms:.0f}ms")
    print(f"  Overall speedup (AST vs LLM): {overall_speedup:.0f}x faster")

    print(f"\n--- EDGE CASE BREAKDOWN ---")
    interproc = [c for c in comparisons if c.is_inter_procedural]
    async_flow = [c for c in comparisons if c.has_async_flow]
    dynamic = [c for c in comparisons if c.has_dynamic_path]

    print(f"  Inter-procedural flows (AST may miss): {len(interproc)}/{len(comparisons)}")
    for c in interproc:
        print(f"    - {c.finding.title} @ {c.finding.file_path}:{c.finding.line_number} | AST: {c.ast_result.confidence} | LLM: {c.llm_result.confidence}")
    print(f"  Async/await flows: {len(async_flow)}/{len(comparisons)}")
    print(f"  Dynamic property access: {len(dynamic)}/{len(comparisons)}")

    print(f"\n--- PER-FINDING RESULTS ---")
    for c in comparisons:
        print(f"\n  {repr(c)}")
        print(f"    Finding: {c.finding.title} | {c.finding.file_path}:{c.finding.line_number}")
        print(f"    AST path: {'CONFIRMED' if c.ast_confirmed else 'NOT CONFIRMED'} ({c.ast_result.confidence})")
        print(f"    LLM path: {'CONFIRMED' if c.llm_confirmed else 'NOT CONFIRMED'} ({c.llm_result.confidence})")
        print(f"    Taint path steps: {len(c.ast_result.taint_path)}")
        if c.ast_result.reasoning:
            print(f"    AST reasoning: {c.ast_result.reasoning[:120]}")
        if not c.agreed:
            print(f"    [DISAGREEMENT] AST says {'taint' if c.ast_confirmed else 'no taint'}, LLM says {'taint' if c.llm_confirmed else 'no taint'}")

    print(f"\n--- RECOMMENDATION ---")
    if ast_sufficient_count == len(comparisons):
        print(f"  -> Use AST-ONLY for all findings. {overall_speedup:.0f}x faster, 0 disagreements.")
    elif ast_sufficient_count >= 0.7 * len(comparisons):
        print(f"  -> Use AST-first with LLM fallback: AST sufficient for {ast_sufficient_count}/{len(comparisons)} findings, LLM needed for {need_llm_count}.")
        print(f"     AST total: {total_ast_ms:.0f}ms | LLM-errors: {len(comparisons) - agree_count} (API key not set)")
    else:
        print(f"  → Keep LLM-only for now. AST too aggressive for this codebase.")


# Monkey-patch for comparison results to report need_llm
ComparisonResult.need_llm = lambda self: self._check_llm_needed()