diff --git a/AGENTS/cmdb_agent.py b/AGENTS/agent_cmdb.py similarity index 91% rename from AGENTS/cmdb_agent.py rename to AGENTS/agent_cmdb.py index e772230..8f31228 100644 --- a/AGENTS/cmdb_agent.py +++ b/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) diff --git a/AGENTS/knowledge_agent.py b/AGENTS/agent_knowledge.py similarity index 98% rename from AGENTS/knowledge_agent.py rename to AGENTS/agent_knowledge.py index 89df54e..7a96742 100644 --- a/AGENTS/knowledge_agent.py +++ b/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") diff --git a/AGENTS/siem_agent.py b/AGENTS/agent_siem.py similarity index 92% rename from AGENTS/siem_agent.py rename to AGENTS/agent_siem.py index 516a844..cfa580e 100644 --- a/AGENTS/siem_agent.py +++ b/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) diff --git a/AGENTS/ti_agent.py b/AGENTS/agent_ti.py similarity index 96% rename from AGENTS/ti_agent.py rename to AGENTS/agent_ti.py index 9edbd10..a6c3819 100644 --- a/AGENTS/ti_agent.py +++ b/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") diff --git a/Automation/Handle/playbook.py b/Automation/Handle/playbook.py index c6abf4a..2d4d53d 100644 --- a/Automation/Handle/playbook.py +++ b/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 diff --git a/DATA/Agent_CMDB/system.md b/DATA/Agent_CMDB/system.md new file mode 100644 index 0000000..02061f5 --- /dev/null +++ b/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. \ No newline at end of file diff --git a/DATA/siem_agent/splunk_datamodels.yml b/DATA/Agent_SIEM/splunk_datamodels.yml similarity index 100% rename from DATA/siem_agent/splunk_datamodels.yml rename to DATA/Agent_SIEM/splunk_datamodels.yml diff --git a/DATA/siem_agent/system_prompt.md b/DATA/Agent_SIEM/system_prompt.md similarity index 100% rename from DATA/siem_agent/system_prompt.md rename to DATA/Agent_SIEM/system_prompt.md diff --git a/DATA/Case_L3_SOC_Analyst_Agent_With_Tools/L3_SOC_Analyst.md b/DATA/Case_L3_SOC_Analyst_Agent_With_Tools/L3_SOC_Analyst.md index 9d370ba..274cbd8 100644 --- a/DATA/Case_L3_SOC_Analyst_Agent_With_Tools/L3_SOC_Analyst.md +++ b/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` 提交报告. diff --git a/DATA/Case_Threat_Hunting_Agent/Analyst_System.md b/DATA/Case_Threat_Hunting_Agent/Analyst_System.md index fab9e48..ee4a3b5 100644 --- a/DATA/Case_Threat_Hunting_Agent/Analyst_System.md +++ b/DATA/Case_Threat_Hunting_Agent/Analyst_System.md @@ -14,13 +14,13 @@ 您拥有以下强大的调查工具： -* **SIEMAgent (安全信息和事件管理系统代理)**： +* **AgentSIEM (安全信息和事件管理系统代理)**： * **用途**：用于搜索**日志、安全事件和活动记录**.例如：特定IP地址的网络连接、用户登录失败尝试、进程执行历史等. * **何时使用**：当问题涉及行为模式、时间线分析或内部系统活动时. -* **CMDBAgent (配置管理数据库代理)**： +* **AgentCMDB (配置管理数据库代理)**： * **用途**：用于查询**内部资产(主机、服务器、用户)的详细信息**.例如：主机操作系统、所属部门、负责人、IP地址归属、已知安全标记等. * **何时使用**：当需要了解内部资产的属性、上下文或识别其重要性时. -* **TIAgent (威胁情报代理)**： +* **AgentTI (威胁情报代理)**： * **用途**：用于检查**外部实体(如公共IP地址、域名、文件哈希值)的威胁信誉**.例如：查询某个IP是否为已知恶意C2服务器、某个哈希值是否关联已知恶意软件. * **何时使用**：当调查涉及外部威胁来源、恶意指标或攻击者基础设施时. diff --git a/DATA/Case_Threat_Hunting_Agent/Analyst_System_EN.md b/DATA/Case_Threat_Hunting_Agent/Analyst_System_EN.md index e358a56..c99c3f3 100644 --- a/DATA/Case_Threat_Hunting_Agent/Analyst_System_EN.md +++ b/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. diff --git a/DATA/cmdb_agent/system.md b/DATA/cmdb_agent/system.md deleted file mode 100644 index 9dcbefc..0000000 --- a/DATA/cmdb_agent/system.md +++ /dev/null @@ -1 +0,0 @@ -你是一个CMDB查询助手,能够根据用户的自然语言查询请求,调用合适的CMDB工具进行查询,并返回JSON格式的结果. \ No newline at end of file diff --git a/Lib/llmapi.py b/Lib/llmapi.py index cca88dd..cb9a7d1 100644 --- a/Lib/llmapi.py +++ b/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 diff --git a/Lib/montior.py b/Lib/montior.py index 233b8f1..cc2274e 100644 --- a/Lib/montior.py +++ b/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}") diff --git a/MODULES/__init__.py b/MODULES/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/MODULES/__init__.py @@ -0,0 +1 @@ + diff --git a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py index e896bd7..e8fa3ae 100644 --- a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py +++ b/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""" diff --git a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent_With_Tools.py b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent_With_Tools.py index 1dc05ae..168cc04 100644 --- a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent_With_Tools.py +++ b/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): diff --git a/PLAYBOOKS/Case_Threat_Hunting_Agent.py b/PLAYBOOKS/Case_Threat_Hunting_Agent.py index de5890c..a219df5 100644 --- a/PLAYBOOKS/Case_Threat_Hunting_Agent.py +++ b/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): diff --git a/PLAYBOOKS/__init__.py b/PLAYBOOKS/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/PLAYBOOKS/__init__.py @@ -0,0 +1 @@ +