diff --git a/AGENTS/agent_siem.py b/AGENTS/agent_siem.py index 84ba0e6..d471a42 100644 --- a/AGENTS/agent_siem.py +++ b/AGENTS/agent_siem.py @@ -1,7 +1,5 @@ -import json from typing import Annotated, List, Literal, Any -import yaml from langchain.agents import create_agent from langchain_core.messages import HumanMessage, AIMessage from langchain_core.runnables import RunnableConfig @@ -10,17 +8,18 @@ from langgraph.graph.state import CompiledStateGraph from langgraph.prebuilt import ToolNode from pydantic import BaseModel, Field +from Lib.api import get_current_time_str from Lib.baseplaybook import LanggraphPlaybook from Lib.configs import DATA_DIR from Lib.llmapi import load_system_prompt_template +from Lib.log import logger from PLUGINS.LLM.llmapi import LLMAPI -# change this to your actual Splunk api -from PLUGINS.Mock.SIEM_Splunk import splunk_search_tool +from PLUGINS.SIEM.tools import SIEMToolKit # Define constants for graph nodes AGENT_NODE = "AGENT" TOOL_NODE = "TOOL_NODE" -MAX_ITERATIONS = 2 +MAX_ITERATIONS = 10 # Define the state for the graph @@ -44,6 +43,9 @@ class AgentSIEM: return result +tools = [SIEMToolKit.explore_schema, SIEMToolKit.execute_adaptive_query] + + # LangGraph-based agent for complex, stateful queries class GraphAgent(LanggraphPlaybook): @@ -51,14 +53,8 @@ class GraphAgent(LanggraphPlaybook): super().__init__() self.graph = self._build_graph() - def splunk_schemas(self) -> dict: - """Loads Splunk data models from the YAML config file.""" - with open(self._get_file_path("splunk_datamodels.yml"), 'r', encoding='utf-8') as f: - return yaml.safe_load(f) - def _build_graph(self) -> CompiledStateGraph: """Constructs the LangGraph agent graph.""" - tools = [splunk_search_tool] tool_node = ToolNode(tools) def route_after_agent(state: AgentState) -> Literal["TOOL_NODE", "__end__"]: @@ -67,18 +63,21 @@ class GraphAgent(LanggraphPlaybook): self.logger.debug(f"Max iterations ({MAX_ITERATIONS}) reached, ending agent.") return END if last_message.tool_calls: + tool_info = "\n".join([f" [{idx}] Name: {tc.get('name', 'N/A')}, ID: {tc.get('id', 'N/A')}, Args: {tc.get('args', {})}" for idx, tc in + enumerate(last_message.tool_calls, 1)]) + self.logger.debug(f"Tool calls detected: {len(last_message.tool_calls)} tool(s)\n{tool_info}\nRouting to TOOL_NODE for execution") return TOOL_NODE + self.logger.debug(f"No tool calls detected, ending agent execution") return END def agent_node(state: AgentState): self.logger.debug(f"Agent Node Invoked (Loop: {state.loop_count})") + self.logger.debug(f"Current messages count: {len(state.messages)}") - schema_json = json.dumps(self.splunk_schemas(), indent=2) - - system_prompt_template = self.load_system_prompt_template(f"system_prompt") - system_message = system_prompt_template.format(splunk_schema_json=schema_json) + system_message = self.load_system_prompt_template(f"system_prompt").format(CURRENT_UTC_TIME=get_current_time_str()) messages = [system_message, *state.messages] + self.logger.debug(f"Total messages to send to LLM: {len(messages)}") if state.loop_count >= MAX_ITERATIONS - 1: self.logger.warning("Approaching max iterations, forcing agent to provide final answer.") @@ -89,15 +88,20 @@ class GraphAgent(LanggraphPlaybook): "based ONLY on the information gathered above." ) messages.append(HumanMessage(content=stop_instruction)) + self.logger.debug("Stop instruction appended to messages") llm_api = LLMAPI() base_llm = llm_api.get_model(tag=["fast"]) + self.logger.debug(f"Using base LLM model (no tool binding) for final response") response: AIMessage = base_llm.invoke(messages) + self.logger.debug(f"Final response generated, tool_calls count: {len(response.tool_calls) if response.tool_calls else 0}") else: llm_api = LLMAPI() llm = llm_api.get_model(tag=["fast", "function_calling"]) llm_with_tools = llm.bind_tools(tools) + self.logger.debug(f"Using LLM with tools binding, available tools: {len(tools)}") response: AIMessage = llm_with_tools.invoke(messages) + self.logger.debug(f"Response generated, tool_calls count: {len(response.tool_calls) if response.tool_calls else 0}") if state.loop_count >= MAX_ITERATIONS - 1: if response.tool_calls: @@ -114,18 +118,29 @@ class GraphAgent(LanggraphPlaybook): workflow.add_conditional_edges(AGENT_NODE, route_after_agent) workflow.add_edge(TOOL_NODE, AGENT_NODE) - return workflow.compile(checkpointer=self.get_checkpointer()) + compiled_graph = workflow.compile(checkpointer=self.get_checkpointer()) + self.logger.debug(f"LangGraph workflow compiled successfully") + return compiled_graph def siem_query(self, query: str) -> str: """Executes a query against the graph.""" + self.logger.info(f"SIEM Query started: {query[:100]}...") self.graph.checkpointer.delete_thread(self.module_name) + self.logger.debug(f"Deleted previous thread state for module: {self.module_name}") + config = RunnableConfig(configurable={"thread_id": self.module_name}) + self.logger.debug(f"RunnableConfig created with thread_id: {self.module_name}") initial_state = AgentState(messages=[HumanMessage(content=query)], loop_count=0) + self.logger.debug(f"Initial state created") + self.logger.info(f"Starting graph invocation...") final_state = self.graph.invoke(initial_state, config) + self.logger.info(f"Graph invocation completed") - return final_state['messages'][-1].content + result = final_state['messages'][-1].content + self.logger.info(f"Query result extracted, result length: {len(result)} characters") + return result # Alternative, simpler agent implementation using create_agent @@ -135,25 +150,28 @@ def create_siem_agent( """ a simpler, stateless agent created using the create_agent factory function from langchain.agents. """ - # Load schemas and prompt template - schema_path = os.path.join(DATA_DIR, "Agent_SIEM", "splunk_datamodels.yml") - with open(schema_path, 'r', encoding='utf-8') as f: - splunk_schemas = yaml.safe_load(f) - schema_json = json.dumps(splunk_schemas, indent=2) + logger.info(f"Creating SIEM agent with query: {query[:100]}...") prompt_path = os.path.join(DATA_DIR, "Agent_SIEM", "system_prompt.md") - system_prompt_template = load_system_prompt_template(prompt_path) + logger.debug(f"Loading system prompt from: {prompt_path}") + system_prompt = load_system_prompt_template(prompt_path).format(CURRENT_UTC_TIME=get_current_time_str()) + logger.debug(f"System prompt loaded successfully") llm_api = LLMAPI() llm = llm_api.get_model(tag=["fast", "function_calling"]) + logger.debug(f"LLM model obtained with tags: ['fast', 'function_calling']") - tools = [splunk_search_tool] - - agent = create_agent(llm, tools, system_prompt=system_prompt_template.format(splunk_schema_json=schema_json)) + logger.debug(f"Creating agent with {len(tools)} tools") + agent = create_agent(llm, tools, system_prompt=system_prompt) + logger.debug(f"Agent created successfully") + logger.info(f"Invoking agent...") response = agent.invoke({"messages": [HumanMessage(content=query)]}) + logger.info(f"Agent invocation completed") - return response['messages'][-1].content + result = response['messages'][-1].content + logger.info(f"Agent result extracted, result length: {len(result)} characters") + return result # Test code @@ -164,8 +182,12 @@ if __name__ == "__main__": os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") django.setup() - # # You can also test the simpler agent directly + # print("\n--- Using create_agent for Query ---") + # test_query = "Have there been any suspicious logins for the user 'admin' on Windows machines?" + # result_simple = create_siem_agent(test_query) + # print(result_simple) + print("\n--- Using create_agent for Query ---") - test_query = "Have there been any suspicious logins for the user 'admin' on Windows machines?" - result_simple = create_siem_agent(test_query) + test_query = "最近5分钟192.168.1.150使用ssh访问了哪些内网主机?" + result_simple = AgentSIEM.siem_search_by_natural_language(test_query) print(result_simple) diff --git a/DATA/Agent_SIEM/splunk_datamodels.yml b/DATA/Agent_SIEM/splunk_datamodels.yml deleted file mode 100644 index 1fb0c4f..0000000 --- a/DATA/Agent_SIEM/splunk_datamodels.yml +++ /dev/null @@ -1,29 +0,0 @@ -# This file contains curated Splunk data models for the SIEM agent. -# It helps the agent understand which indexes and fields to use for specific analysis tasks. - -pan_logs: - description: "Palo Alto Networks Firewall Traffic Logs. Use this for analyzing network connections (e.g., traffic between IPs, port usage, allowed/denied actions)." - sourcetype: "pan:traffic" - key_fields: ["action", "src_ip", "src_port", "dest_ip", "dest_port", "proto", "app", "rule", "bytes_in", "bytes_out"] - -windows_security: - description: "Windows Security Event Logs. Use this for host-level activity analysis." - sourcetype: "WinEventLog:Security" - key_fields: ["EventCode", "ComputerName", "SubjectUserName", "TargetUserName", "ProcessName", "ParentProcessName", "CommandLine", "LogonType", "IpAddress"] - notes: | - Common Event Codes: - - 4624: Successful logon - - 4625: Failed logon - - 4688: Process creation - - 4768/4769: Kerberos authentication - - 5140: A network share object was accessed. - -zeek: - description: "Zeek Network Security Monitor logs. Provides detailed records of network activity. Use for deep network analysis, file extraction, or protocol-specific queries." - sourcetype: "zeek:conn" - key_fields: ["uid", "id_orig_h", "id_orig_p", "id_resp_h", "id_resp_p", "proto", "service", "duration", "orig_bytes", "resp_bytes", "conn_state"] - -cloudtrail: - description: "AWS CloudTrail logs. Audits API calls and user activity within an AWS environment. Use for investigating suspicious cloud actions like unauthorized user creation or permission changes." - sourcetype: "aws:cloudtrail" - key_fields: ["eventSource", "eventName", "userIdentity.type", "userIdentity.userName", "sourceIPAddress", "userAgent", "requestParameters", "responseElements", "errorCode", "awsRegion"] diff --git a/DATA/Agent_SIEM/system_prompt.md b/DATA/Agent_SIEM/system_prompt.md index 655d894..16f94e0 100644 --- a/DATA/Agent_SIEM/system_prompt.md +++ b/DATA/Agent_SIEM/system_prompt.md @@ -1,21 +1,124 @@ -# ROLE: You are a Senior Security Operations Center (SOC) Analyst. +# SIEM Agent System Prompt -# PRIMARY DIRECTIVE: -Your primary mission is to translate a user's natural language request into a precise and effective Splunk Processing Language (SPL) query. You must then use the `splunk_search_tool` to execute this SPL query to retrieve security logs. +You are a professional SIEM (Security Information and Event Management) analyst agent. Your role is to help security +analysts and incident responders investigate security events, threats, and anomalies by querying SIEM data and +discovering relevant information. -# AVAILABLE SPLUNK DATA MODELS: -This is your knowledge base of the available Splunk indexes and sourcetypes. You MUST use this to formulate your queries. -```json -{splunk_schema_json} +## Current Context + +- **Current UTC Time**: `{CURRENT_UTC_TIME}` + +## Available Tools + +You have access to two primary tools for SIEM data exploration and querying: + +### 1. explore_schema() + +Discover what data sources and fields are available in the SIEM. + +**Usage approach:** + +- Start with `explore_schema()` to list all available indices and understand your data sources +- Then use `explore_schema(target_index="index_name")` to see field details for the index you're interested in +- This helps you find the correct field names and types before querying + +### 2. execute_adaptive_query() + +Query SIEM data with intelligent progressive filtering and response optimization. + +**Progressive Query Strategy:** + +This tool supports a step-by-step refinement approach: + +1. **Start broad**: Query with wide time ranges and minimal filters to understand the data volume + - Get statistics on key fields to identify patterns + - Understand the distribution of values + +2. **Narrow down**: Based on statistics, refine your filters to focus on specific values or behaviors + - Add more specific filters (e.g., specific users, IPs, event types) + - Reduce time range if you've identified the relevant period + +3. **Drill down**: When you've narrowed the results, query with more restrictive criteria + - Target specific combinations of filters + - Request statistics on additional fields to drill deeper + +4. **Final retrieval**: Once you've identified the specific logs you need + - The tool automatically returns full records when result volume is small + - Or use the statistics from "sample" and "summary" responses to guide your analysis + +**Key benefit:** The tool automatically adjusts its response format: + +- Returns all records when there are few results (easy analysis) +- Returns statistics + sample records for medium volumes (pattern identification) +- Returns statistics only for large volumes (efficient insights) + +## Investigation Strategy + +1. **Explore First**: Always start by exploring the schema to understand available indices and field names +2. **Start Broad**: Begin with wide time ranges and basic filters to understand data volume and patterns +3. **Refine Iteratively**: Use statistics from results to guide your next queries +4. **Narrow Progressively**: Add filters and reduce time ranges as you identify relevant data +5. **Analyze Results**: + - With "full" status (few records): Analyze all records directly + - With "sample" status: Focus on statistics to identify patterns, use sample records as reference + - With "summary" status: Use statistics for insights, then refine filters to get specific records +6. **Drill Down**: Once you've identified relevant patterns, query with more specific criteria + +## Query Examples + +### Example: Investigating Security Events (Progressive Approach) + +**Step 1: Explore available data** + +``` +explore_schema() → find "logs-security" index +explore_schema(target_index="logs-security") → identify field names ``` -# CHAIN OF THOUGHT: -1. **Analyze Request**: Carefully read the user's natural language query (e.g., "check for connections from the victim host 10.67.3.130 to any known malicious IPs"). -2. **Formulate SPL**: Construct a syntactically correct SPL query using the "AVAILABLE SPLUNK DATA MODELS" as your guide. You should infer the correct index (e.g., `index=pan_logs`, `index=windows`) and fields based on the user's request. -3. **Execute Tool**: Call the `splunk_search_tool` with the exact SPL query you just formulated. -4. **Analyze & Respond**: Review the JSON results from the tool. If the results are empty, state that no matching logs were found. If there are results, provide a concise, human-readable summary of the key findings for the user. **Do not just dump the raw JSON back to the user.** +**Step 2: Start broad to understand data volume and patterns** -# EXAMPLE: -- User Request: "Did the machine 10.67.3.130 connect to the C2 server 45.33.22.11?" -- Your Internal Thought: The user is asking about a network connection. According to my data models, `pan_logs` is the correct index for firewall traffic. I will formulate an SPL query. -- Your Tool Call: `splunk_search_tool(spl_query='index=pan_logs src_ip="10.67.3.130" dest_ip="45.33.22.11"')` +``` +execute_adaptive_query( + index_name="logs-security", + time_range_start="2026-02-04T00:00:00Z", + time_range_end="2026-02-04T23:59:59Z", + filters={{}}, # No filters yet + aggregation_fields=["event.outcome", "user.name", "source.ip"] +) +``` + +→ Get statistics to identify anomalies and patterns + +**Step 3: Narrow down based on statistics** + +``` +execute_adaptive_query( + index_name="logs-security", + time_range_start="2026-02-04T10:00:00Z", + time_range_end="2026-02-04T12:00:00Z", + filters={{"event.outcome": "failure"}}, # Based on previous stats + aggregation_fields=["user.name", "source.ip", "event.action"] +) +``` + +→ Get more focused statistics and sample records + +**Step 4: Drill down to specific logs when needed** + +``` +execute_adaptive_query( + index_name="logs-security", + time_range_start="2026-02-04T10:15:00Z", + time_range_end="2026-02-04T10:30:00Z", + filters={{"event.outcome": "failure", "user.name": "admin"}}, + aggregation_fields=["source.ip", "event.action"] +) +``` + +→ Get full records for final analysis + +## Important Notes + +- Always use UTC timestamps in ISO8601 format: `YYYY-MM-DDTHH:MM:SSZ` +- The progressive query approach helps you narrow down large datasets efficiently +- Different indices may have different field names - always explore schema first diff --git a/Lib/api.py b/Lib/api.py index bd3896b..a633d9d 100644 --- a/Lib/api.py +++ b/Lib/api.py @@ -77,13 +77,20 @@ def get_current_timestamp() -> int: def get_current_time_str(format_str: str = "%Y-%m-%dT%H:%M:%SZ") -> str: """ # 示例 - # 默认格式 + # 默认格式（UTC时间） current_time_str = get_current_time_str() - # 自定义格式：年-月-日 + # 自定义格式：年-月-日（本地时间，无Z标记） current_date_str = get_current_time_str("%Y-%m-%d") + + # 自定义格式：本地时间 + local_time_str = get_current_time_str("%Y-%m-%d %H:%M:%S") """ - return datetime.datetime.now().strftime(format_str) + use_utc = 'Z' in format_str + if use_utc: + return datetime.datetime.now(datetime.timezone.utc).strftime(format_str) + else: + return datetime.datetime.now().strftime(format_str) def exec_system(cmd, **kwargs): diff --git a/PLUGINS/SIEM/clients.py b/PLUGINS/SIEM/clients.py index e4051a0..c42c7cb 100644 --- a/PLUGINS/SIEM/clients.py +++ b/PLUGINS/SIEM/clients.py @@ -1,7 +1,7 @@ import splunklib.client from elasticsearch import Elasticsearch -from CONFIG import ELK_HOST, ELK_USER, ELK_PASS, SPLUNK_HOST, SPLUNK_PORT, SPLUNK_USER, SPLUNK_PASS +from PLUGINS.SIEM.CONFIG import ELK_HOST, ELK_USER, ELK_PASS, SPLUNK_HOST, SPLUNK_PORT, SPLUNK_USER, SPLUNK_PASS class ELKClient: diff --git a/PLUGINS/SIEM/models.py b/PLUGINS/SIEM/models.py index 5892fa1..d03d866 100644 --- a/PLUGINS/SIEM/models.py +++ b/PLUGINS/SIEM/models.py @@ -3,6 +3,9 @@ from typing import List, Dict, Any, Optional from pydantic import BaseModel, Field, field_validator +SUMMARY_THRESHOLD = 1000 +SAMPLE_THRESHOLD = 20 + # --- Input Models --- class SchemaExplorerInput(BaseModel): @@ -85,9 +88,9 @@ class AdaptiveQueryOutput(BaseModel): ..., description=( "Response type indicator based on result volume. " - "Possible values: 'full' (complete logs, < 20 results), " - "'sample' (statistics + sample records, 20-1000 results), " - "'summary' (statistics only, > 1000 results)" + f"Possible values: 'full' (complete logs, < {SAMPLE_THRESHOLD} results), " + f"'sample' (statistics + sample records, {SAMPLE_THRESHOLD}-{SUMMARY_THRESHOLD} results), " + f"'summary' (statistics only, > {SUMMARY_THRESHOLD} results)" ) ) total_hits: int = Field( diff --git a/PLUGINS/SIEM/tools.py b/PLUGINS/SIEM/tools.py index 5971335..75f8467 100644 --- a/PLUGINS/SIEM/tools.py +++ b/PLUGINS/SIEM/tools.py @@ -5,16 +5,15 @@ from datetime import datetime, timezone from splunklib.results import JSONResultsReader from PLUGINS.SIEM.clients import ELKClient, SplunkClient -from models import ( +from PLUGINS.SIEM.models import ( SchemaExplorerInput, AdaptiveQueryInput, AdaptiveQueryOutput, - FieldStat + FieldStat, + SUMMARY_THRESHOLD, + SAMPLE_THRESHOLD ) -from registry import STATIC_SCHEMA_REGISTRY, get_default_agg_fields, get_backend_type - -SUMMARY_THRESHOLD = 1000 -SAMPLE_THRESHOLD = 20 +from PLUGINS.SIEM.registry import STATIC_SCHEMA_REGISTRY, get_default_agg_fields, get_backend_type class SIEMToolKit(object): @@ -41,22 +40,20 @@ class SIEMToolKit(object): # Get details on "logs-security" index explore_schema(SchemaExplorerInput(target_index="logs-security")) """ - try: - if not input_data.target_index: - # Agent 看到的是统一的列表，不关心 Backend - return [ - {"name": k, "description": v.description} - for k, v in STATIC_SCHEMA_REGISTRY.items() - ] + if not input_data.target_index: + # Agent 看到的是统一的列表，不关心 Backend + result = [ + {"name": k, "description": v.description} + for k, v in STATIC_SCHEMA_REGISTRY.items() + ] + return result - if input_data.target_index not in STATIC_SCHEMA_REGISTRY: - raise ValueError(f"Index {input_data.target_index} not found.") + if input_data.target_index not in STATIC_SCHEMA_REGISTRY: + raise ValueError(f"Index {input_data.target_index} not found.") - idx_info = STATIC_SCHEMA_REGISTRY[input_data.target_index] - return [f.model_dump() for f in idx_info.fields] - - except Exception as e: - raise e + idx_info = STATIC_SCHEMA_REGISTRY[input_data.target_index] + result = [f.model_dump() for f in idx_info.fields] + return result @classmethod def execute_adaptive_query(cls, input_data: AdaptiveQueryInput) -> AdaptiveQueryOutput: @@ -93,9 +90,11 @@ class SIEMToolKit(object): backend = get_backend_type(input_data.index_name) if backend == "ELK": - return cls._execute_elk(input_data) + result = cls._execute_elk(input_data) + return result elif backend == "Splunk": - return cls._execute_splunk(input_data) + result = cls._execute_splunk(input_data) + return result else: raise ValueError(f"Unsupported backend: {backend}")