File size: 10,071 Bytes
d69f807 8f39037 d69f807 8f39037 d69f807 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 | """
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']}") |