diff --git a/DATA/PLAYBOOKS/Investigation_Agent/Investigation_System.md b/DATA/PLAYBOOKS/Investigation_Agent/Investigation_System.md deleted file mode 100644 index 39a9869..0000000 --- a/DATA/PLAYBOOKS/Investigation_Agent/Investigation_System.md +++ /dev/null @@ -1,146 +0,0 @@ -你是一名资深 SOC / DFIR 事件调查分析师。你的任务是阅读输入的 Case 全量结构化数据，产出一份严格符合 `InvestigationReport` -Schema 的调查报告。 - -你的职责不是复述原始字段，而是基于证据形成案件判断，说明： - -- 这起 Case 更接近真实安全事件、可疑事件、误报、良性行为还是证据不足。 -- 攻击者或行为主体做了什么，已确认推进到什么阶段。 -- 已获得的权限、控制能力、访问范围或业务影响是什么。 -- 当前最重要的处置动作和仍待补证的不确定点是什么。 - -分析原则： - -1. 只使用输入 Case 中出现的事实、字段、时间、实体、关联对象和原始描述，不得编造不存在的证据。 -2. 可以推理，但所有推理都必须建立在明确证据之上；证据不足时，降低 `confidence`，并在 `unknowns` 中写明缺口。 -3. 区分“已观察到的事实”“基于事实的推论”“无法确认的部分”。不要把怀疑写成既成事实。 -4. 综合使用整个 Case 上下文，包括 `alerts`、`artifacts`、`enrichments`、`tickets`、时间字段、状态字段、文本描述和处置记录。 -5. 多条告警可能只是同一行为的重复观测。去重后再判断，不要把重复观测误写成多个独立攻击步骤。 -6. 没有执行成功、权限获得、持久化建立、横向移动成功或数据访问成功的证据时，不要夸大为“已攻陷”。 -7. `severity`、`impact`、`priority`、`confidence`、`tactic`、`technique`、`sub_technique`、`remediation` - 的已有字段只作为参考；*你需要根据案件整体证据重新评估*。 -8. 报告应优先服务分析和响应，不追求“字段写满”。某个维度没有足够证据时，可以输出空列表或更克制的结论。 - -按以下分析顺序思考： - -1. 先判断案件性质。 - 判断这更像真实入侵、恶意尝试、账号滥用、策略误配、云控制面异常、邮件安全事件、数据访问异常、良性业务行为，还是单点可疑线索。 - -2. 再判断证据强度。 - 优先寻找能形成闭环的证据链：同一主体、同一时间窗口、同一目标、同一行为路径。区分直接证据与间接迹象，评估是否存在合理的良性解释。 - -3. 再重建行为链。 - 仅输出有证据支撑的行为阶段。攻击链不要求完整，也不要求一定覆盖多个 MITRE ATT&CK 阶段；单步恶意活动、账号滥用、策略修改或误报场景同样成立。 - -4. 再判断权限、范围和影响。 - 说明行为主体获得了什么访问能力或控制能力，例如登录成功、命令执行、邮箱规则修改、云 API 调用、策略变更、数据读取、持久化入口、横向访问能力等。 - -5. 最后给出处置和补证建议。 - 优先处理仍在持续、可扩散、可复用、可再次利用的风险；同时明确哪些关键问题尚未确认。 - -字段级输出要求： - -`verdict` - -- 必须明确表达案件最终性质。 -- 优先使用：`True Positive`、`Suspicious`、`False Positive`、`Benign`、`Insufficient Data`。 -- 只有当证据非常明确且更贴近其他枚举语义时，才使用其他 verdict 值。 - -`severity` - -- 反映事件本身的技术和业务严重性，而不是告警来源默认等级。 -- 若已确认真实入侵、执行成功、权限提升、横向移动、关键账号滥用、核心资产受影响，通常不应低于 `High`。 -- 若只是单点异常、弱证据线索或更像待确认噪声，可为 `Low` 或 `Medium`。 - -`impact` - -- 反映实际或潜在影响范围。 -- 若涉及关键业务、身份系统、邮件系统、终端控制、云控制面、敏感数据访问或持续控制能力，应提高影响等级。 - -`priority` - -- 反映响应紧急度，而不只是技术严重性。 -- 若风险仍在持续、可继续扩散、可再次利用、需要立即止血或存在高价值资产暴露，应提高优先级。 - -`confidence` - -- `High`: 多源证据互相印证，关键行为链闭环清晰，替代解释空间小。 -- `Medium`: 主要结论成立，但仍存在关键缺口或部分环节依赖推理。 -- `Low`: 证据弱、歧义大、上下文不足，或良性解释仍然成立。 - -`digest` - -- 写成一段高信息密度的结论性摘要，不要写成字段罗列。 -- 建议用 4 到 6 句完成： -- 先明确 verdict 和案件性质。 -- 再概括已确认的核心行为和已推进到的阶段。 -- 再说明已获得的权限、访问范围或控制能力。 -- 再说明受影响资产、账号、数据或业务风险。 -- 最后说明最强证据以及仍未确认的关键点。 -- 必须让分析员在不看原始 Case 的情况下，也能快速理解案情核心。 - -`affected_assets` - -- 仅列出与案件直接相关、被直接操作、被明确影响，或有较强证据支持的资产。 -- 使用清晰标识，例如主机名、IP、用户名、邮箱地址、资源名、文件路径。 -- 若只是疑似受影响对象，可在语义上体现“疑似”或“潜在目标”，但不要混入大量弱相关对象。 - -`evidence_findings` - -- 这是报告中最关键的“证据发现层”，用于承载支撑结论的高价值发现。 -- 每条发现应围绕一个核心主体或一个关键行为，不要把整案塞进一条 finding。 -- `evidence` 要尽量写出可追溯线索，例如时间、字段名、对象名、告警名、enrichment 结论、处置状态或原始现象。 -- `conclusion` 要说明这条发现对案件判断意味着什么，例如支持真实入侵、支持误报、支持权限获得、支持仍需补证。 -- 允许同时出现支持恶意的发现和支持良性解释的发现；最终由 `verdict` 和 `confidence` 反映综合判断。 - -`attack_chain` - -- 使用 MITRE ATT&CK 阶段描述已确认的行为链。 -- 只输出有证据支撑的阶段，不要为了“完整”补全攻击链。 -- 如果案件不是典型多阶段攻击，输出少量步骤或空列表都可以。 -- 每一步的 `description` 应说明该阶段发生了什么，以及证据依据是什么。 - -`attack_timeline` - -- 按时间顺序输出关键事件。 -- 时间能精确则精确，不能精确则写相对顺序或近似时间。 -- `evidence_field` 尽量写出可追溯的日志字段、告警字段、对象名或记录来源。 -- 只保留能推动案件判断的关键时间点，不要重复同类噪声事件。 - -`ioc_indicators` - -- 仅保留具备排查、封禁、搜索、持续监控价值的 IOC。 -- 若案件没有明确、可复用的 IOC，可以输出空列表。 -- 不要把普通描述、一般性症状或缺乏唯一性的文本误写成 IOC。 -- `context` 要交代该 IOC 在本案中的作用，例如下载地址、C2、恶意登录源、横向移动命令、落地文件等。 - -`remediations` - -- 按响应价值输出具体、可执行的处置建议。 -- 优先考虑遏制，再考虑清除、恢复、验证、加固和持续监控。 -- 高优先级建议应优先覆盖正在进行中的风险、被盗账号、持久化入口、恶意连接、权限滥用和关键资产暴露。 -- 不要写空泛建议，例如“加强监控”或“进一步调查”而没有具体动作。 - -`unknowns` - -- 写出当前仍阻碍案件定性的关键不确定点、缺失证据或待验证事项。 -- 应聚焦最重要的缺口，例如是否登录成功、是否执行成功、是否建立持久化、是否发生数据外传、是否有更多受影响资产。 -- 不要把已经确认的结论重新写成 unknown，也不要写泛泛的“需要进一步调查”。 - -输出纪律： - -1. 输出必须完全符合 `InvestigationReport` 结构，不要添加 schema 之外的字段。 -2. 所有列表都应去重、去噪、保留最高价值内容，避免不同字段之间大段重复。 -3. 允许列表为空，但不允许为了填满结构而编造内容。 -4. 除 `digest` 外，其余字段尽量紧凑、具体、可追溯。 -5. 列表长度控制如下： - -- `affected_assets` 最多 5 条。 -- `evidence_findings` 最多 5 条。 -- `attack_chain` 最多 6 条。 -- `attack_timeline` 最多 8 条。 -- `ioc_indicators` 最多 10 条。 -- `remediations` 最多 6 条。 -- `unknowns` 最多 5 条。 - -6. 若多个字段会表达相同内容，请在最合适的字段表达一次，其他字段只保留必要信息，避免机械重复。 -7. 在信息不足时，优先输出更克制的 verdict、更低的 confidence，以及更明确的 unknowns，而不是补写未经证实的细节。 diff --git a/Lib/baseapi.py b/Lib/baseapi.py index bd8eca0..c57d822 100644 --- a/Lib/baseapi.py +++ b/Lib/baseapi.py @@ -62,6 +62,10 @@ class BaseAPI(ABC): if os.path.isfile(os.path.join(DATA_DIR, fname)): # "ES-Rule-21-Phishing_user_report_mail/senior_phishing_expert.md" template_path = os.path.join(DATA_DIR, fname) + elif os.path.isfile(os.path.join(DATA_DIR, "PLAYBOOKS", self.module_name, fname)): # "ES-Rule-21-Phishing_user_report_mail/senior_phishing_expert.md" + template_path = os.path.join(DATA_DIR, "PLAYBOOKS", self.module_name, fname) + elif os.path.isfile(os.path.join(DATA_DIR, "MODULES",self.module_name, fname)): # "ES-Rule-21-Phishing_user_report_mail/senior_phishing_expert.md" + template_path = os.path.join(DATA_DIR, "MODULES", self.module_name, fname) else: template_path = os.path.join(DATA_DIR, self.module_name, fname) diff --git a/Lib/baseplaybook.py b/Lib/baseplaybook.py index b09b232..8493e43 100644 --- a/Lib/baseplaybook.py +++ b/Lib/baseplaybook.py @@ -1,23 +1,12 @@ -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 BaseAgentState from Lib.log import logger -from PLUGINS.SIRP.sirpapi import Message from PLUGINS.SIRP.sirpapi import Playbook, Notice -from PLUGINS.SIRP.sirpextramodel import MessageType, PlaybookJobStatus, MessageModel, PlaybookModel +from PLUGINS.SIRP.sirpextramodel import PlaybookJobStatus, PlaybookModel class BasePlaybook(BaseAPI): @@ -69,42 +58,6 @@ class LanggraphPlaybook(BasePlaybook): checkpointer = MemorySaver() return checkpointer - def add_message_to_playbook(self, message: BaseMessage | BaseModel, node=None): - - message_model = MessageModel() - message_model.playbook = [self._playbook_model.row_id] - message_model.node = node - - # handle content - if isinstance(message, BaseMessage): - message_model.content = message.content - - if isinstance(message, SystemMessage): - message_model.type = MessageType.SYSTEM - elif isinstance(message, HumanMessage): - message_model.type = MessageType.HUMAN - elif isinstance(message, AIMessage): - if hasattr(message, 'tool_calls') and message.tool_calls: - message_model.type = MessageType.AI - message_model.data = json.dumps(message.tool_calls) - else: - message_model.type = MessageType.AI - 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} - message_model.type = MessageType.TOOL - message_model.data = json.dumps(json_data) - elif isinstance(message, BaseModel): - message_model.type = MessageType.AI - message_model.data = message.model_dump_json(exclude_none=True, exclude_defaults=True, exclude_unset=True) - else: - logger.warning(f"Unknown message type: {message.type}.") - - row_id = Message.create(message_model) - return row_id - # langgraph interface def run_graph(self): self.graph.checkpointer.delete_thread(self.module_name) diff --git a/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py b/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py index e419e90..0a12b10 100644 --- a/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py +++ b/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py @@ -225,8 +225,8 @@ if __name__ == "__main__": # module.run() # 批量测试最早的100条告警 - # module = Module() - # message_ids = module.read_stream_head_ids(2) - # for message_id in message_ids: - # module.debug_message_id = message_id - # module.run() + module = Module() + message_ids = module.read_stream_head_ids(5) + for message_id in message_ids: + module.debug_message_id = message_id + module.run() diff --git a/PLAYBOOKS/Investigation.py b/PLAYBOOKS/Investigation.py new file mode 100644 index 0000000..b9b998b --- /dev/null +++ b/PLAYBOOKS/Investigation.py @@ -0,0 +1,75 @@ +from datetime import datetime +from pathlib import Path + +from langchain_core.messages import HumanMessage + +from Lib.baseplaybook import BasePlaybook +from Lib.configs import DATA_DIR +from PLUGINS.LLM.llmapi import LLMAPI +from PLUGINS.SIRP.analysis import AnalysisRecord, InvestigationReport +from PLUGINS.SIRP.sirpapi import Case +from PLUGINS.SIRP.sirpbasemodel import AI_PROFILE_INVESTIGATION +from PLUGINS.SIRP.sirpcoremodel import CaseModel +from PLUGINS.SIRP.sirpextramodel import PlaybookModel + + +class Playbook(BasePlaybook): + NAME = "Investigation" + DESC = "Investigation" + + def __init__(self): + super().__init__() # do not delete this code + + def run(self): + trigger = f"playbook:{self.NAME}" + case_row_id = self.param_source_row_id + + case = Case.get(case_row_id, lazy_load=False) + if not case: + self.logger.error(f"Case not found. row_id: {case_row_id}") + return + + PROMPT_PATH = Path(DATA_DIR) / "SYSTEM" / "ANALYSIS" / "System_EN.md" + content = case.model_dump_json_for_ai(profile=AI_PROFILE_INVESTIGATION) + system_message = self.load_system_prompt_template(PROMPT_PATH).format() + + llm_api = LLMAPI() + llm = llm_api.get_model(tag="structured_output").with_structured_output(InvestigationReport) + + messages = [ + system_message, + HumanMessage(content=content) + ] + report: InvestigationReport = llm.invoke(messages) + + case_new = CaseModel( + row_id=case_row_id, + verdict_ai=report.verdict, + severity_ai=report.severity, + impact_ai=report.impact, + priority_ai=report.priority, + confidence_ai=report.confidence, + investigation_report_ai_json=AnalysisRecord( + trigger=trigger, + analysis_last_started_at=( + case.analysis_last_started_at.isoformat() if case.analysis_last_started_at else None + ), + analysis_last_completed_at=datetime.now().astimezone().isoformat(), + report=report, + ).model_dump_json(), + ) + Case.update(case_new) + self.logger.info(f"Case analysis completed. row_id: {case_row_id}, trigger: {trigger}") + + +if __name__ == "__main__": + import os + import django + + os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") + django.setup() + model = PlaybookModel(source_row_id='58cac985-341e-4391-814a-c58733a61d35') + module = Playbook() + module._playbook_model = model + + module.run() diff --git a/PLAYBOOKS/Investigation_Agent.py b/PLAYBOOKS/Investigation_Agent.py deleted file mode 100644 index fd85850..0000000 --- a/PLAYBOOKS/Investigation_Agent.py +++ /dev/null @@ -1,114 +0,0 @@ -from typing import List - -from langchain_core.messages import HumanMessage -from pydantic import BaseModel, ConfigDict, Field - -from Lib.baseplaybook import BasePlaybook -from PLUGINS.LLM.llmapi import LLMAPI -from PLUGINS.SIRP.sirpapi import Case -from PLUGINS.SIRP.sirpbasemodel import AI_PROFILE_INVESTIGATION -from PLUGINS.SIRP.sirpcoremodel import AttackStage, CaseModel, CasePriority, CaseVerdict, Confidence, Impact, Severity -from PLUGINS.SIRP.sirpextramodel import PlaybookModel - - -class AffectedAsset(BaseModel): - asset_type: str = Field(description="受影响或被攻击者直接操作的资产类型，例如 Host、IP、User、Mailbox、File、Cloud Resource。") - asset_value: str = Field(description="资产的具体标识，例如主机名、IP、用户名、邮箱地址、文件路径、云资源 ARN。") - - -class EvidenceFinding(BaseModel): - title: str = Field(description="关键发现标题，例如 可疑登录成功后修改邮箱转发规则、主机A出现横向移动痕迹。") - finding_type: str = Field(description="发现类型，例如 Identity、Host、Process、Network、Email、Cloud、Policy、Ticket、Other。") - subject: str = Field(description="该发现围绕的主体，例如某账号、主机、IP、URL、策略名或告警簇。") - evidence: str = Field(description="支撑该发现的核心证据摘要，尽量写出可追溯的字段、对象或现象。") - conclusion: str = Field(description="基于该证据得出的结论，说明它在本案中意味着什么。") - - -class AttackChainStep(BaseModel): - attack_stage: AttackStage = Field(description="MITRE ATT&CK 攻击阶段。") - description: str = Field(description="该阶段发生了什么、攻击者如何实现、证据依据是什么。") - - -class TimelineEvent(BaseModel): - timestamp: str = Field(description="事件发生时间；若无法精确确定，可填相对时间或近似时间。") - attack_behavior: str = Field(description="该时间点发生的关键行为、操作或检测现象。") - evidence_field: str = Field(description="支撑该结论的关键日志字段、原文片段或关联证据。") - - -class IndicatorOfCompromise(BaseModel): - indicator_type: str = Field(description="IOC 类型，只能从 IP、Domain、URL、FileHash、FilePath、Command、RegistryKey 中选择。") - value: str = Field(description="IOC 的具体值。") - context: str = Field(description="该 IOC 在本案中的上下文，例如作为下载地址、C2、落地文件、横向移动命令等。") - - -class Remediation(BaseModel): - action_type: str = Field(description="处置动作类型，例如隔离主机、禁用账号、阻断 URL、删除文件、修复配置。") - description: str = Field(description="可直接执行的处置或加固建议，要求具体。") - priority: CasePriority = Field(description="该处置动作自身的执行优先级。") - - -class InvestigationReport(BaseModel): - model_config = ConfigDict(use_enum_values=False) - - verdict: CaseVerdict = Field(description="AI 对案件最终性质的判断，例如 True Positive、Suspicious、False Positive、Insufficient Data。") - severity: Severity = Field(description="AI 评估的事件严重程度。") - impact: Impact = Field(description="AI 评估的事件影响等级。") - priority: CasePriority = Field(description="AI 评估的响应优先级。") - confidence: Confidence = Field(description="AI 评估的事件置信度。") - digest: str = Field(description="事件综合摘要。") - affected_assets: List[AffectedAsset] = Field(description="受影响资产列表。") - evidence_findings: List[EvidenceFinding] = Field(description="支撑案件结论的关键证据发现列表。") - attack_chain: List[AttackChainStep] = Field(description="基于证据重建的攻击链步骤。") - attack_timeline: List[TimelineEvent] = Field(description="按时间顺序排列的关键事件时间线。") - ioc_indicators: List[IndicatorOfCompromise] = Field(description="可用于排查、封禁、搜索或持续监控的 IOC 列表。") - remediations: List[Remediation] = Field(description="面向分析员的处置与加固建议。") - unknowns: List[str] = Field(description="当前仍无法确认、需要补证或需要进一步排查的不确定点列表。") - - -class Playbook(BasePlaybook): - NAME = "Investigation Agent" - DESC = "Investigation Agent" - - def __init__(self): - super().__init__() # do not delete this code - - def run(self): - case = Case.get(self.param_source_row_id) - content = case.model_dump_json_for_ai(profile=AI_PROFILE_INVESTIGATION) - - system_message = self.load_system_prompt_template("Investigation_System").format() - - llm_api = LLMAPI() - llm = llm_api.get_model(tag="structured_output") - - messages = [ - system_message, - HumanMessage(content=content) - ] - llm = llm.with_structured_output(InvestigationReport) - response: InvestigationReport = llm.invoke(messages) - - case_new = CaseModel( - row_id=self.param_source_row_id, - verdict_ai=response.verdict, - severity_ai=response.severity, - impact_ai=response.impact, - priority_ai=response.priority, - confidence_ai=response.confidence, - investigation_report_ai_json=response.model_dump_json(), - ) - Case.update(case_new) - return - - -if __name__ == "__main__": - import os - import django - - os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") - django.setup() - model = PlaybookModel(source_row_id='91dd3b8f-3143-48d0-b17b-276ff3ae884e') - module = Playbook() - module._playbook_model = model - - module.run() diff --git a/PLAYBOOKS/L3_SOC_Analyst_Agent_With_Tools.py b/PLAYBOOKS/L3_SOC_Analyst_Agent_With_Tools.py index 1513978..805b883 100644 --- a/PLAYBOOKS/L3_SOC_Analyst_Agent_With_Tools.py +++ b/PLAYBOOKS/L3_SOC_Analyst_Agent_With_Tools.py @@ -12,8 +12,8 @@ from Lib.llmapi import BaseAgentState from PLUGINS.LLM.llmapi import LLMAPI from PLUGINS.SIRP.sirpapi import Case, Knowledge from PLUGINS.SIRP.sirpbasemodel import AI_PROFILE_INVESTIGATION -from PLUGINS.SIRP.sirpextramodel import PlaybookJobStatus, PlaybookModel from PLUGINS.SIRP.sirpcoremodel import Severity, AttackStage, Confidence, CaseModel +from PLUGINS.SIRP.sirpextramodel import PlaybookJobStatus, PlaybookModel class AgentState(BaseAgentState): @@ -122,9 +122,6 @@ class Playbook(LanggraphPlaybook): content="You have reached the maximum iterations limit. Based on all the information collected above, provide your final analysis using the AnalyzeResult tool immediately.")) response = llm_with_tools.invoke(messages) - for message in messages: - self.add_message_to_playbook(message, node="analyze_node") - self.add_message_to_playbook(response, node="analyze_node") return {"loop_count": state.loop_count + 1, "messages": [response]} def should_continue(state: AgentState): diff --git a/PLAYBOOKS/Threat_Hunting_Agent.py b/PLAYBOOKS/Threat_Hunting_Agent.py index 39f74d3..75ebc57 100644 --- a/PLAYBOOKS/Threat_Hunting_Agent.py +++ b/PLAYBOOKS/Threat_Hunting_Agent.py @@ -17,8 +17,8 @@ from Lib.baseplaybook import LanggraphPlaybook from PLUGINS.LLM.llmapi import LLMAPI from PLUGINS.SIRP.sirpapi import Case from PLUGINS.SIRP.sirpbasemodel import AI_PROFILE_INVESTIGATION -from PLUGINS.SIRP.sirpextramodel import PlaybookJobStatus, PlaybookModel from PLUGINS.SIRP.sirpcoremodel import CaseModel +from PLUGINS.SIRP.sirpextramodel import PlaybookJobStatus, PlaybookModel MAX_ITERATIONS = 3 MAX_ITERATIONS_OF_FUNCTIONS_CALL = 2 @@ -225,11 +225,6 @@ class Playbook(LanggraphPlaybook): self.logger.info("Stripping hallucinated tool calls in final round.") response.tool_calls = [] - for message in messages: - self.add_message_to_playbook(message, node="analyst_node") - - self.add_message_to_playbook(response, node="analyst_node") - return {"loop_count": state.loop_count + 1, "messages": [response]} # Tool node @@ -276,12 +271,6 @@ class Playbook(LanggraphPlaybook): 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, node="final_answer_node") - - self.add_message_to_playbook(response, node="final_answer_node") - return { "answer": response.answer, "reasoning": response.reasoning, @@ -331,8 +320,9 @@ class Playbook(LanggraphPlaybook): 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.model_dump_json_for_ai(profile=AI_PROFILE_INVESTIGATION), - user_intent=user_intent) + human_message = self.load_human_prompt_template("Intent_Human", lang=PROMPT_LANG).format( + case=case.model_dump_json_for_ai(profile=AI_PROFILE_INVESTIGATION), + user_intent=user_intent) few_shot_examples = [ ] @@ -347,11 +337,6 @@ class Playbook(LanggraphPlaybook): llm = llm_api.get_model(tag="fast") response: AIMessage = llm.invoke(messages) - for message in messages: - self.add_message_to_playbook(message, node="intent_node") - - self.add_message_to_playbook(response, node="intent_node") - node_out = { "case": case, "user_intent": user_intent, @@ -420,10 +405,6 @@ class Playbook(LanggraphPlaybook): self.logger.debug(f"Generated Plan for Round {iteration_count}") - for message in messages: - self.add_message_to_playbook(message, node="planner_node") - self.add_message_to_playbook(response, node="planner_node") - node_out = { "current_plan": current_plan, "iteration_count": iteration_count, @@ -515,10 +496,6 @@ class Playbook(LanggraphPlaybook): case_new = CaseModel(row_id=self.param_source_row_id, threat_hunting_report_ai=response.content) Case.update(case_new) - for message in messages: - self.add_message_to_playbook(message, node="planner_node") - self.add_message_to_playbook(response, node="planner_node") - node_out = {"report": response.content} self.update_playbook_status(PlaybookJobStatus.SUCCESS, "Threat Hunting Agent Finish.") @@ -560,9 +537,9 @@ if __name__ == "__main__": os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") django.setup() model = PlaybookModel( - source_row_id='141a4bd0-f3cf-4e0c-91b6-f8d9fff6f653', - user_input="Has the host in the case been infected", - row_id="401ca83c-4579-4e6f-8329-2e61a6c3405a") + source_row_id='141a4bd0-f3cf-4e0c-91b6-f8d9fff6f653', + user_input="Has the host in the case been infected", + row_id="401ca83c-4579-4e6f-8329-2e61a6c3405a") module = Playbook() module._playbook_model = model diff --git a/PLUGINS/SIRP/sirpapi.py b/PLUGINS/SIRP/sirpapi.py index 9ec43b5..34e61fa 100644 --- a/PLUGINS/SIRP/sirpapi.py +++ b/PLUGINS/SIRP/sirpapi.py @@ -15,7 +15,7 @@ from PLUGINS.SIRP.nocolyapi import WorksheetRow from PLUGINS.SIRP.nocolymodel import Condition, Group, Operator from PLUGINS.SIRP.sirpbasemodel import AutoAccount, BaseSystemModel from PLUGINS.SIRP.sirpcoremodel import Severity, Confidence, EnrichmentModel, TicketModel, ArtifactModel, AlertModel, CaseModel -from PLUGINS.SIRP.sirpextramodel import PlaybookType, PlaybookJobStatus, KnowledgeAction, MessageModel, PlaybookModel, KnowledgeModel +from PLUGINS.SIRP.sirpextramodel import PlaybookType, PlaybookJobStatus, KnowledgeAction, PlaybookModel, KnowledgeModel def model_to_fields(model_instance: BaseModel) -> List[Dict[str, Any]]: @@ -1088,12 +1088,6 @@ class Case(BaseWorksheetEntity[CaseModel]): ) -class Message(BaseWorksheetEntity[MessageModel]): - """Message 实体类""" - WORKSHEET_ID = "message" - MODEL_CLASS = MessageModel - - class Playbook(BaseWorksheetEntity[PlaybookModel]): """PlaybookLoader 实体类""" WORKSHEET_ID = "playbook" diff --git a/PLUGINS/SIRP/sirpextramodel.py b/PLUGINS/SIRP/sirpextramodel.py index b2360e2..ff10731 100644 --- a/PLUGINS/SIRP/sirpextramodel.py +++ b/PLUGINS/SIRP/sirpextramodel.py @@ -1,7 +1,7 @@ from __future__ import annotations from enum import StrEnum -from typing import Optional, List, Union +from typing import Optional, List from pydantic import Field @@ -39,17 +39,6 @@ class KnowledgeAction(StrEnum): DONE = 'Done' -class MessageModel(BaseSystemModel): - playbook: Optional[List[Union[PlaybookModel, str]]] = Field(default="", description="Owning playbook row ID (所属剧本行 ID)") - node: Optional[str] = Field(default="", description="Source node name or ID (源节点名称或 ID)", - json_schema_extra={"ai": [AI_PROFILE_MCP]}) - content: Optional[str] = Field(default="", description="Message text content (消息文本内容)", - json_schema_extra={"ai": [AI_PROFILE_MCP]}) - data: Optional[str] = Field(default="", description="Message JSON payload (消息 JSON 负载)") - type: Optional[MessageType] = Field(default=None, description="Message role type (消息角色类型)", - json_schema_extra={"ai": [AI_PROFILE_MCP]}) - - class PlaybookModel(BaseSystemModel): id: Optional[str] = Field(default=None, init=False, description="Record ID e.g. playbook_000001 (记录 ID e.g. playbook_000001)", json_schema_extra={"ai": [AI_PROFILE_MCP]}) @@ -73,10 +62,6 @@ class PlaybookModel(BaseSystemModel): remark: Optional[str] = Field(default="", description="Execution remark (执行备注)", json_schema_extra={"ai": [AI_PROFILE_MCP]}) - # 关联表 - messages: Optional[List[Union[MessageModel, str]]] = Field(default=None, description="Execution message history (执行消息历史)", - json_schema_extra={"ai": [AI_PROFILE_MCP]}) - class KnowledgeModel(BaseSystemModel): id: Optional[str] = Field(default=None, init=False, description="Record ID e.g. knowledge_000001 (记录 ID e.g. knowledge_000001)",