update api

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

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

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
 
 

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
 
 

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

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
 
         # --- 构建主图 ---

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":

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},