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

rename

Browse Source
This commit is contained in:
rootkit
2026-01-27 02:57:23 +08:00
parent 65f22329bc
commit 517f86f057
11 changed files with 138 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Lib/basemodule.py
+2 -2
View File
@@ -4,7 +4,7 @@ from langgraph.graph.state import CompiledStateGraph
from Lib.baseapi import BaseAPI
from Lib.configs import REDIS_CONSUMER_GROUP
from Lib.llmapi import AgentState
from Lib.llmapi import BaseAgentState
from PLUGINS.Redis.redis_stream_api import RedisStreamAPI
@@ -43,7 +43,7 @@ class LanggraphModule(BaseModule):
config = RunnableConfig()
config["configurable"] = {"thread_id": self.module_name}
if self.agent_state is None:
self.agent_state = AgentState()
self.agent_state = BaseAgentState()
for event in self.graph.stream(self.agent_state, config, stream_mode="values"):
self.logger.debug(event)
self.logger.debug(f"{self.module_name} finished processing.")
Lib/baseplaybook.py
+2 -2
View File
@@ -13,7 +13,7 @@ from langgraph.graph.state import CompiledStateGraph
from pydantic import BaseModel
from Lib.baseapi import BaseAPI
from Lib.llmapi import AgentState
from Lib.llmapi import BaseAgentState
from Lib.log import logger
from PLUGINS.SIRP.sirpapi import Message
from PLUGINS.SIRP.sirpapi import Playbook, Notice
@@ -116,7 +116,7 @@ class LanggraphPlaybook(BasePlaybook):
config = RunnableConfig()
config["configurable"] = {"thread_id": self.module_name}
if self.agent_state is None:
self.agent_state = AgentState()
self.agent_state = BaseAgentState()
for event in self.graph.stream(self.agent_state, config, stream_mode="values"):
self.logger.debug(event)
Lib/llmapi.py
+1 -1
View File
@@ -8,7 +8,7 @@ from Lib.log import logger
from PLUGINS.SIRP.sirpmodel import CaseModel, AlertModel, ArtifactModel
class AgentState(BaseModel):
class BaseAgentState(BaseModel):
messages: Annotated[List[Any], add_messages] = []
case: CaseModel = None
alert: AlertModel = None
MODULES/ES-Rule-21-Phishing-User-Report-Mail.py
+16 -13
View File
@@ -10,7 +10,7 @@ from pydantic import BaseModel, Field
from Lib.api import string_to_string_time, get_current_time_str
from Lib.basemodule import LanggraphModule
from Lib.llmapi import AgentState
from Lib.llmapi import BaseAgentState
from PLUGINS.LLM.llmapi import LLMAPI
from PLUGINS.SIRP.sirpapi import Alert
from PLUGINS.SIRP.sirpmodel import AlertModel, ArtifactModel, ArtifactType, ArtifactRole, Severity, AlertStatus, AlertAnalyticType, ProductCategory, Confidence
@@ -18,11 +18,15 @@ from PLUGINS.SIRP.sirpmodel import AlertModel, ArtifactModel, ArtifactType, Arti
class AnalyzeResult(BaseModel):
"""Structure for extracting phishing analysis result"""
is_phishing: bool = Field(description="Whether it is a phishing email, True or False")
confidence: Confidence = Field(description="Confidence level assessment")
is_phishing: bool = Field(description="Whether it is a phishing email, True or False", default=False)
confidence: Confidence = Field(description="Confidence level assessment", default=Confidence.UNKNOWN)
reasoning: Optional[Union[str, Dict[str, Any]]] = Field(description="Reasoning process", default=None)
class AgentState(BaseAgentState):
analyze_result: AnalyzeResult = None
class Module(LanggraphModule):
THREAD_NUM = 2
@@ -99,7 +103,6 @@ class Module(LanggraphModule):
))
alert_model.artifacts = artifacts
return {"alert": alert_model}
def alert_analyze_node(state: AgentState):
@@ -228,7 +231,7 @@ class Module(LanggraphModule):
]
few_shot_examples = [
HumanMessage(content=legitimate_alert.model_dump_json()),
HumanMessage(content=legitimate_alert.model_dump_json_for_ai()),
AIMessage(
content=str(AnalyzeResult(
is_phishing=False,
@@ -236,7 +239,7 @@ class Module(LanggraphModule):
reasoning="The email is from a known colleague within the same organization, discussing a legitimate project. SPF and authentication checks pass."
).model_dump())
),
HumanMessage(content=phishing_alert.model_dump_json()),
HumanMessage(content=phishing_alert.model_dump_json_for_ai()),
AIMessage(
content=str(AnalyzeResult(
is_phishing=True,
@@ -250,7 +253,7 @@ class Module(LanggraphModule):
messages = [
system_message,
*few_shot_examples,
HumanMessage(content=json.dumps(alert.model_dump_for_ai())),
HumanMessage(content=json.dumps(alert.model_dump_json_for_ai())),
]
llm_api = LLMAPI()
@@ -258,15 +261,15 @@ class Module(LanggraphModule):
llm_structured = llm.with_structured_output(AnalyzeResult)
response: AnalyzeResult = llm_structured.invoke(messages)
state.analyze_result = response.model_dump()
state.analyze_result = response
return state
def alert_output_node(state: AgentState):
"""
Save analysis result to AlertModel and persist using SIRP API.
"""
alert_model: AlertModel = state.alert_model
analyze_result: AnalyzeResult = AnalyzeResult(**state.analyze_result)
alert_model: AlertModel = state.alert
analyze_result: AnalyzeResult = state.analyze_result
if analyze_result.is_phishing and analyze_result.confidence in [Confidence.HIGH, Confidence.MEDIUM]:
alert_model.severity = Severity.HIGH
@@ -276,10 +279,10 @@ class Module(LanggraphModule):
alert_model.summary_ai = str(analyze_result.reasoning)
alert_model.confidence = analyze_result.confidence
tags = list(alert_model.tags) if alert_model.tags else []
labels = list(alert_model.labels) if alert_model.labels else []
if analyze_result.is_phishing:
tags.append("confirmed-phishing")
alert_model.tags = tags
labels.append("confirmed-phishing")
alert_model.labels = labels
alert_model.uid = f"phishing-{get_current_time_str()}"
MODULES/NDR-Rule-05-Suspect-C2-Communication.py
+5 -1
View File
@@ -10,7 +10,7 @@ from pydantic import BaseModel, Field, ConfigDict
from Lib.api import get_current_time_str
from Lib.basemodule import LanggraphModule
from Lib.llmapi import AgentState
from Lib.llmapi import BaseAgentState
from PLUGINS.LLM.llmapi import LLMAPI
from PLUGINS.SIRP.grouprule import GroupRule
from PLUGINS.SIRP.sirpapi import create_alert_with_group_rule, InputAlert, Case
@@ -45,6 +45,10 @@ class AnalyzeResult(BaseModel):
recommended_actions: Optional[Union[str, Dict[str, Any]]] = Field(description="e.g., 'Isolate host 10.1.1.5'", default=None)
class AgentState(BaseAgentState):
analyze_result: AnalyzeResult = None
class Module(LanggraphModule):
THREAD_NUM = 2
PLAYBOOKS/Case_L3_SOC_Analyst_Agent.py
+8 -6
View File
@@ -1,4 +1,3 @@
import json
from typing import Any
from langchain_core.messages import HumanMessage
@@ -7,7 +6,7 @@ from langgraph.graph.state import CompiledStateGraph
from pydantic import BaseModel, Field, ConfigDict
from Lib.baseplaybook import LanggraphPlaybook
from Lib.llmapi import AgentState
from Lib.llmapi import BaseAgentState
from PLUGINS.LLM.llmapi import LLMAPI
from PLUGINS.SIRP.sirpapi import Case
from PLUGINS.SIRP.sirpmodel import PlaybookJobStatus, PlaybookModel, CaseModel
@@ -27,6 +26,10 @@ class AnalyzeResult(BaseModel):
recommended_actions: str | dict[str, Any] | None = Field(description="e.g., 'Isolate host 10.1.1.5'", default=None)
class AgentState(BaseAgentState):
analyze_result: AnalyzeResult = None
class Playbook(LanggraphPlaybook):
TYPE = "CASE"
NAME = "L3 SOC Analyst Agent"
@@ -63,18 +66,17 @@ class Playbook(LanggraphPlaybook):
messages = [
system_message,
*few_shot_examples,
HumanMessage(content=json.dumps(state.case.model_dump_for_ai()))
HumanMessage(content=state.case.model_dump_json_for_ai())
]
llm = llm.with_structured_output(AnalyzeResult)
response: AnalyzeResult = llm.invoke(messages)
analyze_result = response.model_dump()
self.logger.debug(f"Analyze result: {response.model_dump()}")
return {"analyze_result": analyze_result}
return {"analyze_result": response}
def output_node(state: AgentState):
"""Process analysis results"""
analyze_result: AnalyzeResult = AnalyzeResult(**state.analyze_result)
analyze_result: AnalyzeResult = state.analyze_result
case_new = CaseModel(rowid=self.param_source_rowid,
severity_ai=analyze_result.new_severity,
PLAYBOOKS/Case_L3_SOC_Analyst_Agent_With_Tools.py
+1 -2
View File
@@ -1,4 +1,3 @@
import json
from typing import Annotated, Any, Dict, List
from langchain_core.messages import HumanMessage
@@ -103,7 +102,7 @@ class Playbook(LanggraphPlaybook):
def init(self):
def preprocess_node(state: AgentState):
case = Case.get(self.param_source_rowid)
content = f"Current Case Data (includes latest alert): {json.dumps(case.model_dump_for_ai())}"
content = f"Current Case Data (includes latest alert): {case.model_dump_json_for_ai()}"
return {"case": case, "messages": [HumanMessage(content=content)]}
def analyze_node(state: AgentState):
PLAYBOOKS/Case_Threat_Hunting_Agent.py
+2 -1
View File
@@ -355,7 +355,8 @@ 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_for_ai(), user_intent=user_intent)
human_message = self.load_human_prompt_template("Intent_Human", lang=PROMPT_LANG).format(case=case.model_dump_json_for_ai(),
user_intent=user_intent)
# Construct few-shot examples
few_shot_examples = [
PLUGINS/SIRP/sirpapi.py
+1 -7
View File
@@ -1,4 +1,4 @@
from typing import List, Dict, Union
from typing import List, Union
import requests
@@ -115,12 +115,6 @@ class Case(BaseWorksheetEntity[CaseModel]):
model.tickets = Ticket.batch_update_or_create(model.tickets)
return model
@classmethod
def get_ai_friendly_data(cls, rowid: str) -> Dict:
"""获取LLM友好的原始数据"""
model: CaseModel = cls.get(rowid, include_system_fields=True)
return model.model_dump_for_ai()
class Message(BaseWorksheetEntity[MessageModel]):
"""Message 实体类"""
PLUGINS/SIRP/sirpbase.py
+2 -1
View File
@@ -158,7 +158,8 @@ class BaseWorksheetEntity(ABC, Generic[T]):
Returns:
新创建的记录ID
"""
model = cls._prepare_for_save(model)
fields = model_to_fields(model)
rowid = WorksheetRow.create(cls.WORKSHEET_ID, fields)
return rowid
PLUGINS/SIRP/sirpmodel.py
+98 -7
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import json
from datetime import datetime
from enum import StrEnum
from typing import List, Optional, Any, Union, ClassVar
@@ -358,28 +359,118 @@ class BaseSystemModel(BaseModel):
return v.strftime("%Y-%m-%dT%H:%M:%SZ")
return v
def model_dump_for_ai(self) -> dict[str, Any]:
def model_dump_json_for_ai(
self,
*,
exclude_none: bool = True,
exclude_unset: bool = True,
exclude_default: bool = True,
) -> str:
"""
递归序列化模型为 AI 友好的 JSON 字符串格式。
在序列化前处理嵌套对象，确保每层都能应用自己的 ai_exclude_fields。
Args:
exclude_none: 是否排除值为None的字段
exclude_unset: 是否排除未被显式设置的字段
exclude_default: 是否排除值为默认值的字段
"""
dict_representation = self.model_dump_for_ai(
exclude_none=exclude_none,
exclude_unset=exclude_unset,
exclude_default=exclude_default
)
return json.dumps(dict_representation, ensure_ascii=False)
def model_dump_for_ai(
self,
*,
exclude_none: bool = True,
exclude_unset: bool = True,
exclude_default: bool = True,
) -> dict[str, Any]:
"""
递归序列化模型为 AI 友好的字典格式。
在序列化前处理嵌套对象，确保每层都能应用自己的 ai_exclude_fields。
Args:
exclude_none: 是否排除值为None的字段
exclude_unset: 是否排除未被显式设置的字段
exclude_default: 是否排除值为默认值的字段
"""
result = {}
fields_set = self.__pydantic_fields_set__ if hasattr(self, '__pydantic_fields_set__') else set()
for field_name, field_value in self.__dict__.items():
if field_name in self.ai_exclude_fields:
continue
result[field_name] = self._process_value_before_dump(field_value)
if self._should_exclude_field(
field_name, field_value, fields_set, exclude_none, exclude_unset, exclude_default
):
continue
result[field_name] = self._process_value_before_dump(
field_value, exclude_none, exclude_unset, exclude_default
)
return result
def _process_value_before_dump(self, value: Any) -> Any:
def _should_exclude_field(
self,
field_name: str,
field_value: Any,
fields_set: set,
exclude_none: bool,
exclude_unset: bool,
exclude_default: bool
) -> bool:
"""
在序列化前处理值，支持递归调用嵌套模型的 model_dump_for_ai()。
判断是否应该排除该字段。
"""
if exclude_none and field_value is None:
return True
if exclude_unset and field_name not in fields_set:
return True
if exclude_default:
model_fields = self.model_fields
if field_name in model_fields:
field_info = model_fields[field_name]
default_value = field_info.default
if default_value is not None and field_value == default_value:
return True
return False
def _process_value_before_dump(
self,
value: Any,
exclude_none: bool = False,
exclude_unset: bool = False,
exclude_default: bool = False
) -> Any:
"""
在序列化前处理值，支持递归调用嵌套模型的 model_dump_json_for_ai()。
"""
if isinstance(value, BaseSystemModel):
return value.model_dump_for_ai()
return value.model_dump_for_ai(
exclude_none=exclude_none,
exclude_unset=exclude_unset,
exclude_default=exclude_default
)
elif isinstance(value, list):
return [self._process_value_before_dump(item) for item in value]
return [
self._process_value_before_dump(item, exclude_none, exclude_unset, exclude_default)
for item in value
]
elif isinstance(value, dict):
return {k: self._process_value_before_dump(v) for k, v in value.items()}
return {
k: self._process_value_before_dump(v, exclude_none, exclude_unset, exclude_default)
for k, v in value.items()
}
else:
return self._serialize_value(value)