tmp

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)

DATA/Agent_SIEM/splunk_datamodels.yml

@@ -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"]

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

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):

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:

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(

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}")