crewai 知识库针对信息安全应急演练的定位和使用

知识库针对信息安全应急演练的定位和使用

在信息安全应急演练（以下简称“应急演练”）场景中，CrewAI的核心定位是**“应急演练全流程智能协作中枢与知识支撑平台”**。其核心价值在于通过多智能体（Agent）协同配合，结合专属知识库的精准赋能，实现应急演练从“场景触发-事件处置-复盘优化”的全流程自动化、标准化与智能化，既替代部分人工重复工作（如信息查询、流程校验），又通过知识沉淀与复用持续提升演练真实性和应急响应能力，为真实网络安全事件处置积累实战经验。

CrewAI的使用需以专属知识库为核心支撑，知识库的建设与应用直接决定应急演练的效率、精度与规范性。以下从“建设内容、建设方式、作用发挥”三大核心维度，梳理CrewAI知识库在应急演练场景中的完整落地逻辑。

一、建设内容

应急演练的核心目标是模拟真实网络安全事件（如数据泄露、勒索攻击、DDoS攻击等），检验团队响应能力，其知识库内容需严格围绕“演练准备-事件检测-应急处置-复盘优化”全流程，筛选“可支撑决策、可指导操作、可验证结果、可沉淀经验”的核心信息，杜绝冗余无效内容。结合应急演练流程特点，知识库核心建设内容分为四大模块，同时遵循明确的内容筛选原则。

根据应急演练的流程，数据库主要需要四个部分。

1.领域/专家知识 (Domain Knowledge)

该模块是保障应急演练合规性、专业性的基础，为CrewAI智能体提供权威的判断标准、技术支撑和行业规范，确保演练处置流程不偏离行业准则和技术常识。

核心内容包括：一是行业标准与法规，如《网络安全法》《数据安全法》中关于应急响应的强制性要求、GB/T 22240《信息安全技术 网络安全应急响应指南》、NIST SP 800-61应急响应框架、等保2.0对应急演练的合规考核指标等；二是技术专家手册，如网络安全设备（防火墙、WAF、IDS/IPS）的应急配置手册、漏洞库（CVE、CNVD）的应急修复方案、常见攻击类型（勒索病毒、SQL注入、供应链攻击）的技术原理与识别特征；三是合规性文件，如企业内部网络安全政策、应急演练合规检查表、第三方审计要求（如SOC 2、ISO 27001对应急响应的考核标准）。

该模块内容主要用于智能体在演练中快速核验处置方案的合规性、判断攻击类型的技术特征、匹配漏洞修复的标准流程，避免出现“处置不合规、技术不专业”的问题。

2.运营/事实数据 (Operational Data)

该模块以结构化数据为主，是智能体快速定位问题、明确责任边界、开展协同处置的关键，需确保数据的实时性、准确性和完整性，支持智能体快速查询与关联分析。

核心内容包括：一是网络资产信息，如企业内网IP地址段分配表（含网段归属部门、负责人、设备类型）、服务器资产清单（主机名、操作系统、部署业务、开放端口、管理员联系方式）、数据库资产信息（库名、存储位置、敏感数据类型、访问权限）；二是联系人与权限数据，如应急响应团队成员通讯录（姓名、岗位、电话、负责模块）、第三方服务商联系方式（安全厂商、律师事务所、公关公司）、权限审批流程表（如紧急停机、数据备份恢复的审批人及审批流程）；三是环境配置数据，如演练环境拓扑图（标注关键节点、网络链路）、生产环境与演练环境的隔离配置、应急工具部署路径（如漏洞扫描工具、日志分析平台的访问地址与账号）。

该模块内容主要用于智能体在演练中快速查询资产归属、联系相关责任人、获取工具访问权限，为应急处置的快速落地提供精准数据支撑。

3.流程与规则 (Procedures & Rules)

该模块是规范智能体处置步骤、确保演练流程标准化的核心，需具备强可执行性，明确“遇到什么情况、该做什么、怎么做、注意什么”，避免智能体出现操作混乱或流程偏差。

核心内容包括：一是应急响应SOP（标准作业流程），按攻击类型分类的标准化处置流程（如勒索病毒应急响应SOP：发现告警→隔离受感染主机→备份关键数据→分析病毒溯源→修复漏洞→恢复业务→审计复盘）、关键操作步骤的先后顺序（如“必须先关闭受影响端口，再排查漏洞”）；二是应急预案，针对不同场景的专项预案（数据泄露应急预案、DDoS攻击应急预案、核心系统故障应急预案）、预案触发条件（如“单台服务器CPU占用率持续10分钟≥95%，触发资源耗尽应急处置”）；三是演练规则，如演练评分标准（响应时长、处置准确率、合规性得分）、禁止操作清单（禁止在演练中修改生产环境配置、禁止泄露真实用户数据）、演练场景参数（模拟攻击的IP范围、漏洞利用路径）。

