app/tools.py

"""
Tool implementations for VulnGraph ReAct agent

Each tool has clear function:
Input: typed parameters
Output: ToolResult with status ,data and error

Tools:
search_knowledge_base - semantic search over CWE/OWASP/Bandit docs
get_file_context- read code around specific line number
query_attack_graph - query Neo4j for finding details and relationships
generate_patch- use LLM to generate a code fix for finding
get_finding_explanation - fetch existing LLM explanation from Neo4j
"""
import os
import json
from pathlib import Path
from typing import Optional
from dataclasses import dataclass,field
from dotenv import load_dotenv
from neo4j import GraphDatabase

load_dotenv()

#Config
NEO4J_URI      = os.getenv("NEO4J_URI", "bolt://localhost:7687")
NEO4J_USER     = os.getenv("NEO4J_USER", "neo4j")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "vulngraph123")
BASE_DIR       = Path(__file__).parent.parent.resolve()
CONTEXT_LINES  = 10  # lines of code to show around the vulnerable line

#ToolResult
@dataclass
class ToolResult:
    """
    Standard return type for all
    The agent uses status to decide whether to retry or move on.
    Data is what gets injected into agent's observation
    """
    tool_name: str
    status: str
    data: str =""
    metadata: dict=field(default_factory=dict)
    error: str=""

    def to_observation(self)->str:
        """Format result as observation string for agent's reasoning loop"""
        if self.status=="error":
            return f"[{self.tool_name}] ERROR: {self.error}"
        if self.status=="empty":
            return f"[{self.tool_name}] No results found. "
        return  f"[{self.tool_name}] \n {self.data}"
    
def search_knowledge_base(query:str,top_k: int=3)->ToolResult:
    """
    Semantic search over the VulnGraph knowledge base (CWE, OWASP, Bandit docs).
    Uses the same ChromaDB collection as the RAG pipeline.

    Args:
        query: natural language search query e.g. "subprocess command injection"
        top_k: number of documents to retrieve (default 3)

    Returns:
        ToolResult with formatted context string ready for agent observation
    """

    try:
        from rag import retrieve_context
        context =retrieve_context(
            finding_id=query,
            find_txt=query,
            top_k=top_k
        )
        if not context:
            return ToolResult(
                tool_name="search_knowledge_base",
                status="empty",
                data="",
                metadata={"query":query}
            )
        return ToolResult(
            tool_name="search_knowledge_base",
            status="success",
            data=context,
            metadata={"query":query, "top_k":top_k}
        )
    except Exception as e:
        return ToolResult(
            tool_name="search_knowledge_base",
            status="error",
            error=str(e),
            metadata={"query":query}
        )
    
def get_file_context(
    file_path: str,
    line_number: int,
    context_lines: int = CONTEXT_LINES,
    target_dir: Optional[str] = None
) -> ToolResult:
    """
    Read code around a specific line number from a file.
    The agent uses this to see the actual vulnerable code before generating a patch.

    Args:
        file_path:     relative path to file (as stored in Neo4j)
        line_number:   the vulnerable line number
        context_lines: how many lines before/after to include
        target_dir:    override base directory (for scanning external repos)

    Returns:
        ToolResult with code snippet and line numbers
    """
    try:
        base=Path(target_dir) if target_dir else BASE_DIR
        full_path=base/file_path

        if not full_path.exists():
            full_path=Path(file_path)
            if not full_path.exists():
                return ToolResult(
                tool_name="get_file_context",
                status="error",
                error=f"File not found: {file_path}",
                metadata={"file_path":file_path,"line_number":line_number}
            )
        with open(full_path,"r",encoding="utf-8",errors="replace") as f:
            lines=f.readlines()

        total_lines=len(lines)
        start=max(0,line_number-context_lines-1)
        end=min(total_lines,line_number+context_lines)

        snippet_lines=[]
        for i , line in enumerate(lines[start:end],start=start+1):
            marker=">>>"if i==line_number else "   "
            snippet_lines.append(f"{marker}{i:4d} | {line.rstrip()}")

        snippet= "\n".join(snippet_lines)
        data=f"File: {file_path}\n Vulnerable Line: {line_number}\n\n{snippet}"

        return ToolResult(
                tool_name="get_file_context",
                status="success",
                data=data,
                metadata={"file_path":file_path, "line_number":line_number, "start_line": start+1,"end_line":end,"total_lines":total_lines}
        )
    except Exception as e:
        return ToolResult(
                tool_name="get_file_context",
                status="error",
                error=str(e),
                metadata={"file_path":file_path, "line_number":line_number}
        )
    
