diff --git a/Lib/baseplaybook.py b/Lib/baseplaybook.py index 3ec1b75..4f7d1a1 100644 --- a/Lib/baseplaybook.py +++ b/Lib/baseplaybook.py @@ -15,7 +15,8 @@ from pydantic import BaseModel from Lib.baseapi import BaseAPI from Lib.llmapi import AgentState from Lib.log import logger -from PLUGINS.SIRP.sirpapi import PlaybookMessage +from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook, Notice +from PLUGINS.SIRP.sirpapi import PlaybookMessage, PlaybookStatusType class BasePlaybook(BaseAPI): @@ -31,6 +32,35 @@ class BasePlaybook(BaseAPI): def param(self, key, default=None): return self._params.get(key, default) + # 定义内部参数 + @property + def param_playbook_rowid(self): + return self.param("playbook_rowid") + + @property + def param_rowid(self): + return self.param("rowid") + + @property + def param_worksheet(self): + return self.param("worksheet") + + @property + def param_user(self): + return self.param("user") + + @property + def param_user_input(self): + return self.param("user_input") + + def update_playbook(self, status: PlaybookStatusType, remark: str): + rowid = SIRPPlaybook.update_status_and_remark(self.param_playbook_rowid, status, remark) + return rowid + + def send_notice(self, title: str, body: str): + result = Notice.send(self.param_user, title, body) + return result + class LanggraphPlaybook(BasePlaybook): def __init__(self): @@ -43,16 +73,10 @@ class LanggraphPlaybook(BasePlaybook): checkpointer = MemorySaver() return checkpointer - def run_graph(self): - self.graph.checkpointer.delete_thread(self.module_name) - config = RunnableConfig() - config["configurable"] = {"thread_id": self.module_name} - if self.agent_state is None: - self.agent_state = AgentState() - for event in self.graph.stream(self.agent_state, config, stream_mode="values"): - self.logger.debug(event) - def add_message_to_playbook(self, message: BaseMessage | BaseModel, playbook_rowid=None, node=None): + if playbook_rowid is None: + playbook_rowid = self.param_playbook_rowid + if isinstance(message, SystemMessage): fields = [ {"id": "type", "value": "SystemMessage"}, @@ -117,6 +141,16 @@ class LanggraphPlaybook(BasePlaybook): row_id = PlaybookMessage.create(fields) return row_id + # langgraph interface + def run_graph(self): + self.graph.checkpointer.delete_thread(self.module_name) + config = RunnableConfig() + config["configurable"] = {"thread_id": self.module_name} + if self.agent_state is None: + self.agent_state = AgentState() + for event in self.graph.stream(self.agent_state, config, stream_mode="values"): + self.logger.debug(event) + def run(self): self.run_graph() return self.agent_state diff --git a/PLAYBOOKS/Alert_Analysis_Agent.py b/PLAYBOOKS/Alert_Analysis_Agent.py index 16efe12..a34422d 100644 --- a/PLAYBOOKS/Alert_Analysis_Agent.py +++ b/PLAYBOOKS/Alert_Analysis_Agent.py @@ -8,10 +8,7 @@ from pydantic import BaseModel from Lib.baseplaybook import LanggraphPlaybook from PLUGINS.LLM.llmapi import LLMAPI -from PLUGINS.SIRP.nocolyapi import WorksheetRow -from PLUGINS.SIRP.sirpapi import Alert, Artifact -from PLUGINS.SIRP.sirpapi import Notice -from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook +from PLUGINS.SIRP.sirpapi import Alert class AgentState(BaseModel): @@ -33,12 +30,7 @@ class Playbook(LanggraphPlaybook): def preprocess_node(state: AgentState): """预处理数据""" # worksheet = self.param("worksheet") - rowid = self.param("rowid") - worksheet = self.param("worksheet") - alert = WorksheetRow.get(worksheet, rowid, include_system_fields=False) - artifacts = WorksheetRow.relations(Alert.WORKSHEET_ID, alert.get("rowId"), "artifact", relation_worksheet_id=Artifact.WORKSHEET_ID, - include_system_fields=False) - alert["artifact"] = artifacts + alert = Alert.get(self.param_rowid) state.alert = alert return state @@ -83,19 +75,16 @@ class Playbook(LanggraphPlaybook): def output_node(state: AgentState): """处理分析结果""" - suggestion = state.suggestion fields = [ {"id": "suggestion_ai", "value": suggestion}, ] - rowid = self.param("rowid") - WorksheetRow.update(Alert.WORKSHEET_ID, rowid, fields) + Alert.update(self.param_rowid, fields) + + self.send_notice("Alert_Suggestion_Gen_By_LLM output_node Finish", f"rowid：{self.param('rowid')}") + self.update_playbook("Success", "Get suggestion by ai agent completed.") self.agent_state = state - - Notice.send(self.param("user"), "Alert_Suggestion_Gen_By_LLM output_node Finish", f"rowid：{self.param('rowid')}") - - SIRPPlaybook.update_status_and_remark(self.param("playbook_rowid"), "Success", "Get suggestion by ai agent completed.") # Success/Failed return state # 编译graph diff --git a/PLAYBOOKS/Artifact_TI_Enrichment_By_AlienVaultOTX.py b/PLAYBOOKS/Artifact_TI_Enrichment_By_AlienVaultOTX.py index 1bbdbfe..22d3227 100644 --- a/PLAYBOOKS/Artifact_TI_Enrichment_By_AlienVaultOTX.py +++ b/PLAYBOOKS/Artifact_TI_Enrichment_By_AlienVaultOTX.py @@ -3,8 +3,7 @@ import json from Lib.api import is_ipaddress from Lib.baseplaybook import BasePlaybook from PLUGINS.AlienVaultOTX.alienvaultotx import AlienVaultOTX -from PLUGINS.SIRP.nocolyapi import WorksheetRow -from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook +from PLUGINS.SIRP.sirpapi import Artifact class Playbook(BasePlaybook): @@ -16,10 +15,7 @@ class Playbook(BasePlaybook): def run(self): try: - worksheet = self.param("worksheet") - rowid = self.param("rowid") - - artifact = WorksheetRow.get(worksheet, rowid, include_system_fields=False) + artifact = Artifact.get(self.param_rowid) self.logger.info(f"Querying threat intelligence for : {artifact}") if "ip" in artifact.get("type"): @@ -34,12 +30,12 @@ class Playbook(BasePlaybook): ti_result = {"error": "Unsupported type. Please use 'ip', 'vm_ip', or 'hash'."} fields = [{"id": "enrichment", "value": json.dumps(ti_result)}] - WorksheetRow.update(worksheet, rowid, fields) - SIRPPlaybook.update_status_and_remark(self.param("playbook_rowid"), "Success", "Threat intelligence enrichment completed.") # Success/Failed + Artifact.update(self.param_rowid, fields) + self.update_playbook("Success", "Threat intelligence enrichment completed.") except Exception as e: self.logger.exception(e) - SIRPPlaybook.update_status_and_remark(self.param("playbook_rowid"), "Failed", f"Error during TI enrichment: {e}") # Success/Failed + self.update_playbook("Failed", f"Error during TI enrichment: {e}") return diff --git a/PLAYBOOKS/Artifact_TI_Enrichment_By_Mock.py b/PLAYBOOKS/Artifact_TI_Enrichment_By_Mock.py index ea1f36d..d11c7e2 100644 --- a/PLAYBOOKS/Artifact_TI_Enrichment_By_Mock.py +++ b/PLAYBOOKS/Artifact_TI_Enrichment_By_Mock.py @@ -2,8 +2,7 @@ import json import time from Lib.baseplaybook import BasePlaybook -from PLUGINS.SIRP.nocolyapi import WorksheetRow -from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook +from PLUGINS.SIRP.sirpapi import Artifact class Playbook(BasePlaybook): @@ -15,10 +14,7 @@ class Playbook(BasePlaybook): def run(self): try: - worksheet = self.param("worksheet") - rowid = self.param("rowid") - - artifact = WorksheetRow.get(worksheet, rowid, include_system_fields=False) + artifact = Artifact.get(self.param_rowid) self.logger.info(f"Querying threat intelligence for : {artifact}") # 模拟查询威胁情报数据库,在实际应用中，这里应该调用外部API或数据库进行查询 @@ -30,12 +26,11 @@ class Playbook(BasePlaybook): "last_seen": "2024-10-01T12:34:56Z"} fields = [{"id": "enrichment", "value": json.dumps(ti_result)}] - WorksheetRow.update(worksheet, rowid, fields) - - SIRPPlaybook.update_status_and_remark(self.param("playbook_rowid"), "Success", "Threat intelligence enrichment completed.") # Success/Failed + Artifact.update(self.param_rowid, fields) + self.update_playbook("Success", "Threat intelligence enrichment completed.") except Exception as e: self.logger.exception(e) - SIRPPlaybook.update_status_and_remark(self.param("playbook_rowid"), "Failed", f"Error during TI enrichment: {e}") # Success/Failed + self.update_playbook("Failed", f"Error during TI enrichment: {e}") return diff --git a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py index 9c3a720..2531594 100644 --- a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py +++ b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py @@ -10,11 +10,7 @@ from pydantic import BaseModel, Field, ConfigDict from Lib.baseplaybook import LanggraphPlaybook from Lib.llmapi import AgentState from PLUGINS.LLM.llmapi import LLMAPI -from PLUGINS.SIRP.nocolyapi import WorksheetRow -from PLUGINS.SIRP.sirpapi import Alert, Artifact from PLUGINS.SIRP.sirpapi import Case -from PLUGINS.SIRP.sirpapi import Notice -from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook class ConfidenceLevel(str, Enum): @@ -57,17 +53,7 @@ class Playbook(LanggraphPlaybook): def init(self): def preprocess_node(state: AgentState): """预处理数据""" - # worksheet = self.param("worksheet") - rowid = self.param("rowid") - - case = WorksheetRow.get(Case.WORKSHEET_ID, rowid, include_system_fields=False) - - alerts = WorksheetRow.relations(Case.WORKSHEET_ID, rowid, "alert", relation_worksheet_id=Alert.WORKSHEET_ID, include_system_fields=False) - for alert in alerts: - artifacts = WorksheetRow.relations(Alert.WORKSHEET_ID, alert.get("rowId"), "artifact", relation_worksheet_id=Artifact.WORKSHEET_ID, - include_system_fields=False) - alert["artifact"] = artifacts - case["alert"] = alerts + case = Case.get_raw_data(self.param_rowid) state.case = case return state @@ -119,8 +105,6 @@ class Playbook(LanggraphPlaybook): analyze_result: AnalyzeResult = AnalyzeResult(**state.analyze_result) - case_row_id = self.param("rowid") - case_field = [ {"id": "severity", "value": analyze_result.new_severity}, {"id": "confidence_ai", "value": analyze_result.confidence}, @@ -128,12 +112,10 @@ class Playbook(LanggraphPlaybook): {"id": "attack_stage_ai", "value": analyze_result.current_attack_stage}, {"id": "recommended_actions_ai", "value": analyze_result.recommended_actions}, ] + Case.update(self.param_rowid, case_field) - Case.update(case_row_id, case_field) - - Notice.send(self.param("user"), "Case_L3_SOC_Analyst_Agent Finish", f"rowid：{self.param('rowid')}") - - SIRPPlaybook.update_status_and_remark(self.param("playbook_rowid"), "Success", "Get suggestion by ai agent completed.") # Success/Failed + self.send_notice("Case_L3_SOC_Analyst_Agent Finish", f"rowid：{self.param_rowid}") + self.update_playbook("Success", "Get suggestion by ai agent completed.") return state # 编译graph diff --git a/PLAYBOOKS/Case_Threat_Hunting_Agent.py b/PLAYBOOKS/Case_Threat_Hunting_Agent.py index 1765613..3e7c2c4 100644 --- a/PLAYBOOKS/Case_Threat_Hunting_Agent.py +++ b/PLAYBOOKS/Case_Threat_Hunting_Agent.py @@ -16,7 +16,6 @@ from AGENTS.ti_agent import TIAgent from Lib.baseplaybook import LanggraphPlaybook from PLUGINS.LLM.llmapi import LLMAPI from PLUGINS.SIRP.sirpapi import Case -from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook MAX_ITERATIONS = 3 PROMPT_LANG = None @@ -173,7 +172,7 @@ class MainState(BaseModel): ) planning_history: Annotated[ - List[PlanningRecord], # 这里改为 PlanningRecord + List[PlanningRecord], operator.add ] = Field( default_factory=list, @@ -227,9 +226,9 @@ class Playbook(LanggraphPlaybook): # update record for message in messages: - self.add_message_to_playbook(message, self.param("playbook_rowid"), node="analyst_node") + self.add_message_to_playbook(message, node="analyst_node") - self.add_message_to_playbook(response, self.param("playbook_rowid"), node="analyst_node") + self.add_message_to_playbook(response, node="analyst_node") # 返回更新的消息列表，LangGraph 会自动追加到 state.messages return {"messages": [response]} @@ -281,9 +280,9 @@ class Playbook(LanggraphPlaybook): # update record for message in messages: - self.add_message_to_playbook(message, self.param("playbook_rowid"), node="final_answer_node") + self.add_message_to_playbook(message, node="final_answer_node") - self.add_message_to_playbook(response, self.param("playbook_rowid"), node="final_answer_node") + self.add_message_to_playbook(response, node="final_answer_node") return { "answer": response.answer, @@ -324,20 +323,19 @@ class Playbook(LanggraphPlaybook): """意图识别：确定总目标""" self.logger.info("Intent Node Invoked") - # 获取数据 - rowid = self.param("rowid") + # 获取 Case 数据 + case = Case.get_raw_data(rowid=self.param_rowid) - case = Case.get_raw_data(rowid=rowid) - - user_intent = self.param("user_input") + # 用户意图 + user_intent = self.param_user_input if not user_intent: user_intent = "None (Auto-Pilot Mode)" # 加载system prompt system_prompt_template = self.load_system_prompt_template("Intent_System", lang=PROMPT_LANG) - system_message = system_prompt_template.format() + human_message = self.load_human_prompt_template("Intent_Human", lang=PROMPT_LANG).format(case=case, user_intent=user_intent) # 构建few-shot示例 @@ -358,9 +356,9 @@ class Playbook(LanggraphPlaybook): # update record for message in messages: - self.add_message_to_playbook(message, self.param("playbook_rowid"), node="intent_node") + self.add_message_to_playbook(message, node="intent_node") - self.add_message_to_playbook(response, self.param("playbook_rowid"), node="intent_node") + self.add_message_to_playbook(response, node="intent_node") node_out = { "case": case, @@ -438,8 +436,8 @@ class Playbook(LanggraphPlaybook): # update record for message in messages: - self.add_message_to_playbook(message, self.param("playbook_rowid"), node="planner_node") - self.add_message_to_playbook(response, self.param("playbook_rowid"), node="planner_node") + self.add_message_to_playbook(message, node="planner_node") + self.add_message_to_playbook(response, node="planner_node") node_out = { "current_plan": current_plan, @@ -534,25 +532,22 @@ class Playbook(LanggraphPlaybook): llm = llm_api.get_model(tag=["powerful"]) response = llm.invoke(messages) - case_row_id = self.param("rowid") - case_field = [ {"id": "threat_hunting_report", "value": response.content}, {"id": "threat_hunting_tool_calls", "value": json.dumps(total_tool_calls)}, ] - Case.update(case_row_id, case_field) + Case.update(self.param_rowid, case_field) # update record for message in messages: - self.add_message_to_playbook(message, self.param("playbook_rowid"), node="planner_node") - self.add_message_to_playbook(response, self.param("playbook_rowid"), node="planner_node") + self.add_message_to_playbook(message, node="planner_node") + self.add_message_to_playbook(response, node="planner_node") node_out = {"report": response.content} # update playbook status - SIRPPlaybook.update_status_and_remark(self.param("playbook_rowid"), "Success", "Get suggestion by ai agent completed.") # Success/Failed - + self.update_playbook("Success", "Threat Hunting Agent Finish.") return node_out # --- 构建主图 --- diff --git a/PLAYBOOKS/TI_Artifact_query_by_AlienVaultOTX.py b/PLAYBOOKS/TI_Artifact_query_by_AlienVaultOTX.py index 2f0f0cc..318d139 100644 --- a/PLAYBOOKS/TI_Artifact_query_by_AlienVaultOTX.py +++ b/PLAYBOOKS/TI_Artifact_query_by_AlienVaultOTX.py @@ -2,7 +2,7 @@ import json from Lib.baseplaybook import BasePlaybook from PLUGINS.AlienVaultOTX.alienvaultotx import AlienVaultOTX -from PLUGINS.SIRP.nocolyapi import WorksheetRow +from PLUGINS.SIRP.sirpapi import Artifact class Playbook(BasePlaybook): @@ -12,9 +12,8 @@ class Playbook(BasePlaybook): super().__init__() # do not delete this code def run(self): - worksheet = self.param("worksheet") - rowid = self.param("rowid") - artifact = WorksheetRow.get(worksheet, rowid) + artifact = Artifact.get(self.param_rowid) + self.logger.info(f"Querying threat intelligence for : {artifact}") if artifact.get("type") == "ip" or artifact.get("type") == "vm_ip": diff --git a/PLUGINS/SIRP/sirpapi.py b/PLUGINS/SIRP/sirpapi.py index 17a5d4a..ed0a706 100644 --- a/PLUGINS/SIRP/sirpapi.py +++ b/PLUGINS/SIRP/sirpapi.py @@ -1,6 +1,6 @@ import json import os -from typing import TypedDict, List, Optional, Union, Dict, Any, NotRequired +from typing import TypedDict, List, Optional, Union, Dict, Any, NotRequired, Literal import requests @@ -65,6 +65,11 @@ class Artifact(object): def __init__(self): pass + @staticmethod + def get(rowid, include_system_fields=False): + artifact = WorksheetRow.get(Artifact.WORKSHEET_ID, rowid, include_system_fields=include_system_fields) + return artifact + @staticmethod def list(filter: dict): result = WorksheetRow.list(Artifact.WORKSHEET_ID, filter) @@ -103,13 +108,18 @@ class Alert(object): pass @staticmethod - def get(rowid): - alert = WorksheetRow.get(Alert.WORKSHEET_ID, rowid, include_system_fields=False) + def get(rowid, include_system_fields=False): + alert = WorksheetRow.get(Alert.WORKSHEET_ID, rowid, include_system_fields=include_system_fields) artifacts = WorksheetRow.relations(Alert.WORKSHEET_ID, rowid, Alert.ARTIFACT_FIELD_ID, relation_worksheet_id=Artifact.WORKSHEET_ID, include_system_fields=False) alert[Alert.ARTIFACT_FIELD_ID] = artifacts return alert + @staticmethod + def update(rowid, fields: list): + row_id = WorksheetRow.update(Alert.WORKSHEET_ID, rowid, fields) + return row_id + @staticmethod def create(alert: InputAlert): artifact_rowid_list = [] @@ -219,6 +229,7 @@ class Case(object): @staticmethod def get_raw_data(rowid, include_system_fields=False) -> Dict: + """获取案件及其关联告警和工单的原始数据，并只保留对LLM有用字段""" case = WorksheetRow.get(Case.WORKSHEET_ID, rowid, include_system_fields=include_system_fields) useful_case_fields = ["rowId", "title", 'case_status', 'created_date', 'tags', 'severity', 'type', 'description', 'close_reason', 'alert_date', @@ -331,6 +342,9 @@ class Case(object): return f.read() +PlaybookStatusType = Literal["Success", "Failed", "Pending"] + + class Playbook(object): WORKSHEET_ID = "playbook" @@ -348,7 +362,7 @@ class Playbook(object): return row_id @staticmethod - def update_status_and_remark(row_id, status, remark): + def update_status_and_remark(row_id, status: StatusType, remark): fields = [ {"id": "job_status", "value": status}, {"id": "remark", "value": remark},