该模块内容主要用于智能体在演练中明确操作流程、遵循处置规则、把控演练边界，确保应急处置动作标准化、规范化。

4.历史与经验 (History & Experience)

该模块由CrewAI框架自动管理，无需人工手动维护，核心是沉淀过往演练的关键信息，让智能体“从历史演练中学习”，持续优化处置效率和精度，实现演练能力的迭代升级。

核心内容包括：一是历史演练结果，如过往演练的场景类型、处置过程记录、得分情况、存在的问题（如“2024年Q2勒索病毒演练中，漏洞修复耗时过长”）；二是经验教训总结，如成功处置案例、失败案例复盘、优化建议；三是偏好与适配调整，如根据多次演练结果记录的企业处置偏好、不同场景下的智能体协作模式（“DDoS攻击演练中，由网络Agent先处置，再移交安全Agent溯源”）。

该模块内容主要用于智能体在演练中复用成功经验、规避过往错误、适配企业专属处置习惯，同时为演练复盘优化提供核心依据。

内容筛选原则

为确保知识库内容的“精准性、可用性、安全性”，建设过程中需严格遵循三大筛选原则：

1. 必要性原则：仅保留与应急演练直接相关、智能体完成任务必须的信息，坚决剔除冗余内容，如过时的旧版应急预案、与演练无关的日常运维数据、重复的漏洞说明文档、测试用的无效数据等，避免占用存储资源、影响检索效率。

2. 可解析性原则：优先选择结构化（CSV、Excel、JSON）或半结构化（Markdown、带标题的Word）数据，确保智能体可快速提取关键字段；对于非结构化文本（如长PDF手册），需按“场景-步骤-注意事项”分层排版（如应急预案按“事件类型→处置流程→风险提示”分章节），便于智能体进行文本切片和精准检索，避免因格式混乱导致信息无法解析。

3. 安全性原则：针对敏感信息（如服务器管理员密码、真实用户数据、核心网络拓扑图、核心算法），必须进行加密存储或脱敏处理（如密码加密、用户数据替换为测试数据、拓扑图隐藏关键节点地址）；同时严格限制智能体对敏感数据的访问权限（如仅允许指定智能体读取权限审批表），避免演练过程中出现信息泄露，保障演练环境和企业核心信息安全。

二、建设方式

应急演练的数据源呈现“结构化数据与非结构化文档并存、小规模配置文件与大规模手册共生”的特征，不同类型、规模的数据适配不同的建设路径。结合CrewAI框架特性，知识库建设主要有三种核心路径，可单独使用或混合搭配，核心目标是平衡“开发成本、检索效率、使用便捷性”，确保知识库能高效支撑应急演练需求。

路径一：工具挂载 (Tool-based)

1. 适用场景

核心适配运营/事实数据中的小规模结构化文件，以及简易版流程规则表格，具体满足：数据类型为CSV、Excel、JSON等结构化格式；单文件体积≤10MB，总文件数量≤20个；使用需求为快速查询字段、提取关键信息，无需复杂语义解析（如“查询某IP的归属部门”“提取应急联系人电话”）。

典型应用：挂载网络资产清单（CSV）、应急响应团队通讯录（Excel）、演练评分标准（Excel）等。

2. 建设方法

无需复杂开发，直接使用CrewAI框架内置的FileReadTool（单文件读取工具）和DirectoryReadTool（目录读取工具），按以下步骤完成建设：

第一步，数据预处理：对结构化文件进行清理，确保字段清晰，删除空行和无效数据，统一文件编码，避免读取乱码；

第二步，工具配置：针对单个关键文件，使用FileReadTool配置文件路径、编码格式，并添加工具描述（如“用于查询网络资产的IP地址、归属部门、负责人等信息”），便于智能体识别工具用途；针对存放多个规则表格的目录，使用DirectoryReadTool配置目录路径，通过valid_extensions参数限定仅读取CSV、Excel、JSON格式文件，避免无关文件干扰；

第三步，工具绑定与测试：将配置好的工具绑定至负责“信息查询”的智能体（如对应的Agent），测试工具调用是否正常，确保智能体可精准提取目标字段。

3. 核心优势与局限

优势：零开发成本、部署速度快、检索效率高，适配演练中快速查询的刚需；

局限：仅支持结构化数据和小规模文件，不具备语义检索能力，无法适配大规模非结构化文档的检索需求。

路径二：自定义 RAG (Custom RAG)

1. 适用场景

核心适配领域/专家知识、流程与规则中的大规模非结构化文档，具体满足：数据类型为PDF手册、长文本Markdown、复杂Word文档；单文档体积≥50MB、总文档数量≥50个；使用需求为精准语义检索（如“查询勒索病毒应急响应中的数据备份步骤”）、需精细化控制文本切片逻辑（如按章节、场景拆分文档）。

