app/agent.py

"""
agent.py - ReAct agent for VulnGraph

The loop:
    1. THINK  — agent reasons about what it knows and what it needs
    2. ACT    — agent calls a tool with specific parameters
    3. OBSERVE — agent reads the tool result
    4. REPEAT  — until agent has enough to give a final answer

For VulnGraph, the agent's job is:
    Given a finding ID → understand it → see the code → generate a patch
"""
import json
import os
import argparse
from datetime import datetime
from typing import Optional
from dotenv import load_dotenv
from pathlib import Path
BASE_DIR = Path(__file__).parent.parent.resolve()

from tools import(
    TOOL_REGISTRY,
    ToolResult
)

load_dotenv()

#Config
OLLAMA_MODEL  = os.getenv("OLLAMA_MODEL", "llama3.2:3b")


#Tool executor
def execute_tool(tool_name:str,parameters:dict)->ToolResult:
    """
    Execute a tool by name with given parameters.
    Looks up the tool in TOOL_REGISTRY and calls it.
    """
    if tool_name not in TOOL_REGISTRY:
        return ToolResult(
            tool_name=tool_name,
            status="error",
            error=f"Unknown tool '{tool_name}'Available: {list(TOOL_REGISTRY.keys())}"
        )
    tool_fn=TOOL_REGISTRY[tool_name]["function"]
    req=TOOL_REGISTRY[tool_name].get("required",[])

    #check required params
    missing=[p for p in req if p not in parameters]
    if missing:
        return ToolResult(
            tool_name=tool_name,
            status="error",
            error=f"Missing required parameters :{missing} "

        )
    try:
        return tool_fn(**parameters)
    except TypeError as e:
        return ToolResult(
            tool_name=tool_name,
            status="error",
            error=f"Parameter error:{e}"
        )
def _build_final_answer(finding_id, graph_result, kb_result,explanation_result, file_result, patch_result) -> str:
    parts = [f"# Security Analysis: {finding_id}\n"]

    if graph_result.status == "success":
        parts.append(f"## Finding Details\n{graph_result.data}\n")

    if explanation_result.status == "success":
        parts.append(f"## AI Explanation\n{explanation_result.data}\n")

    if kb_result.status == "success":
        parts.append(f"## Security References\n{kb_result.data[:500]}\n")

    if file_result and file_result.status == "success":
        parts.append(f"## Vulnerable Code\n{file_result.data}\n")

    if patch_result and patch_result.status == "success":
        parts.append(f"## Generated Patch\n{patch_result.data}\n")
    else:
        parts.append("## Patch\nPatch generation failed or was skipped.\n")

    return "\n".join(parts)