def query_attack_graph(finding_id:str)-> ToolResult:
    """
    Query Neo4j for a finding's details, connected files, and attack path context.
    Gives the agent graph-level context — not just the finding in isolation.

    Args:
        finding_id: vulnerability/rule ID e.g. "B404", "CVE-2025-1234"

    Returns:
        ToolResult with finding details, connected files, and severity
    """
    try:
        driver=GraphDatabase.driver(NEO4J_URI,auth=(NEO4J_USER,NEO4J_PASSWORD))
        with driver.session() as session:
            res=session.run("""
                        MATCH (f:File)-[r]->(n)
                        WHERE (n:Vulnerability OR n:Secret)
                        AND (n.id =$fid OR n.rule=$fid)
                        RETURN f.path AS file_path,
                                type(r) AS relationship,
                                labels(n)[0] AS node_type,
                                CASE WHEN n:Vulnerability THEN n.id ELSE n.rule END AS id,
                                CASE WHEN n:Vulnerability THEN coalesce(n.severity,"UNKNOWN")
                                ELSE 'SECRET' END AS severity,
                                coalesce(n.source,'unknown') AS source,
                                coalesce(n.text,n.title,'') AS description,
                                n.explanation AS existing_explanation,
                                n.cwe AS cwe
                            """, fid=finding_id)
            rows=[dict(r) for r in res]
            driver.close()
        
        if not rows:
            return ToolResult(
                tool_name="query_attack_graph",
                status="empty",
                metadata={"finding_id": finding_id}
            )
        # Format as readable summary
        lines = [f"Finding: {finding_id}"]
        lines.append(f"Type: {rows[0]['node_type']}")
        lines.append(f"Severity: {rows[0]['severity']}")
        lines.append(f"Source: {rows[0]['source']}")
        if rows[0].get("cwe"):
            lines.append(f"CWE: {rows[0]['cwe']}")
        if rows[0].get("description"):
            lines.append(f"Description: {rows[0]['description'][:200]}")
        if rows[0].get("existing_explanation"):
            lines.append(f"Existing explanation: {rows[0]['existing_explanation'][:200]}")

        lines.append(f"\n Affected files ({len(rows)})")
        for row in rows:
            lines.append(f" -{row['file_path']} [{row['relationship']}]")

        return ToolResult(
            tool_name="query_attack_graph",
            status="success",
            data="\n".join(lines),
            metadata={
                "finding_id": finding_id,
                "affected_files": [r["file_path"] for r in rows],
                "severity": rows[0]["severity"],
                "source": rows[0]["source"],
                "node_type": rows[0]["node_type"],
                "description": rows[0].get("description", "")
            }
        )
    except Exception as e:
        return ToolResult(
            tool_name="query_attack_graph",
            status="error",
            error=str(e),
            metadata={"finding_id": finding_id}
        )
    