典型应用：挂载GB/T 22240标准PDF、NIST应急响应框架手册、详细版应急预案（Markdown）、复杂攻击类型处置手册等。

2. 建设方法

需基于向量库和嵌入模型构建检索增强生成（RAG）系统，封装为CrewAI自定义工具，核心步骤如下：

第一步，选型配置：向量库优先选择轻量级ChromaDB（适合私有部署、演练环境隔离，成本低、部署便捷），大规模场景可选用Milvus；嵌入模型选择SentenceTransformer的all-MiniLM-L6-v2，高精度需求可选用OpenAI的text-embedding-ada-002；

第二步，文本预处理与切片：对非结构化文档进行清理，按“场景-章节-核心内容”分层切片，为每个切片添加元数据标签，便于后续溯源和精准检索；

第三步，向量库构建：将切片文本、元数据一同传入嵌入模型生成向量，存入向量库并建立索引，设置定期更新机制；

第四步，自定义工具封装：基于CrewAI的BaseTool类，封装检索逻辑，明确检索规则，确保智能体可快速获取相关内容；

第五步，工具绑定与调试：将自定义RAG工具绑定至负责“专业解读、流程指导”的智能体，调试检索精度，确保查询意图与检索结果匹配度≥80%。

3. 核心优势与局限

优势：检索精度高（支持语义检索，可精准匹配复杂查询需求）、支持大规模文档存储、切片和检索规则可自定义，适配演练中复杂场景的专业知识查询需求；

局限：开发成本较高、需维护向量库、本地部署需适配硬件资源。

路径三：原生 Memory 系统

1. 适用场景

核心适配历史与经验数据的沉淀，以及跨任务关联信息的传递，具体满足：数据类型为过往演练结果、经验教训、跨任务协作数据（如前序智能体的查询结果）；使用需求为无需手动维护、智能体自主记忆与复用信息、支持多智能体协作演练（如资产查询Agent→处置执行Agent→复盘优化Agent的信息传递）。

典型应用：存储历史演练复盘报告、跨智能体协作数据、智能体处置偏好等。

2. 建设方法

无需手动构建知识库结构，直接依托CrewAI框架原生Memory系统，按以下步骤完成配置：

第一步，基础配置：在Crew（智能体团队）初始化时，开启memory=True参数，指定记忆存储路径，避免默认临时目录导致记忆丢失；

第二步，嵌入模型适配：演练环境为私有部署时，选择Ollama的nomic-embed-text模型；可访问公网时，选用OpenAI Embedding提升记忆检索精度；

第三步，记忆类型配置：根据演练需求灵活开启记忆类型，跨任务协作时开启短期记忆（Short-term Memory）+实体记忆（Entity Memory），经验沉淀时开启长期记忆（Long-term Memory），减少资源占用；

第四步，通用知识挂载：将通用型应急知识（如基础应急响应流程、常见攻击识别特征）直接通过knowledge_sources参数挂载，框架自动处理文本切片、向量化与存储，无需人工干预。

3. 核心优势与局限

优势：开箱即用、自动管理记忆生命周期、支持跨任务信息复用、可自动沉淀历史经验；

局限：记忆存储路径配置灵活度低、切片和检索逻辑不可自定义、大规模数据下性能略有下降。

三、发挥作用

知识库的核心价值不在于“存得多”，而在于“用得好”。需将知识库深度融入应急演练“准备-处置-复盘”全环节，通过“检索增强、决策约束、跨任务记忆”三大核心机制，让智能体自主调用知识、规范执行流程、高效协同协作，实现知识的动态赋能。

（一）检索增强 (RAG)

Agent 在思考时检索相关知识，将其注入 Prompt，使其能基于特定文档回答问题，减少幻觉。

1. 核心逻辑

智能体在执行演练任务时，自动检索知识库中与当前场景匹配的专业知识、流程规则和数据信息，将检索结果注入Prompt，让智能体基于权威、准确的知识做出决策，避免因知识储备不足或主观判断导致的“幻觉”。

2. 典型应用场景

1. 事件识别阶段：当演练触发“服务器异常告警”，处置智能体自动调用自定义RAG工具，检索“常见服务器异常对应的攻击类型”，结合告警特征快速判断为DDoS攻击或勒索病毒入侵，避免误判；

2. 处置执行阶段：执行“漏洞修复”步骤时，智能体调用领域/专家知识模块，检索对应漏洞的应急修复方案，严格按手册中的操作步骤执行，确保操作专业合规；

3. 合规验证阶段：演练结束后，复盘智能体调用RAG工具检索行业标准和企业合规要求，验证处置流程是否符合响应时长、操作规范等考核指标，生成合规性报告。

3. 优化技巧

在工具描述（description）中明确检索场景，引导智能体在特定场景下精准调用对应知识库工具，提升检索效率和匹配精度；同时设置检索结果过滤规则，仅返回与当前任务高度相关的信息，避免信息过载。

