change name

AGENTS/agent_cmdb.py

@@ -1,3 +1,4 @@
+import os
 from typing import Annotated, List, Literal, Any
 
 from langchain.agents import create_agent
@@ -9,13 +10,15 @@ from langgraph.prebuilt import ToolNode
 from pydantic import BaseModel, Field
 
 from Lib.baseplaybook import LanggraphPlaybook
+from Lib.configs import DATA_DIR
+from Lib.llmapi import load_system_prompt_template
 from PLUGINS.LLM.llmapi import LLMAPI
 # Modify the following functions to the actual CMDB API
 from PLUGINS.Mock.CMDB import get_ci_context_tool, fuzzy_search_ci_tool, get_cis_by_software_tool, get_cis_by_port_tool, get_cis_by_service_tool, \
     get_cis_by_user_tool
 
 
-class CMDBAgent(object):
+class AgentCMDB(object):
 
     @staticmethod
     def cmdb_query_asset(
@@ -28,7 +31,7 @@ class CMDBAgent(object):
         # result = cmdb_query(query)
         # return result
 
-        agent = GraphAgent()
+        agent = AgentGraphCMDB()
         result = agent.cmdb_query(query)
         return result
 
@@ -42,7 +45,7 @@ class AgentState(BaseModel):
 
 
 # Use langgraph to create a CMDB query agent for finer-grained control
-class GraphAgent(LanggraphPlaybook):
+class AgentGraphCMDB(LanggraphPlaybook):
 
     def __init__(self):
         super().__init__()  # do not delete this code
@@ -124,10 +127,12 @@ def cmdb_query(
         get_cis_by_service_tool,
         get_cis_by_user_tool,
     ]
+    prompt_path = os.path.join(DATA_DIR, "Agent_CMDB", "system.md")
+    system_prompt_template = load_system_prompt_template(prompt_path)
     agent = create_agent(
         model=llm,
         tools=CMDB_AGENT_TOOLS,
-        system_prompt="You are a CMDB query assistant. You can call appropriate CMDB tools to query based on the user's natural language query request and return the results in JSON format.",
+        system_prompt=system_prompt_template.format(),
     )
 
     response = agent.invoke({"messages": [HumanMessage(content=query)]})
@@ -141,6 +146,6 @@ if __name__ == "__main__":
 
     # result = cmdb_query(query)
 
-    agent = GraphAgent()
+    agent = AgentGraphCMDB()
     result = agent.cmdb_query(query)
     print(result)

AGENTS/agent_knowledge.py

@@ -13,7 +13,7 @@ if MEM_ZERO_USE:
 from langchain_core.tools import tool
 
 
-class KnowledgeAgent(object):
+class AgentKnowledge(object):
 
     @staticmethod
     @tool("internal_knowledge_base_search")

AGENTS/agent_siem.py

@@ -13,6 +13,7 @@ from pydantic import BaseModel, Field
 
 from Lib.baseplaybook import LanggraphPlaybook
 from Lib.configs import DATA_DIR
+from Lib.llmapi import load_system_prompt_template
 from PLUGINS.LLM.llmapi import LLMAPI
 from PLUGINS.Mock.SIEM_Splunk import splunk_search_tool
 
@@ -27,7 +28,7 @@ class AgentState(BaseModel):
 
 
 # Main class for the SIEM Agent, serving as the public interface
-class SIEMAgent:
+class AgentSIEM:
     def search(
             self,
             natural_query: Annotated[str, "A natural language query for SIEM. (e.g., 'Find connections from 10.10.10.10 to any malicious IP')"]
@@ -110,23 +111,20 @@ 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, "siem_agent", "splunk_datamodels.yml")
+    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)
-
-    prompt_path = os.path.join(DATA_DIR, "siem_agent", "system_prompt.md")
-    with open(prompt_path, 'r', encoding='utf-8') as f:
-        system_prompt_template = f.read()
-
     schema_json = json.dumps(splunk_schemas, indent=2)
