diff --git a/.run/ASP.run.xml b/.run/ASP.run.xml deleted file mode 100644 index c8ba7ca..0000000 --- a/.run/ASP.run.xml +++ /dev/null @@ -1,37 +0,0 @@ - - - - - \ No newline at end of file diff --git a/ASP/settings.py b/ASP/settings.py index 7f7c08a..44139b8 100644 --- a/ASP/settings.py +++ b/ASP/settings.py @@ -83,6 +83,7 @@ LOGGING = { 'class': 'logging.FileHandler', 'formatter': 'standard', 'filename': os.path.join(settings.BASE_DIR, 'Docker', 'log', 'django.log'), + 'encoding': 'utf-8', }, }, diff --git a/Lib/baseplaybook.py b/Lib/baseplaybook.py index 1fa2de6..3ec1b75 100644 --- a/Lib/baseplaybook.py +++ b/Lib/baseplaybook.py @@ -1,10 +1,21 @@ +import json + +from langchain_core.messages import ( + BaseMessage, + SystemMessage, + HumanMessage, + AIMessage, + ToolMessage +) from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.memory import MemorySaver from langgraph.graph.state import CompiledStateGraph +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 class BasePlaybook(BaseAPI): @@ -41,6 +52,71 @@ class LanggraphPlaybook(BasePlaybook): 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 isinstance(message, SystemMessage): + fields = [ + {"id": "type", "value": "SystemMessage"}, + {"id": "node", "value": node}, + {"id": "playbook_rowid", "value": playbook_rowid}, + {"id": "content", "value": message.content}, + {"id": "json", "value": None}, + ] + elif isinstance(message, HumanMessage): + fields = [ + {"id": "type", "value": "HumanMessage"}, + {"id": "node", "value": node}, + {"id": "playbook_rowid", "value": playbook_rowid}, + {"id": "content", "value": message.content}, + {"id": "json", "value": None}, + ] + elif isinstance(message, AIMessage): + if hasattr(message, 'tool_calls') and message.tool_calls: + fields = [ + {"id": "type", "value": "AIMessage"}, + {"id": "node", "value": node}, + {"id": "playbook_rowid", "value": playbook_rowid}, + {"id": "content", "value": message.content}, + {"id": "json", "value": json.dumps(message.tool_calls)}, + ] + else: + fields = [ + {"id": "type", "value": "AIMessage"}, + {"id": "node", "value": node}, + {"id": "playbook_rowid", "value": playbook_rowid}, + {"id": "content", "value": message.content}, + {"id": "json", "value": None}, + ] + elif isinstance(message, ToolMessage): + try: + json_data = {"name": message.name, "tool_call_id": message.tool_call_id, "result": json.loads(message.content)} + except json.decoder.JSONDecodeError: + json_data = {"name": message.name, "tool_call_id": message.tool_call_id, "result": message.content} + + fields = [ + {"id": "type", "value": "ToolMessage"}, + {"id": "node", "value": node}, + {"id": "playbook_rowid", "value": playbook_rowid}, + {"id": "json", "value": json.dumps(json_data)}, + ] + elif isinstance(message, BaseModel): + fields = [ + {"id": "type", "value": "AIMessage"}, + {"id": "node", "value": node}, + {"id": "playbook_rowid", "value": playbook_rowid}, + {"id": "content", "value": None}, + {"id": "json", "value": message.model_dump_json()}, + ] + else: + fields = [ + {"id": "role", "value": message.type}, + {"id": "node", "value": node}, + {"id": "playbook_rowid", "value": playbook_rowid}, + {"id": "content", "value": message.content}, + {"id": "json", "value": None}, + ] + row_id = PlaybookMessage.create(fields) + return row_id + def run(self): self.run_graph() return self.agent_state diff --git a/PLAYBOOKS/Case_Threat_Hunting_Agent.py b/PLAYBOOKS/Case_Threat_Hunting_Agent.py index 48d979f..1765613 100644 --- a/PLAYBOOKS/Case_Threat_Hunting_Agent.py +++ b/PLAYBOOKS/Case_Threat_Hunting_Agent.py @@ -15,9 +15,8 @@ from AGENTS.siem_agent import SIEMAgent from AGENTS.ti_agent import TIAgent 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 Case +from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook MAX_ITERATIONS = 3 PROMPT_LANG = None @@ -195,7 +194,6 @@ class Playbook(LanggraphPlaybook): def __init__(self): super().__init__() # do not delete this code self.analyst_graph: CompiledStateGraph - self.main_graph: CompiledStateGraph self.max_iterations = MAX_ITERATIONS self.build_analyst_graph() self.build_main_graph() @@ -225,7 +223,13 @@ 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.query_asset, TIAgent.lookup]) - response = llm_with_tools.invoke(messages) + response: AIMessage = llm_with_tools.invoke(messages) + + # update record + for message in messages: + self.add_message_to_playbook(message, self.param("playbook_rowid"), node="analyst_node") + + self.add_message_to_playbook(response, self.param("playbook_rowid"), node="analyst_node") # 返回更新的消息列表，LangGraph 会自动追加到 state.messages return {"messages": [response]} @@ -255,9 +259,6 @@ class Playbook(LanggraphPlaybook): # get answer reasoning last_message = state.messages[-1] - llm_api = LLMAPI() - formatter_llm = llm_api.get_model(tag=["cheap", "structured_output"]) - structured_llm = formatter_llm.with_structured_output(AnalystOutput) system_prompt_template = self.load_system_prompt_template("Analyst_Final_System", lang=PROMPT_LANG) system_message = system_prompt_template.format() @@ -273,9 +274,17 @@ class Playbook(LanggraphPlaybook): human_message ] - response = structured_llm.invoke(messages) - self.logger.info(f"[final_answer_node] Question: {state.question}") - self.logger.info(f"[final_answer_node] Answer: {response.answer}") + llm_api = LLMAPI() + formatter_llm = llm_api.get_model(tag=["cheap", "structured_output"]) + structured_llm = formatter_llm.with_structured_output(AnalystOutput) + response: AnalystOutput = structured_llm.invoke(messages) + + # 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(response, self.param("playbook_rowid"), node="final_answer_node") + return { "answer": response.answer, "reasoning": response.reasoning, @@ -314,18 +323,11 @@ class Playbook(LanggraphPlaybook): def intent_node(state: MainState): """意图识别：确定总目标""" self.logger.info("Intent Node Invoked") + # 获取数据 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 = json.dumps(case, indent=2) + case = Case.get_raw_data(rowid=rowid) user_intent = self.param("user_input") @@ -342,11 +344,6 @@ class Playbook(LanggraphPlaybook): few_shot_examples = [ ] - # 运行 - llm_api = LLMAPI() - - llm = llm_api.get_model(tag="fast") - # 构建消息列表 messages = [ system_message, @@ -354,12 +351,21 @@ class Playbook(LanggraphPlaybook): human_message ] - resp = llm.invoke(messages) + # 运行 + llm_api = LLMAPI() + llm = llm_api.get_model(tag="fast") + response: AIMessage = llm.invoke(messages) + + # update record + for message in messages: + self.add_message_to_playbook(message, self.param("playbook_rowid"), node="intent_node") + + self.add_message_to_playbook(response, self.param("playbook_rowid"), node="intent_node") node_out = { "case": case, "user_intent": user_intent, - "hunting_objective": resp.content, + "hunting_objective": response.content, "iteration_count": 0, "findings": [] } @@ -419,18 +425,21 @@ class Playbook(LanggraphPlaybook): ] llm = llm.with_structured_output(HuntingPlan) - plan_result: HuntingPlan = llm.invoke(messages) + response: HuntingPlan = llm.invoke(messages) current_record = PlanningRecord( iteration=iteration_count, - rationale=plan_result.rationale, - plan=plan_result.current_plan + rationale=response.rationale, + plan=response.current_plan ) - current_plan = plan_result.current_plan + current_plan = response.current_plan self.logger.info(f"Generated Plan for Round {iteration_count}") - for plan in current_plan: - self.logger.info(f" - {plan}") + + # 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") node_out = { "current_plan": current_plan, @@ -513,11 +522,6 @@ class Playbook(LanggraphPlaybook): few_shot_examples = [ ] - # 运行 - llm_api = LLMAPI() - - llm = llm_api.get_model(tag=["powerful"]) - # 构建消息列表 messages = [ system_message, @@ -525,18 +529,30 @@ class Playbook(LanggraphPlaybook): human_message ] - resp = llm.invoke(messages) + # 运行 + llm_api = LLMAPI() + llm = llm_api.get_model(tag=["powerful"]) + response = llm.invoke(messages) case_row_id = self.param("rowid") case_field = [ - {"id": "threat_hunting_report", "value": resp.content}, + {"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) - node_out = {"report": resp.content} + # 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") + + 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 + return node_out # --- 构建主图 --- @@ -580,7 +596,12 @@ if __name__ == "__main__": os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") django.setup() - params_debug = {'rowid': 'e07f7b79-caac-43aa-b7d3-d99d2354c5c9', 'worksheet': 'case', "user_input": "Has the host in the case been infected"} + params_debug = { + 'rowid': '47da1d00-c9bf-4b5f-8ab8-8877ec292b98', + 'worksheet': 'case', + "user_input": "Has the host in the case been infected", + "playbook_rowid": "9fb4a3e1-6ae7-47b2-9b15-95264272dff5" + } module = Playbook() module._params = params_debug module.run() diff --git a/PLUGINS/SIRP/sirpapi.py b/PLUGINS/SIRP/sirpapi.py index 546a3a6..17a5d4a 100644 --- a/PLUGINS/SIRP/sirpapi.py +++ b/PLUGINS/SIRP/sirpapi.py @@ -217,6 +217,38 @@ class Case(object): case[Case.ALERT_FIELD_ID] = alerts return case + @staticmethod + def get_raw_data(rowid, include_system_fields=False) -> Dict: + 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', + 'case_id', + 'respond_time', 'note', 'acknowledged_date'] + + case_clean = {key: case[key] for key in useful_case_fields if key in case} + + # alert id + alerts = WorksheetRow.relations(Case.WORKSHEET_ID, rowid, Case.ALERT_FIELD_ID, relation_worksheet_id=Alert.WORKSHEET_ID, + include_system_fields=include_system_fields) + alerts_clean = [] + for alert in alerts: + useful_alert_fields = ["rowId", 'severity', 'rule_id', 'rule_name', 'id'] + alert_clean = {key: alert[key] for key in useful_alert_fields if key in alert} + + artifacts = WorksheetRow.relations(Alert.WORKSHEET_ID, alert.get("rowId"), Alert.ARTIFACT_FIELD_ID, relation_worksheet_id=Artifact.WORKSHEET_ID, + include_system_fields=include_system_fields) + artifacts_clean = [] + for artifact in artifacts: + useful_artifact_fields = ["rowId", "type", "value", "enrichment", 'is_whitelisted', 'is_evidence'] + artifact_clean = {key: artifact[key] for key in useful_artifact_fields if key in artifact} + artifacts_clean.append(artifact_clean) + + alert_clean[Alert.ARTIFACT_FIELD_ID] = artifacts_clean + alerts_clean.append(alert_clean) + + case_clean[Case.ALERT_FIELD_ID] = alerts_clean + return case_clean + @staticmethod def create(case: InputCase): case_fields = [ @@ -343,7 +375,6 @@ class PlaybookMessage(object): class Notice(object): - @staticmethod def send(user, title, body=None): result = requests.post(SIRP_NOTICE_WEBHOOK, json={"title": title, "body": body, "user": user}) diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..562f934 --- /dev/null +++ b/TODO.md @@ -0,0 +1,9 @@ +## 文档 + +- 插件: LLM Mock Redis +- Agents +- 安装文档 + +## 代码 + +- Agent 接口都抽离成Agent \ No newline at end of file