（二）决策约束

通过检索到的规则，强制 Agent 遵循业务逻辑。

1. 核心逻辑

通过知识库中的流程规则、禁止操作清单、优先级排序等内容，为智能体的决策和操作设置“刚性边界”，强制智能体遵循业务逻辑和演练规范，避免出现违规操作、流程偏差等问题。

2. 典型应用场景

1. 权限约束：当智能体尝试执行“关闭核心业务服务器”“修改数据库访问权限”等高危操作时，自动检索流程与规则模块中的“权限审批流程”和“禁止操作清单”，若未获取审批权限，智能体将暂停操作并发起权限申请，确保操作合规；

2. 流程约束：处置智能体在执行“数据恢复”步骤时，检索到应急预案中的明确规则“数据恢复前必须先完成漏洞修复和病毒清除”，因此自动调整操作顺序，先执行漏洞修复和病毒清除，再进行数据恢复，避免恢复后系统再次被攻击；

3. 优先级约束：当同时出现多个处置任务，智能体检索“任务优先级规则”，按优先级顺序执行任务，确保核心目标达成。

3. 实现方式

在自定义RAG工具或文件工具中添加“规则校验逻辑”，智能体检索到相关规则后，自动将其作为操作前提条件；若检测到违反规则的操作意图，智能体将拒绝执行并提示原因，同时引导其执行合规流程。

（三）跨任务记忆

利用 Short-term Memory，让后续任务的 Agent 直接利用前序 Agent 的查证结果，无需重复劳动。

1. 核心逻辑

利用CrewAI的短期记忆（Short-term Memory）和实体记忆（Entity Memory），实现跨任务、跨智能体的信息无缝流转，前序任务的执行结果自动传递给后续任务的智能体，无需重复查询；同时通过长期记忆（Long-term Memory）沉淀历史经验，持续优化演练流程和智能体处置能力。

2. 典型应用场景

1. 跨智能体信息复用：资产查询智能体完成“受攻击IP的归属查询”，处置执行智能体直接从短期记忆中获取该结果，无需重新调用工具查询，快速联系受害ip负责人确认服务器状态并开展协同处置；

2. 多步骤任务信息传递：在“攻击检测→资产定位→漏洞修复→业务恢复→复盘优化”的全流程演练中，每个步骤的执行结果自动存入记忆，后续步骤的智能体可随时调取，确保流程连贯、信息不丢失；

3. 历史经验复用：复盘优化智能体从长期记忆中提取“往年演练的经验教训”，结合本次演练结果，生成针对性优化建议，并自动更新至知识库，供后续演练使用。

3. 优化技巧

设置记忆有效期（如短期记忆保留24小时，适配单场演练的跨任务协作；长期记忆永久保留，适配多场演练的经验沉淀），定期清理无效记忆（如重复的查询结果、过时的演练数据），避免占用过多存储资源；同时通过实体记忆提取关键信息（如IP地址、漏洞编号、负责人姓名），确保记忆内容精准、可用。

四、方案总结

在信息安全应急演练场景中，CrewAI的定位是“全流程智能协作中枢与知识支撑平台”，其核心价值实现依赖于专属知识库的精准建设与高效应用。知识库需围绕“领域/专家知识、运营/事实数据、流程与规则、历史与经验”四大核心内容构建，遵循“必要性、可解析性、安全性”三大筛选原则；建设过程中需结合数据特征，采用“工具挂载+自定义RAG+原生Memory系统”的混合建设路径，平衡开发成本与使用需求；通过“检索增强、决策约束、跨任务记忆”三大机制，实现知识在演练全流程的动态赋能，最终达成“演练流程标准化、决策判断精准化、协同协作高效化、能力迭代自动化”的核心目标，为真实网络安全事件的处置积累实战经验，提升企业整体应急响应能力。

五、具体实施

想要让crewai在应急响应中发挥作用，首先需要有庞大的知识库作为支撑。这个知识库不仅能指导ai在演练被红队攻击时处理事务的优先级，还能规范ai对应急的响应操作，确保其迅速有效。此外，我们还可以把公司的资产以及对应负责人的信息放进csv文件中，方便crewai使用readtools对其进行排查上报。

1、数据生成

import os
import csv
import random
from datetime import datetime

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
KNOWLEDGE_DIR = os.path.join(BASE_DIR, "knowledge")

def ensure_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)