-    system_prompt = system_prompt_template.format(splunk_schema_json=schema_json)
+
+    prompt_path = os.path.join(DATA_DIR, "Agent_SIEM", "system_prompt.md")
+    system_prompt_template = load_system_prompt_template(prompt_path)
 
     llm_api = LLMAPI()
     llm = llm_api.get_model(tag=["fast", "function_calling"])
 
     tools = [splunk_search_tool]
 
-    agent = create_agent(llm, tools, system_prompt=system_prompt)
+    agent = create_agent(llm, tools, system_prompt=system_prompt_template.format(splunk_schema_json=schema_json))
 
     response = agent.invoke({"messages": [HumanMessage(content=query)]})
 
@@ -135,13 +133,13 @@ a simpler, stateless agent created using the create_agent factory function from
 
 # Test code
 if __name__ == "__main__":
-    # siem_agent = SIEMAgent()
+    # Agent_SIEM = SIEMAgent()
     #
     # # Example query that requires the agent to formulate an SPL query
     # test_query = "Have there been any suspicious logins for the user 'admin' on Windows machines?"
     #
     # print(f"--- Using GraphAgent for Query: '{test_query}' ---")
-    # result = siem_agent.search(test_query)
+    # result = Agent_SIEM.search(test_query)
     # print("\n--- Final Answer ---")
     # print(result)
     #
@@ -150,7 +148,7 @@ if __name__ == "__main__":
     # # A more complex query
     # test_query_2 = "check for connections from the victim host 10.67.3.130 to any known malicious IPs, like 45.33.22.11"
     # print(f"--- Using GraphAgent for Query: '{test_query_2}' ---")
-    # result_2 = siem_agent.search(test_query_2)
+    # result_2 = Agent_SIEM.search(test_query_2)
     # print("\n--- Final Answer ---")
     # print(result_2)
 

AGENTS/agent_ti.py

@@ -6,7 +6,7 @@ from langchain_core.tools import tool
 from PLUGINS.Mock.TI import TI
 
 
-class TIAgent(object):
+class AgentTI(object):
 
     @staticmethod
     @tool("ti_lookup")

Automation/Handle/playbook.py

@@ -23,7 +23,7 @@ class Playbook(object):
             module_config = Xcache.get_module_config_by_name_and_type(type, name)
             if module_config is None:
                 # try again to load all module config
-                Playbook.load_all_module_config()
+                Playbook.load_all_playbook_config()
                 module_config = Xcache.get_module_config_by_name_and_type(type, name)
             if module_config is None:
                 context = data_return(305, {"status": "Failed", "job_id": None}, Playbook_MSG_ZH.get(305), Playbook_MSG_EN.get(305))
@@ -63,7 +63,7 @@ class Playbook(object):
                 return context
 
     @staticmethod
-    def get_module_intent(modulename, module_files_dir):
+    def get_playbook_intent(modulename, module_files_dir):
         if modulename == "__init__" or modulename == "__pycache__" or modulename == '':  # Special handling for __init__.py
             return None
         try:
@@ -75,8 +75,8 @@ class Playbook(object):
             return None
 
     @staticmethod
-    def gen_module_config(modulename, module_files_dir="PLAYBOOKS"):
-        module_intent = Playbook.get_module_intent(modulename, module_files_dir)
+    def gen_playbook_config(modulename, module_files_dir="PLAYBOOKS"):
+        module_intent = Playbook.get_playbook_intent(modulename, module_files_dir)
 
         if module_intent is None:
             return None
@@ -96,14 +96,14 @@ class Playbook(object):
             return None
 
     @staticmethod
-    def load_all_module_config():
+    def load_all_playbook_config():
         all_modules_config = []
         # post module
         module_count = 0
         module_filenames = os.listdir(os.path.join(settings.BASE_DIR, 'PLAYBOOKS'))
         for module_filename in module_filenames:
             module_name = module_filename.split(".")[0]
-            one_module_config = Playbook.gen_module_config(module_name, 'PLAYBOOKS')
+            one_module_config = Playbook.gen_playbook_config(module_name, 'PLAYBOOKS')
             if one_module_config is not None:
                 all_modules_config.append(one_module_config)
                 module_count += 1

