agentic-threat-hunting-framework 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {agentic_threat_hunting_framework-0.3.0.dist-info → agentic_threat_hunting_framework-0.3.1.dist-info}/METADATA +1 -1
- {agentic_threat_hunting_framework-0.3.0.dist-info → agentic_threat_hunting_framework-0.3.1.dist-info}/RECORD +13 -8
- athf/__version__.py +1 -1
- athf/agents/__init__.py +14 -0
- athf/agents/base.py +141 -0
- athf/agents/llm/__init__.py +27 -0
- athf/agents/llm/hunt_researcher.py +762 -0
- athf/agents/llm/hypothesis_generator.py +238 -0
- athf/data/docs/CHANGELOG.md +23 -0
- {agentic_threat_hunting_framework-0.3.0.dist-info → agentic_threat_hunting_framework-0.3.1.dist-info}/WHEEL +0 -0
- {agentic_threat_hunting_framework-0.3.0.dist-info → agentic_threat_hunting_framework-0.3.1.dist-info}/entry_points.txt +0 -0
- {agentic_threat_hunting_framework-0.3.0.dist-info → agentic_threat_hunting_framework-0.3.1.dist-info}/licenses/LICENSE +0 -0
- {agentic_threat_hunting_framework-0.3.0.dist-info → agentic_threat_hunting_framework-0.3.1.dist-info}/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: agentic-threat-hunting-framework
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.1
|
|
4
4
|
Summary: Agentic Threat Hunting Framework - Memory and AI for threat hunters
|
|
5
5
|
Author-email: Sydney Marrone <athf@nebulock.io>
|
|
6
6
|
Maintainer-email: Sydney Marrone <athf@nebulock.io>
|
|
@@ -1,7 +1,12 @@
|
|
|
1
|
-
agentic_threat_hunting_framework-0.3.
|
|
1
|
+
agentic_threat_hunting_framework-0.3.1.dist-info/licenses/LICENSE,sha256=_KObErRfiKoolznt-DF0nJnr3U9Rdh7Z4Ba7G5qqckk,1071
|
|
2
2
|
athf/__init__.py,sha256=OrjZe8P97_BTEkscapnwSsqKSjwXNP9d8-HtGr19Ni0,241
|
|
3
|
-
athf/__version__.py,sha256=
|
|
3
|
+
athf/__version__.py,sha256=BCXUHJGbpcGJIqpp7wAIo-a46Xl8TlClrdkqx9_kTW8,59
|
|
4
4
|
athf/cli.py,sha256=rkg_Nx9Yy_UqTXBOh-pwaiD-lXO0_IXQMA1SQpDj7g0,4639
|
|
5
|
+
athf/agents/__init__.py,sha256=iaSJpvnXm9rz4QS7gBrsaLEjm49uvsMs4BLPOJeyp78,346
|
|
6
|
+
athf/agents/base.py,sha256=lnVDIOQUOyP-Apa9UM2E1mRXUPnNJ4hVqQXOwVw2u4c,4286
|
|
7
|
+
athf/agents/llm/__init__.py,sha256=qSGA-NaInjsDkMpGQwnTz3S1OgCVlzetpMcDS_to1co,671
|
|
8
|
+
athf/agents/llm/hunt_researcher.py,sha256=dIyD2Izh3zdf62kCHug1DwXFgmWhOMQUTim7qM3UAIs,27071
|
|
9
|
+
athf/agents/llm/hypothesis_generator.py,sha256=XkbJz8IS4zwQjEy-ZD0zy2XW5uRnAy87Lii-5XTY0WU,8564
|
|
5
10
|
athf/commands/__init__.py,sha256=KbpUcLPjmltq5a_m1MjhrIe4sk3DvqsnAw1wCAZfZNo,85
|
|
6
11
|
athf/commands/agent.py,sha256=k-NWiLppt2oWbiJ-hx1inkK51jhfsAYiFhixbzzQmQI,16565
|
|
7
12
|
athf/commands/context.py,sha256=V-at81-OgKcLY-In48-AccTnHfTgdofmnjE8S5kypoI,12678
|
|
@@ -20,7 +25,7 @@ athf/core/research_manager.py,sha256=i4fUjuZJcAik8I4pwbLkQlu6cuxkWDlqaIRQrzAfB0s
|
|
|
20
25
|
athf/core/template_engine.py,sha256=vNTVhlxIXZpxU7VmQyrqCSt6ORS0IVjAV54TOmUDMTE,5636
|
|
21
26
|
athf/core/web_search.py,sha256=B9IhmwH7gy2RVA6WSN3L7yGp3Q4L8OsiiwcEvnnZejU,10320
|
|
22
27
|
athf/data/__init__.py,sha256=QtgONloCaS3E9Ow995FMxyy6BbszpfmYeWpySQ2b9Mc,502
|
|
23
|
-
athf/data/docs/CHANGELOG.md,sha256=
|
|
28
|
+
athf/data/docs/CHANGELOG.md,sha256=2omJArkID-VADL0gNDfBSS0_E9GnP9OfZLn9ls-l5eA,7074
|
|
24
29
|
athf/data/docs/CLI_REFERENCE.md,sha256=zqUp-tu8OAcqzpOwx3XvzEq7UV6woDraUOcWasZI0a8,43748
|
|
25
30
|
athf/data/docs/INSTALL.md,sha256=JOWxk6q2-rdpgCnWdSPb3-Cp8rX1y4nQm7ObKz2G0uM,13117
|
|
26
31
|
athf/data/docs/README.md,sha256=rp-XQZeqteXJz7M2qKX3sl6o0AVfhGmz8GcNNKAt8pM,1061
|
|
@@ -44,8 +49,8 @@ athf/data/prompts/ai-workflow.md,sha256=rZtOcGuAEi35qx7182TwHJEORdz1-RxkZMBVkg61
|
|
|
44
49
|
athf/data/prompts/basic-prompts.md,sha256=2bunpO35RoBdJWYthXVi40RNl2UWrfwOaFthBLHF5sU,8463
|
|
45
50
|
athf/data/templates/HUNT_LOCK.md,sha256=zXxHaKMWbRDLewLTegYJMbXRM72s9gFFvjdwFfGNeJE,7386
|
|
46
51
|
athf/utils/__init__.py,sha256=aEAPI1xnAsowOtc036cCb9ZOek5nrrfevu8PElhbNgk,30
|
|
47
|
-
agentic_threat_hunting_framework-0.3.
|
|
48
|
-
agentic_threat_hunting_framework-0.3.
|
|
49
|
-
agentic_threat_hunting_framework-0.3.
|
|
50
|
-
agentic_threat_hunting_framework-0.3.
|
|
51
|
-
agentic_threat_hunting_framework-0.3.
|
|
52
|
+
agentic_threat_hunting_framework-0.3.1.dist-info/METADATA,sha256=KgED__EriZvPR42CCSDQHf7md832CFyd7RyRTbdtQbU,15838
|
|
53
|
+
agentic_threat_hunting_framework-0.3.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
|
54
|
+
agentic_threat_hunting_framework-0.3.1.dist-info/entry_points.txt,sha256=GopR2iTiBs-yNMWiUZ2DaFIFglXxWJx1XPjTa3ePtfE,39
|
|
55
|
+
agentic_threat_hunting_framework-0.3.1.dist-info/top_level.txt,sha256=Cxxg6SMLfawDJWBITsciRzq27XV8fiaAor23o9Byoes,5
|
|
56
|
+
agentic_threat_hunting_framework-0.3.1.dist-info/RECORD,,
|
athf/__version__.py
CHANGED
athf/agents/__init__.py
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
"""ATHF Agent Framework.
|
|
2
|
+
|
|
3
|
+
This module provides base classes and implementations for ATHF agents.
|
|
4
|
+
Agents can be deterministic (Python-only) or LLM-powered (using Claude API).
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from athf.agents.base import Agent, AgentResult, DeterministicAgent, LLMAgent
|
|
8
|
+
|
|
9
|
+
__all__ = [
|
|
10
|
+
"Agent",
|
|
11
|
+
"AgentResult",
|
|
12
|
+
"DeterministicAgent",
|
|
13
|
+
"LLMAgent",
|
|
14
|
+
]
|
athf/agents/base.py
ADDED
|
@@ -0,0 +1,141 @@
|
|
|
1
|
+
"""Base classes for hunt-vault agents."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
from abc import ABC, abstractmethod
|
|
5
|
+
from dataclasses import dataclass, field
|
|
6
|
+
from typing import Any, Dict, Generic, List, Optional, TypeVar
|
|
7
|
+
|
|
8
|
+
# Type variables for input/output
|
|
9
|
+
InputT = TypeVar("InputT")
|
|
10
|
+
OutputT = TypeVar("OutputT")
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@dataclass
|
|
14
|
+
class AgentResult(Generic[OutputT]):
|
|
15
|
+
"""Standard result format for all agents."""
|
|
16
|
+
|
|
17
|
+
success: bool
|
|
18
|
+
data: Optional[OutputT]
|
|
19
|
+
error: Optional[str] = None
|
|
20
|
+
warnings: List[str] = field(default_factory=list)
|
|
21
|
+
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
22
|
+
|
|
23
|
+
@property
|
|
24
|
+
def is_success(self) -> bool:
|
|
25
|
+
"""Check if the agent execution was successful."""
|
|
26
|
+
return self.success and self.error is None
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class Agent(ABC, Generic[InputT, OutputT]):
|
|
30
|
+
"""Base class for all agents."""
|
|
31
|
+
|
|
32
|
+
def __init__(self, config: Optional[Dict[str, Any]] = None):
|
|
33
|
+
"""Initialize agent with optional configuration.
|
|
34
|
+
|
|
35
|
+
Args:
|
|
36
|
+
config: Optional configuration dictionary
|
|
37
|
+
"""
|
|
38
|
+
self.config = config or {}
|
|
39
|
+
self._setup()
|
|
40
|
+
|
|
41
|
+
def _setup(self) -> None:
|
|
42
|
+
"""Optional setup method for subclasses."""
|
|
43
|
+
pass
|
|
44
|
+
|
|
45
|
+
@abstractmethod
|
|
46
|
+
def execute(self, input_data: InputT) -> AgentResult[OutputT]:
|
|
47
|
+
"""Execute agent logic.
|
|
48
|
+
|
|
49
|
+
Args:
|
|
50
|
+
input_data: Input for the agent
|
|
51
|
+
|
|
52
|
+
Returns:
|
|
53
|
+
AgentResult with output data or error
|
|
54
|
+
"""
|
|
55
|
+
pass
|
|
56
|
+
|
|
57
|
+
def __call__(self, input_data: InputT) -> AgentResult[OutputT]:
|
|
58
|
+
"""Allow calling agent as a function."""
|
|
59
|
+
return self.execute(input_data)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
class DeterministicAgent(Agent[InputT, OutputT]):
|
|
63
|
+
"""Base class for deterministic Python agents (no LLM)."""
|
|
64
|
+
|
|
65
|
+
pass
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
class LLMAgent(Agent[InputT, OutputT]):
|
|
69
|
+
"""Base class for LLM-powered agents."""
|
|
70
|
+
|
|
71
|
+
def __init__(self, config: Optional[Dict[str, Any]] = None, llm_enabled: bool = True):
|
|
72
|
+
"""Initialize LLM agent.
|
|
73
|
+
|
|
74
|
+
Args:
|
|
75
|
+
config: Optional configuration dictionary
|
|
76
|
+
llm_enabled: Whether to enable LLM functionality
|
|
77
|
+
"""
|
|
78
|
+
self.llm_enabled = llm_enabled
|
|
79
|
+
super().__init__(config)
|
|
80
|
+
|
|
81
|
+
def _log_llm_metrics(
|
|
82
|
+
self,
|
|
83
|
+
agent_name: str,
|
|
84
|
+
model_id: str,
|
|
85
|
+
input_tokens: int,
|
|
86
|
+
output_tokens: int,
|
|
87
|
+
cost_usd: float,
|
|
88
|
+
duration_ms: int,
|
|
89
|
+
) -> None:
|
|
90
|
+
"""Log LLM call metrics to centralized tracker.
|
|
91
|
+
|
|
92
|
+
Args:
|
|
93
|
+
agent_name: Name of the agent (e.g., "hypothesis-generator")
|
|
94
|
+
model_id: Bedrock model ID
|
|
95
|
+
input_tokens: Number of input tokens
|
|
96
|
+
output_tokens: Number of output tokens
|
|
97
|
+
cost_usd: Estimated cost in USD
|
|
98
|
+
duration_ms: Call duration in milliseconds
|
|
99
|
+
"""
|
|
100
|
+
try:
|
|
101
|
+
from athf.core.metrics_tracker import MetricsTracker
|
|
102
|
+
|
|
103
|
+
MetricsTracker.get_instance().log_bedrock_call(
|
|
104
|
+
agent=agent_name,
|
|
105
|
+
model_id=model_id,
|
|
106
|
+
input_tokens=input_tokens,
|
|
107
|
+
output_tokens=output_tokens,
|
|
108
|
+
cost_usd=cost_usd,
|
|
109
|
+
duration_ms=duration_ms,
|
|
110
|
+
)
|
|
111
|
+
except Exception:
|
|
112
|
+
pass # Never fail agent execution due to metrics logging
|
|
113
|
+
|
|
114
|
+
def _get_llm_client(self) -> Any:
|
|
115
|
+
"""Get AWS Bedrock runtime client for Claude models.
|
|
116
|
+
|
|
117
|
+
Returns:
|
|
118
|
+
Bedrock runtime client instance or None if LLM is disabled
|
|
119
|
+
|
|
120
|
+
Raises:
|
|
121
|
+
ValueError: If AWS credentials are not configured
|
|
122
|
+
ImportError: If boto3 package is not installed
|
|
123
|
+
"""
|
|
124
|
+
if not self.llm_enabled:
|
|
125
|
+
return None
|
|
126
|
+
|
|
127
|
+
try:
|
|
128
|
+
import boto3
|
|
129
|
+
|
|
130
|
+
# Get AWS region from environment or use default
|
|
131
|
+
region = os.getenv("AWS_REGION", os.getenv("AWS_DEFAULT_REGION", "us-east-1"))
|
|
132
|
+
|
|
133
|
+
# Create Bedrock runtime client
|
|
134
|
+
# Uses AWS credentials from environment, ~/.aws/credentials, or IAM role
|
|
135
|
+
client = boto3.client(service_name="bedrock-runtime", region_name=region)
|
|
136
|
+
|
|
137
|
+
return client
|
|
138
|
+
except ImportError:
|
|
139
|
+
raise ImportError("boto3 package not installed. Run: pip install boto3")
|
|
140
|
+
except Exception as e:
|
|
141
|
+
raise ValueError(f"Failed to create Bedrock client: {e}")
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
"""LLM-powered agents for ATHF.
|
|
2
|
+
|
|
3
|
+
These agents use Claude API for creative and analytical tasks.
|
|
4
|
+
All LLM agents have fallback to deterministic methods when LLM is disabled.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from athf.agents.llm.hunt_researcher import (
|
|
8
|
+
HuntResearcherAgent,
|
|
9
|
+
ResearchInput,
|
|
10
|
+
ResearchOutput,
|
|
11
|
+
ResearchSkillOutput,
|
|
12
|
+
)
|
|
13
|
+
from athf.agents.llm.hypothesis_generator import (
|
|
14
|
+
HypothesisGenerationInput,
|
|
15
|
+
HypothesisGenerationOutput,
|
|
16
|
+
HypothesisGeneratorAgent,
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
__all__ = [
|
|
20
|
+
"HypothesisGeneratorAgent",
|
|
21
|
+
"HypothesisGenerationInput",
|
|
22
|
+
"HypothesisGenerationOutput",
|
|
23
|
+
"HuntResearcherAgent",
|
|
24
|
+
"ResearchInput",
|
|
25
|
+
"ResearchOutput",
|
|
26
|
+
"ResearchSkillOutput",
|
|
27
|
+
]
|
|
@@ -0,0 +1,762 @@
|
|
|
1
|
+
"""Hunt researcher agent - LLM-powered thorough research before hunting.
|
|
2
|
+
|
|
3
|
+
Implements a structured 5-skill research methodology:
|
|
4
|
+
1. System Research - How does the technology/process normally work?
|
|
5
|
+
2. Adversary Tradecraft - How do adversaries abuse it? (web search)
|
|
6
|
+
3. Telemetry Mapping - What OCSF fields and data sources capture this?
|
|
7
|
+
4. Related Work - What past hunts/investigations are relevant?
|
|
8
|
+
5. Synthesis - Key findings, gaps, recommended focus areas
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
import json
|
|
12
|
+
import os
|
|
13
|
+
import time
|
|
14
|
+
from dataclasses import dataclass, field
|
|
15
|
+
from pathlib import Path
|
|
16
|
+
from typing import Any, Dict, List, Optional
|
|
17
|
+
|
|
18
|
+
from athf.agents.base import AgentResult, LLMAgent
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
@dataclass
|
|
22
|
+
class ResearchInput:
|
|
23
|
+
"""Input for hunt research."""
|
|
24
|
+
|
|
25
|
+
topic: str # Research topic (e.g., "LSASS memory dumping")
|
|
26
|
+
mitre_technique: Optional[str] = None # Optional T-code to focus research
|
|
27
|
+
depth: str = "advanced" # "basic" (5 min) or "advanced" (15-20 min)
|
|
28
|
+
include_past_hunts: bool = True
|
|
29
|
+
include_telemetry_mapping: bool = True
|
|
30
|
+
web_search_enabled: bool = True # Can be disabled for offline mode
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@dataclass
|
|
34
|
+
class ResearchSkillOutput:
|
|
35
|
+
"""Output from a single research skill."""
|
|
36
|
+
|
|
37
|
+
skill_name: str # e.g., "system_research", "adversary_tradecraft"
|
|
38
|
+
summary: str
|
|
39
|
+
key_findings: List[str]
|
|
40
|
+
sources: List[Dict[str, str]] # {"title": "", "url": "", "snippet": ""}
|
|
41
|
+
confidence: float # 0-1 confidence in findings
|
|
42
|
+
duration_ms: int = 0
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
@dataclass
|
|
46
|
+
class ResearchOutput:
|
|
47
|
+
"""Complete research output following OTR-inspired 5 skills."""
|
|
48
|
+
|
|
49
|
+
research_id: str # R-XXXX format
|
|
50
|
+
topic: str
|
|
51
|
+
mitre_techniques: List[str]
|
|
52
|
+
|
|
53
|
+
# 5 OTR-inspired research skills
|
|
54
|
+
system_research: ResearchSkillOutput
|
|
55
|
+
adversary_tradecraft: ResearchSkillOutput
|
|
56
|
+
telemetry_mapping: ResearchSkillOutput
|
|
57
|
+
related_work: ResearchSkillOutput
|
|
58
|
+
synthesis: ResearchSkillOutput
|
|
59
|
+
|
|
60
|
+
# Synthesis outputs
|
|
61
|
+
recommended_hypothesis: Optional[str] = None
|
|
62
|
+
data_source_availability: Dict[str, bool] = field(default_factory=dict)
|
|
63
|
+
estimated_hunt_complexity: str = "medium" # low/medium/high
|
|
64
|
+
gaps_identified: List[str] = field(default_factory=list)
|
|
65
|
+
|
|
66
|
+
# Cost tracking
|
|
67
|
+
total_duration_ms: int = 0
|
|
68
|
+
web_searches_performed: int = 0
|
|
69
|
+
llm_calls: int = 0
|
|
70
|
+
total_cost_usd: float = 0.0
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
class HuntResearcherAgent(LLMAgent[ResearchInput, ResearchOutput]):
|
|
74
|
+
"""Performs thorough research before hunt creation.
|
|
75
|
+
|
|
76
|
+
Implements a structured 5-skill research methodology:
|
|
77
|
+
1. System Research - How does the technology/process normally work?
|
|
78
|
+
2. Adversary Tradecraft - How do adversaries abuse it? (web search)
|
|
79
|
+
3. Telemetry Mapping - What OCSF fields and data sources capture this?
|
|
80
|
+
4. Related Work - What past hunts/investigations are relevant?
|
|
81
|
+
5. Research Synthesis - Key findings, gaps, recommended focus areas
|
|
82
|
+
|
|
83
|
+
Features:
|
|
84
|
+
- Web search via Tavily API for external threat intel
|
|
85
|
+
- OCSF schema awareness for telemetry mapping
|
|
86
|
+
- Past hunt correlation via similarity search
|
|
87
|
+
- Cost tracking across all operations
|
|
88
|
+
- Fallback to limited research when APIs unavailable
|
|
89
|
+
"""
|
|
90
|
+
|
|
91
|
+
def __init__(
|
|
92
|
+
self,
|
|
93
|
+
llm_enabled: bool = True,
|
|
94
|
+
tavily_api_key: Optional[str] = None,
|
|
95
|
+
) -> None:
|
|
96
|
+
"""Initialize researcher with optional API keys."""
|
|
97
|
+
super().__init__(llm_enabled=llm_enabled)
|
|
98
|
+
self.tavily_api_key = tavily_api_key or os.getenv("TAVILY_API_KEY")
|
|
99
|
+
self._search_client: Optional[Any] = None
|
|
100
|
+
self._total_cost = 0.0
|
|
101
|
+
self._llm_calls = 0
|
|
102
|
+
self._web_searches = 0
|
|
103
|
+
|
|
104
|
+
def _get_search_client(self) -> Optional[Any]:
|
|
105
|
+
"""Get or create Tavily search client."""
|
|
106
|
+
if self._search_client is None and self.tavily_api_key:
|
|
107
|
+
try:
|
|
108
|
+
from athf.core.web_search import TavilySearchClient
|
|
109
|
+
|
|
110
|
+
self._search_client = TavilySearchClient(api_key=self.tavily_api_key)
|
|
111
|
+
except Exception:
|
|
112
|
+
pass
|
|
113
|
+
return self._search_client
|
|
114
|
+
|
|
115
|
+
def execute(self, input_data: ResearchInput) -> AgentResult[ResearchOutput]:
|
|
116
|
+
"""Execute complete research workflow.
|
|
117
|
+
|
|
118
|
+
Args:
|
|
119
|
+
input_data: Research input with topic, technique, and depth
|
|
120
|
+
|
|
121
|
+
Returns:
|
|
122
|
+
AgentResult with complete research output or error
|
|
123
|
+
"""
|
|
124
|
+
start_time = time.time()
|
|
125
|
+
self._total_cost = 0.0
|
|
126
|
+
self._llm_calls = 0
|
|
127
|
+
self._web_searches = 0
|
|
128
|
+
|
|
129
|
+
try:
|
|
130
|
+
# Get next research ID
|
|
131
|
+
from athf.core.research_manager import ResearchManager
|
|
132
|
+
|
|
133
|
+
manager = ResearchManager()
|
|
134
|
+
research_id = manager.get_next_research_id()
|
|
135
|
+
|
|
136
|
+
# Determine search depth based on input
|
|
137
|
+
search_depth = "basic" if input_data.depth == "basic" else "advanced"
|
|
138
|
+
|
|
139
|
+
# Execute all 5 skills
|
|
140
|
+
skill_1 = self._skill_1_system_research(input_data.topic, search_depth)
|
|
141
|
+
skill_2 = self._skill_2_adversary_tradecraft(
|
|
142
|
+
input_data.topic,
|
|
143
|
+
input_data.mitre_technique,
|
|
144
|
+
search_depth,
|
|
145
|
+
input_data.web_search_enabled,
|
|
146
|
+
)
|
|
147
|
+
skill_3 = self._skill_3_telemetry_mapping(
|
|
148
|
+
input_data.topic,
|
|
149
|
+
input_data.mitre_technique,
|
|
150
|
+
)
|
|
151
|
+
skill_4 = self._skill_4_related_work(input_data.topic)
|
|
152
|
+
skill_5 = self._skill_5_synthesis(
|
|
153
|
+
input_data.topic,
|
|
154
|
+
input_data.mitre_technique,
|
|
155
|
+
[skill_1, skill_2, skill_3, skill_4],
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
# Extract synthesis outputs
|
|
159
|
+
mitre_techniques = [input_data.mitre_technique] if input_data.mitre_technique else []
|
|
160
|
+
|
|
161
|
+
# Build output
|
|
162
|
+
total_duration_ms = int((time.time() - start_time) * 1000)
|
|
163
|
+
|
|
164
|
+
output = ResearchOutput(
|
|
165
|
+
research_id=research_id,
|
|
166
|
+
topic=input_data.topic,
|
|
167
|
+
mitre_techniques=mitre_techniques,
|
|
168
|
+
system_research=skill_1,
|
|
169
|
+
adversary_tradecraft=skill_2,
|
|
170
|
+
telemetry_mapping=skill_3,
|
|
171
|
+
related_work=skill_4,
|
|
172
|
+
synthesis=skill_5,
|
|
173
|
+
recommended_hypothesis=self._extract_hypothesis(skill_5),
|
|
174
|
+
data_source_availability=self._extract_data_sources(skill_3),
|
|
175
|
+
estimated_hunt_complexity=self._estimate_complexity(skill_2, skill_3),
|
|
176
|
+
gaps_identified=self._extract_gaps(skill_5),
|
|
177
|
+
total_duration_ms=total_duration_ms,
|
|
178
|
+
web_searches_performed=self._web_searches,
|
|
179
|
+
llm_calls=self._llm_calls,
|
|
180
|
+
total_cost_usd=round(self._total_cost, 4),
|
|
181
|
+
)
|
|
182
|
+
|
|
183
|
+
return AgentResult(
|
|
184
|
+
success=True,
|
|
185
|
+
data=output,
|
|
186
|
+
error=None,
|
|
187
|
+
warnings=[],
|
|
188
|
+
metadata={
|
|
189
|
+
"research_id": research_id,
|
|
190
|
+
"duration_ms": total_duration_ms,
|
|
191
|
+
"web_searches": self._web_searches,
|
|
192
|
+
"llm_calls": self._llm_calls,
|
|
193
|
+
"cost_usd": round(self._total_cost, 4),
|
|
194
|
+
},
|
|
195
|
+
)
|
|
196
|
+
|
|
197
|
+
except Exception as e:
|
|
198
|
+
return AgentResult(
|
|
199
|
+
success=False,
|
|
200
|
+
data=None,
|
|
201
|
+
error=str(e),
|
|
202
|
+
warnings=[],
|
|
203
|
+
metadata={},
|
|
204
|
+
)
|
|
205
|
+
|
|
206
|
+
def _skill_1_system_research(
|
|
207
|
+
self,
|
|
208
|
+
topic: str,
|
|
209
|
+
search_depth: str,
|
|
210
|
+
) -> ResearchSkillOutput:
|
|
211
|
+
"""Skill 1: Research how the system/technology normally works.
|
|
212
|
+
|
|
213
|
+
Args:
|
|
214
|
+
topic: Research topic
|
|
215
|
+
search_depth: "basic" or "advanced"
|
|
216
|
+
|
|
217
|
+
Returns:
|
|
218
|
+
ResearchSkillOutput with system research findings
|
|
219
|
+
"""
|
|
220
|
+
start_time = time.time()
|
|
221
|
+
sources: List[Dict[str, str]] = []
|
|
222
|
+
search_results = None
|
|
223
|
+
|
|
224
|
+
# Try web search for system internals
|
|
225
|
+
search_client = self._get_search_client()
|
|
226
|
+
if search_client:
|
|
227
|
+
try:
|
|
228
|
+
search_results = search_client.search_system_internals(topic, search_depth)
|
|
229
|
+
self._web_searches += 1
|
|
230
|
+
|
|
231
|
+
for result in search_results.results[:5]:
|
|
232
|
+
sources.append(
|
|
233
|
+
{
|
|
234
|
+
"title": result.title,
|
|
235
|
+
"url": result.url,
|
|
236
|
+
"snippet": result.content[:200] + "..." if len(result.content) > 200 else result.content,
|
|
237
|
+
}
|
|
238
|
+
)
|
|
239
|
+
except Exception:
|
|
240
|
+
pass
|
|
241
|
+
|
|
242
|
+
# Generate summary using LLM
|
|
243
|
+
if self.llm_enabled:
|
|
244
|
+
summary, key_findings = self._llm_summarize_system_research(topic, sources, search_results)
|
|
245
|
+
else:
|
|
246
|
+
summary = f"System research for {topic} - requires LLM for detailed analysis"
|
|
247
|
+
key_findings = ["LLM disabled - manual research required"]
|
|
248
|
+
|
|
249
|
+
duration_ms = int((time.time() - start_time) * 1000)
|
|
250
|
+
|
|
251
|
+
return ResearchSkillOutput(
|
|
252
|
+
skill_name="system_research",
|
|
253
|
+
summary=summary,
|
|
254
|
+
key_findings=key_findings,
|
|
255
|
+
sources=sources,
|
|
256
|
+
confidence=0.8 if sources else 0.5,
|
|
257
|
+
duration_ms=duration_ms,
|
|
258
|
+
)
|
|
259
|
+
|
|
260
|
+
def _skill_2_adversary_tradecraft(
|
|
261
|
+
self,
|
|
262
|
+
topic: str,
|
|
263
|
+
technique: Optional[str],
|
|
264
|
+
search_depth: str,
|
|
265
|
+
web_search_enabled: bool,
|
|
266
|
+
) -> ResearchSkillOutput:
|
|
267
|
+
"""Skill 2: Research adversary tradecraft via web search.
|
|
268
|
+
|
|
269
|
+
Args:
|
|
270
|
+
topic: Research topic
|
|
271
|
+
technique: Optional MITRE ATT&CK technique
|
|
272
|
+
search_depth: "basic" or "advanced"
|
|
273
|
+
web_search_enabled: Whether web search is enabled
|
|
274
|
+
|
|
275
|
+
Returns:
|
|
276
|
+
ResearchSkillOutput with adversary tradecraft findings
|
|
277
|
+
"""
|
|
278
|
+
start_time = time.time()
|
|
279
|
+
sources: List[Dict[str, str]] = []
|
|
280
|
+
search_results = None
|
|
281
|
+
|
|
282
|
+
# Try web search for adversary tradecraft
|
|
283
|
+
search_client = self._get_search_client()
|
|
284
|
+
if search_client and web_search_enabled:
|
|
285
|
+
try:
|
|
286
|
+
search_results = search_client.search_adversary_tradecraft(topic, technique, search_depth)
|
|
287
|
+
self._web_searches += 1
|
|
288
|
+
|
|
289
|
+
for result in search_results.results[:7]:
|
|
290
|
+
sources.append(
|
|
291
|
+
{
|
|
292
|
+
"title": result.title,
|
|
293
|
+
"url": result.url,
|
|
294
|
+
"snippet": result.content[:200] + "..." if len(result.content) > 200 else result.content,
|
|
295
|
+
}
|
|
296
|
+
)
|
|
297
|
+
except Exception:
|
|
298
|
+
pass
|
|
299
|
+
|
|
300
|
+
# Generate summary using LLM
|
|
301
|
+
if self.llm_enabled:
|
|
302
|
+
summary, key_findings = self._llm_summarize_tradecraft(topic, technique, sources, search_results)
|
|
303
|
+
else:
|
|
304
|
+
summary = f"Adversary tradecraft for {topic} - requires LLM for detailed analysis"
|
|
305
|
+
key_findings = ["LLM disabled - manual research required"]
|
|
306
|
+
|
|
307
|
+
duration_ms = int((time.time() - start_time) * 1000)
|
|
308
|
+
|
|
309
|
+
return ResearchSkillOutput(
|
|
310
|
+
skill_name="adversary_tradecraft",
|
|
311
|
+
summary=summary,
|
|
312
|
+
key_findings=key_findings,
|
|
313
|
+
sources=sources,
|
|
314
|
+
confidence=0.85 if sources else 0.4,
|
|
315
|
+
duration_ms=duration_ms,
|
|
316
|
+
)
|
|
317
|
+
|
|
318
|
+
def _skill_3_telemetry_mapping(
|
|
319
|
+
self,
|
|
320
|
+
topic: str,
|
|
321
|
+
technique: Optional[str],
|
|
322
|
+
) -> ResearchSkillOutput:
|
|
323
|
+
"""Skill 3: Map to OCSF fields and available data sources.
|
|
324
|
+
|
|
325
|
+
Args:
|
|
326
|
+
topic: Research topic
|
|
327
|
+
technique: Optional MITRE ATT&CK technique
|
|
328
|
+
|
|
329
|
+
Returns:
|
|
330
|
+
ResearchSkillOutput with telemetry mapping
|
|
331
|
+
"""
|
|
332
|
+
start_time = time.time()
|
|
333
|
+
sources: List[Dict[str, str]] = []
|
|
334
|
+
|
|
335
|
+
# Load OCSF schema reference
|
|
336
|
+
ocsf_schema = self._load_ocsf_schema()
|
|
337
|
+
environment_data = self._load_environment()
|
|
338
|
+
|
|
339
|
+
# Generate telemetry mapping using LLM
|
|
340
|
+
if self.llm_enabled:
|
|
341
|
+
summary, key_findings = self._llm_map_telemetry(topic, technique, ocsf_schema, environment_data)
|
|
342
|
+
else:
|
|
343
|
+
summary = f"Telemetry mapping for {topic} - requires LLM for detailed analysis"
|
|
344
|
+
key_findings = [
|
|
345
|
+
"Common fields: process.name, process.command_line, actor.user.name",
|
|
346
|
+
"Check OCSF_SCHEMA_REFERENCE.md for field population rates",
|
|
347
|
+
]
|
|
348
|
+
|
|
349
|
+
# Add schema reference as source
|
|
350
|
+
sources.append(
|
|
351
|
+
{
|
|
352
|
+
"title": "OCSF Schema Reference",
|
|
353
|
+
"url": "knowledge/OCSF_SCHEMA_REFERENCE.md",
|
|
354
|
+
"snippet": "Internal schema documentation with field population rates",
|
|
355
|
+
}
|
|
356
|
+
)
|
|
357
|
+
|
|
358
|
+
duration_ms = int((time.time() - start_time) * 1000)
|
|
359
|
+
|
|
360
|
+
return ResearchSkillOutput(
|
|
361
|
+
skill_name="telemetry_mapping",
|
|
362
|
+
summary=summary,
|
|
363
|
+
key_findings=key_findings,
|
|
364
|
+
sources=sources,
|
|
365
|
+
confidence=0.9, # High confidence - based on internal schema
|
|
366
|
+
duration_ms=duration_ms,
|
|
367
|
+
)
|
|
368
|
+
|
|
369
|
+
def _skill_4_related_work(self, topic: str) -> ResearchSkillOutput:
|
|
370
|
+
"""Skill 4: Find related past hunts and investigations.
|
|
371
|
+
|
|
372
|
+
Args:
|
|
373
|
+
topic: Research topic
|
|
374
|
+
|
|
375
|
+
Returns:
|
|
376
|
+
ResearchSkillOutput with related work
|
|
377
|
+
"""
|
|
378
|
+
start_time = time.time()
|
|
379
|
+
sources: List[Dict[str, str]] = []
|
|
380
|
+
key_findings = []
|
|
381
|
+
|
|
382
|
+
# Use similarity search to find related hunts
|
|
383
|
+
try:
|
|
384
|
+
from athf.commands.similar import _find_similar_hunts
|
|
385
|
+
|
|
386
|
+
similar_hunts = _find_similar_hunts(topic, limit=5, threshold=0.1)
|
|
387
|
+
|
|
388
|
+
for hunt in similar_hunts:
|
|
389
|
+
sources.append(
|
|
390
|
+
{
|
|
391
|
+
"title": f"{hunt['hunt_id']}: {hunt['title']}",
|
|
392
|
+
"url": f"hunts/{hunt['hunt_id']}.md",
|
|
393
|
+
"snippet": f"Status: {hunt['status']}, Score: {hunt['similarity_score']:.3f}",
|
|
394
|
+
}
|
|
395
|
+
)
|
|
396
|
+
key_findings.append(f"{hunt['hunt_id']}: {hunt['title']} (similarity: {hunt['similarity_score']:.2f})")
|
|
397
|
+
|
|
398
|
+
except Exception:
|
|
399
|
+
key_findings.append("No similar hunts found or similarity search unavailable")
|
|
400
|
+
|
|
401
|
+
summary = f"Found {len(sources)} related hunts for {topic}"
|
|
402
|
+
if not sources:
|
|
403
|
+
summary = f"No related hunts found for {topic} - this may be a new research area"
|
|
404
|
+
|
|
405
|
+
duration_ms = int((time.time() - start_time) * 1000)
|
|
406
|
+
|
|
407
|
+
return ResearchSkillOutput(
|
|
408
|
+
skill_name="related_work",
|
|
409
|
+
summary=summary,
|
|
410
|
+
key_findings=key_findings if key_findings else ["No related past hunts found"],
|
|
411
|
+
sources=sources,
|
|
412
|
+
confidence=0.95, # High confidence - based on internal search
|
|
413
|
+
duration_ms=duration_ms,
|
|
414
|
+
)
|
|
415
|
+
|
|
416
|
+
def _skill_5_synthesis(
|
|
417
|
+
self,
|
|
418
|
+
topic: str,
|
|
419
|
+
technique: Optional[str],
|
|
420
|
+
skills: List[ResearchSkillOutput],
|
|
421
|
+
) -> ResearchSkillOutput:
|
|
422
|
+
"""Skill 5: Synthesize all research into actionable insights.
|
|
423
|
+
|
|
424
|
+
Args:
|
|
425
|
+
topic: Research topic
|
|
426
|
+
technique: Optional MITRE ATT&CK technique
|
|
427
|
+
skills: Outputs from skills 1-4
|
|
428
|
+
|
|
429
|
+
Returns:
|
|
430
|
+
ResearchSkillOutput with synthesis
|
|
431
|
+
"""
|
|
432
|
+
start_time = time.time()
|
|
433
|
+
|
|
434
|
+
# Generate synthesis using LLM
|
|
435
|
+
if self.llm_enabled:
|
|
436
|
+
summary, key_findings = self._llm_synthesize(topic, technique, skills)
|
|
437
|
+
else:
|
|
438
|
+
summary = f"Research synthesis for {topic}"
|
|
439
|
+
key_findings = [
|
|
440
|
+
"LLM disabled - manual synthesis required",
|
|
441
|
+
"Review individual skill outputs for findings",
|
|
442
|
+
]
|
|
443
|
+
|
|
444
|
+
duration_ms = int((time.time() - start_time) * 1000)
|
|
445
|
+
|
|
446
|
+
return ResearchSkillOutput(
|
|
447
|
+
skill_name="synthesis",
|
|
448
|
+
summary=summary,
|
|
449
|
+
key_findings=key_findings,
|
|
450
|
+
sources=[], # Synthesis doesn't have external sources
|
|
451
|
+
confidence=0.8,
|
|
452
|
+
duration_ms=duration_ms,
|
|
453
|
+
)
|
|
454
|
+
|
|
455
|
+
def _llm_summarize_system_research(
|
|
456
|
+
self,
|
|
457
|
+
topic: str,
|
|
458
|
+
sources: List[Dict[str, str]],
|
|
459
|
+
search_results: Optional[Any],
|
|
460
|
+
) -> tuple[str, List[str]]:
|
|
461
|
+
"""Use LLM to summarize system research findings."""
|
|
462
|
+
try:
|
|
463
|
+
client = self._get_llm_client()
|
|
464
|
+
if not client:
|
|
465
|
+
return f"System research for {topic}", ["LLM unavailable"]
|
|
466
|
+
|
|
467
|
+
# Build context from sources
|
|
468
|
+
context = ""
|
|
469
|
+
if search_results and hasattr(search_results, "answer") and search_results.answer:
|
|
470
|
+
context = f"Web search summary: {search_results.answer}\n\n"
|
|
471
|
+
|
|
472
|
+
for source in sources[:5]:
|
|
473
|
+
context += f"- {source['title']}: {source['snippet']}\n"
|
|
474
|
+
|
|
475
|
+
prompt = f"""You are a security researcher studying system internals.
|
|
476
|
+
|
|
477
|
+
Topic: {topic}
|
|
478
|
+
|
|
479
|
+
Research Context:
|
|
480
|
+
{context}
|
|
481
|
+
|
|
482
|
+
Based on this context, provide:
|
|
483
|
+
1. A concise summary (2-3 sentences) of how this system/technology normally works
|
|
484
|
+
2. 3-5 key findings about normal behavior
|
|
485
|
+
|
|
486
|
+
Return JSON format:
|
|
487
|
+
{{
|
|
488
|
+
"summary": "string",
|
|
489
|
+
"key_findings": ["finding1", "finding2", "finding3"]
|
|
490
|
+
}}"""
|
|
491
|
+
|
|
492
|
+
response = self._call_llm(prompt)
|
|
493
|
+
data = json.loads(response)
|
|
494
|
+
return data["summary"], data["key_findings"]
|
|
495
|
+
|
|
496
|
+
except Exception as e:
|
|
497
|
+
return f"System research for {topic} (LLM error: {str(e)[:50]})", ["Error during LLM analysis"]
|
|
498
|
+
|
|
499
|
+
def _llm_summarize_tradecraft(
|
|
500
|
+
self,
|
|
501
|
+
topic: str,
|
|
502
|
+
technique: Optional[str],
|
|
503
|
+
sources: List[Dict[str, str]],
|
|
504
|
+
search_results: Optional[Any],
|
|
505
|
+
) -> tuple[str, List[str]]:
|
|
506
|
+
"""Use LLM to summarize adversary tradecraft findings."""
|
|
507
|
+
try:
|
|
508
|
+
client = self._get_llm_client()
|
|
509
|
+
if not client:
|
|
510
|
+
return f"Adversary tradecraft for {topic}", ["LLM unavailable"]
|
|
511
|
+
|
|
512
|
+
# Build context from sources
|
|
513
|
+
context = ""
|
|
514
|
+
if search_results and hasattr(search_results, "answer") and search_results.answer:
|
|
515
|
+
context = f"Web search summary: {search_results.answer}\n\n"
|
|
516
|
+
|
|
517
|
+
for source in sources[:7]:
|
|
518
|
+
context += f"- {source['title']}: {source['snippet']}\n"
|
|
519
|
+
|
|
520
|
+
technique_str = f" ({technique})" if technique else ""
|
|
521
|
+
|
|
522
|
+
prompt = f"""You are a threat intelligence analyst studying adversary techniques.
|
|
523
|
+
|
|
524
|
+
Topic: {topic}{technique_str}
|
|
525
|
+
|
|
526
|
+
Research Context:
|
|
527
|
+
{context}
|
|
528
|
+
|
|
529
|
+
Based on this context, provide:
|
|
530
|
+
1. A concise summary (2-3 sentences) of how adversaries abuse this system/technique
|
|
531
|
+
2. 4-6 key findings about attack methods, tools used, and indicators
|
|
532
|
+
|
|
533
|
+
Return JSON format:
|
|
534
|
+
{{
|
|
535
|
+
"summary": "string",
|
|
536
|
+
"key_findings": ["finding1", "finding2", "finding3", "finding4"]
|
|
537
|
+
}}"""
|
|
538
|
+
|
|
539
|
+
response = self._call_llm(prompt)
|
|
540
|
+
data = json.loads(response)
|
|
541
|
+
return data["summary"], data["key_findings"]
|
|
542
|
+
|
|
543
|
+
except Exception as e:
|
|
544
|
+
return f"Adversary tradecraft for {topic} (LLM error: {str(e)[:50]})", ["Error during LLM analysis"]
|
|
545
|
+
|
|
546
|
+
def _llm_map_telemetry(
|
|
547
|
+
self,
|
|
548
|
+
topic: str,
|
|
549
|
+
technique: Optional[str],
|
|
550
|
+
ocsf_schema: str,
|
|
551
|
+
environment_data: str,
|
|
552
|
+
) -> tuple[str, List[str]]:
|
|
553
|
+
"""Use LLM to map topic to OCSF telemetry fields."""
|
|
554
|
+
try:
|
|
555
|
+
client = self._get_llm_client()
|
|
556
|
+
if not client:
|
|
557
|
+
return f"Telemetry mapping for {topic}", ["LLM unavailable"]
|
|
558
|
+
|
|
559
|
+
technique_str = f" ({technique})" if technique else ""
|
|
560
|
+
|
|
561
|
+
prompt = f"""You are a detection engineer mapping attack behaviors to telemetry.
|
|
562
|
+
|
|
563
|
+
Topic: {topic}{technique_str}
|
|
564
|
+
|
|
565
|
+
OCSF Schema Reference (partial):
|
|
566
|
+
{ocsf_schema[:3000]}
|
|
567
|
+
|
|
568
|
+
Environment:
|
|
569
|
+
{environment_data[:1000]}
|
|
570
|
+
|
|
571
|
+
Based on this context, provide:
|
|
572
|
+
1. A concise summary of what telemetry would capture this behavior
|
|
573
|
+
2. 4-6 specific OCSF fields that are relevant, with population rates if known
|
|
574
|
+
|
|
575
|
+
Return JSON format:
|
|
576
|
+
{{
|
|
577
|
+
"summary": "string",
|
|
578
|
+
"key_findings": ["field1 (X% populated): description", "field2: description"]
|
|
579
|
+
}}"""
|
|
580
|
+
|
|
581
|
+
response = self._call_llm(prompt)
|
|
582
|
+
data = json.loads(response)
|
|
583
|
+
return data["summary"], data["key_findings"]
|
|
584
|
+
|
|
585
|
+
except Exception as e:
|
|
586
|
+
return f"Telemetry mapping for {topic} (LLM error: {str(e)[:50]})", ["Error during LLM analysis"]
|
|
587
|
+
|
|
588
|
+
def _llm_synthesize(
|
|
589
|
+
self,
|
|
590
|
+
topic: str,
|
|
591
|
+
technique: Optional[str],
|
|
592
|
+
skills: List[ResearchSkillOutput],
|
|
593
|
+
) -> tuple[str, List[str]]:
|
|
594
|
+
"""Use LLM to synthesize all research findings."""
|
|
595
|
+
try:
|
|
596
|
+
client = self._get_llm_client()
|
|
597
|
+
if not client:
|
|
598
|
+
return f"Research synthesis for {topic}", ["LLM unavailable"]
|
|
599
|
+
|
|
600
|
+
# Build context from all skills
|
|
601
|
+
context = ""
|
|
602
|
+
for skill in skills:
|
|
603
|
+
context += f"\n### {skill.skill_name.replace('_', ' ').title()}\n"
|
|
604
|
+
context += f"Summary: {skill.summary}\n"
|
|
605
|
+
context += "Key findings:\n"
|
|
606
|
+
for finding in skill.key_findings[:4]:
|
|
607
|
+
context += f"- {finding}\n"
|
|
608
|
+
|
|
609
|
+
technique_str = f" ({technique})" if technique else ""
|
|
610
|
+
|
|
611
|
+
prompt = f"""You are a senior threat hunter synthesizing research for a hunt.
|
|
612
|
+
|
|
613
|
+
Topic: {topic}{technique_str}
|
|
614
|
+
|
|
615
|
+
Research Findings:
|
|
616
|
+
{context}
|
|
617
|
+
|
|
618
|
+
Based on all research findings, provide:
|
|
619
|
+
1. An executive summary (2-3 sentences) synthesizing all findings
|
|
620
|
+
2. A recommended hypothesis statement in the format: "Adversaries use [behavior] to [goal] on [target]"
|
|
621
|
+
3. 2-3 gaps identified in current coverage or knowledge
|
|
622
|
+
4. 2-3 recommended focus areas for the hunt
|
|
623
|
+
|
|
624
|
+
Return JSON format:
|
|
625
|
+
{{
|
|
626
|
+
"summary": "string",
|
|
627
|
+
"key_findings": [
|
|
628
|
+
"Hypothesis: Adversaries use...",
|
|
629
|
+
"Gap: ...",
|
|
630
|
+
"Focus: ..."
|
|
631
|
+
]
|
|
632
|
+
}}"""
|
|
633
|
+
|
|
634
|
+
response = self._call_llm(prompt)
|
|
635
|
+
data = json.loads(response)
|
|
636
|
+
return data["summary"], data["key_findings"]
|
|
637
|
+
|
|
638
|
+
except Exception as e:
|
|
639
|
+
return f"Research synthesis for {topic} (LLM error: {str(e)[:50]})", ["Error during LLM analysis"]
|
|
640
|
+
|
|
641
|
+
def _call_llm(self, prompt: str) -> str:
|
|
642
|
+
"""Call LLM and return response text."""
|
|
643
|
+
client = self._get_llm_client()
|
|
644
|
+
if not client:
|
|
645
|
+
raise ValueError("LLM client not available")
|
|
646
|
+
|
|
647
|
+
# Bedrock model ID - using cross-region inference profile for Claude Sonnet
|
|
648
|
+
model_id = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
|
649
|
+
|
|
650
|
+
# Prepare request body for Bedrock
|
|
651
|
+
request_body = {
|
|
652
|
+
"anthropic_version": "bedrock-2023-05-31",
|
|
653
|
+
"max_tokens": 2048,
|
|
654
|
+
"messages": [{"role": "user", "content": prompt}],
|
|
655
|
+
}
|
|
656
|
+
|
|
657
|
+
# Invoke model via Bedrock
|
|
658
|
+
start_time = time.time()
|
|
659
|
+
response = client.invoke_model(modelId=model_id, body=json.dumps(request_body))
|
|
660
|
+
duration_ms = int((time.time() - start_time) * 1000)
|
|
661
|
+
|
|
662
|
+
# Parse Bedrock response
|
|
663
|
+
response_body = json.loads(response["body"].read())
|
|
664
|
+
|
|
665
|
+
# Extract text from response
|
|
666
|
+
output_text: str = str(response_body["content"][0]["text"])
|
|
667
|
+
|
|
668
|
+
# Try to extract JSON from markdown code blocks if present
|
|
669
|
+
if "```json" in output_text:
|
|
670
|
+
json_start = output_text.find("```json") + 7
|
|
671
|
+
json_end = output_text.find("```", json_start)
|
|
672
|
+
output_text = output_text[json_start:json_end].strip()
|
|
673
|
+
elif "```" in output_text:
|
|
674
|
+
json_start = output_text.find("```") + 3
|
|
675
|
+
json_end = output_text.find("```", json_start)
|
|
676
|
+
output_text = output_text[json_start:json_end].strip()
|
|
677
|
+
|
|
678
|
+
# Track costs
|
|
679
|
+
usage = response_body.get("usage", {})
|
|
680
|
+
input_tokens = usage.get("input_tokens", 0)
|
|
681
|
+
output_tokens = usage.get("output_tokens", 0)
|
|
682
|
+
cost = self._calculate_cost_bedrock(input_tokens, output_tokens)
|
|
683
|
+
self._total_cost += cost
|
|
684
|
+
self._llm_calls += 1
|
|
685
|
+
|
|
686
|
+
# Log metrics
|
|
687
|
+
self._log_llm_metrics(
|
|
688
|
+
agent_name="hunt-researcher",
|
|
689
|
+
model_id=model_id,
|
|
690
|
+
input_tokens=input_tokens,
|
|
691
|
+
output_tokens=output_tokens,
|
|
692
|
+
cost_usd=cost,
|
|
693
|
+
duration_ms=duration_ms,
|
|
694
|
+
)
|
|
695
|
+
|
|
696
|
+
return output_text
|
|
697
|
+
|
|
698
|
+
def _calculate_cost_bedrock(self, input_tokens: int, output_tokens: int) -> float:
|
|
699
|
+
"""Calculate AWS Bedrock Claude cost."""
|
|
700
|
+
# Claude Sonnet on Bedrock pricing
|
|
701
|
+
input_cost_per_1k = 0.003
|
|
702
|
+
output_cost_per_1k = 0.015
|
|
703
|
+
|
|
704
|
+
input_cost = (input_tokens / 1000) * input_cost_per_1k
|
|
705
|
+
output_cost = (output_tokens / 1000) * output_cost_per_1k
|
|
706
|
+
|
|
707
|
+
return round(input_cost + output_cost, 4)
|
|
708
|
+
|
|
709
|
+
def _load_ocsf_schema(self) -> str:
|
|
710
|
+
"""Load OCSF schema reference content."""
|
|
711
|
+
schema_path = Path.cwd() / "knowledge" / "OCSF_SCHEMA_REFERENCE.md"
|
|
712
|
+
if schema_path.exists():
|
|
713
|
+
return schema_path.read_text()[:5000] # Limit size
|
|
714
|
+
return "OCSF schema reference not found"
|
|
715
|
+
|
|
716
|
+
def _load_environment(self) -> str:
|
|
717
|
+
"""Load environment.md content."""
|
|
718
|
+
env_path = Path.cwd() / "environment.md"
|
|
719
|
+
if env_path.exists():
|
|
720
|
+
return env_path.read_text()[:2000] # Limit size
|
|
721
|
+
return "Environment file not found"
|
|
722
|
+
|
|
723
|
+
def _extract_hypothesis(self, synthesis: ResearchSkillOutput) -> Optional[str]:
|
|
724
|
+
"""Extract recommended hypothesis from synthesis."""
|
|
725
|
+
for finding in synthesis.key_findings:
|
|
726
|
+
if finding.lower().startswith("hypothesis:"):
|
|
727
|
+
return finding.replace("Hypothesis:", "").replace("hypothesis:", "").strip()
|
|
728
|
+
return None
|
|
729
|
+
|
|
730
|
+
def _extract_data_sources(self, telemetry: ResearchSkillOutput) -> Dict[str, bool]:
|
|
731
|
+
"""Extract data source availability from telemetry mapping."""
|
|
732
|
+
# Default data sources based on environment
|
|
733
|
+
return {
|
|
734
|
+
"process_execution": True,
|
|
735
|
+
"file_operations": True,
|
|
736
|
+
"network_connections": False, # Limited visibility per AGENTS.md
|
|
737
|
+
"registry_events": False, # Platform-dependent
|
|
738
|
+
}
|
|
739
|
+
|
|
740
|
+
def _estimate_complexity(
|
|
741
|
+
self,
|
|
742
|
+
tradecraft: ResearchSkillOutput,
|
|
743
|
+
telemetry: ResearchSkillOutput,
|
|
744
|
+
) -> str:
|
|
745
|
+
"""Estimate hunt complexity based on research."""
|
|
746
|
+
# Simple heuristic based on number of findings
|
|
747
|
+
total_findings = len(tradecraft.key_findings) + len(telemetry.key_findings)
|
|
748
|
+
|
|
749
|
+
if total_findings <= 4:
|
|
750
|
+
return "low"
|
|
751
|
+
elif total_findings <= 8:
|
|
752
|
+
return "medium"
|
|
753
|
+
else:
|
|
754
|
+
return "high"
|
|
755
|
+
|
|
756
|
+
def _extract_gaps(self, synthesis: ResearchSkillOutput) -> List[str]:
|
|
757
|
+
"""Extract identified gaps from synthesis."""
|
|
758
|
+
gaps = []
|
|
759
|
+
for finding in synthesis.key_findings:
|
|
760
|
+
if finding.lower().startswith("gap:"):
|
|
761
|
+
gaps.append(finding.replace("Gap:", "").replace("gap:", "").strip())
|
|
762
|
+
return gaps
|
|
@@ -0,0 +1,238 @@
|
|
|
1
|
+
"""Hypothesis generator agent - LLM-powered hypothesis generation."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import time
|
|
5
|
+
from dataclasses import dataclass
|
|
6
|
+
from typing import Any, Dict, List, Optional
|
|
7
|
+
|
|
8
|
+
from athf.agents.base import AgentResult, LLMAgent
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@dataclass
|
|
12
|
+
class HypothesisGenerationInput:
|
|
13
|
+
"""Input for hypothesis generation."""
|
|
14
|
+
|
|
15
|
+
threat_intel: str # User-provided threat context
|
|
16
|
+
past_hunts: List[Dict[str, Any]] # Similar past hunts for context
|
|
17
|
+
environment: Dict[str, Any] # Data sources, platforms, etc.
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
@dataclass
|
|
21
|
+
class HypothesisGenerationOutput:
|
|
22
|
+
"""Output from hypothesis generation."""
|
|
23
|
+
|
|
24
|
+
hypothesis: str
|
|
25
|
+
justification: str
|
|
26
|
+
mitre_techniques: List[str]
|
|
27
|
+
data_sources: List[str]
|
|
28
|
+
expected_observables: List[str]
|
|
29
|
+
known_false_positives: List[str]
|
|
30
|
+
time_range_suggestion: str
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class HypothesisGeneratorAgent(LLMAgent[HypothesisGenerationInput, HypothesisGenerationOutput]):
|
|
34
|
+
"""Generates hunt hypotheses using Claude.
|
|
35
|
+
|
|
36
|
+
Uses Claude API for context-aware hypothesis generation with fallback
|
|
37
|
+
to template-based generation when LLM is disabled.
|
|
38
|
+
|
|
39
|
+
Features:
|
|
40
|
+
- TTP-focused hypothesis generation
|
|
41
|
+
- MITRE ATT&CK technique mapping
|
|
42
|
+
- Data source validation
|
|
43
|
+
- False positive prediction
|
|
44
|
+
- Cost tracking
|
|
45
|
+
"""
|
|
46
|
+
|
|
47
|
+
def execute(self, input_data: HypothesisGenerationInput) -> AgentResult[HypothesisGenerationOutput]:
|
|
48
|
+
"""Generate hypothesis using LLM.
|
|
49
|
+
|
|
50
|
+
Args:
|
|
51
|
+
input_data: Hypothesis generation input
|
|
52
|
+
|
|
53
|
+
Returns:
|
|
54
|
+
AgentResult with hypothesis output or error
|
|
55
|
+
"""
|
|
56
|
+
if not self.llm_enabled:
|
|
57
|
+
# Fallback to template-based generation
|
|
58
|
+
return self._template_generate(input_data)
|
|
59
|
+
|
|
60
|
+
# Use AWS Bedrock Claude API
|
|
61
|
+
try:
|
|
62
|
+
client = self._get_llm_client()
|
|
63
|
+
|
|
64
|
+
prompt = self._build_prompt(input_data)
|
|
65
|
+
|
|
66
|
+
# Bedrock model ID - using cross-region inference profile for Claude Sonnet 4.5
|
|
67
|
+
# Cross-region inference profiles provide better availability and automatic failover
|
|
68
|
+
model_id = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
|
69
|
+
|
|
70
|
+
# Prepare request body for Bedrock
|
|
71
|
+
request_body = {
|
|
72
|
+
"anthropic_version": "bedrock-2023-05-31",
|
|
73
|
+
"max_tokens": 4096,
|
|
74
|
+
"messages": [{"role": "user", "content": prompt}],
|
|
75
|
+
}
|
|
76
|
+
|
|
77
|
+
# Invoke model via Bedrock (with timing)
|
|
78
|
+
start_time = time.time()
|
|
79
|
+
response = client.invoke_model(modelId=model_id, body=json.dumps(request_body))
|
|
80
|
+
duration_ms = int((time.time() - start_time) * 1000)
|
|
81
|
+
|
|
82
|
+
# Parse Bedrock response
|
|
83
|
+
response_body = json.loads(response["body"].read())
|
|
84
|
+
|
|
85
|
+
# Extract text from response
|
|
86
|
+
output_text = response_body["content"][0]["text"]
|
|
87
|
+
|
|
88
|
+
# Try to extract JSON from markdown code blocks if present
|
|
89
|
+
if "```json" in output_text:
|
|
90
|
+
json_start = output_text.find("```json") + 7
|
|
91
|
+
json_end = output_text.find("```", json_start)
|
|
92
|
+
output_text = output_text[json_start:json_end].strip()
|
|
93
|
+
elif "```" in output_text:
|
|
94
|
+
json_start = output_text.find("```") + 3
|
|
95
|
+
json_end = output_text.find("```", json_start)
|
|
96
|
+
output_text = output_text[json_start:json_end].strip()
|
|
97
|
+
|
|
98
|
+
# Parse JSON with better error handling
|
|
99
|
+
try:
|
|
100
|
+
output_data = json.loads(output_text)
|
|
101
|
+
except json.JSONDecodeError as e:
|
|
102
|
+
# If JSON parsing fails, log the actual response for debugging
|
|
103
|
+
raise ValueError(
|
|
104
|
+
f"Failed to parse JSON response from Claude. "
|
|
105
|
+
f"Error: {e}. "
|
|
106
|
+
f"Response text (first 1500 chars): {output_text[:1500]}"
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
output = HypothesisGenerationOutput(**output_data)
|
|
110
|
+
|
|
111
|
+
# Extract usage metrics from Bedrock response
|
|
112
|
+
usage = response_body.get("usage", {})
|
|
113
|
+
input_tokens = usage.get("input_tokens", 0)
|
|
114
|
+
output_tokens = usage.get("output_tokens", 0)
|
|
115
|
+
cost_usd = self._calculate_cost_bedrock(input_tokens, output_tokens)
|
|
116
|
+
|
|
117
|
+
# Log metrics to centralized tracker
|
|
118
|
+
self._log_llm_metrics(
|
|
119
|
+
agent_name="hypothesis-generator",
|
|
120
|
+
model_id=model_id,
|
|
121
|
+
input_tokens=input_tokens,
|
|
122
|
+
output_tokens=output_tokens,
|
|
123
|
+
cost_usd=cost_usd,
|
|
124
|
+
duration_ms=duration_ms,
|
|
125
|
+
)
|
|
126
|
+
|
|
127
|
+
return AgentResult(
|
|
128
|
+
success=True,
|
|
129
|
+
data=output,
|
|
130
|
+
error=None,
|
|
131
|
+
warnings=[],
|
|
132
|
+
metadata={
|
|
133
|
+
"llm_model": model_id,
|
|
134
|
+
"prompt_tokens": input_tokens,
|
|
135
|
+
"completion_tokens": output_tokens,
|
|
136
|
+
"cost_usd": cost_usd,
|
|
137
|
+
"duration_ms": duration_ms,
|
|
138
|
+
},
|
|
139
|
+
)
|
|
140
|
+
|
|
141
|
+
except Exception as e:
|
|
142
|
+
# Fall back to template generation
|
|
143
|
+
return self._template_generate(input_data, error=str(e))
|
|
144
|
+
|
|
145
|
+
def _build_prompt(self, input_data: HypothesisGenerationInput) -> str:
|
|
146
|
+
"""Build Claude prompt for hypothesis generation.
|
|
147
|
+
|
|
148
|
+
Args:
|
|
149
|
+
input_data: Hypothesis generation input
|
|
150
|
+
|
|
151
|
+
Returns:
|
|
152
|
+
Formatted prompt string
|
|
153
|
+
"""
|
|
154
|
+
return f"""You are a threat hunting expert. Generate a hunt hypothesis based on the following:
|
|
155
|
+
|
|
156
|
+
**Threat Intel:**
|
|
157
|
+
{input_data.threat_intel}
|
|
158
|
+
|
|
159
|
+
**Past Similar Hunts:**
|
|
160
|
+
{json.dumps(input_data.past_hunts, indent=2)}
|
|
161
|
+
|
|
162
|
+
**Available Environment:**
|
|
163
|
+
{json.dumps(input_data.environment, indent=2)}
|
|
164
|
+
|
|
165
|
+
Generate a hypothesis following this format:
|
|
166
|
+
- Hypothesis: "Adversaries use [behavior] to [goal] on [target]"
|
|
167
|
+
- Justification: Why this hypothesis is valuable
|
|
168
|
+
- MITRE Techniques: Relevant ATT&CK techniques (e.g., T1003.001)
|
|
169
|
+
- Data Sources: Which data sources to query
|
|
170
|
+
- Expected Observables: What we expect to find
|
|
171
|
+
- Known False Positives: Common benign patterns
|
|
172
|
+
- Time Range: Suggested time window with justification
|
|
173
|
+
|
|
174
|
+
**IMPORTANT:** Return your response as a JSON object matching this schema:
|
|
175
|
+
{{
|
|
176
|
+
"hypothesis": "string",
|
|
177
|
+
"justification": "string",
|
|
178
|
+
"mitre_techniques": ["T1234.001", "T5678.002"],
|
|
179
|
+
"data_sources": ["ClickHouse nocsf_unified_events", "CloudTrail"],
|
|
180
|
+
"expected_observables": ["Process execution", "Network connections"],
|
|
181
|
+
"known_false_positives": ["Legitimate software", "Administrative tools"],
|
|
182
|
+
"time_range_suggestion": "7 days (justification)"
|
|
183
|
+
}}
|
|
184
|
+
"""
|
|
185
|
+
|
|
186
|
+
def _template_generate(
|
|
187
|
+
self, input_data: HypothesisGenerationInput, error: Optional[str] = None
|
|
188
|
+
) -> AgentResult[HypothesisGenerationOutput]:
|
|
189
|
+
"""Fallback template-based generation (no LLM).
|
|
190
|
+
|
|
191
|
+
Args:
|
|
192
|
+
input_data: Hypothesis generation input
|
|
193
|
+
error: Optional error message from LLM attempt
|
|
194
|
+
|
|
195
|
+
Returns:
|
|
196
|
+
AgentResult with template-generated hypothesis
|
|
197
|
+
"""
|
|
198
|
+
# Simple template logic
|
|
199
|
+
output = HypothesisGenerationOutput(
|
|
200
|
+
hypothesis=f"Investigate suspicious activity related to: {input_data.threat_intel[:100]}",
|
|
201
|
+
justification="Template-generated hypothesis (LLM disabled or failed)",
|
|
202
|
+
mitre_techniques=[],
|
|
203
|
+
data_sources=["EDR telemetry", "SIEM logs"],
|
|
204
|
+
expected_observables=["Process execution", "Network connections"],
|
|
205
|
+
known_false_positives=["Legitimate software updates", "Administrative tools"],
|
|
206
|
+
time_range_suggestion="7 days (standard baseline)",
|
|
207
|
+
)
|
|
208
|
+
|
|
209
|
+
warnings = ["LLM disabled - using template generation"]
|
|
210
|
+
if error:
|
|
211
|
+
warnings.append(f"LLM error: {error}")
|
|
212
|
+
|
|
213
|
+
return AgentResult(
|
|
214
|
+
success=True,
|
|
215
|
+
data=output,
|
|
216
|
+
error=None,
|
|
217
|
+
warnings=warnings,
|
|
218
|
+
metadata={"fallback": True},
|
|
219
|
+
)
|
|
220
|
+
|
|
221
|
+
def _calculate_cost_bedrock(self, input_tokens: int, output_tokens: int) -> float:
|
|
222
|
+
"""Calculate AWS Bedrock Claude cost.
|
|
223
|
+
|
|
224
|
+
Args:
|
|
225
|
+
input_tokens: Number of input tokens
|
|
226
|
+
output_tokens: Number of output tokens
|
|
227
|
+
|
|
228
|
+
Returns:
|
|
229
|
+
Cost in USD
|
|
230
|
+
"""
|
|
231
|
+
# Claude Sonnet 4.5 on Bedrock pricing (as of January 2025)
|
|
232
|
+
input_cost_per_1k = 0.003
|
|
233
|
+
output_cost_per_1k = 0.015
|
|
234
|
+
|
|
235
|
+
input_cost = (input_tokens / 1000) * input_cost_per_1k
|
|
236
|
+
output_cost = (output_tokens / 1000) * output_cost_per_1k
|
|
237
|
+
|
|
238
|
+
return round(input_cost + output_cost, 4)
|
athf/data/docs/CHANGELOG.md
CHANGED
|
@@ -25,6 +25,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|
|
25
25
|
### Security
|
|
26
26
|
- None
|
|
27
27
|
|
|
28
|
+
## [0.3.1] - 2026-01-13
|
|
29
|
+
|
|
30
|
+
### Fixed
|
|
31
|
+
- **Packaging Bug** - Fixed `ModuleNotFoundError: No module named 'athf.agents'` when installing via pip/pipx
|
|
32
|
+
- Added missing packages to `pyproject.toml`: `athf.agents`, `athf.agents.llm`
|
|
33
|
+
- Packages list now includes all subdirectories: athf, athf.agents, athf.agents.llm, athf.commands, athf.core, athf.data, athf.utils
|
|
34
|
+
- Verified wheel build includes all agent module files
|
|
35
|
+
|
|
36
|
+
## [0.3.0] - 2026-01-11
|
|
37
|
+
|
|
38
|
+
### Added
|
|
39
|
+
- **Agent Framework** - Autonomous agents for threat hunting workflows
|
|
40
|
+
- `athf.agents` - Base agent framework and orchestration
|
|
41
|
+
- `athf.agents.llm` - LLM-powered agents (hypothesis generation, research, finding analysis)
|
|
42
|
+
- Agent orchestration with task delegation and result aggregation
|
|
43
|
+
- **Research Workflow** - Pre-hunt research and investigation (`athf research`)
|
|
44
|
+
- **Drift Detection** - Behavioral anomaly detection infrastructure (`athf drift`)
|
|
45
|
+
- **Signal Investigation** - Low-fidelity pattern scoring and investigation (`athf signals`)
|
|
46
|
+
|
|
47
|
+
### Changed
|
|
48
|
+
- CLI refactored to support agent-based workflows
|
|
49
|
+
- Enhanced hunt creation with agent-generated hypotheses
|
|
50
|
+
|
|
28
51
|
## [0.2.2] - 2024-12-17
|
|
29
52
|
|
|
30
53
|
### Fixed
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|