File size: 16,043 Bytes
e4932aa |
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 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 |
"""
Database Agent - A specialized ReAct agent for MITRE ATT&CK technique retrieval
This agent provides semantic search capabilities over the MITRE ATT&CK knowledge base
with support for filtered searches by tactics, platforms, and other metadata.
"""
import os
import json
import sys
import time
from typing import List, Dict, Any, Optional, Literal
from pathlib import Path
# LangGraph and LangChain imports
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, AIMessage
from langchain.chat_models import init_chat_model
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_text_splitters import TokenTextSplitter
from langgraph.prebuilt import create_react_agent
# LangSmith imports
from langsmith import traceable, Client, get_current_run_tree
# Import prompts from the separate file
from src.agents.database_agent.prompts import DATABASE_AGENT_SYSTEM_PROMPT
# Import the cyber knowledge base
try:
from src.knowledge_base.cyber_knowledge_base import CyberKnowledgeBase
except Exception as e:
print(
f"[WARNING] Could not import CyberKnowledgeBase. Please adjust import paths. {e}"
)
sys.exit(1)
ls_client = Client(api_key=os.getenv("LANGSMITH_API_KEY"))
def truncate_to_tokens(text: str, max_tokens: int) -> str:
"""
Truncate text to a maximum number of tokens using LangChain's TokenTextSplitter.
Args:
text: The text to truncate
max_tokens: Maximum number of tokens
Returns:
Truncated text within the token limit
"""
if not text:
return ""
# Clean the text by replacing newlines with spaces
cleaned_text = text.replace("\n", " ")
# Use TokenTextSplitter to split by tokens
splitter = TokenTextSplitter(
encoding_name="o200k_base", chunk_size=max_tokens, chunk_overlap=0
)
chunks = splitter.split_text(cleaned_text)
return chunks[0] if chunks else ""
class DatabaseAgent:
"""
A specialized ReAct agent for MITRE ATT&CK technique retrieval and search.
This agent provides intelligent search capabilities over the MITRE ATT&CK knowledge base,
including semantic search, filtered search, and multi-query search with RRF fusion.
"""
def __init__(
self,
kb_path: str = "./cyber_knowledge_base",
llm_client: BaseChatModel = None,
):
"""
Initialize the Database Agent.
Args:
kb_path: Path to the cyber knowledge base directory
llm_client: LLM model to use for the agent
"""
self.kb_path = kb_path
self.kb = self._init_knowledge_base()
if llm_client:
self.llm = llm_client
else:
self.llm = init_chat_model(
"google_genai:gemini-2.0-flash",
temperature=0.1,
)
print(
f"[INFO] Database Agent: Using default LLM model: google_genai:gemini-2.0-flash"
)
# Create tools
self.tools = self._create_tools()
# Create ReAct agent
self.agent = self._create_react_agent()
@traceable(name="database_agent_init_kb")
def _init_knowledge_base(self) -> CyberKnowledgeBase:
"""Initialize and load the cyber knowledge base."""
kb = CyberKnowledgeBase()
if kb.load_knowledge_base(self.kb_path):
print("[SUCCESS] Database Agent: Loaded existing knowledge base")
return kb
else:
print(
f"[ERROR] Database Agent: Could not load knowledge base from {self.kb_path}"
)
print("Please ensure the knowledge base is built and available.")
raise RuntimeError("Knowledge base not available")
@traceable(name="database_agent_format_results")
def _format_results_as_json(self, results) -> List[Dict[str, Any]]:
"""Format search results as structured JSON."""
output = []
for doc in results:
technique_info = {
"attack_id": doc.metadata.get("attack_id", "Unknown"),
"name": doc.metadata.get("name", "Unknown"),
"tactics": [
t.strip()
for t in doc.metadata.get("tactics", "").split(",")
if t.strip()
],
"platforms": [
p.strip()
for p in doc.metadata.get("platforms", "").split(",")
if p.strip()
],
"description": truncate_to_tokens(doc.page_content, 300),
"relevance_score": doc.metadata.get("relevance_score", None),
"rrf_score": doc.metadata.get("rrf_score", None),
"mitigation_count": doc.metadata.get("mitigation_count", 0),
# "mitigations": truncate_to_tokens(
# doc.metadata.get("mitigations", ""), 50
# ),
}
output.append(technique_info)
return output
def _log_search_metrics(
self,
search_type: str,
query: str,
results_count: int,
execution_time: float,
success: bool,
):
"""Log search performance metrics to LangSmith."""
try:
current_run = get_current_run_tree()
if current_run:
ls_client.create_feedback(
run_id=current_run.id,
key="database_search_performance",
score=1.0 if success else 0.0,
value={
"search_type": search_type,
"query": query,
"results_count": results_count,
"execution_time": execution_time,
"success": success,
},
)
except Exception as e:
print(f"Failed to log search metrics: {e}")
def _log_agent_performance(
self, query: str, message_count: int, execution_time: float, success: bool
):
"""Log overall agent performance metrics."""
try:
current_run = get_current_run_tree()
if current_run:
ls_client.create_feedback(
run_id=current_run.id,
key="database_agent_performance",
score=1.0 if success else 0.0,
value={
"query": query,
"message_count": message_count,
"execution_time": execution_time,
"success": success,
"agent_type": "database_search",
},
)
except Exception as e:
print(f"Failed to log agent metrics: {e}")
def _create_tools(self):
"""Create the search tools for the Database Agent."""
@tool
@traceable(name="database_search_techniques")
def search_techniques(query: str, top_k: int = 5) -> str:
"""
Search for MITRE ATT&CK techniques using semantic search.
Args:
query: Search query string
top_k: Number of results to return (default: 5, max: 20)
Returns:
JSON string with search results containing technique details
"""
start_time = time.time()
try:
# Limit top_k for performance
top_k = min(max(top_k, 1), 20) # Ensure top_k is between 1 and 20
# Single query search
results = self.kb.search(query, top_k=top_k)
techniques = self._format_results_as_json(results)
execution_time = time.time() - start_time
self._log_search_metrics(
"single_query", query, len(techniques), execution_time, True
)
return json.dumps(
{
"search_type": "single_query",
"query": query,
"techniques": techniques,
"total_results": len(techniques),
},
indent=2,
)
except Exception as e:
execution_time = time.time() - start_time
self._log_search_metrics(
"single_query", query, 0, execution_time, False
)
return json.dumps(
{
"error": str(e),
"techniques": [],
"message": "Error occurred during search",
},
indent=2,
)
@tool
@traceable(name="database_search_techniques_filtered")
def search_techniques_filtered(
query: str,
top_k: int = 5,
filter_tactics: Optional[List[str]] = None,
filter_platforms: Optional[List[str]] = None,
) -> str:
"""
Search for MITRE ATT&CK techniques with metadata filters.
Args:
query: Search query string
top_k: Number of results to return (default: 5, max: 20)
filter_tactics: Filter by specific tactics (e.g., ['defense-evasion', 'privilege-escalation'])
filter_platforms: Filter by platforms (e.g., ['Windows', 'Linux'])
Returns:
JSON string with filtered search results
Examples of tactics: initial-access, execution, persistence, privilege-escalation,
defense-evasion, credential-access, discovery, lateral-movement, collection,
command-and-control, exfiltration, impact
Examples of platforms: Windows, macOS, Linux, AWS, Azure, GCP, SaaS, Network,
Containers, Android, iOS
"""
start_time = time.time()
try:
# Limit top_k for performance
top_k = min(max(top_k, 1), 20)
# Single query search with filters
results = self.kb.search(
query,
top_k=top_k,
filter_tactics=filter_tactics,
filter_platforms=filter_platforms,
)
techniques = self._format_results_as_json(results)
execution_time = time.time() - start_time
self._log_search_metrics(
"filtered_query", query, len(techniques), execution_time, True
)
return json.dumps(
{
"search_type": "single_query_filtered",
"query": query,
"filters": {
"tactics": filter_tactics,
"platforms": filter_platforms,
},
"techniques": techniques,
"total_results": len(techniques),
},
indent=2,
)
except Exception as e:
execution_time = time.time() - start_time
self._log_search_metrics(
"filtered_query", query, 0, execution_time, False
)
return json.dumps(
{
"error": str(e),
"techniques": [],
"message": "Error occurred during filtered search",
},
indent=2,
)
# return [search_techniques, search_techniques_filtered]
return [search_techniques]
def _create_react_agent(self):
"""Create the ReAct agent with the search tools using the prompt from prompts.py."""
return create_react_agent(
model=self.llm,
tools=self.tools,
prompt=DATABASE_AGENT_SYSTEM_PROMPT,
name="database_agent",
)
@traceable(name="database_agent_search")
def search(self, query: str, **kwargs) -> Dict[str, Any]:
"""
Search for techniques using the agent's capabilities.
Args:
query: The search query or question
**kwargs: Additional parameters passed to the agent
Returns:
Dictionary with the agent's response
"""
start_time = time.time()
try:
messages = [HumanMessage(content=query)]
response = self.agent.invoke({"messages": messages}, **kwargs)
execution_time = time.time() - start_time
self._log_agent_performance(
query, len(response.get("messages", [])), execution_time, True
)
return {
"success": True,
"messages": response["messages"],
"final_response": (
response["messages"][-1].content if response["messages"] else ""
),
}
except Exception as e:
execution_time = time.time() - start_time
self._log_agent_performance(query, 0, execution_time, False)
return {
"success": False,
"error": str(e),
"messages": [],
"final_response": f"Error during search: {str(e)}",
}
@traceable(name="database_agent_stream_search")
def stream_search(self, query: str, **kwargs):
"""
Stream the agent's search process for real-time feedback.
Args:
query: The search query or question
**kwargs: Additional parameters passed to the agent
Yields:
Streaming responses from the agent
"""
try:
messages = [HumanMessage(content=query)]
for chunk in self.agent.stream({"messages": messages}, **kwargs):
yield chunk
except Exception as e:
yield {"error": str(e)}
@traceable(name="database_agent_test")
def test_database_agent():
"""Test function to demonstrate Database Agent capabilities."""
print("Testing Database Agent...")
# Initialize agent
try:
agent = DatabaseAgent()
print("Database Agent initialized successfully")
except Exception as e:
print(f"Failed to initialize Database Agent: {e}")
return
# Test queries
test_queries = [
"Find techniques related to credential dumping and LSASS memory access",
"What are Windows-specific privilege escalation techniques?",
"Search for defense evasion techniques that work on Linux platforms",
"Find lateral movement techniques involving SMB or WMI",
"What techniques are used for persistence on macOS systems?",
]
for i, query in enumerate(test_queries, 1):
print(f"\n--- Test Query {i} ---")
print(f"Query: {query}")
print("-" * 50)
# Test regular search
result = agent.search(query)
if result["success"]:
print("Search completed successfully")
# Print last AI message (the summary)
for msg in reversed(result["messages"]):
if isinstance(msg, AIMessage) and not hasattr(msg, "tool_calls"):
print(f"Response: {msg.content[:300]}...")
break
else:
print(f"Search failed: {result['error']}")
if __name__ == "__main__":
test_database_agent()
|