DATA/Agent_CMDB/system.md

@@ -0,0 +1,2 @@
+You are a CMDB query assistant. You can call appropriate CMDB tools to query based on the user's natural language query request and return the results in JSON
+format.

DATA/Case_L3_SOC_Analyst_Agent_With_Tools/L3_SOC_Analyst.md

@@ -10,14 +10,14 @@
 
 你拥有两个核心工具,请根据分析进度选择使用：
 
-1. **KnowledgeAgent.search**: 【探测工具】.如果你发现 Case 中的 Artifacts(IP, 域名, 哈希, 文件名)没有背景情报,或者你对某种攻击技术不熟悉,**必须**
+1. **AgentKnowledge.search**: 【探测工具】.如果你发现 Case 中的 Artifacts(IP, 域名, 哈希, 文件名)没有背景情报,或者你对某种攻击技术不熟悉,**必须**
    调用此工具进行搜索.你可以进行多轮搜索以完善证据链.
 2. **AnalyzeResult**: 【终结工具】.当你认为证据已经充分,可以做出最终判定时,**必须且只能**通过调用此工具提交结果.调用此工具代表任务结束.
 
 # 工作流 (Workflow)
 
 1. **证据盘点**: 检查 Case 中所有告警的关联性,识别关键 Artifacts.
-2. **情报补全**: 针对未知的恶意指标或行为,调用 `KnowledgeAgent.search`.
+2. **情报补全**: 针对未知的恶意指标或行为,调用 `AgentKnowledge.search`.
 3. **攻击链映射**: 结合搜索结果,判断攻击处于 MITRE ATT&CK 的哪个阶段.
 4. **提交结论**: 调用 `AnalyzeResult` 提交报告.
 

DATA/Case_Threat_Hunting_Agent/Analyst_System.md

@@ -14,13 +14,13 @@
 
 您拥有以下强大的调查工具：
 
-*   **SIEMAgent (安全信息和事件管理系统代理)**：
+*   **AgentSIEM (安全信息和事件管理系统代理)**：
     *   **用途**：用于搜索**日志、安全事件和活动记录**.例如：特定IP地址的网络连接、用户登录失败尝试、进程执行历史等.
     *   **何时使用**：当问题涉及行为模式、时间线分析或内部系统活动时.
-*   **CMDBAgent (配置管理数据库代理)**：
+*   **AgentCMDB (配置管理数据库代理)**：
     *   **用途**：用于查询**内部资产(主机、服务器、用户)的详细信息**.例如：主机操作系统、所属部门、负责人、IP地址归属、已知安全标记等.
     *   **何时使用**：当需要了解内部资产的属性、上下文或识别其重要性时.
-*   **TIAgent (威胁情报代理)**：
+*   **AgentTI (威胁情报代理)**：
     *   **用途**：用于检查**外部实体(如公共IP地址、域名、文件哈希值)的威胁信誉**.例如：查询某个IP是否为已知恶意C2服务器、某个哈希值是否关联已知恶意软件.
     *   **何时使用**：当调查涉及外部威胁来源、恶意指标或攻击者基础设施时.
 

DATA/Case_Threat_Hunting_Agent/Analyst_System_EN.md

@@ -14,13 +14,13 @@ Based on the assigned "Investigation Question" and "Case Background", you must:
 
 You have the following powerful investigation tools:
 
-*   **SIEMAgent (Security Information and Event Management System Agent)**:
+*   **AgentSIEM (Security Information and Event Management System Agent)**:
     *   **Purpose**: Used to search for **logs, security events, and activity records**. For example: network connections from a specific IP address, user login failure attempts, process execution history, etc.
     *   **When to use**: When the question involves behavioral patterns, timeline analysis, or internal system activities.
-*   **CMDBAgent (Configuration Management Database Agent)**:
+*   **AgentCMDB (Configuration Management Database Agent)**:
     *   **Purpose**: Used to query for **detailed information about internal assets (hosts, servers, users)**. For example: host operating system, department, owner, IP address attribution, known security tags, etc.
     *   **When to use**: When you need to understand the properties, context, or identify the importance of internal assets.
