From da012def9fb049270368c362f4ec72ae37cba5d7 Mon Sep 17 00:00:00 2001 From: rootkit Date: Mon, 26 Jan 2026 02:55:38 +0800 Subject: [PATCH] update case api and playbook --- .gemini/GEMINI.md | 5 +- CLAUDE.md | 5 +- Lib/llmapi.py | 7 +- PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py | 66 ++++++++----------- PLUGINS/SIRP/sirpapi.py | 17 +++-- PLUGINS/SIRP/sirpmodel.py | 90 ++++++++++++++++++++------ 6 files changed, 118 insertions(+), 72 deletions(-) diff --git a/.gemini/GEMINI.md b/.gemini/GEMINI.md index 2b15403..bfde4d8 100644 --- a/.gemini/GEMINI.md +++ b/.gemini/GEMINI.md @@ -1,5 +1,6 @@ # 必须遵守的规则 - 每次生成代码前,需要描述方案和代码结构,等待用户确认后再生成代码 -- 除非用户明确要求,否则代码中不要添加任何注释,不要生成任何测试代码,不要执行任何单元测试 -- 除非用户明确要求,否则不要生成和添加任何文档 \ No newline at end of file +- 除非用户明确要求,否则代码中不要添加任何注释 +- 除非用户明确要求,否则不要生成和添加任何文档 +- 不要运行任何测试,不要添加任何测试脚本和代码,不要进行语法检查 \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md index 2b15403..bfde4d8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,5 +1,6 @@ # 必须遵守的规则 - 每次生成代码前,需要描述方案和代码结构,等待用户确认后再生成代码 -- 除非用户明确要求,否则代码中不要添加任何注释,不要生成任何测试代码,不要执行任何单元测试 -- 除非用户明确要求,否则不要生成和添加任何文档 \ No newline at end of file +- 除非用户明确要求,否则代码中不要添加任何注释 +- 除非用户明确要求,否则不要生成和添加任何文档 +- 不要运行任何测试,不要添加任何测试脚本和代码,不要进行语法检查 \ No newline at end of file diff --git a/Lib/llmapi.py b/Lib/llmapi.py index cb9a7d1..5f66baf 100644 --- a/Lib/llmapi.py +++ b/Lib/llmapi.py @@ -5,13 +5,14 @@ from langgraph.graph import add_messages from pydantic import BaseModel from Lib.log import logger +from PLUGINS.SIRP.sirpmodel import CaseModel, AlertModel, ArtifactModel class AgentState(BaseModel): messages: Annotated[List[Any], add_messages] = [] - case: Dict[str, Any] = {} - alert: Dict[str, Any] = {} - artifact: Dict[str, Any] = {} + case: CaseModel = None + alert: AlertModel = None + artifact: ArtifactModel = None temp_data: Dict[str, Any] = {} analyze_result: Dict[str, Any] = {} diff --git a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py index bac0ba6..6272578 100644 --- a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py +++ b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py @@ -1,5 +1,4 @@ import json -from enum import Enum from typing import Any from langchain_core.messages import HumanMessage @@ -11,23 +10,8 @@ from Lib.baseplaybook import LanggraphPlaybook from Lib.llmapi import AgentState from PLUGINS.LLM.llmapi import LLMAPI from PLUGINS.SIRP.sirpapi import Case -from PLUGINS.SIRP.sirpmodel import PlaybookJobStatus - - -class ConfidenceLevel(str, Enum): - """Confidence Level""" - LOW = "Low" - MEDIUM = "Medium" - HIGH = "High" - - -class Severity(str, Enum): - """Severity Level""" - INFO = "Info" - LOW = "Low" - MEDIUM = "Medium" - HIGH = "High" - CRITICAL = "Critical" +from PLUGINS.SIRP.sirpmodel import PlaybookJobStatus, PlaybookModel, CaseModel +from PLUGINS.SIRP.sirpmodel import Severity, Confidence class AnalyzeResult(BaseModel): @@ -37,9 +21,9 @@ class AnalyzeResult(BaseModel): original_severity: Severity = Field(description="Original alert severity") new_severity: Severity = Field(description="Recommended new severity level") - confidence: ConfidenceLevel = Field(description="Confidence score, only one of 'Low', 'Medium', or 'High'") + confidence: Confidence = Field(description="Confidence score, only one of 'Low', 'Medium', or 'High'") analysis_rationale: str | None = Field(description="Analysis process and reasons", default=None) - current_attack_stage: str | dict[str, Any] | None = Field(description="e.g., 'T1059 - Command and Control', 'Lateral Movement'", default=None) + attack_stage: str | dict[str, Any] | None = Field(description="e.g., 'T1059 - Command and Control', 'Lateral Movement'", default=None) recommended_actions: str | dict[str, Any] | None = Field(description="e.g., 'Isolate host 10.1.1.5'", default=None) @@ -54,9 +38,8 @@ class Playbook(LanggraphPlaybook): def init(self): def preprocess_node(state: AgentState): """Preprocess data""" - case = Case.get_ai_friendly_data(self.param_source_rowid) - state.case = case - return state + case = Case.get(self.param_source_rowid) + return {"case": case} # Define node def analyze_node(state: AgentState): @@ -80,30 +63,27 @@ class Playbook(LanggraphPlaybook): messages = [ system_message, *few_shot_examples, - HumanMessage(content=json.dumps(state.case)) + HumanMessage(content=json.dumps(state.case.model_dump_for_ai())) ] llm = llm.with_structured_output(AnalyzeResult) response: AnalyzeResult = llm.invoke(messages) - state.analyze_result = response.model_dump() - - # response = llm.invoke(messages) - # response = LLMAPI.extract_think(response) # Temporary solution for langchain chatollama bug - # state.analyze_result = json.loads(response.content) - return state + analyze_result = response.model_dump() + self.logger.debug(f"Analyze result: {response.model_dump()}") + return {"analyze_result": analyze_result} def output_node(state: AgentState): """Process analysis results""" analyze_result: AnalyzeResult = AnalyzeResult(**state.analyze_result) - case_field = [ - {"id": "severity", "value": analyze_result.new_severity}, - {"id": "confidence_ai", "value": analyze_result.confidence}, - {"id": "analysis_rationale_ai", "value": analyze_result.analysis_rationale}, - {"id": "attack_stage_ai", "value": analyze_result.current_attack_stage}, - {"id": "recommended_actions_ai", "value": analyze_result.recommended_actions}, - ] - Case.update(self.param_source_rowid, case_field) + case_new = CaseModel(rowid=self.param_source_rowid, + severity_ai=analyze_result.new_severity, + confidence_ai=analyze_result.confidence, + analysis_rationale_ai=analyze_result.analysis_rationale, + attack_stage_ai=analyze_result.attack_stage, + recommended_actions_ai=analyze_result.recommended_actions, + ) + Case.update(case_new) self.send_notice("Case_L3_SOC_Analyst_Agent Finish", f"rowid:{self.param_source_rowid}") self.update_playbook_status(PlaybookJobStatus.SUCCESS, "Get suggestion by ai agent completed.") @@ -130,7 +110,13 @@ class Playbook(LanggraphPlaybook): if __name__ == "__main__": - params_debug = {'source_rowid': '47da1d00-c9bf-4b5f-8ab8-8877ec292b98', 'source_worksheet': 'case'} + import os + import django + + os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") + django.setup() + model = PlaybookModel(source_worksheet='case', source_rowid='141a4bd0-f3cf-4e0c-91b6-f8d9fff6f653') module = Playbook() - # module._params = params_debug + module._playbook_model = model + module.run() diff --git a/PLUGINS/SIRP/sirpapi.py b/PLUGINS/SIRP/sirpapi.py index 22531c3..db02ce8 100644 --- a/PLUGINS/SIRP/sirpapi.py +++ b/PLUGINS/SIRP/sirpapi.py @@ -2,6 +2,7 @@ from typing import List, Dict, Union import requests +from Lib.log import logger from PLUGINS.SIRP.CONFIG import SIRP_NOTICE_WEBHOOK from PLUGINS.SIRP.nocolymodel import AccountModel, Condition, Group, Operator from PLUGINS.SIRP.sirpbase import BaseWorksheetEntity @@ -117,12 +118,18 @@ class Case(BaseWorksheetEntity[CaseModel]): @classmethod def get_ai_friendly_data(cls, rowid: str) -> Dict: """获取LLM友好的原始数据""" - model = cls.get(rowid, include_system_fields=True) + model: CaseModel = cls.get(rowid, include_system_fields=True) - # TODO : 这里可以根据需要添加清理的字段 + # 这里可以根据需要添加清理AI不需要的字段 + model.workbook = None + model.analysis_rationale_ai = None + model.recommended_actions_ai = None + model.attack_stage_ai = None + model.severity_ai = None + model.confidence_ai = None model.threat_hunting_report_ai = None - data = model.model_dump(mode='json', exclude_unset=True, exclude_none=True, exclude_defaults=True) + data = model.model_dump(mode='json', exclude_unset=True, exclude_none=True, exclude_defaults=True, exclude={"threat_hunting_report_ai", }) return data @@ -204,8 +211,8 @@ class Notice(object): elif isinstance(user, list): users = user else: - raise ValueError("user 参数必须是 AccountModel 实例或 AccountModel 实例列表") - + logger.error("user 参数必须是 AccountModel 实例或 AccountModel 实例列表") + return False for user in users: result = requests.post(SIRP_NOTICE_WEBHOOK, json={"title": title, "body": body, "user": user.fullname}) return True diff --git a/PLUGINS/SIRP/sirpmodel.py b/PLUGINS/SIRP/sirpmodel.py index e13dd33..822e614 100644 --- a/PLUGINS/SIRP/sirpmodel.py +++ b/PLUGINS/SIRP/sirpmodel.py @@ -2,7 +2,7 @@ from __future__ import annotations from datetime import datetime from enum import StrEnum -from typing import List, Optional, Any, Union +from typing import List, Optional, Any, Union, ClassVar from pydantic import BaseModel, Field, field_validator, ConfigDict, field_serializer @@ -125,7 +125,7 @@ class ArtifactReputationScore(StrEnum): OTHER = 'Other' -class SeverityLevel(StrEnum): +class Severity(StrEnum): UNKNOWN = "Unknown" INFORMATIONAL = "Informational" LOW = "Low" @@ -186,7 +186,7 @@ class AlertAction(StrEnum): OTHER = "Other" -class ConfidenceLevel(StrEnum): +class Confidence(StrEnum): UNKNOWN = "Unknown" LOW = "Low" MEDIUM = "Medium" @@ -312,9 +312,11 @@ class KnowledgeAction(StrEnum): class BaseSystemModel(BaseModel): model_config = ConfigDict(populate_by_name=True) + ai_exclude_fields: ClassVar[set[str]] = set() + rowid: Optional[str] = Field(default=None, description="数据的唯一行ID") - ownerid: Optional[Union[List[AccountModel], AccountModel]] = Field(default=None, description="数据的所有者/创建者") - caid: Optional[AccountModel] = Field(default=None, description="当前处理人/负责人") + ownerid: Optional[AccountModel] = Field(default=None, description="数据的所有者") + caid: Optional[AccountModel] = Field(default=None, description="创建人") ctime: Optional[Union[datetime, str]] = Field(default=None, description="创建时间") utime: Optional[Union[datetime, str]] = Field(default=None, description="最后更新时间") uaid: Optional[AccountModel] = Field(default=None, description="最后更新人") @@ -356,6 +358,42 @@ class BaseSystemModel(BaseModel): return v.strftime("%Y-%m-%dT%H:%M:%SZ") return v + def model_dump_for_ai(self) -> dict[str, Any]: + """ + 递归序列化模型为 AI 友好的字典格式。 + 在序列化前处理嵌套对象，确保每层都能应用自己的 ai_exclude_fields。 + """ + result = {} + for field_name, field_value in self.__dict__.items(): + if field_name in self.ai_exclude_fields: + continue + result[field_name] = self._process_value_before_dump(field_value) + return result + + def _process_value_before_dump(self, value: Any) -> Any: + """ + 在序列化前处理值，支持递归调用嵌套模型的 model_dump_for_ai()。 + """ + if isinstance(value, BaseSystemModel): + return value.model_dump_for_ai() + elif isinstance(value, list): + return [self._process_value_before_dump(item) for item in value] + elif isinstance(value, dict): + return {k: self._process_value_before_dump(v) for k, v in value.items()} + else: + return self._serialize_value(value) + + def _serialize_value(self, value: Any) -> Any: + """ + 序列化特殊类型的值（datetime、枚举等）。 + """ + if isinstance(value, datetime): + return value.strftime("%Y-%m-%dT%H:%M:%SZ") + elif isinstance(value, StrEnum): + return value.value + else: + return value + class MessageModel(BaseSystemModel): playbook: Optional[List[Union[PlaybookModel, str]]] = Field(default="", description="所属Playbook的唯一行ID") @@ -376,7 +414,7 @@ class PlaybookModel(BaseSystemModel): name: Optional[str] = Field(default="", description="执行的Playbook的名称") user_input: Optional[str] = Field(default="", description="用户对Playbook的初始输入或后续指令") - user: Optional[List[AccountModel]] = Field(default=None, description="发起Playbook的用户") + user: Optional[Union[List[AccountModel], AccountModel, str]] = Field(default=None, description="发起Playbook的用户") # 关联表 messages: Optional[List[Union[MessageModel, str]]] = Field(default=None, description="Playbook执行过程中的所有消息记录，构成对话历史") @@ -392,6 +430,8 @@ class KnowledgeModel(BaseSystemModel): class EnrichmentModel(BaseSystemModel): + ai_exclude_fields: ClassVar[set[str]] = {'ownerid', 'caid', 'uaid'} + name: Optional[str] = Field(default="", description="富化信息的名称或标题") type: Optional[str] = Field(default="Other", description="富化信息的类型", json_schema_extra={"type": 2}) provider: Optional[str] = Field(default="Other", description="富化信息的提供方，例如威胁情报厂商", json_schema_extra={"type": 2}) @@ -402,6 +442,8 @@ class EnrichmentModel(BaseSystemModel): class TicketModel(BaseSystemModel): + ai_exclude_fields: ClassVar[set[str]] = {'ownerid', 'caid', 'uaid'} + status: Optional[TicketStatus] = Field( default=None, description="外部工单系统中的状态") type: Optional[TicketType] = Field(default=None, description="外部工单系统的类型", @@ -412,6 +454,8 @@ class TicketModel(BaseSystemModel): class ArtifactModel(BaseSystemModel): + ai_exclude_fields: ClassVar[set[str]] = {'ownerid', 'caid', 'uaid'} + name: Optional[str] = Field(default="", description="实体（Artifact）的名称，通常与值相同或为其描述") type: Optional[ArtifactType] = Field( default=None, description="实体的类型, 例如: IP地址, 主机名, 文件哈希等") @@ -432,16 +476,18 @@ class ArtifactModel(BaseSystemModel): class AlertModel(BaseSystemModel): - severity: Optional[SeverityLevel] = Field(default=None, - description="告警的严重性，由源安全产品定义") + ai_exclude_fields: ClassVar[set[str]] = {'ownerid', 'caid', 'uaid', "summary_ai", "case"} + + severity: Optional[Severity] = Field(default=None, + description="告警的严重性，由源安全产品定义") title: Optional[str] = Field(default="", description="告警的标题") impact: Optional[ImpactLevel] = Field(default=None, description="告警可能造成的影响范围") disposition: Optional[AlertDisposition] = Field( default=None, description="安全产品对该活动的处置结果, 如'Blocked', 'Allowed等") action: Optional[AlertAction] = Field(default=None, description="检测到的原始行为, 如'Allowed', 'Denied等") - confidence: Optional[ConfidenceLevel] = Field(default=None, - description="告警的置信度，表示该告警为真阳性的可能性") + confidence: Optional[Confidence] = Field(default=None, + description="告警的置信度，表示该告警为真阳性的可能性") uid: Optional[str] = Field(default="", description="告警的唯一标识符") labels: Optional[List[str]] = Field(default=[], description="为告警打上的标签", json_schema_extra={"type": 2}) desc: Optional[str] = Field(default="", description="对告警的详细描述") @@ -510,13 +556,17 @@ class AlertModel(BaseSystemModel): class CaseModel(BaseSystemModel): + ai_exclude_fields: ClassVar[set[str]] = {'ownerid', 'caid', 'uaid', "workbook", "analysis_rationale_ai", "recommended_actions_ai", "attack_stage_ai", + "severity_ai", "confidence_ai", + "threat_hunting_report_ai"} + title: Optional[str] = Field(default="", description="安全事件的标题, 应能简明扼要地概括事件的核心内容") - severity: Optional[SeverityLevel] = Field(default=None, - description="由分析师评估或重新定义的事件严重性") + severity: Optional[Severity] = Field(default=None, + description="由分析师评估或重新定义的事件严重性") impact: Optional[ImpactLevel] = Field(default=None, description="由分析师评估的事件实际影响") priority: Optional[CasePriority] = Field(default=None, description="事件的处置优先级") src_url: Optional[str] = Field(default="", description="在源系统中查看此事件的URL") - confidence: Optional[ConfidenceLevel] = Field(default=None, description="由分析师评估的事件置信度") + confidence: Optional[Confidence] = Field(default=None, description="由分析师评估的事件置信度") description: Optional[str] = Field(default="", description="对安全事件的详细描述") category: Optional[ProductCategory] = Field(default=None, @@ -525,13 +575,13 @@ class CaseModel(BaseSystemModel): status: Optional[CaseStatus] = Field(default=None, description="安全事件的处理状态") - assignee_l1: Optional[AccountModel] = Field(default=None, description="分配给L1一线分析师") + assignee_l1: Optional[Union[List[AccountModel], AccountModel, str]] = Field(default=None, description="分配给L1一线分析师") acknowledged_time: Optional[Union[datetime, str]] = Field(default=None, description="L1分析师首次确认接收事件的时间") comment: Optional[str] = Field(default="", description="分析师对整个事件的评论或处置记录") - attachments: Optional[List[Union[AttachmentModel, str]]] = Field(default=[], description="与事件相关的附件列表") + attachments: Optional[Union[List[AttachmentModel], str]] = Field(default=[], description="与事件相关的附件列表") - assignee_l2: Optional[AccountModel] = Field(default=None, description="分配或升级给L2二线分析师") - assignee_l3: Optional[AccountModel] = Field(default=None, description="分配或升级给L3专家分析师") + assignee_l2: Optional[Union[List[AccountModel], AccountModel, str]] = Field(default=None, description="分配或升级给L2二线分析师") + assignee_l3: Optional[Union[List[AccountModel], AccountModel, str]] = Field(default=None, description="分配或升级给L3专家分析师") closed_time: Optional[Union[datetime, str]] = Field(default=None, description="事件关闭的时间") verdict: Optional[CaseVerdict] = Field( default=None, description="对事件的最终裁定结论") @@ -545,9 +595,9 @@ class CaseModel(BaseSystemModel): analysis_rationale_ai: Optional[str] = Field(default="", description="AI对事件的分析基本原理和逻辑") recommended_actions_ai: Optional[str] = Field(default="", description="AI推荐的下一步操作或修复建议") attack_stage_ai: Optional[str] = Field(default="", description="AI评估的攻击阶段") - severity_ai: Optional[SeverityLevel] = Field(default=None, - description="AI评估的事件严重性") - confidence_ai: Optional[ConfidenceLevel] = Field(default=None, description="AI评估的事件置信度") + severity_ai: Optional[Severity] = Field(default=None, + description="AI评估的事件严重性") + confidence_ai: Optional[Confidence] = Field(default=None, description="AI评估的事件置信度") threat_hunting_report_ai: Optional[str] = Field(default="", description="AI生成的与此事件相关的威胁狩猎报告")