mirror of
https://github.com/FunnyWolf/agentic-soc-platform.git
synced 2026-08-22 13:12:56 +02:00
tmp
This commit is contained in:
funnywolf
@@ -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，而不是补写未经证实的细节。
|
||||
@@ -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)
|
||||
|
||||
|
||||
+1
-48
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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)",
|
||||
|
||||
Reference in New Issue
Block a user