-*   **TIAgent (Threat Intelligence Agent)**:
+*   **AgentTI (Threat Intelligence Agent)**:
     *   **Purpose**: Used to check the **threat reputation of external entities (like public IP addresses, domains, file hashes)**. For example: querying if an IP is a known malicious C2 server, or if a hash is associated with known malware.
     *   **When to use**: When the investigation involves external threat sources, malicious indicators, or attacker infrastructure.
 
@@ -36,7 +36,7 @@ You have the following powerful investigation tools:
     *   Think about how to construct the parameters for the tool calls to ensure query precision and effectiveness.
 3.  **Iterative Investigation**:
     *   If the first tool call does not completely resolve the issue or generates new leads, **continue to call tools** for the next exploration.
-    *   For example: You first use SIEMAgent to discover a suspicious external IP connection, then you should immediately use TIAgent to query the reputation of that IP.
+    *   For example: You first use AgentSIEM to discover a suspicious external IP connection, then you should immediately use AgentTI to query the reputation of that IP.
     *   **Note**: After each tool call, you will receive the tool's output. Please adjust your next action based on the output.
 4.  **Synthesis & Conclusion**:
     *   When you believe you have collected enough evidence to answer the "Investigation Question", stop calling tools.

DATA/cmdb_agent/system.md

@@ -1 +0,0 @@
-你是一个CMDB查询助手,能够根据用户的自然语言查询请求,调用合适的CMDB工具进行查询,并返回JSON格式的结果.

Lib/llmapi.py

@@ -1,8 +1,11 @@
 from typing import Annotated, Any, Dict, List
 
+from langchain_core.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate
 from langgraph.graph import add_messages
 from pydantic import BaseModel
 
+from Lib.log import logger
+
 
 class AgentState(BaseModel):
     messages: Annotated[List[Any], add_messages] = []
@@ -11,3 +14,26 @@ class AgentState(BaseModel):
     artifact: Dict[str, Any] = {}
     temp_data: Dict[str, Any] = {}
     analyze_result: Dict[str, Any] = {}
+
+
+def load_system_prompt_template(template_path):
+    """Load system prompt template"""
+    try:
+        with open(template_path, 'r', encoding='utf-8') as f:
+            system_prompt_template: SystemMessagePromptTemplate = SystemMessagePromptTemplate.from_template(f.read())
+            logger.debug(f"Loaded system prompt template from: {template_path}")
+            return system_prompt_template
+    except Exception as e:
+        logger.warning(f"Failed to load prompt template {template_path}: {str(e)}")
+        raise e
+
+
+def load_human_prompt_template(template_path):
+    try:
+        with open(template_path, 'r', encoding='utf-8') as f:
+            human_prompt_template: HumanMessagePromptTemplate = HumanMessagePromptTemplate.from_template(f.read())
+            logger.debug(f"Loaded human prompt template from: {template_path}")
+            return human_prompt_template
+    except Exception as e:
+        logger.warning(f"Failed to load prompt template {template_path}: {str(e)}")
+        raise e

Lib/montior.py

@@ -91,7 +91,7 @@ class MainMonitor(object):
         Xcache.set_token_user(ASP_REST_API_TOKEN, api_usr, None)
 
         logger.info("Load Playbook module config")
-        Playbook.load_all_module_config()
+        Playbook.load_all_playbook_config()
 
         # self.MainScheduler.add_job(func=self.subscribe_clean_thread,
         #                            max_instances=1,
@@ -124,7 +124,7 @@ class MainMonitor(object):
             row_id = one_record.get("rowId")
             module_config = Xcache.get_module_config_by_name_and_type(type, name)
             if module_config is None:
-                Playbook.load_all_module_config()
+                Playbook.load_all_playbook_config()
                 module_config = Xcache.get_module_config_by_name_and_type(type, name)
             if module_config is None:
                 logger.error(f"Playbook module config not found: {type} - {name}")

MODULES/__init__.py

@@ -0,0 +1 @@
+

PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py

