diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 0000000..4dd34f4 --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,9 @@ +{ + "permissions": { + "allow": [ + "Bash(python -m py_compile:*)" + ], + "deny": [], + "ask": [] + } +} diff --git a/.gemini/GEMINI.md b/.gemini/GEMINI.md index e12ab7b..2b15403 100644 --- a/.gemini/GEMINI.md +++ b/.gemini/GEMINI.md @@ -1,4 +1,5 @@ # 必须遵守的规则 - 每次生成代码前,需要描述方案和代码结构,等待用户确认后再生成代码 -- 除非用户明确要求,否则代码中不要添加任何注释,不要生成任何测试代码,不要执行任何单元测试 \ No newline at end of file +- 除非用户明确要求,否则代码中不要添加任何注释,不要生成任何测试代码,不要执行任何单元测试 +- 除非用户明确要求,否则不要生成和添加任何文档 \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..2b15403 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,5 @@ +# 必须遵守的规则 + +- 每次生成代码前,需要描述方案和代码结构,等待用户确认后再生成代码 +- 除非用户明确要求,否则代码中不要添加任何注释,不要生成任何测试代码,不要执行任何单元测试 +- 除非用户明确要求,否则不要生成和添加任何文档 \ No newline at end of file diff --git a/Lib/baseplaybook.py b/Lib/baseplaybook.py index afe3e76..f489c35 100644 --- a/Lib/baseplaybook.py +++ b/Lib/baseplaybook.py @@ -16,7 +16,8 @@ from Lib.baseapi import BaseAPI from Lib.llmapi import AgentState from Lib.log import logger from PLUGINS.SIRP.sirpapi import Message -from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook, Notice, PlaybookStatus +from PLUGINS.SIRP.sirpapi import Playbook, Notice +from PLUGINS.SIRP.sirptype import PlaybookModel, PlaybookJobStatus class BasePlaybook(BaseAPI): @@ -26,35 +27,35 @@ class BasePlaybook(BaseAPI): def __init__(self): super().__init__() - self._params = {} self.logger = logger - - def param(self, key, default=None): - return self._params.get(key, default) + # noinspection PyTypeChecker + self._playbook_model: PlaybookModel = None # 定义内部参数 @property def param_rowid(self): - return self.param("rowid") + return self._playbook_model.rowid @property def param_source_rowid(self): - return self.param("source_rowid") + return self._playbook_model.source_rowid @property def param_source_worksheet(self): - return self.param("source_worksheet") + return self._playbook_model.source_worksheet @property def param_user(self): - return self.param("user") + return self._playbook_model.user @property def param_user_input(self): - return self.param("user_input") + return self._playbook_model.user_input - def update_playbook(self, status: PlaybookStatus, remark: str): - rowid = SIRPPlaybook.update_status_and_remark(self.param_rowid, status, remark) + def update_playbook_status(self, status: PlaybookJobStatus, remark: str): + self._playbook_model.job_status = status + self._playbook_model.remark = remark + rowid = Playbook.update_or_create(self._playbook_model) return rowid def send_notice(self, title: str, body: str): diff --git a/Lib/llmfunc.py b/Lib/llmfunc.py index ffc388b..cc743c6 100644 --- a/Lib/llmfunc.py +++ b/Lib/llmfunc.py @@ -14,17 +14,17 @@ def function_call_debug( return f"This-is-a-test-function-to-debug_function_call-The-magic-number-is-{magic_num * 10}." -def get_case_by_case_id(case_id: Annotated[str, "Case ID"]): +def get_case_by_rowid(rowid: Annotated[str, "Case Rowid"]): """ - Retrieve a security case by its unique Case ID. + Retrieve a security case by its unique Case Rowid. This tool allows you to look up full details of a specific case when you have its ID. Useful for retrieving context, status, or artifacts associated with a known case identifier. Args: - case_id: The unique string identifier of the case (e.g., 'CASE-00000313'). + rowid: The unique string identifier of the case (e.g., '2101ff98-f52e-4f38-b107-fe53f7f77b5c'). Returns: The Case object containing all case details if found, otherwise None. """ - return Case.get_by_case_id(case_id) + return Case.get(rowid) diff --git a/Lib/montior.py b/Lib/montior.py index ebca75a..e67da02 100644 --- a/Lib/montior.py +++ b/Lib/montior.py @@ -12,12 +12,13 @@ from Lib.apsmodule import aps_module from Lib.baseplaybook import BasePlaybook from Lib.engine import Engine from Lib.log import logger -from Lib.playbook import Playbook +from Lib.playbookloader import PlaybookLoader from Lib.xcache import Xcache from PLUGINS.Embeddings.embeddings_qdrant import embedding_api_singleton_qdrant from PLUGINS.Mem0.CONFIG import USE as MEM_ZERO_USE from PLUGINS.Redis.redis_stream_api import RedisStreamAPI -from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook, Knowledge, KnowledgeAction, PlaybookStatus +from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook, Knowledge, KnowledgeAction +from PLUGINS.SIRP.sirptype import PlaybookJobStatus if MEM_ZERO_USE: from PLUGINS.Mem0.mem_zero import mem_zero_singleton @@ -90,15 +91,8 @@ class MainMonitor(object): Xcache.set_token_user(ASP_REST_API_TOKEN, api_usr, None) - logger.info("Load Playbook module config") - Playbook.load_all_playbook_config() - - # self.MainScheduler.add_job(func=self.subscribe_clean_thread, - # max_instances=1, - # trigger='interval', - # hours=1, - # id='subscribe_clean_thread') - # self.MainScheduler.start() + logger.info("Load PlaybookLoader module config") + PlaybookLoader.load_all_playbook_config() delay_time = 3 @@ -112,54 +106,44 @@ class MainMonitor(object): @staticmethod def subscribe_pending_playbook(): - records = SIRPPlaybook.get_pending_playbooks() - for one_record in records: - name = one_record.get("name") - type = one_record.get("type") - row_id = one_record.get("rowid") - module_config = Xcache.get_module_config_by_name_and_type(type, name) - if module_config is None: - Playbook.load_all_playbook_config() - module_config = Xcache.get_module_config_by_name_and_type(type, name) - if module_config is None: - logger.error(f"Playbook module config not found: {type} - {name}") + models = SIRPPlaybook.list_pending_playbooks() - SIRPPlaybook.update_status_and_remark(row_id, PlaybookStatus.FAILED, f"Playbook module config not found: {type} - {name}") + for model in models: + module_config = Xcache.get_module_config_by_name_and_type(model.type, model.name) + if module_config is None: + PlaybookLoader.load_all_playbook_config() # try again + module_config = Xcache.get_module_config_by_name_and_type(model.type, model.name) + if module_config is None: + logger.error(f"PlaybookLoader module config not found: {model.type} - {model.name}") + model.job_status = PlaybookJobStatus.FAILED + model.remark = f"PlaybookLoader module config not found: {model.type} - {model.name}" + SIRPPlaybook.update_or_create(model) continue + load_path = module_config.get("load_path") - if one_record.get("user"): - user = one_record.get("user")[0].get("fullname") - else: - user = None - - params = { - "rowid": row_id, - "source_worksheet": one_record.get("type").lower(), - "source_rowid": one_record.get("source_rowid"), - "user_input": one_record.get("user_input"), - "user": user, - } - try: class_intent = importlib.import_module(load_path) playbook_intent: BasePlaybook = class_intent.Playbook() - playbook_intent._params = params + playbook_intent._playbook_model = model except Exception as E: logger.exception(E) - SIRPPlaybook.update_status_and_remark(row_id, PlaybookStatus.FAILED, f"{E}") + model.job_status = PlaybookJobStatus.FAILED + model.remark = str(E) + SIRPPlaybook.update_or_create(model) continue job_id = aps_module.putin_post_python_module_queue(playbook_intent) - if job_id: - logger.info(f"Create playbook job success: {job_id}") - fields = [ - {"id": "job_status", "value": "Running"}, - {"id": "job_id", "value": job_id}, - ] - SIRPPlaybook.update(row_id, fields) + if not job_id: + model.job_status = PlaybookJobStatus.FAILED + model.remark = "Failed to create playbook job." + SIRPPlaybook.update_or_create(model) + continue else: - SIRPPlaybook.update_status_and_remark(row_id, PlaybookStatus.FAILED, f"Failed to create playbook job.") + logger.info(f"Create playbook job success: {job_id}") + model.job_status = PlaybookJobStatus.RUNNING + model.job_id = job_id + SIRPPlaybook.update_or_create(model) @staticmethod def subscribe_knowledge_action(): diff --git a/Lib/playbook.py b/Lib/playbookloader.py similarity index 89% rename from Lib/playbook.py rename to Lib/playbookloader.py index 37345b5..2a6539c 100644 --- a/Lib/playbook.py +++ b/Lib/playbookloader.py @@ -6,7 +6,7 @@ from Lib.log import logger from Lib.xcache import Xcache -class Playbook(object): +class PlaybookLoader(object): """Task Adder""" def __init__(self): @@ -26,7 +26,7 @@ class Playbook(object): @staticmethod def gen_playbook_config(modulename, module_files_dir="PLAYBOOKS"): - module_intent = Playbook.get_playbook_intent(modulename, module_files_dir) + module_intent = PlaybookLoader.get_playbook_intent(modulename, module_files_dir) if module_intent is None: return None @@ -53,7 +53,7 @@ class Playbook(object): module_filenames = os.listdir(os.path.join(BASE_DIR, 'PLAYBOOKS')) for module_filename in module_filenames: module_name = module_filename.split(".")[0] - one_module_config = Playbook.gen_playbook_config(module_name, 'PLAYBOOKS') + one_module_config = PlaybookLoader.gen_playbook_config(module_name, 'PLAYBOOKS') if one_module_config is not None: all_modules_config.append(one_module_config) module_count += 1 diff --git a/MODULES/NDR-Rule-05-Suspect-C2-Communication.py b/MODULES/NDR-Rule-05-Suspect-C2-Communication.py index 67ed77c..53579cc 100644 --- a/MODULES/NDR-Rule-05-Suspect-C2-Communication.py +++ b/MODULES/NDR-Rule-05-Suspect-C2-Communication.py @@ -115,7 +115,7 @@ class Module(LanggraphModule): case_row_id = state.temp_data.get("case_row_id") case = Case.get(case_row_id) - fields_useful_to_llm = ['title', 'alert', 'tags', 'severity', 'type', 'description', 'close_reason', 'alert_date', 'case_id', 'note', 'summary', + fields_useful_to_llm = ['title', 'alert', 'tags', 'severity', 'type', 'description', 'close_reason', 'alert_date', 'rowid', 'note', 'summary', 'attachment'] case_for_llm = {} for key in case: diff --git a/PLAYBOOKS/Alert_Analysis_Agent.py b/PLAYBOOKS/Alert_Analysis_Agent.py index b3a0f0a..715083f 100644 --- a/PLAYBOOKS/Alert_Analysis_Agent.py +++ b/PLAYBOOKS/Alert_Analysis_Agent.py @@ -8,7 +8,8 @@ from pydantic import BaseModel from Lib.baseplaybook import LanggraphPlaybook from PLUGINS.LLM.llmapi import LLMAPI -from PLUGINS.SIRP.sirpapi import Alert, PlaybookStatus +from PLUGINS.SIRP.sirpapi import Alert +from PLUGINS.SIRP.sirptype import PlaybookJobStatus class AgentState(BaseModel): @@ -20,7 +21,7 @@ class AgentState(BaseModel): class Playbook(LanggraphPlaybook): TYPE = "ALERT" # Classification tag - NAME = "Alert Analysis Agent" # Playbook name + NAME = "Alert Analysis Agent" # PlaybookLoader name def __init__(self): super().__init__() # do not delete this code @@ -81,7 +82,7 @@ class Playbook(LanggraphPlaybook): Alert.update(self.param_source_rowid, fields) self.send_notice("Alert_Suggestion_Gen_By_LLM output_node Finish", f"rowid:{self.param_source_rowid}") - self.update_playbook(PlaybookStatus.SUCCESS, "Get suggestion by ai agent completed.") + self.update_playbook_status(PlaybookJobStatus.SUCCESS, "Get suggestion by ai agent completed.") self.agent_state = state return state @@ -109,5 +110,5 @@ class Playbook(LanggraphPlaybook): if __name__ == "__main__": params_debug = {'source_rowid': '13782e0a-2423-4fc3-9b16-7f2eb15ae83f', 'source_worksheet': 'alert'} module = Playbook() - module._params = params_debug + # module._params = params_debug module.run() diff --git a/PLAYBOOKS/Artifact_TI_Enrichment_By_AlienVaultOTX.py b/PLAYBOOKS/Artifact_TI_Enrichment_By_AlienVaultOTX.py index 1fbde3d..b7d6da8 100644 --- a/PLAYBOOKS/Artifact_TI_Enrichment_By_AlienVaultOTX.py +++ b/PLAYBOOKS/Artifact_TI_Enrichment_By_AlienVaultOTX.py @@ -3,7 +3,8 @@ import json from Lib.api import is_ipaddress from Lib.baseplaybook import BasePlaybook from PLUGINS.AlienVaultOTX.alienvaultotx import AlienVaultOTX -from PLUGINS.SIRP.sirpapi import Artifact, PlaybookStatus +from PLUGINS.SIRP.sirpapi import Artifact +from PLUGINS.SIRP.sirptype import PlaybookJobStatus class Playbook(BasePlaybook): @@ -32,15 +33,15 @@ class Playbook(BasePlaybook): fields = [{"id": "enrichment", "value": json.dumps(ti_result)}] Artifact.update(self.param_source_rowid, fields) - self.update_playbook(PlaybookStatus.SUCCESS, "Threat intelligence enrichment completed.") + self.update_playbook_status(PlaybookJobStatus.SUCCESS, "Threat intelligence enrichment completed.") except Exception as e: self.logger.exception(e) - self.update_playbook(PlaybookStatus.FAILED, f"Error during TI enrichment: {e}") + self.update_playbook_status(PlaybookJobStatus.FAILED, f"Error during TI enrichment: {e}") return if __name__ == "__main__": params_debug = {'source_rowid': '54725ee3-c85d-49e7-ac09-4cb982dab957', 'source_worksheet': 'Artifact'} module = Playbook() - module._params = params_debug + # module._params = params_debug module.run() diff --git a/PLAYBOOKS/Artifact_TI_Enrichment_By_Mock.py b/PLAYBOOKS/Artifact_TI_Enrichment_By_Mock.py index 42d4758..21bdb41 100644 --- a/PLAYBOOKS/Artifact_TI_Enrichment_By_Mock.py +++ b/PLAYBOOKS/Artifact_TI_Enrichment_By_Mock.py @@ -2,7 +2,8 @@ import json import time from Lib.baseplaybook import BasePlaybook -from PLUGINS.SIRP.sirpapi import Artifact, PlaybookStatus +from PLUGINS.SIRP.sirpapi import Artifact +from PLUGINS.SIRP.sirptype import PlaybookJobStatus class Playbook(BasePlaybook): @@ -27,15 +28,15 @@ class Playbook(BasePlaybook): fields = [{"id": "enrichment", "value": json.dumps(ti_result)}] Artifact.update(self.param_source_rowid, fields) - self.update_playbook(PlaybookStatus.SUCCESS, "Threat intelligence enrichment completed.") + self.update_playbook_status(PlaybookJobStatus.SUCCESS, "Threat intelligence enrichment completed.") except Exception as e: self.logger.exception(e) - self.update_playbook(PlaybookStatus.FAILED, f"Error during TI enrichment: {e}") + self.update_playbook_status(PlaybookJobStatus.FAILED, f"Error during TI enrichment: {e}") return if __name__ == "__main__": params_debug = {'source_rowid': 'a966036e-b29e-4449-be48-23293bacac5d', 'source_worksheet': 'Artifact'} module = Playbook() - module._params = params_debug + # module._params = params_debug module.run() diff --git a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py index 945fd2a..e372ff3 100644 --- a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py +++ b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py @@ -10,7 +10,8 @@ from pydantic import BaseModel, Field, ConfigDict from Lib.baseplaybook import LanggraphPlaybook from Lib.llmapi import AgentState from PLUGINS.LLM.llmapi import LLMAPI -from PLUGINS.SIRP.sirpapi import Case, PlaybookStatus +from PLUGINS.SIRP.sirpapi import Case +from PLUGINS.SIRP.sirptype import PlaybookJobStatus class ConfidenceLevel(str, Enum): @@ -53,7 +54,7 @@ class Playbook(LanggraphPlaybook): def init(self): def preprocess_node(state: AgentState): """Preprocess data""" - case = Case.get_raw_data(self.param_source_rowid) + case = Case.get_ai_friendly_data(self.param_source_rowid) state.case = case return state @@ -105,7 +106,7 @@ class Playbook(LanggraphPlaybook): Case.update(self.param_source_rowid, case_field) self.send_notice("Case_L3_SOC_Analyst_Agent Finish", f"rowid:{self.param_source_rowid}") - self.update_playbook(PlaybookStatus.SUCCESS, "Get suggestion by ai agent completed.") + self.update_playbook_status(PlaybookJobStatus.SUCCESS, "Get suggestion by ai agent completed.") return state # Compile graph @@ -131,5 +132,5 @@ class Playbook(LanggraphPlaybook): if __name__ == "__main__": params_debug = {'source_rowid': '47da1d00-c9bf-4b5f-8ab8-8877ec292b98', 'source_worksheet': 'case'} module = Playbook() - module._params = params_debug + # module._params = params_debug module.run() diff --git a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent_With_Tools.py b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent_With_Tools.py index 86040e9..ceca8d9 100644 --- a/PLAYBOOKS/Case_L3_SOC_Analyst_Agent_With_Tools.py +++ b/PLAYBOOKS/Case_L3_SOC_Analyst_Agent_With_Tools.py @@ -11,7 +11,8 @@ from pydantic import BaseModel, Field, ConfigDict from AGENTS.agent_knowledge import AgentKnowledge from Lib.baseplaybook import LanggraphPlaybook from PLUGINS.LLM.llmapi import LLMAPI -from PLUGINS.SIRP.sirpapi import Case, PlaybookStatus +from PLUGINS.SIRP.sirpapi import Case +from PLUGINS.SIRP.sirptype import PlaybookJobStatus class AgentState(BaseModel): @@ -96,7 +97,7 @@ class Playbook(LanggraphPlaybook): def init(self): def preprocess_node(state: AgentState): - case = Case.get_raw_data(self.param_source_rowid) + case = Case.get_ai_friendly_data(self.param_source_rowid) content = f"Current Case Data (includes latest alert): {json.dumps(case)}" return {"case": case, "messages": [HumanMessage(content=content)]} @@ -154,7 +155,7 @@ class Playbook(LanggraphPlaybook): Case.update(self.param_source_rowid, case_field) self.send_notice("Case_L3_SOC_Analyst_Agent Finish", f"rowid:{self.param_source_rowid}") - self.update_playbook(PlaybookStatus.SUCCESS, "SOC analysis completed with potential tool-assisted enrichment.") + self.update_playbook_status(PlaybookJobStatus.SUCCESS, "SOC analysis completed with potential tool-assisted enrichment.") return {"analyze_result": result_data} workflow = StateGraph(AgentState) @@ -189,5 +190,5 @@ class Playbook(LanggraphPlaybook): if __name__ == "__main__": params_debug = {'source_rowid': 'f0189cf8-44af-4c46-90c7-988a159bb34c', 'source_worksheet': 'case'} module = Playbook() - module._params = params_debug + # module._params = params_debug module.run() diff --git a/PLAYBOOKS/Case_Threat_Hunting_Agent.py b/PLAYBOOKS/Case_Threat_Hunting_Agent.py index a88645d..5ce8dc3 100644 --- a/PLAYBOOKS/Case_Threat_Hunting_Agent.py +++ b/PLAYBOOKS/Case_Threat_Hunting_Agent.py @@ -16,7 +16,8 @@ from AGENTS.agent_ti import AgentTI from Lib.api import get_current_time_str from Lib.baseplaybook import LanggraphPlaybook from PLUGINS.LLM.llmapi import LLMAPI -from PLUGINS.SIRP.sirpapi import Case, PlaybookStatus +from PLUGINS.SIRP.sirpapi import Case +from PLUGINS.SIRP.sirptype import PlaybookJobStatus MAX_ITERATIONS = 3 PROMPT_LANG = None @@ -345,7 +346,7 @@ class Playbook(LanggraphPlaybook): self.logger.debug("Intent Node Invoked") # Get Case data - case = Case.get_raw_data(rowid=self.param_source_rowid) + case = Case.get_ai_friendly_data(rowid=self.param_source_rowid) # User intent user_intent = self.param_user_input @@ -566,7 +567,7 @@ class Playbook(LanggraphPlaybook): node_out = {"report": response.content} # update playbook status - self.update_playbook(PlaybookStatus.SUCCESS, "Threat Hunting Agent Finish.") + self.update_playbook_status(PlaybookJobStatus.SUCCESS, "Threat Hunting Agent Finish.") return node_out # --- Build the main graph --- @@ -617,5 +618,5 @@ if __name__ == "__main__": "playbook_rowid": "9fb4a3e1-6ae7-47b2-9b15-95264272dff5" } module = Playbook() - module._params = params_debug + # module._params = params_debug module.run() diff --git a/PLUGINS/Dify/Sentinel.yml b/PLUGINS/Dify/Sentinel.yml index c9fe968..12608e6 100644 --- a/PLUGINS/Dify/Sentinel.yml +++ b/PLUGINS/Dify/Sentinel.yml @@ -35,8 +35,8 @@ model_config: provider_id: sirp_new provider_name: SIRP provider_type: mcp - tool_label: get_case_by_case_id - tool_name: get_case_by_case_id + tool_label: get_case_by_rowid + tool_name: get_case_by_rowid tool_parameters: case_id: '' annotation_reply: diff --git a/PLUGINS/MCP/mcpserver.py b/PLUGINS/MCP/mcpserver.py index c416fb8..cc2ed55 100644 --- a/PLUGINS/MCP/mcpserver.py +++ b/PLUGINS/MCP/mcpserver.py @@ -14,7 +14,7 @@ if __name__ == "__main__": import django django.setup() - from Lib.llmfunc import function_call_debug, get_case_by_case_id + from Lib.llmfunc import function_call_debug, get_case_by_rowid # Define UUID file path uuid_file_path = os.path.join(BASE_DIR, "Docker", "mcp_uuid") @@ -40,6 +40,6 @@ if __name__ == "__main__": # add tools mcp.add_tool(function_call_debug) - mcp.add_tool(get_case_by_case_id) + mcp.add_tool(get_case_by_rowid) print(f"mcp server url: http://your_server_ip:{port}/{uuid_str}/sse") mcp.run(transport="sse") diff --git a/PLUGINS/SIRP/base_entity.py b/PLUGINS/SIRP/base_entity.py new file mode 100644 index 0000000..06091b9 --- /dev/null +++ b/PLUGINS/SIRP/base_entity.py @@ -0,0 +1,293 @@ +from abc import ABC +from typing import TypeVar, Generic, Type, List, Dict, Union, Any + +from pydantic import BaseModel + +from PLUGINS.SIRP.nocolyapi import WorksheetRow, Group, Condition, Operator +from PLUGINS.SIRP.sirptype import BaseSystemModel + + +def model_to_fields(model_instance: BaseModel) -> List[Dict[str, Any]]: + fields = [] + model_data = model_instance.model_dump(mode='json', exclude_unset=True) + for key, value in model_data.items(): + field_info = model_instance.model_fields.get(key) + field_item = { + 'id': key, + 'value': value + } + if field_info and field_info.json_schema_extra: + field_item.update(field_info.json_schema_extra) + fields.append(field_item) + return fields + + +# 定义泛型类型 +T = TypeVar('T', bound=BaseSystemModel) + + +class BaseWorksheetEntity(ABC, Generic[T]): + """通用工作表实体基类 - 支持泛型和关联加载""" + + WORKSHEET_ID: str + MODEL_CLASS: Type[T] + + @classmethod + def get( + cls, + rowid: str, + include_system_fields: bool = True, + lazy_load: bool = False + ) -> T: + """获取单条记录 + + Args: + rowid: 记录ID + include_system_fields: 是否包含系统字段 + lazy_load: 是否延迟加载关联数据（True时不加载关联） + + Returns: + 模型实例 + """ + result = WorksheetRow.get( + cls.WORKSHEET_ID, + rowid, + include_system_fields=include_system_fields + ) + model = cls.MODEL_CLASS(**result) + + if not lazy_load: + model = cls._load_relations(model, include_system_fields) + + return model + + @classmethod + def list( + cls, + model: Group, + include_system_fields: bool = True, + lazy_load: bool = False + ) -> List[T]: + """按过滤条件列表查询 + + Args: + model: 过滤条件Group对象 + include_system_fields: 是否包含系统字段 + lazy_load: 是否延迟加载关联数据（True时不加载关联） + + Returns: + 模型实例列表 + """ + filter_dict = model.model_dump() + result = WorksheetRow.list( + cls.WORKSHEET_ID, + filter_dict, + include_system_fields=include_system_fields + ) + + model_list = [] + for item in result: + model_obj = cls.MODEL_CLASS(**item) + if not lazy_load: + model_obj = cls._load_relations(model_obj, include_system_fields) + model_list.append(model_obj) + + return model_list + + @classmethod + def list_by_rowids( + cls, + rowids: Union[List[str], None], + include_system_fields: bool = True, + lazy_load: bool = False + ) -> Union[List[T], List[str], None]: + """按ID列表查询 + + Args: + rowids: 记录ID列表 + include_system_fields: 是否包含系统字段 + lazy_load: 是否延迟加载关联数据 + + Returns: + 模型实例列表或原始rowids列表 + """ + if rowids is not None and rowids != []: + filter_model = Group( + logic="AND", + children=[ + Condition( + field="rowid", + operator=Operator.IN, + value=rowids + ) + ] + ) + return cls.list(filter_model, include_system_fields=include_system_fields, lazy_load=lazy_load) + return rowids + + @classmethod + def create(cls, model: T) -> str: + """创建记录 + + Args: + model: 模型实例 + + Returns: + 新创建的记录ID + """ + + fields = model_to_fields(model) + rowid = WorksheetRow.create(cls.WORKSHEET_ID, fields) + return rowid + + @classmethod + def update(cls, model: T) -> str: + """更新记录 + + Args: + model: 模型实例（必须包含rowid） + + Returns: + 更新的记录ID + + Raises: + ValueError: 当rowid为None时 + """ + if model.rowid is None: + raise ValueError(f"{cls.__name__} rowid is None, cannot update.") + + fields = model_to_fields(model) + rowid = WorksheetRow.update(cls.WORKSHEET_ID, model.rowid, fields) + return rowid + + @classmethod + def update_or_create(cls, model: T) -> str: + """更新或创建记录 + + Args: + model: 模型实例 + + Returns: + 记录ID + """ + model = cls._prepare_for_save(model) + + fields = model_to_fields(model) + + if model.rowid is None: + rowid = WorksheetRow.create(cls.WORKSHEET_ID, fields) + else: + rowid = WorksheetRow.update(cls.WORKSHEET_ID, model.rowid, fields) + + return rowid + + @classmethod + def batch_update(cls, model_list: List[Union[T, str]]) -> Union[List[str], None]: + """批量更新 + + Args: + model_list: 模型实例或ID字符串的列表 + + Returns: + 更新后的记录ID列表 + + Raises: + TypeError: 当列表中包含不支持的类型时 + """ + if model_list is None: + return model_list + + rowids = [] + for model in model_list: + if isinstance(model, str): + rowids.append(model) # just link + elif isinstance(model, cls.MODEL_CLASS): + rowid = cls.update_or_create(model) + rowids.append(rowid) + else: + raise TypeError( + f"Unsupported {cls.__name__} data type: {type(model).__name__}. " + f"Expected str or {cls.MODEL_CLASS.__name__}" + ) + + return rowids + + @classmethod + def _load_relations(cls, model: T, include_system_fields: bool = True) -> T: + """加载关联数据（子类可覆盖） + + Args: + model: 模型实例 + include_system_fields: 是否包含系统字段 + + Returns: + 加载了关联数据的模型实例 + """ + return model + + @classmethod + def _prepare_for_save(cls, model: T) -> T: + """保存前准备（子类可覆盖） + + Args: + model: 模型实例 + + Returns: + 准备好的模型实例 + """ + return model + + +class BaseSimpleEntity(ABC): + """简化的工作表实体基类（不使用模型）""" + + WORKSHEET_ID: str + + @classmethod + def list(cls, filter_dict: dict) -> List[Dict]: + """列表查询 + + Args: + filter_dict: 过滤条件字典 + + Returns: + 字典列表 + """ + return WorksheetRow.list(cls.WORKSHEET_ID, filter_dict, include_system_fields=False) + + @classmethod + def get(cls, rowid: str) -> Dict: + """获取单条记录 + + Args: + rowid: 记录ID + + Returns: + 字典 + """ + return WorksheetRow.get(cls.WORKSHEET_ID, rowid, include_system_fields=False) + + @classmethod + def create(cls, fields: List[Dict]) -> str: + """创建记录 + + Args: + fields: 字段列表 + + Returns: + 新创建的记录ID + """ + return WorksheetRow.create(cls.WORKSHEET_ID, fields) + + @classmethod + def update(cls, rowid: str, fields: List[Dict]) -> str: + """更新记录 + + Args: + rowid: 记录ID + fields: 字段列表 + + Returns: + 更新的记录ID + """ + return WorksheetRow.update(cls.WORKSHEET_ID, rowid, fields) diff --git a/PLUGINS/SIRP/sirpapi.py b/PLUGINS/SIRP/sirpapi.py index 8fbcec9..1181a5d 100644 --- a/PLUGINS/SIRP/sirpapi.py +++ b/PLUGINS/SIRP/sirpapi.py @@ -1,535 +1,150 @@ from enum import StrEnum -from typing import List, Dict, Any, Literal, Union +from typing import List, Dict, Literal, Union import requests -from pydantic import BaseModel -from Lib.log import logger from PLUGINS.SIRP.CONFIG import SIRP_NOTICE_WEBHOOK -from PLUGINS.SIRP.nocolyapi import WorksheetRow, OptionSet, Group, Condition, Operator -from PLUGINS.SIRP.sirptype import EnrichmentModel, ArtifactModel, AlertModel, CaseModel, TicketModel +from PLUGINS.SIRP.base_entity import BaseWorksheetEntity, BaseSimpleEntity +from PLUGINS.SIRP.nocolyapi import OptionSet, Group, Condition, Operator +from PLUGINS.SIRP.sirptype import EnrichmentModel, ArtifactModel, AlertModel, CaseModel, TicketModel, MessageModel, PlaybookModel, PlaybookJobStatus, \ + AccountModel -def model_to_fields(model_instance: BaseModel) -> List[Dict[str, Any]]: - fields = [] - model_data = model_instance.model_dump(mode='json', exclude_unset=True) - for key, value in model_data.items(): - field_info = model_instance.model_fields.get(key) - field_item = { - 'id': key, - 'value': value - } - if field_info and field_info.json_schema_extra: - field_item.update(field_info.json_schema_extra) - fields.append(field_item) - return fields - - -class Enrichment(object): +class Enrichment(BaseWorksheetEntity[EnrichmentModel]): + """Enrichment 实体类""" WORKSHEET_ID = "enrichment" - - def __init__(self): - pass - - @staticmethod - def get(rowid, include_system_fields=True) -> EnrichmentModel: - result = WorksheetRow.get(Enrichment.WORKSHEET_ID, rowid, include_system_fields=include_system_fields) - model = EnrichmentModel(**result) - return model - - @staticmethod - def list(model: Group, include_system_fields=True) -> List[EnrichmentModel]: - filter = model.model_dump() - result = WorksheetRow.list(Enrichment.WORKSHEET_ID, filter, include_system_fields=include_system_fields) - model_list = [] - for one in result: - model_list.append(EnrichmentModel(**one)) - return model_list - - @staticmethod - def list_by_rowids(rowids: Union[List[str], None], include_system_fields=True) -> Union[List[EnrichmentModel], List[str], None]: - if rowids is not None and rowids != []: - filter_model = Group( - logic="AND", - children=[ - Condition( - field="rowid", - operator=Operator.IN, - value=rowids - ) - ] - ) - enrichment_list = Enrichment.list(filter_model, include_system_fields=include_system_fields) - return enrichment_list - else: - return rowids - - @staticmethod - def update(model: EnrichmentModel) -> str: - if model.rowid is not None: - fields = model_to_fields(model) - rowid = WorksheetRow.update(Enrichment.WORKSHEET_ID, model.rowid, fields) - else: - raise Exception("Enrichment rowid is None, cannot update.") - return rowid - - @staticmethod - def batch_update(model_list: List[Union[EnrichmentModel, str]]) -> Union[List[str], None]: - if model_list is not None: - rowids = [] - for model in model_list: - if isinstance(model, str): - rowids.append(model) # just link - continue - elif isinstance(model, EnrichmentModel): - rowid = Enrichment.update_or_create(model) # update or create record - rowids.append(rowid) - else: - raise Exception("Unsupported enrichment data type") - - return rowids - else: - return model_list - - @staticmethod - def create(model: EnrichmentModel) -> str: - fields = model_to_fields(model) - rowid = WorksheetRow.create(Enrichment.WORKSHEET_ID, fields) - return rowid - - @staticmethod - def update_or_create(model: EnrichmentModel) -> str: - fields = model_to_fields(model) - if model.rowid is None: - rowid = WorksheetRow.create(Enrichment.WORKSHEET_ID, fields) - else: - rowid = WorksheetRow.update(Enrichment.WORKSHEET_ID, model.rowid, fields) - return rowid + MODEL_CLASS = EnrichmentModel -class Ticket(object): +class Ticket(BaseWorksheetEntity[TicketModel]): + """Ticket 实体类""" WORKSHEET_ID = "ticket" - - def __init__(self): - pass - - @staticmethod - def get(rowid, include_system_fields=True) -> TicketModel: - result = WorksheetRow.get(Ticket.WORKSHEET_ID, rowid, include_system_fields=include_system_fields) - model = TicketModel(**result) - return model - - @staticmethod - def list(model: Group, include_system_fields=True) -> List[TicketModel]: - filter = model.model_dump() - result = WorksheetRow.list(Ticket.WORKSHEET_ID, filter, include_system_fields=include_system_fields) - model_list = [] - for one in result: - model_list.append(TicketModel(**one)) - return model_list - - @staticmethod - def list_by_rowids(rowids: Union[List[str], None], include_system_fields=True) -> Union[List[TicketModel], List[str], None]: - if rowids is not None and rowids != []: - filter_model = Group( - logic="AND", - children=[ - Condition( - field="rowid", - operator=Operator.IN, - value=rowids - ) - ] - ) - ticket_list = Ticket.list(filter_model, include_system_fields=include_system_fields) - return ticket_list - else: - return rowids - - @staticmethod - def update(model: TicketModel) -> str: - if model.rowid is not None: - fields = model_to_fields(model) - rowid = WorksheetRow.update(Ticket.WORKSHEET_ID, model.rowid, fields) - else: - raise Exception("Ticket rowid is None, cannot update.") - return rowid - - @staticmethod - def batch_update(model_list: List[Union[TicketModel, str]]) -> Union[List[str], None]: - if model_list is not None: - rowids = [] - for model in model_list: - if isinstance(model, str): - rowids.append(model) # just link - continue - elif isinstance(model, TicketModel): - rowid = Ticket.update_or_create(model) # update or create record - rowids.append(rowid) - else: - raise Exception("Unsupported ticket data type") - - return rowids - else: - return model_list - - @staticmethod - def create(model: TicketModel) -> str: - fields = model_to_fields(model) - rowid = WorksheetRow.create(Ticket.WORKSHEET_ID, fields) - return rowid - - @staticmethod - def update_or_create(model: TicketModel) -> str: - fields = model_to_fields(model) - if model.rowid is None: - rowid = WorksheetRow.create(Ticket.WORKSHEET_ID, fields) - else: - rowid = WorksheetRow.update(Ticket.WORKSHEET_ID, model.rowid, fields) - return rowid + MODEL_CLASS = TicketModel -class Artifact(object): +class Artifact(BaseWorksheetEntity[ArtifactModel]): + """Artifact 实体类 - 关联 Enrichment""" WORKSHEET_ID = "artifact" + MODEL_CLASS = ArtifactModel - def __init__(self): - pass - - @staticmethod - def get(rowid, include_system_fields=True) -> ArtifactModel: - result = WorksheetRow.get(Artifact.WORKSHEET_ID, rowid, include_system_fields=include_system_fields) - model = ArtifactModel(**result) - - # enrichments - model.enrichments = Enrichment.list_by_rowids(model.enrichments) - + @classmethod + def _load_relations(cls, model: ArtifactModel, include_system_fields: bool = True) -> ArtifactModel: + """加载关联的enrichments""" + model.enrichments = Enrichment.list_by_rowids( + model.enrichments, + include_system_fields=include_system_fields, + lazy_load=False + ) return model - @staticmethod - def list(model: Group, include_system_fields=True) -> List[ArtifactModel]: - filter = model.model_dump() - result = WorksheetRow.list(Artifact.WORKSHEET_ID, filter, include_system_fields=include_system_fields) - artifact_list = [] - for artifact_data in result: - artifact_model = ArtifactModel(**artifact_data) - - # enrichments - artifact_model.enrichments = Enrichment.list_by_rowids(artifact_model.enrichments) - - artifact_list.append(artifact_model) - return artifact_list - - @staticmethod - def list_by_rowids(rowids: Union[List[str], None], include_system_fields=True) -> Union[List[ArtifactModel], List[str], None]: - if rowids is not None and rowids != []: - filter_model = Group( - logic="AND", - children=[ - Condition( - field="rowid", - operator=Operator.IN, - value=rowids - ) - ] - ) - artifact_list = Artifact.list(filter_model, include_system_fields=include_system_fields) - return artifact_list - else: - return rowids - - @staticmethod - def update_or_create(model: ArtifactModel) -> str: - # enrichments + @classmethod + def _prepare_for_save(cls, model: ArtifactModel) -> ArtifactModel: + """保存前处理关联数据""" model.enrichments = Enrichment.batch_update(model.enrichments) - - fields = model_to_fields(model) - if model.rowid is None: - rowid = WorksheetRow.create(Artifact.WORKSHEET_ID, fields) - else: - rowid = WorksheetRow.update(Artifact.WORKSHEET_ID, model.rowid, fields) - return rowid - - @staticmethod - def batch_update(model_list: List[Union[ArtifactModel, str]]) -> Union[List[str], None]: - if model_list is not None: - rowids = [] - for model in model_list: - if isinstance(model, str): - rowids.append(model) # just link - continue - elif isinstance(model, ArtifactModel): - rowid = Artifact.update_or_create(model) # update or create record - rowids.append(rowid) - else: - raise Exception("Unsupported enrichment data type") - - return rowids - else: - return model_list + return model -class Alert(object): +class Alert(BaseWorksheetEntity[AlertModel]): + """Alert 实体类 - 关联 Artifact 和 Enrichment""" WORKSHEET_ID = "alert" + MODEL_CLASS = AlertModel - def __init__(self): - pass - - @staticmethod - def get(rowid, include_system_fields=True) -> AlertModel: - result = WorksheetRow.get(Alert.WORKSHEET_ID, rowid, include_system_fields=include_system_fields) - model = AlertModel(**result) - - # artifacts - model.artifacts = Artifact.list_by_rowids(model.artifacts) - - # enrichments - model.enrichments = Enrichment.list_by_rowids(model.enrichments) - + @classmethod + def _load_relations(cls, model: AlertModel, include_system_fields: bool = True) -> AlertModel: + """加载关联的artifacts和enrichments""" + model.artifacts = Artifact.list_by_rowids( + model.artifacts, + include_system_fields=include_system_fields, + lazy_load=False + ) + model.enrichments = Enrichment.list_by_rowids( + model.enrichments, + include_system_fields=include_system_fields, + lazy_load=False + ) return model - @staticmethod - def list(model: Group, include_system_fields=True) -> List[AlertModel]: - filter = model.model_dump() - result = WorksheetRow.list(Alert.WORKSHEET_ID, filter, include_system_fields=include_system_fields) - alert_list = [] - for alert_data in result: - alert_model = AlertModel(**alert_data) - - # artifacts - alert_model.artifacts = Artifact.list_by_rowids(alert_model.artifacts) - - # enrichments - alert_model.enrichments = Enrichment.list_by_rowids(alert_model.enrichments) - - alert_list.append(alert_model) - return alert_list - - @staticmethod - def list_by_rowids(rowids: Union[List[str], None], include_system_fields=True) -> Union[List[AlertModel], List[str], None]: - if rowids is not None and rowids != []: - filter_model = Group( - logic="AND", - children=[ - Condition( - field="rowid", - operator=Operator.IN, - value=rowids - ) - ] - ) - model_list = Alert.list(filter_model, include_system_fields=include_system_fields) - return model_list - else: - return rowids - - @staticmethod - def update_or_create(model: AlertModel) -> str: - - # artifacts + @classmethod + def _prepare_for_save(cls, model: AlertModel) -> AlertModel: + """保存前处理关联数据""" model.artifacts = Artifact.batch_update(model.artifacts) - - # enrichments model.enrichments = Enrichment.batch_update(model.enrichments) - - fields = model_to_fields(model) - if model.rowid is None: - rowid = WorksheetRow.create(Alert.WORKSHEET_ID, fields) - else: - rowid = WorksheetRow.update(Alert.WORKSHEET_ID, model.rowid, fields) - return rowid - - @staticmethod - def batch_update(model_list: List[Union[AlertModel, str]]) -> Union[List[str], None]: - if model_list is not None: - rowids = [] - for model in model_list: - if isinstance(model, str): - rowids.append(model) # just link - continue - elif isinstance(model, AlertModel): - rowid = Alert.update_or_create(model) # update or create record - rowids.append(rowid) - else: - raise Exception("Unsupported enrichment data type") - - return rowids - else: - return model_list - - -class Case(object): - WORKSHEET_ID = "case" - - def __init__(self): - pass - - @staticmethod - def get(rowid, include_system_fields=True) -> CaseModel: - result = WorksheetRow.get(Case.WORKSHEET_ID, rowid, include_system_fields=include_system_fields) - model = CaseModel(**result) - - # alerts - model.alerts = Alert.list_by_rowids(model.alerts) - - # enrichments - model.enrichments = Enrichment.list_by_rowids(model.enrichments) - - # tickets - model.tickets = Ticket.list_by_rowids(model.tickets) - return model - @staticmethod - def list(model: Group, include_system_fields=True) -> List[CaseModel]: - filter = model.model_dump() - result = WorksheetRow.list(Case.WORKSHEET_ID, filter, include_system_fields=include_system_fields) - case_list = [] - for case_data in result: - case_model = CaseModel(**case_data) - # alerts - case_model.alerts = Alert.list_by_rowids(case_model.alerts) +class Case(BaseWorksheetEntity[CaseModel]): + """Case 实体类 - 关联 Alert、Enrichment 和 Ticket""" + WORKSHEET_ID = "case" + MODEL_CLASS = CaseModel - # enrichments - case_model.enrichments = Enrichment.list_by_rowids(case_model.enrichments) + @classmethod + def _load_relations(cls, model: CaseModel, include_system_fields: bool = True) -> CaseModel: + """加载所有关联数据""" + model.alerts = Alert.list_by_rowids( + model.alerts, + include_system_fields=include_system_fields, + lazy_load=False + ) + model.enrichments = Enrichment.list_by_rowids( + model.enrichments, + include_system_fields=include_system_fields, + lazy_load=False + ) + model.tickets = Ticket.list_by_rowids( + model.tickets, + include_system_fields=include_system_fields, + lazy_load=False + ) + return model - # tickets - case_model.tickets = Ticket.list_by_rowids(case_model.tickets) - - case_list.append(case_model) - return case_list - - @staticmethod - def get_raw_data(rowid, include_system_fields=False) -> Dict: - """Get the raw data of the case and its associated alarms and work orders, and only keep the fields useful for LLM""" - case = WorksheetRow.get(Case.WORKSHEET_ID, rowid, include_system_fields=include_system_fields) - - useful_case_fields = ["rowid", "title", 'case_status', 'created_date', 'tags', 'severity', 'type', 'description', 'close_reason', 'alert_date', - 'case_id', - 'respond_time', 'note', 'acknowledged_date'] - - case_clean = {key: case[key] for key in useful_case_fields if key in case} - - # alert id - alerts = WorksheetRow.relations(Case.WORKSHEET_ID, rowid, "alerts", relation_worksheet_id=Alert.WORKSHEET_ID, - include_system_fields=include_system_fields) - alerts_clean = [] - for alert in alerts: - useful_alert_fields = ["rowid", 'severity', 'rule_id', 'rule_name', 'id'] - alert_clean = {key: alert[key] for key in useful_alert_fields if key in alert} - - artifacts = WorksheetRow.relations(Alert.WORKSHEET_ID, alert.get("rowid"), "artifacts", relation_worksheet_id=Artifact.WORKSHEET_ID, - include_system_fields=include_system_fields) - artifacts_clean = [] - for artifact in artifacts: - useful_artifact_fields = ["rowid", "type", "value", "enrichment", 'is_whitelisted', 'is_evidence'] - artifact_clean = {key: artifact[key] for key in useful_artifact_fields if key in artifact} - artifacts_clean.append(artifact_clean) - - alert_clean["artifacts"] = artifacts_clean - alerts_clean.append(alert_clean) - - case_clean["alerts"] = alerts_clean - return case_clean - - @staticmethod - def update_or_create(model: CaseModel) -> str: - - # alerts + @classmethod + def _prepare_for_save(cls, model: CaseModel) -> CaseModel: + """保存前处理关联数据""" model.alerts = Alert.batch_update(model.alerts) - - # enrichments model.enrichments = Enrichment.batch_update(model.enrichments) - - # tickets model.tickets = Ticket.batch_update(model.tickets) + return model - fields = model_to_fields(model) - if model.rowid is None: - rowid = WorksheetRow.create(Case.WORKSHEET_ID, fields) - else: - rowid = WorksheetRow.update(Case.WORKSHEET_ID, model.rowid, fields) - return rowid + @classmethod + def get_ai_friendly_data(cls, rowid: str) -> Dict: + """获取LLM友好的原始数据""" + model = cls.get(rowid, include_system_fields=True) - @staticmethod - def get_by_case_id(case_id: str): - filter = { - "type": "group", - "logic": "AND", - "children": [ - { - "type": "condition", - "field": "case_id", - "operator": "eq", - "value": case_id - }, - ] - } - rows = WorksheetRow.list(Case.WORKSHEET_ID, filter) - if rows: - if len(rows) > 1: - logger.warning(f"found multiple rows with case_id {case_id}") - return Case.get(rows[0]['rowid']) - else: - return None + # TODO : 这里可以根据需要添加清理的字段 + model.threat_hunting_report_ai = None + + data = model.model_dump(mode='json', exclude_unset=True, exclude_none=True, exclude_defaults=True) + return data -class PlaybookStatus(StrEnum): - SUCCESS = 'Success' - FAILED = 'Failed' - PENDING = 'Pending' - RUNNING = 'Running' +class Message(BaseWorksheetEntity[MessageModel]): + """Message 实体类""" + WORKSHEET_ID = "message" + MODEL_CLASS = MessageModel -class Playbook(object): +class Playbook(BaseWorksheetEntity[PlaybookModel]): + """PlaybookLoader 实体类""" WORKSHEET_ID = "playbook" + MODEL_CLASS = PlaybookModel - def __init__(self): - pass + @classmethod + def list_pending_playbooks(cls) -> List[PlaybookModel]: + """获取待处理的playbooks""" - @staticmethod - def list(filter: dict): - result = WorksheetRow.list(Playbook.WORKSHEET_ID, filter, include_system_fields=False) - return result - - @staticmethod - def create(fields: List): - row_id = WorksheetRow.create(Playbook.WORKSHEET_ID, fields) - return row_id - - @staticmethod - def update(row_id, fields: List): - row_id = WorksheetRow.update(Playbook.WORKSHEET_ID, row_id, fields) - return row_id - - @staticmethod - def update_status_and_remark(row_id, status: PlaybookStatus, remark): - fields = [ - {"id": "job_status", "value": status}, - {"id": "remark", "value": remark}, - ] - row_id = WorksheetRow.update(Playbook.WORKSHEET_ID, row_id, fields) - return row_id - - @staticmethod - def get_pending_playbooks(): - pending_option_value = OptionSet.get_option_key_by_name_and_value("playbook_status", "Pending") - artifact_filter = { - "type": "group", - "logic": "AND", - "children": [ - { - "type": "condition", - "field": "job_status", - "operator": "in", - "value": [ - pending_option_value - ] - } + # pending_option_value = OptionSet.get_option_key_by_name_and_value("playbook__status", PlaybookJobStatus.PENDING) + filter_model = Group( + logic="AND", + children=[ + Condition( + field="job_status", + operator=Operator.IN, + value=[PlaybookJobStatus.PENDING] + ) ] - } - result = Playbook.list(artifact_filter) - return result + ) + + return cls.list(filter_model, lazy_load=True) class KnowledgeAction(StrEnum): @@ -541,47 +156,27 @@ class KnowledgeAction(StrEnum): KnowledgeUsing = Literal[0, 1] -class Knowledge(object): +class Knowledge(BaseSimpleEntity): + """Knowledge 实体类""" WORKSHEET_ID = "knowledge" COLLECTION_NAME = "sirp_knowledge" - def __init__(self): - pass - - @staticmethod - def list(filter: dict): - result = WorksheetRow.list(Knowledge.WORKSHEET_ID, filter, include_system_fields=False) - return result - - @staticmethod - def create(fields: list): - row_id = WorksheetRow.create(Knowledge.WORKSHEET_ID, fields) - return row_id - - @staticmethod - def update(row_id, fields: list): - row_id = WorksheetRow.update(Knowledge.WORKSHEET_ID, row_id, fields) - return row_id - - @staticmethod - def update_action_and_using(row_id, action: KnowledgeAction, using: KnowledgeUsing): + @classmethod + def update_action_and_using(cls, row_id: str, action: KnowledgeAction, using: KnowledgeUsing) -> str: + """更新action和using""" fields = [ {"id": "action", "value": action}, {"id": "using", "value": using}, ] - row_id = WorksheetRow.update(Knowledge.WORKSHEET_ID, row_id, fields) - return row_id + return cls.update(row_id, fields) - @staticmethod - def get_undone_actions(): + @classmethod + def get_undone_actions(cls) -> List[Dict]: + """获取未完成的actions""" options = OptionSet.get("knowledge_action") + action_list = [opt.get("key") for opt in options if opt.get("value") != "Done"] - action_list = [] - for option in options: - if option.get("value") != "Done": - action_list.append(option.get("key")) - - artifact_filter = { + filter_dict = { "type": "group", "logic": "AND", "children": [ @@ -593,34 +188,19 @@ class Knowledge(object): } ] } - result = Knowledge.list(artifact_filter) - return result - - -class Message(object): - WORKSHEET_ID = "message" - - def __init__(self): - pass - - @staticmethod - def list(filter: dict): - result = WorksheetRow.list(Message.WORKSHEET_ID, filter, include_system_fields=False) - return result - - @staticmethod - def create(fields: list): - row_id = WorksheetRow.create(Message.WORKSHEET_ID, fields) - return row_id - - @staticmethod - def update(row_id, fields: list): - row_id = WorksheetRow.update(Message.WORKSHEET_ID, row_id, fields) - return row_id + return cls.list(filter_dict) class Notice(object): @staticmethod - def send(user, title, body=None): - result = requests.post(SIRP_NOTICE_WEBHOOK, json={"title": title, "body": body, "user": user}) - return result + def send(user: Union[AccountModel, List[AccountModel]], title, body=None): + if isinstance(user, AccountModel): + users = [user] + elif isinstance(user, list): + users = user + else: + raise ValueError("user 参数必须是 AccountModel 实例或 AccountModel 实例列表") + + 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/sirptype.py b/PLUGINS/SIRP/sirptype.py index 53ac8a8..c015c7d 100644 --- a/PLUGINS/SIRP/sirptype.py +++ b/PLUGINS/SIRP/sirptype.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import datetime +from enum import StrEnum from typing import List, Optional, Literal, Any, Union from pydantic import BaseModel, Field, field_validator, ConfigDict, field_serializer @@ -95,7 +96,7 @@ class BaseSystemModel(BaseModel): class MessageModel(BaseSystemModel): - playbook_rowid: str = Field(..., description="所属Playbook的唯一行ID") + playbook: List[Union[PlaybookModel, str]] = Field(..., description="所属Playbook的唯一行ID") node: Optional[str] = Field(default="", description="消息来源的节点名称或ID") content: Optional[str] = Field(default="", description="消息的文本内容") data: Optional[str] = Field(default="", description="消息的JSON格式内容，通常用于工具调用和返回") @@ -103,17 +104,24 @@ class MessageModel(BaseSystemModel): description="消息类型，用于区分不同角色的发言") +class PlaybookJobStatus(StrEnum): + SUCCESS = 'Success' + FAILED = 'Failed' + PENDING = 'Pending' + RUNNING = 'Running' + + class PlaybookModel(BaseSystemModel): source_worksheet: Optional[str] = Field(default="", description="Playbook触发源所在的工作表名称") source_rowid: Optional[str] = Field(default="", description="Playbook触发源的行ID，例如具体的告警ID或事件ID") job_id: Optional[str] = Field(default="", description="执行Playbook的后台任务ID") - job_status: Optional[Literal["Pending", "Running", "Success", "Failed", None]] = Field(default=None, description="Playbook执行任务的状态") + job_status: Optional[PlaybookJobStatus] = Field(default=None, description="Playbook执行任务的状态") remark: Optional[str] = Field(default="", description="关于Playbook执行的备注信息") type: Optional[Literal["CASE", "ALERT", "ARTIFACT", None]] = Field(default=None, description="Playbook关联的对象类型") name: Optional[str] = Field(default="", description="执行的Playbook的名称") user_input: Optional[str] = Field(default="", description="用户对Playbook的初始输入或后续指令") - user: Optional[AccountModel] = Field(default=None, description="发起或与Playbook交互的用户") + user: Optional[List[AccountModel]] = Field(default=None, description="发起Playbook的用户") # 关联表 messages: Optional[List[Union[MessageModel, str]]] = Field(default=None, description="Playbook执行过程中的所有消息记录，构成对话历史") @@ -214,7 +222,7 @@ class AlertModel(BaseSystemModel): product_category: Optional[Literal["DLP", "Email", "OT", "Proxy", "UEBA", "TI", "IAM", "EDR", "NDR", "Cloud", "Other", None]] = Field(default=None, description="产生告警的安全产品类别") - product_vender: Optional[str] = Field(default=None, description="安全产品的厂商") + product_vendor: Optional[str] = Field(default=None, description="安全产品的厂商") product_name: Optional[str] = Field(default=None, description="安全产品的名称") product_feature: Optional[str] = Field(default=None, description="产生告警的产品具体功能模块") diff --git a/PLUGINS/SIRP/test_sirptype.py b/PLUGINS/SIRP/test_sirptype.py index caf1274..fb734f2 100644 --- a/PLUGINS/SIRP/test_sirptype.py +++ b/PLUGINS/SIRP/test_sirptype.py @@ -2,7 +2,7 @@ import json from datetime import datetime, timedelta, timezone from PLUGINS.SIRP.nocolyapi import Group, Condition, Operator -from PLUGINS.SIRP.sirpapi import Enrichment, Artifact, Alert +from PLUGINS.SIRP.sirpapi import Enrichment, Artifact, Alert, Playbook from PLUGINS.SIRP.sirptype import CaseModel, AlertModel, ArtifactModel, EnrichmentModel, TicketModel now = datetime.now(timezone.utc) @@ -58,7 +58,7 @@ def generate_test_cases(): verdict=None, summary="", correlation_uid="CORR-PHISH-XYZ-123", - workbook="### Phishing Investigation Playbook\n1. Analyze headers (`done`)\n2. Detonate URL/Attachment (`done`)\n3. Identify recipients (`in-progress`)\n4. Purge emails from mailboxes\n5. Reset compromised user passwords\n", + workbook="### Phishing Investigation PlaybookLoader\n1. Analyze headers (`done`)\n2. Detonate URL/Attachment (`done`)\n3. Identify recipients (`in-progress`)\n4. Purge emails from mailboxes\n5. Reset compromised user passwords\n", analysis_rationale_ai="The email originates from an external, un-reputable domain and uses urgent language, a common phishing tactic. The URL leads to a non-standard login page with a self-signed certificate. The attachment hash matches known malware.", recommended_actions_ai="- Block sender domain 'evil-domain.com'\n- Reset passwords for all users who clicked the link\n- Scan all endpoints for the malware hash 'a1b2c3d4e5f6...'", attack_stage_ai="Initial Access, Execution", @@ -118,7 +118,7 @@ def generate_test_cases(): sub_technique="", mitigation="User Training, Email Filtering", product_category="Email", - product_vender="Microsoft", + product_vendor="Microsoft", product_name="Outlook", product_feature="Phishing Report Add-in", policy_name="", @@ -185,7 +185,7 @@ def generate_test_cases(): sub_technique="", mitigation="Email Attachment Sandboxing, Threat Intelligence Feed Integration", product_category="Email", - product_vender="SecureMail Inc.", + product_vendor="SecureMail Inc.", product_name="SecureMail Gateway", product_feature="AV-Scan-Module", policy_name="Inbound-Malware-Policy", @@ -240,7 +240,7 @@ def generate_test_cases(): verdict="True Positive", summary="Attacker compromised DC01 and moved to WS-FINANCE-05. Both hosts have been isolated and are pending reimaging. All domain admin credentials have been rotated.", correlation_uid="CORR-LAT-MOV-456", - workbook="### Lateral Movement Playbook\n1. Isolate source and destination (`done`)\n2. Dump memory from hosts (`done`)\n3. Analyze for persistence (`done`)\n4. Rotate credentials (`done`)", + workbook="### Lateral Movement PlaybookLoader\n1. Isolate source and destination (`done`)\n2. Dump memory from hosts (`done`)\n3. Analyze for persistence (`done`)\n4. Rotate credentials (`done`)", analysis_rationale_ai="PsExec execution from a domain controller to a workstation is highly anomalous. The initial compromise vector on DC01 appears to be related to a credential dumping alert moments before the lateral movement.", recommended_actions_ai="- Isolate both DC01 and WS-FINANCE-05 immediately.\n- Investigate DC01 for initial compromise.\n- Rotate all privileged credentials.", attack_stage_ai="Lateral Movement", @@ -294,7 +294,7 @@ def generate_test_cases(): sub_technique="", mitigation="Restrict Service Creation, Network Segmentation", product_category="EDR", - product_vender="CrowdStrike", + product_vendor="CrowdStrike", product_name="Falcon", product_feature="Behavioral-Detection-Engine", policy_name="Default Workstation Policy", @@ -359,7 +359,7 @@ def generate_test_cases(): sub_technique="", mitigation="Credential Guard, LSA Protection", product_category="EDR", - product_vender="CrowdStrike", + product_vendor="CrowdStrike", product_name="Falcon", product_feature="Credential-Theft-Protection", policy_name="Domain Controller Policy", @@ -413,7 +413,7 @@ def generate_test_cases(): verdict="Suspicious", summary="", correlation_uid="CORR-DNS-TUN-789", - workbook="### DNS Tunneling Playbook\n1. Analyze query patterns (TXT/NULL record types, query length)\n2. Check domain reputation\n3. Perform packet capture on host\n4. Compare against baseline DNS traffic", + workbook="### DNS Tunneling PlaybookLoader\n1. Analyze query patterns (TXT/NULL record types, query length)\n2. Check domain reputation\n3. Perform packet capture on host\n4. Compare against baseline DNS traffic", analysis_rationale_ai="The high volume of TXT queries to a single, non-business related domain is a strong indicator of DNS tunneling. The query payloads appear to be encoded.", recommended_actions_ai="- Place the host in a sinkhole network to observe C2 traffic safely.\n- Do not block immediately to gather more intelligence on the attacker's infrastructure.", attack_stage_ai="Command and Control", @@ -459,7 +459,7 @@ def generate_test_cases(): sub_technique="", mitigation="DNS Sinkholing, Egress Traffic Filtering", product_category="NDR", - product_vender="Vectra", + product_vendor="Vectra", product_name="Cognito", product_feature="DNS-Analytics", policy_name="", @@ -524,7 +524,7 @@ def generate_test_cases(): sub_technique="", mitigation="Egress DNS Filtering", product_category="Cloud", - product_vender="Palo Alto", + product_vendor="Palo Alto", product_name="PA-Series Firewall", product_feature="DNS-Security", policy_name="Default-DNS-Allow", @@ -604,15 +604,10 @@ def test_enrichment(): Enrichment.list(filter_model) -if __name__ == "__main__": - import os - import django +def test_alert(): + alert = Alert.get("ae83212e-5064-42dd-9e3f-f95b0aeded2d") - os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") - django.setup() - rowid = Alert.get("0949a2df-7592-44f6-ac29-73994152aaa6") - - rowid = Artifact.get("0e4527f9-a0b9-4d71-a805-95a7d8d3267e") + artifact = Artifact.get("0e4527f9-a0b9-4d71-a805-95a7d8d3267e") artifact_model = ArtifactModel( rowid="0e4527f9-a0b9-4d71-a805-95a7d8d3267e", @@ -650,3 +645,13 @@ if __name__ == "__main__": ) rowid = Artifact.update_or_create(artifact_model) print(rowid) + + +if __name__ == "__main__": + import os + import django + + os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings") + django.setup() + models = Playbook.list_pending_playbooks() + print(models) \ No newline at end of file diff --git a/README.md b/README.md index 58e74a8..e9904b7 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ ASP processes security alerts and incidents through a simplified multi-stage pro Agents), enrich data, and determine outcomes. 5. **SIRP Platform**: The output of the modules (now formatted into standardized security records) is sent to the **SIRP ** platform, where cases, alerts, and artifacts are created or updated. -6. **Playbook Engine**: Analysts can trigger **playbooks** from the SIRP user interface against cases, alerts, or +6. **PlaybookLoader Engine**: Analysts can trigger **playbooks** from the SIRP user interface against cases, alerts, or artifacts to perform further automated actions, such as threat intelligence enrichment or remediation. ![img_1.webp](Docker/IMG/img_20.png)