External/agentic-soc-platform
1
0
Fork 0
mirror of https://github.com/FunnyWolf/agentic-soc-platform.git synced 2026-08-22 13:12:56 +02:00
Code Issues Packages Projects Releases Wiki Activity

update case api and playbook

Browse Source
This commit is contained in:
rootkit
2026-01-26 02:55:38 +08:00
parent 8e2803bdd2
commit da012def9f
6 changed files with 118 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gemini/GEMINI.md
+3 -2
View File
@@ -1,5 +1,6 @@
# 必须遵守的规则
- 每次生成代码前,需要描述方案和代码结构,等待用户确认后再生成代码
- 除非用户明确要求,否则代码中不要添加任何注释,不要生成任何测试代码,不要执行任何单元测试
- 除非用户明确要求,否则不要生成和添加任何文档
- 除非用户明确要求,否则代码中不要添加任何注释
- 除非用户明确要求,否则不要生成和添加任何文档
- 不要运行任何测试,不要添加任何测试脚本和代码,不要进行语法检查
CLAUDE.md
+3 -2
View File
@@ -1,5 +1,6 @@
# 必须遵守的规则
- 每次生成代码前,需要描述方案和代码结构,等待用户确认后再生成代码
- 除非用户明确要求,否则代码中不要添加任何注释,不要生成任何测试代码,不要执行任何单元测试
- 除非用户明确要求,否则不要生成和添加任何文档
- 除非用户明确要求,否则代码中不要添加任何注释
- 除非用户明确要求,否则不要生成和添加任何文档
- 不要运行任何测试,不要添加任何测试脚本和代码,不要进行语法检查
Lib/llmapi.py
+4 -3
View File
@@ -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] = {}
PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py
+26 -40
View File
@@ -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()
PLUGINS/SIRP/sirpapi.py
+12 -5
View File
@@ -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
PLUGINS/SIRP/sirpmodel.py
+70 -20
View File
@@ -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生成的与此事件相关的威胁狩猎报告")