@@ -7,14 +7,11 @@ from langgraph.graph import StateGraph
 from langgraph.graph.state import CompiledStateGraph
 from pydantic import BaseModel, Field, ConfigDict
 
-from AGENTS.knowledge_agent import KnowledgeAgent
 from Lib.baseplaybook import LanggraphPlaybook
 from Lib.llmapi import AgentState
 from PLUGINS.LLM.llmapi import LLMAPI
 from PLUGINS.SIRP.sirpapi import Case
 
-tools = [KnowledgeAgent.search]
-
 
 class ConfidenceLevel(str, Enum):
     """Confidence Level"""

PLAYBOOKS/Case_L3_SOC_Analyst_Agent_With_Tools.py

@@ -8,7 +8,7 @@ from langgraph.graph.state import CompiledStateGraph
 from langgraph.prebuilt import ToolNode
 from pydantic import BaseModel, Field, ConfigDict
 
-from AGENTS.knowledge_agent import KnowledgeAgent
+from AGENTS.agent_knowledge import AgentKnowledge
 from Lib.baseplaybook import LanggraphPlaybook
 from Lib.llmapi import AgentState
 from PLUGINS.LLM.llmapi import LLMAPI
@@ -78,7 +78,7 @@ NODE_TOOLS = "tools"
 NODE_OUTPUT = "output_node"
 
 FINAL_TOOL_NAME = AnalyzeResult.__name__
-SEARCH_TOOL = KnowledgeAgent.search
+SEARCH_TOOL = AgentKnowledge.search
 
 
 class Playbook(LanggraphPlaybook):

PLAYBOOKS/Case_Threat_Hunting_Agent.py

@@ -10,9 +10,9 @@ from langgraph.prebuilt import ToolNode
 from langgraph.types import Send
 from pydantic import BaseModel, Field, ConfigDict
 
-from AGENTS.cmdb_agent import CMDBAgent
-from AGENTS.siem_agent import SIEMAgent
-from AGENTS.ti_agent import TIAgent
+from AGENTS.agent_cmdb import AgentCMDB
+from AGENTS.agent_siem import AgentSIEM
+from AGENTS.agent_ti import AgentTI
 from Lib.api import get_current_time_str
 from Lib.baseplaybook import LanggraphPlaybook
 from PLUGINS.LLM.llmapi import LLMAPI
@@ -44,7 +44,8 @@ class AnalystOutput(BaseModel):
 # Define the structured output of the Planner
 class HuntingPlan(BaseModel):
     # Allows generating multiple tasks at once, or an empty list to indicate the end
-    current_plan: List[str] = Field(description="A list of specific questions to be investigated in parallel next. Returns an empty list if there are no more questions.")
+    current_plan: List[str] = Field(
+        description="A list of specific questions to be investigated in parallel next. Returns an empty list if there are no more questions.")
     rationale: str = Field(description="The reason for making this plan")
 
 
@@ -210,7 +211,7 @@ class Playbook(LanggraphPlaybook):
 
             llm_api = LLMAPI()
             base_llm = llm_api.get_model(tag=["powerful", "function_calling"])
-            llm_with_tools = base_llm.bind_tools([SIEMAgent.search, CMDBAgent.cmdb_query_asset, TIAgent.lookup])
+            llm_with_tools = base_llm.bind_tools([AgentSIEM.search, AgentCMDB.cmdb_query_asset, AgentTI.lookup])
             response: AIMessage = llm_with_tools.invoke(messages)
 
             # update record
@@ -223,7 +224,7 @@ class Playbook(LanggraphPlaybook):
             return {"messages": [response]}
 
         # Tool node
-        tool_node = ToolNode([SIEMAgent.search, CMDBAgent.cmdb_query_asset, TIAgent.lookup])
+        tool_node = ToolNode([AgentSIEM.search, AgentCMDB.cmdb_query_asset, AgentTI.lookup])
 
         # Result generation node: when there is no tool call, it is responsible for converting the last message into a structured output
         def final_answer_node(state: AnalystState):

PLAYBOOKS/__init__.py

@@ -0,0 +1 @@
+