def generate_assets():
    print("Generating Asset Inventory...")
    assets_dir = os.path.join(KNOWLEDGE_DIR, "assets")
    ensure_dir(assets_dir)
    
    # 1. Server Inventory
    departments = ["Finance", "HR", "R&D", "Marketing", "Operations", "IT_Security"]
    os_types = ["Ubuntu 22.04", "CentOS 7", "Windows Server 2019", "Windows Server 2022", "RedHat 8"]
    services = ["Web Server", "Database", "File Share", "Domain Controller", "App Server", "Backup Server"]
    
    with open(os.path.join(assets_dir, "server_inventory.csv"), "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["Asset_ID", "IP_Address", "Hostname", "OS", "Service", "Department", "Owner", "Criticality", "Last_Patched"])
        
        for i in range(1, 201):  # 200 Servers
            dept = random.choice(departments)
            ip = f"192.168.{random.randint(10, 50)}.{random.randint(1, 254)}"
            hostname = f"{dept[:3].upper()}-SRV-{i:03d}"
            os_ver = random.choice(os_types)
            svc = random.choice(services)
            owner = f"admin_{dept.lower()}@company.com"
            criticality = "High" if svc in ["Database", "Domain Controller"] else random.choice(["Medium", "Low"])
            last_patched = datetime(2025, random.randint(1, 12), random.randint(1, 28)).strftime("%Y-%m-%d")
            
            writer.writerow([f"ASSET-{i:04d}", ip, hostname, os_ver, svc, dept, owner, criticality, last_patched])

    # 2. Emergency Contacts
    with open(os.path.join(assets_dir, "emergency_contacts.csv"), "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["Role", "Name", "Phone", "Email", "Availability"])
        roles = [
            ("CISO", "Zhang Wei", "+86-13800000001", "ciso@company.com", "24/7"),
            ("Network Lead", "Li Qiang", "+86-13800000002", "net_lead@company.com", "On-Call"),
            ("Server Lead", "Wang Fang", "+86-13800000003", "srv_lead@company.com", "9am-6pm"),
            ("Legal Counsel", "Zhao Min", "+86-13800000004", "legal@company.com", "Business Hours"),
            ("PR Manager", "Sun Li", "+86-13800000005", "pr@company.com", "On-Call")
        ]
        for r in roles:
            writer.writerow(r)

def generate_policies():
    print("Generating Security Policies...")
    policies_dir = os.path.join(KNOWLEDGE_DIR, "policies")
    ensure_dir(policies_dir)
    
    # 1. Network Security Policy (Long Text)
    content = "# Enterprise Network Security Policy\n\n"
    content += "## Chapter 1: General Provisions\n"
    content += "This policy dictates the requirements for network security within the organization.\n\n"
    
    for i in range(1, 51):
        content += f"### Article 1.{i}: Network Access Control\n"
        content += f"All network access must be authenticated. Rule 1.{i} states that access to VLAN {100+i} is restricted to authorized personnel only. "
        content += "Users must use MFA for VPN access. Port security must be enabled on all edge switches.\n\n"

    for i in range(1, 51):
        content += f"### Article 2.{i}: Firewall Configuration\n"
        content += f"Firewall rule set {i} must default to DENY ALL. Only ports 80, 443, and 53 are allowed outbound for user segments. "
        content += "Management interfaces must not be exposed to the internet.\n\n"

    with open(os.path.join(policies_dir, "network_security_policy.md"), "w", encoding="utf-8") as f:
        f.write(content)

    # 2. Data Protection Standard
    content = "# Data Protection & Privacy Standard\n\n"
    content += "## Classification of Data\n"
    levels = ["Public", "Internal", "Confidential", "Restricted"]
    for level in levels:
        content += f"### {level} Data\n"
        content += f"Data classified as {level} requires specific handling controls. Encryption at rest is {'mandatory' if level in ['Confidential', 'Restricted'] else 'optional'}.\n\n"
    
    with open(os.path.join(policies_dir, "data_protection.md"), "w", encoding="utf-8") as f:
        f.write(content)

def generate_sops():
    print("Generating SOPs...")
    sops_dir = os.path.join(KNOWLEDGE_DIR, "sops")
    ensure_dir(sops_dir)
    
    # 1. Ransomware SOP
    content = "# Ransomware Incident Response SOP\n\n"
    phases = [
        ("Preparation", "Ensure backups are immutable. Update antivirus signatures."),
        ("Detection", "Monitor for mass file renames, encryption extensions (e.g., .locked), and ransom notes."),
        ("Containment", "1. Disconnect infected host from network immediately.\n2. Isolate the VLAN.\n3. Do NOT power off if RAM analysis is needed, but hibernate."),
        ("Eradication", "1. Identify the strain of ransomware.\n2. Locate the entry point (RDP, Phishing).\n3. Reimage the machine from gold image."),
        ("Recovery", "1. Restore data from offline backups.\n2. Verify data integrity.\n3. Phased network reconnection."),
        ("Lessons Learned", "Conduct post-incident review within 48 hours.")
    ]
    
    for phase, action in phases:
        content += f"## Phase: {phase}\n{action}\n\n"
        content += "### Detailed Checkpoints\n"
        for i in range(1, 6):
            content += f"- [ ] Checkpoint {phase[0]}-{i}: Verify step {i} of {phase} is complete.\n"
        content += "\n"

    with open(os.path.join(sops_dir, "ransomware_sop.md"), "w", encoding="utf-8") as f:
        f.write(content)

    # 2. DDoS SOP
    content = "# DDoS Incident Response SOP\n\n"
    content += "## Immediate Actions\n1. Verify if traffic is legitimate.\n2. Contact ISP for upstream blocking.\n3. Enable WAF under attack mode.\n\n"
    content += "## Analysis\nCheck logs for IP patterns. Identify if Layer 3 (Volumetric) or Layer 7 (App).\n"
    
    with open(os.path.join(sops_dir, "ddos_sop.md"), "w", encoding="utf-8") as f:
        f.write(content)

if __name__ == "__main__":
    generate_assets()
    generate_policies()
    generate_sops()
    print("Knowledge Base Generation Complete.")

因为是演练，所以我们用随机数+循环函数生成了大量的数据来模拟ai在日志中找出疑点，此外在生成一些应急报告、操作指南以及相关的联系人员名单等

2、rag_tools封装

还是依照之前的方法，通过猴子补丁把embedder的调用路径给替换成本地的模型库进行切片分隔

import os
from typing import Type, Any
from pydantic import BaseModel, Field, PrivateAttr
from crewai.tools import BaseTool
import chromadb
from sentence_transformers import SentenceTransformer

# ==============================================================================
# Global Singleton Resources (Module Level)
# To avoid Pydantic serialization issues and re-initialization
# ==============================================================================
_RAG_CLIENT = None
_RAG_COLLECTION = None
_RAG_EMBED_MODEL = None
_LOCAL_EMBEDDING_MODEL_PATH = os.getenv(
    "LOCAL_EMBEDDING_MODEL",
    "/root/.cache/huggingface/hub/models--sentence-transformers--all-MiniLM-L6-v2/snapshots/c9745ed1d9f207416be6d2e6f8de32d1f16199bf"
)

def _get_or_init_resources(knowledge_path: str):
    global _RAG_CLIENT, _RAG_COLLECTION, _RAG_EMBED_MODEL, _LOCAL_EMBEDDING_MODEL_PATH
    
    if _RAG_CLIENT is None:
        print("Initializing RAG System Resources...")
        
        # 1. Setup ChromaDB
        db_path = os.path.join(os.path.dirname(knowledge_path), "rag_vector_db")
        _RAG_CLIENT = chromadb.PersistentClient(path=db_path)
        
        # 2. Setup Embedding Model
        print(f"Loading Embedding Model from: {_LOCAL_EMBEDDING_MODEL_PATH}")
        try:
            _RAG_EMBED_MODEL = SentenceTransformer(_LOCAL_EMBEDDING_MODEL_PATH)
        except Exception as e:
            print(f"Failed to load local model from {_LOCAL_EMBEDDING_MODEL_PATH}, falling back to 'all-MiniLM-L6-v2'")
            _RAG_EMBED_MODEL = SentenceTransformer('all-MiniLM-L6-v2')

        # 3. Get Collection
        _RAG_COLLECTION = _RAG_CLIENT.get_or_create_collection(name="security_knowledge")
        
        # 4. Check ingestion
        if _RAG_COLLECTION.count() == 0:
            _ingest_knowledge(knowledge_path)
        else:
            print(f"RAG DB already contains {_RAG_COLLECTION.count()} documents.")

def _ingest_knowledge(knowledge_path: str):
    global _RAG_COLLECTION, _RAG_EMBED_MODEL
    print(f"Ingesting knowledge from {knowledge_path}...")
    docs = []
    ids = []
    metadatas = []
    
    for root, _, files in os.walk(knowledge_path):
        for file in files:
            if file.endswith(".md") or file.endswith(".txt"):
                file_path = os.path.join(root, file)
                with open(file_path, "r", encoding="utf-8") as f:
                    text = f.read()
                
                # Simple chunking
                chunks = [c.strip() for c in text.split('\n\n') if c.strip()]
                
                for i, chunk in enumerate(chunks):
                    if len(chunk) < 10: continue
                    docs.append(chunk)
                    ids.append(f"{file}_{i}")
                    metadatas.append({"source": file, "category": os.path.basename(root)})

    if docs:
        print(f"Embedding {len(docs)} chunks...")
        embeddings = _RAG_EMBED_MODEL.encode(docs).tolist()
        
        batch_size = 100
        for i in range(0, len(docs), batch_size):
            end = min(i + batch_size, len(docs))
            _RAG_COLLECTION.add(
                ids=ids[i:end],
                documents=docs[i:end],
                embeddings=embeddings[i:end],
                metadatas=metadatas[i:end]
            )
        print("Ingestion Complete.")


# Define Input Schema for the Tool
class SecurityKnowledgeInput(BaseModel):
    query: str = Field(..., description="The search query regarding security policies, SOPs, or standards.")

class SecurityKnowledgeRAGTool(BaseTool):
    name: str = "Security Knowledge Base Search"
    description: str = (
        "Useful for searching security policies, compliance standards, and incident response SOPs. "
        "Input should be a specific question or keyword related to the incident."
    )
    args_schema: Type[BaseModel] = SecurityKnowledgeInput
    
    def __init__(self, knowledge_path: str = "./knowledge", **kwargs):
        super().__init__(**kwargs)
        _get_or_init_resources(knowledge_path)

    def _run(self, query: str) -> str:
        global _RAG_COLLECTION, _RAG_EMBED_MODEL
        print(f"Searching RAG for: {query}")
        
        if _RAG_EMBED_MODEL is None or _RAG_COLLECTION is None:
             return "Error: RAG System not initialized."

        query_embedding = _RAG_EMBED_MODEL.encode([query]).tolist()
        
        results = _RAG_COLLECTION.query(
            query_embeddings=query_embedding,
            n_results=3,
            include=['documents', 'metadatas']
        )
        
        output = "### Search Results:\n"
        if results['documents'] and results['documents'][0]:
            for i, doc in enumerate(results['documents'][0]):
                meta = results['metadatas'][0][i]
                output += f"**Source:** {meta['source']} ({meta['category']})\n"
                output += f"**Content:** {doc}\n\n"
        else:
            output += "No relevant documents found."
            
        return output

3.main函数执行

最后按照我们直接的建设内容进行人员任务划分，并接入llm，最后将输出结果写到md当中去。

import os
import sys
from typing import List, Any
from sentence_transformers import SentenceTransformer

# ==============================================================================
# 1. Environment Configuration & Patching
# ==============================================================================

# LLM Configuration
LLM_API_KEY = os.getenv("LLM_API_KEY", "YOUR_LLM_API_KEY_HERE")
LLM_BASE_URL = os.getenv("LLM_BASE_URL", "https://YOUR_LLM_BASE_URL/api/v1")
LLM_MODEL = os.getenv("LLM_MODEL", "deepseek-v3")

# Important: Set OPENAI_API_KEY to satisfy validation checks in CrewAI/ChromaDB
# even though we will intercept the actual calls.
os.environ["OPENAI_API_KEY"] = LLM_API_KEY

# Local Embedding Model Path
LOCAL_EMBEDDING_MODEL = os.getenv(
    "LOCAL_EMBEDDING_MODEL",
    "{YOUR_LOCAL_EMBEDDING_MODEL_CACHE_PATH}",
)

print(f"🔧 [Patch] Loading local embedding model: {LOCAL_EMBEDDING_MODEL}")
try:
    _local_embed_model = SentenceTransformer(LOCAL_EMBEDDING_MODEL)
except Exception:
    print(f"⚠️ [Patch] Failed to load from {LOCAL_EMBEDDING_MODEL}, falling back to 'all-MiniLM-L6-v2'")
    _local_embed_model = SentenceTransformer('all-MiniLM-L6-v2')

# --- MONKEY PATCH: Intercept OpenAI Embedding Calls ---
import openai.resources.embeddings
from openai.types import CreateEmbeddingResponse, Embedding
from openai.types.create_embedding_response import Usage

def local_create_embedding(self, *args, **kwargs):
    """
    Mock function to replace openai.resources.embeddings.Embeddings.create
    Intercepts calls from CrewAI/ChromaDB and uses local SentenceTransformer.
    """
    input_text = kwargs.get('input')
    if not input_text:
        return CreateEmbeddingResponse(
            data=[], 
            model="local-model", 
            object="list", 
            usage=Usage(prompt_tokens=0, total_tokens=0)
        )

    if isinstance(input_text, str):
        input_text = [input_text]
    
    # Generate embeddings locally
    vectors = _local_embed_model.encode(input_text).tolist()
    
    # Construct response object mimicking OpenAI API
    data_items = []
    for i, vec in enumerate(vectors):
        data_items.append(Embedding(
            index=i,
            embedding=vec,
            object="embedding"
        ))
    
    return CreateEmbeddingResponse(
        data=data_items,
        model="text-embedding-3-small-local",
        object="list",
        usage=Usage(prompt_tokens=len(input_text), total_tokens=len(input_text))
    )

print("🔧 [Patch] Replacing openai.resources.embeddings.Embeddings.create with local implementation")
openai.resources.embeddings.Embeddings.create = local_create_embedding

# ==============================================================================
# 2. Application Logic
# ==============================================================================

from crewai import Agent, Task, Crew, Process, LLM
from crewai_tools import FileReadTool
from rag_tool import SecurityKnowledgeRAGTool

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
KNOWLEDGE_DIR = os.path.join(BASE_DIR, "knowledge")
ASSETS_DIR = os.path.join(KNOWLEDGE_DIR, "assets")

# Initialize LLM
my_llm = LLM(
    model=f"openai/{LLM_MODEL}",
    base_url=LLM_BASE_URL,
    api_key=LLM_API_KEY
)

# Tools
asset_file_tool = FileReadTool(
    file_path=os.path.join(ASSETS_DIR, "server_inventory.csv"),
    description="读取服务器资产清单，用于查找IP地址对应的归属部门、负责人和操作系统信息。"
)

contact_file_tool = FileReadTool(
    file_path=os.path.join(ASSETS_DIR, "emergency_contacts.csv"),
    description="读取应急联系人列表。"
)

rag_tool = SecurityKnowledgeRAGTool(knowledge_path=KNOWLEDGE_DIR)

# Agents
analyst = Agent(
    role='高级安全分析师',
    goal='分析安全告警，识别攻击类型，并从知识库中检索正确的应急响应SOP。',
    backstory=(
        "你是网络安全取证专家。你的职责是根据告警信息，"
        "查询知识库中的SOP和安全策略，确定技术层面的处置步骤。"
    ),
    tools=[rag_tool],
    verbose=True,
    memory=True,
    llm=my_llm
)

asset_manager = Agent(
    role='IT资产管理员',
    goal='定位受影响的系统及其负责人。',
    backstory=(
        "你管理着企业的CMDB。当提供IP地址或主机名时，"
        "你能快速查出该资产属于哪个部门、谁是负责人以及其重要性等级。"
    ),
    tools=[asset_file_tool, contact_file_tool],
    verbose=True,
    memory=True,
    llm=my_llm
)

commander = Agent(
    role='应急演练指挥官',
    goal='制定最终的应急响应方案并确保合规。',
    backstory=(
        "你是本次演练的总指挥。你需要汇总分析师的技术方案和资产管理员的背景信息，"
        "生成一份可执行的、符合公司规定的《应急响应行动计划》。"
    ),
    verbose=True,
    allow_delegation=True,
    llm=my_llm
)

# Tasks
INCIDENT_SERVER_IP = "{INFECTED_SERVER_IP}"
incident_input = f"告警：在服务器 {INCIDENT_SERVER_IP} 上检测到大量文件被重命名为 .locked 后缀，疑似勒索病毒感染。"

task_analyze = Task(
    description=(
        f"分析以下告警信息：'{incident_input}'。\n"
        "1. 使用RAG工具搜索知识库，判断这是什么类型的攻击。\n"
        "2. 检索并总结该攻击类型对应的 SOP (标准作业程序) 的关键步骤。\n"
        "3. 找出任何相关的安全策略条款。"
    ),
    expected_output="一份事故分析报告，包含攻击类型、SOP关键摘要和相关策略引用。",
    agent=analyst
)

task_asset_context = Task(
    description=(
        f"基于告警IP ({INCIDENT_SERVER_IP})，查询资产详情。\n"
        "1. 在资产清单中查找该IP的主机名、归属部门和负责人。\n"
        "2. 确定该系统的关键性等级 (Criticality)。\n"
        "3. 在联系人列表中找到相关负责人的电话。"
    ),
    expected_output="受影响资产画像，包含负责人联系方式和业务重要性。",
    agent=asset_manager
)

task_response_plan = Task(
    description=(
        "综合分析报告和资产画像，制定最终的《应急响应行动计划》。\n"
        "计划必须包含：\n"
        "1. 事故摘要 (Executive Summary)。\n"
        "2. 通知对象 (谁需要被叫醒)。\n"
        "3. 详细处置步骤 (基于SOP，例如是否断网、是否关机等)。\n"
        "4. 合规性说明 (引用具体的策略条款)。"
    ),
    expected_output="一份专业的、可立即执行的应急响应行动计划。",
    agent=commander,
    context=[task_analyze, task_asset_context],
    output_file='incident_response_plan.md'
)

# Crew
security_drill_crew = Crew(
    agents=[analyst, asset_manager, commander],
    tasks=[task_analyze, task_asset_context, task_response_plan],
    process=Process.sequential,
    memory=True, 
    verbose=True,
    embedder={
        "provider": "openai",
        "config": {
            "model": "text-embedding-3-small"
        }
    }
)

if __name__ == "__main__":
    print("## 信息安全应急演练系统启动 ##")
    print(f"当前演练场景: {incident_input}")
    print("------------------------------------------------")
    
    result = security_drill_crew.kickoff()
    
    print("\n\n########################")
    print("## 最终响应方案 ##")
    print("########################\n")
    print(result)

    # 确保保存到文件（双重保险）
    with open("final_response.md", "w", encoding="utf-8") as f:
        f.write("# 应急响应演练最终报告\n\n")
        f.write(str(result))
    print("\n\n✅ 结果已保存至: incident_response_plan.md 和 final_response.md")