add playbook message function in Threat Hunting Agent

2026-08-22 13:12:56 +02:00 · 2025-12-05 03:33:56 +08:00
parent 7c385da415
commit d590aeb0f1
6 changed files with 180 additions and 79 deletions
@@ -1,37 +0,0 @@
-<component name="ProjectRunConfigurationManager">
-  <configuration default="false" name="ASP" type="Python.DjangoServer" factoryName="Django server">
-    <module name="agentic-soc-platform" />
-    <option name="ENV_FILES" value="" />
-    <option name="INTERPRETER_OPTIONS" value="" />
-    <option name="PARENT_ENVS" value="true" />
-    <envs>
-      <env name="PYTHONUNBUFFERED" value="1" />
-    </envs>
-    <option name="SDK_HOME" value="" />
-    <option name="SDK_NAME" value="uv (agentic-soc-platform)" />
-    <option name="WORKING_DIRECTORY" value="" />
-    <option name="IS_MODULE_SDK" value="false" />
-    <option name="ADD_CONTENT_ROOTS" value="true" />
-    <option name="ADD_SOURCE_ROOTS" value="true" />
-    <EXTENSION ID="net.ashald.envfile">
-      <option name="IS_ENABLED" value="false" />
-      <option name="IS_SUBST" value="false" />
-      <option name="IS_PATH_MACRO_SUPPORTED" value="false" />
-      <option name="IS_IGNORE_MISSING_FILES" value="false" />
-      <option name="IS_ENABLE_EXPERIMENTAL_INTEGRATIONS" value="false" />
-      <ENTRIES>
-        <ENTRY IS_ENABLED="true" PARSER="runconfig" IS_EXECUTABLE="false" />
-      </ENTRIES>
-    </EXTENSION>
-    <option name="launchJavascriptDebuger" value="false" />
-    <option name="port" value="7000" />
-    <option name="host" value="0.0.0.1" />
-    <option name="additionalOptions" value="" />
-    <option name="browserUrl" value="" />
-    <option name="runTestServer" value="false" />
-    <option name="runNoReload" value="false" />
-    <option name="useCustomRunCommand" value="false" />
-    <option name="customRunCommand" value="" />
-    <method v="2" />
-  </configuration>
-</component>
@@ -83,6 +83,7 @@ LOGGING = {
            'class': 'logging.FileHandler',
            'formatter': 'standard',
            'filename': os.path.join(settings.BASE_DIR, 'Docker', 'log', 'django.log'),
+            'encoding': 'utf-8',
        },

    },
@@ -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
@@ -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()
@@ -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})
@@ -0,0 +1,9 @@
+## 文档
+
+- 插件: LLM Mock Redis
+- Agents
+- 安装文档
+
+## 代码
+
+- Agent 接口都抽离成Agent