#Agent run
class Agent:
    """
    ReAct agent that analyzes vulnerability findings and generates patches.

    The agent maintains:
    - messages: full conversation history (system + user + assistant turns)
    - steps: count of reasoning steps taken
    - observations: list of tool results for final summary
    """
    def __init__(self,verbose:bool=True):
        self.verbose=verbose
        self.observations=[]
        self.tools_called=[]

    def _log(self,msg:str,prefix:str=""):
        if self.verbose:
            print(f"{prefix}{msg}")

    def run(self, finding_id: str, file_path: Optional[str] = None) -> dict:
        start_time = datetime.now()
        self._log(f"\n{'='*60}")
        self._log(f"VulnGraph Agent starting for finding: {finding_id}")
        self._log(f"{'='*60}\n")

        # Step 1 — Query attack graph
        self._log("[Step 1/5] Querying attack graph...")
        graph_result = execute_tool("query_attack_graph", {"finding_id": finding_id})
        self._log(f"Status: {graph_result.status}")
        self.observations.append({
            "step": 1,
            "tool": "query_attack_graph",
            "status": graph_result.status,
            "observation": graph_result.data[:200]
        })
        # Step 2 — Search knowledge base
        self._log("[Step 2/5] Searching knowledge base...")
        kb_query = finding_id
        if graph_result.status == "success":
            desc = graph_result.metadata.get("description", "")
            severity = graph_result.metadata.get("severity", "")
            source = graph_result.metadata.get("source", "")
            kb_query = f"{finding_id} {desc} {severity} {source}".strip()
            self._log(f"KB query: {kb_query}")

        kb_result = execute_tool("search_knowledge_base", {"query": kb_query})
        self._log(f"Status: {kb_result.status}")
        self.observations.append({
            "step": 2,
            "tool": "search_knowledge_base",
            "status": kb_result.status,
            "observation": kb_result.data[:200]
        })
        # Step 3 — Get existing explanation
        self._log("[Step 3/5] Fetching existing explanation...")
        explanation_result = execute_tool("get_finding_explanation", {"finding_id": finding_id})
        self._log(f"Status: {explanation_result.status}")
        self.observations.append({
            "step": 3,
            "tool": "get_finding_explanation",
            "status": explanation_result.status,
            "observation": explanation_result.data[:200]
        })
        # Step 4 — Get file context
        # Extract file path from graph result if not provided
        actual_file = file_path
        if not actual_file and graph_result.status == "success":
            files = graph_result.metadata.get("affected_files", [])
            for f in files:
                # Try the path as-is first
                if (BASE_DIR / f).exists():
                    actual_file = f
                    self._log(f"Discovered file from graph: {actual_file}")
                    break
                # Try stripping leading path separators
                stripped = f.lstrip("/\\")
                if (BASE_DIR / stripped).exists():
                    actual_file = stripped
                    self._log(f"Discovered file from graph: {actual_file}")
                    break

        file_result = None
        if actual_file:
            self._log(f"[Step 4/5] Getting file context for {actual_file}...")
            file_result = execute_tool("get_file_context", {
                "file_path": actual_file,
                "line_number": 1
            })
            self._log(f"Status: {file_result.status}")
            self.observations.append({
            "step": 4,
            "tool": "get_file_context",
            "status": file_result.status,
            "observation": file_result.data[:200] if file_result.status== "success" else file_result.error
            })
        else:
            self._log("[Step 4/5] No file path available — skipping file context")

        # Step 5 — Generate patch with all gathered context
        self._log("[Step 5/5] Generating patch...")
        patch_result = None
        if actual_file and file_result and file_result.status == "success":
            patch_result = execute_tool("generate_patch", {
                "file_path": actual_file,
                "line_number": graph_result.metadata.get("line_number", 1) or 1,
                "finding_id": finding_id,
                "code_context": file_result.data,
                "vulnerability_description": graph_result.data,
                "knowledge_context": kb_result.data if kb_result.status == "success" else ""
            })
            self._log(f"Status: {patch_result.status if patch_result else 'skipped'}")
            self.observations.append({
            "step": 5,
            "tool": "generate_patch",
            "status": patch_result.status if patch_result else "skipped",
            "observation": patch_result.data[:200] if patch_result and patch_result.status=="success" else ""
            })
        # Build final answer from all gathered context
        final_answer = _build_final_answer(
            finding_id, graph_result, kb_result,
            explanation_result, file_result, patch_result
        )

        duration = (datetime.now() - start_time).total_seconds()
        tools_called = ["query_attack_graph", "search_knowledge_base",
                        "get_finding_explanation"]
        if file_result:
            tools_called.append("get_file_context")
        if patch_result:
            tools_called.append("generate_patch")

        self._log(f"\nAgent completed in {duration:.1f}s")
        self._log(f"Tools used: {' → '.join(tools_called)}")
        
        return {
            "finding_id": finding_id,
            "final_answer": final_answer,
            "patch": patch_result.metadata.get("patch", {}) if patch_result and patch_result.status == "success" else {},
            "steps_taken": 5,
            "tools_called": tools_called,
            "observations": self.observations,
            "duration_sec": round(duration, 2),
            "model": OLLAMA_MODEL,
            "timestamp": datetime.now().isoformat()
        }
    
def run_agent_for_finding(finding_id:str, file_path:Optional[str]=None)->dict:
    """
    For calling agent using FASTAPI
    """
    agent=Agent(verbose=False)
    return agent.run(finding_id,file_path)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="VulnGraph ReAct Agent")
    parser.add_argument("--finding", required=True, help="Finding ID to analyze e.g. B404")
    parser.add_argument("--file", default=None, help="File path hint (optional)")
    parser.add_argument("--quiet", action="store_true", help="Suppress verbose output")
    args = parser.parse_args()

    agent = Agent(verbose=not args.quiet)
    result = agent.run(args.finding, args.file)

    print("\n" + "="*60)
    print("FINAL RESULT")
    print("="*60)
    print(f"Finding:      {result['finding_id']}")
    print(f"Steps taken:  {result['steps_taken']}")
    print(f"Tools called: {' → '.join(result['tools_called'])}")
    print(f"Duration:     {result['duration_sec']}s")
    if result.get("patch"):
        print(f"\nPatch generated: {result['patch'].get('patch_description', 'N/A')}")
    print(f"\nFinal Answer:\n{result['final_answer']}")