def generate_patch(
    file_path: str,
    line_number: int,
    finding_id: str,
    code_context: str,
    vulnerability_description: str,
    knowledge_context: str = ""
) -> ToolResult:
    """
    Use the LLM to generate a concrete code patch for a vulnerability.
    This is the most complex tool — it synthesizes all previous observations.

    Args:
        file_path:                  path to the vulnerable file
        line_number:                line number of the vulnerability
        finding_id:                    vulnerability ID (B404, CVE-xxx, etc.)
        code_context:                   code snippet from get_file_context
        vulnerability_description:         what the vulnerability is
        knowledge_context:              RAG-retrieved security references

    Returns:
        ToolResult with the generated patch as a unified diff
    """
    import requests as http_requests

    ollama_url =os.getenv("OLLAMA_URL","http://localhost:11434")
    ollama_model = os.getenv("OLLAMA_MODEL","llama3.2:3b")

    prompt = f"""You are a senior security engineer. Generate a concrete code patch to fix this vulnerability.

        Vulnerability: {finding_id}
        File: {file_path} (line {line_number})
        Description: {vulnerability_description}
        
        Current code:
        {code_context}
        
        Security references:
        {knowledge_context if knowledge_context else "No additional context available."}
        
        Respond ONLY with a JSON object — no preamble, no markdown:
        {{
        "patch_description": "One sentence describing what the patch does",
        "patched_code": "The fixed version of the vulnerable code section",
        "explanation": "Why this patch fixes the vulnerability",
        "breaking_changes": "Any breaking changes or migration notes, or null if none"
    }}"""
    try:
        response=http_requests.post(
            f"{ollama_url}/api/generate",
            json={
                "model": ollama_model,
                "prompt":prompt,
                "stream":False,
                "format":"json",
                "options":{"temperature":0.1,"num_predict":600}
            },
            timeout=120
        )
        response.raise_for_status()
        raw= response.json().get("response","").strip()

        #Strip markdown fences if model ignores format=json
        if raw.startswith("```"):
            raw = raw.split("```")[1]
            if raw.startswith("json"):
                raw = raw[4:]
            raw = raw.strip()

        parsed = json.loads(raw)
        required = {"patch_description", "patched_code", "explanation"}
        if not required.issubset(parsed.keys()):
            return ToolResult(
                tool_name="generate_patch",
                status="error",
                error=f"LLM response missing required fields: {parsed.keys()}",
                metadata={"finding_id": finding_id}
            )
        data= (
            f"Patch for {finding_id} in {file_path}:{line_number}\n\n"
            f"Description: {parsed['patch_description']}\n\n"
            f"Fixed code:\n{parsed['patched_code']}\n\n"
            f"Why this fixes it: {parsed['explanation']}\n"
        )
        if parsed.get("breaking_changes"):
            data+=f"\nBreaking changes: {parsed['breaking_changes']}"

        return ToolResult(
            tool_name="generate_patch",
            status="success",
            data=data,
            metadata={
                "finding_id": finding_id,
                "file_path": file_path,
                "line_number": line_number,
                "patch": parsed
            }
        )
    except json.JSONDecodeError as e:
        return ToolResult(
            tool_name="generate_patch",
            status="error",
            error=f"JSON parse failed: {e}. Raw: {raw[:200]}",
            metadata={
                "finding_id": finding_id
            }
        )
    except Exception as e:
        return ToolResult(
            tool_name="generate_patch",
            status="error",
            error=str(e),
            metadata={
                "finding_id": finding_id
            }
        )

def get_finding_explanation(finding_id: str) -> ToolResult:
    """
    Fetch the existing LLM-generated explanation for a finding from Neo4j.
    Use this to check if a finding already has an explanation before generating
    a new one — avoids redundant LLM calls.

    Args:
        finding_id: vulnerability/rule ID e.g. "B404", "CVE-2024-1234"

    Returns:
        ToolResult with explanation, why_dangerous, fix, and cwe fields
    """
    try:
        driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
        with driver.session() as session:
            result = session.run("""
                MATCH (n)
                WHERE (n:Vulnerability OR n:Secret)
                AND (n.id = $fid OR n.rule = $fid)
                AND n.explanation IS NOT NULL
                RETURN  n.explanation    AS explanation,
                        n.why_dangerous  AS why_dangerous,
                        n.fix            AS fix,
                        n.cwe            AS cwe,
                        n.llm_model      AS llm_model,
                        n.llm_version    AS prompt_version
                LIMIT 1
            """, fid=finding_id).single()
        driver.close()

        if not result:
            return ToolResult(
                tool_name="get_finding_explanation",
                status="empty",
                metadata={"finding_id": finding_id}
            )

        data = (
            f"Existing explanation for {finding_id}:\n\n"
            f"WHAT: {result['explanation']}\n\n"
            f"RISK: {result['why_dangerous']}\n\n"
            f"FIX: {result['fix']}\n\n"
            f"CWE: {result.get('cwe') or 'N/A'}\n"
            f"Model: {result.get('llm_model') or 'unknown'}"
        )

        return ToolResult(
            tool_name="get_finding_explanation",
            status="success",
            data=data,
            metadata={
                "finding_id": finding_id,
                "explanation": result["explanation"],
                "why_dangerous": result["why_dangerous"],
                "fix": result["fix"],
                "cwe": result.get("cwe"),
                "llm_model": result.get("llm_model")
            }
        )
    except Exception as e:
        return ToolResult(
            tool_name="get_finding_explanation",
            status="error",
            error=str(e),
            metadata={"finding_id": finding_id}
        )
#Tool Registry
# The agent uses this to discover available tools at runtime.
# The MCP server also reads this registry to expose tools over the protocol.
TOOL_REGISTRY = {
    "search_knowledge_base": {
        "function": search_knowledge_base,
        "description": "Search the security knowledge base (CWE, OWASP, Bandit) for information about a vulnerability type. Use this first to understand what a finding means.",
        "parameters": {
            "query": "Search query string e.g. 'subprocess command injection'",
            "top_k": "Number of documents to retrieve (default 3)"
        },
        "required": ["query"]
    },
    "get_file_context": {
        "function": get_file_context,
        "description": "Read the code around a specific line number in a file. Use this to see the actual vulnerable code before generating a patch.",
        "parameters": {
            "file_path": "Relative path to the file as stored in Neo4j",
            "line_number": "The vulnerable line number",
            "context_lines": "Lines of context around the vulnerable line (default 10)"
        },
        "required": ["file_path", "line_number"]
    },
    "query_attack_graph": {
        "function": query_attack_graph,
        "description": "Query the Neo4j attack graph for a finding's details, severity, and which files are affected. Use this to understand the scope of a vulnerability.",
        "parameters": {
            "finding_id": "Vulnerability or rule ID e.g. 'B404', 'CVE-2025-1234'"
        },
        "required": ["finding_id"]
    },
    "generate_patch": {
        "function": generate_patch,
        "description": "Generate a concrete code patch to fix a vulnerability. Use this last, after you have retrieved knowledge context and file context.",
        "parameters": {
            "file_path": "Path to the vulnerable file",
            "line_number": "Line number of the vulnerability",
            "finding_id": "Vulnerability ID",
            "code_context": "Code snippet from get_file_context",
            "vulnerability_description": "Description of the vulnerability",
            "knowledge_context": "RAG context from search_knowledge_base (optional)"
        },
        "required": ["file_path", "line_number", "finding_id", "code_context", "vulnerability_description"]
    },
    "get_finding_explanation": {
    "function": get_finding_explanation,
    "description": "Fetch the existing LLM explanation for a finding from Neo4j. Use this before generate_patch to check if an explanation already exists and understand the vulnerability context.",
    "parameters": {
        "finding_id": "Vulnerability or rule ID e.g. 'B404', 'CVE-2024-1234'"
    },
    "required": ["finding_id"]
    }
}
def get_tool_description()-> str:
    """
    Format all tool descriptions for injection into the agent's system prompt.
    The agent reads this to know what tools are available.
    """
    lines = ["Available tools:\n"]
    for name, info in TOOL_REGISTRY.items():
        lines.append(f"- {name}: {info['description']}")
        required = info.get("required", [])
        lines.append(f"  Required parameters: {', '.join(required)}")
    return "\n".join(lines)

# ── CLI test ──────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    print("Testing VulnGraph agent tools\n")

    # Test 1: Knowledge base search
    print("=" * 50)
    print("Test 1: search_knowledge_base")
    result = search_knowledge_base("subprocess command injection B404")
    print(result.to_observation()[:300])

    # Test 2: File context
    print("\n" + "=" * 50)
    print("Test 2: get_file_context")
    result = get_file_context("app/scanner.py", 30)
    print(result.to_observation()[:300])

    # Test 3: Attack graph query
    print("\n" + "=" * 50)
    print("Test 3: query_attack_graph")
    result = query_attack_graph("B404")
    print(result.to_observation())

    # Test 4: Tool registry
    print("\n" + "=" * 50)
    print("Test 4: Tool descriptions")
    print(get_tool_description())

    # Test 5: Get existing explanation
    print("\n" + "=" * 50)
    print("Test 5: get_finding_explanation")
    result = get_finding_explanation("B404")
    print(result.to_observation()[:300])