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

update the performance of redis

Browse Source
This commit is contained in:
funnywolf
2025-12-26 21:34:22 +08:00
parent 6e37a3c17c
commit 08cd582e44
11 changed files with 111 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ASP/settings.py
+8
View File
@@ -96,6 +96,14 @@ LOGGING = {
'level': 'INFO',
'propagate': True
},
'asp': {
'handlers': [
'console',
'file',
],
'level': 'DEBUG',
'propagate': True,
},
}
}
Lib/basemodule.py
+1
View File
@@ -46,6 +46,7 @@ class LanggraphModule(BaseModule):
self.agent_state = AgentState()
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.")
def run(self):
self.run_graph()
Lib/log.py
+1 -1
View File
@@ -1,3 +1,3 @@
import logging
logger = logging.getLogger("django")
logger = logging.getLogger("asp")
Lib/montior.py
+1 -7
View File
@@ -15,10 +15,9 @@ from Lib.engine import Engine
from Lib.log import logger
from Lib.xcache import Xcache
from PLUGINS.Embeddings.embeddings_qdrant import embedding_api_singleton_qdrant
from PLUGINS.Redis.CONFIG import REDIS_STREAM_STORE_DAYS
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
from PLUGINS.Mem0.CONFIG import USE as MEM_ZERO_USE
if MEM_ZERO_USE:
from PLUGINS.Mem0.mem_zero import mem_zero_singleton
@@ -101,10 +100,8 @@ class MainMonitor(object):
# self.MainScheduler.start()
delay_time = 3
delay_time_clean_thread = 60 * 60
# Start background tasks
self.start_background_task(self.subscribe_clean_thread, "subscribe_clean_thread", delay_time_clean_thread)
self.start_background_task(self.subscribe_pending_playbook, "subscribe_pending_playbook", delay_time)
self.start_background_task(self.subscribe_knowledge_action, "subscribe_knowledge_action", delay_time)
@@ -112,9 +109,6 @@ class MainMonitor(object):
self.engine.start()
logger.info("Background services started.")
def subscribe_clean_thread(self):
self.redis_stream_api.clean_redis_stream(max_age_days=REDIS_STREAM_STORE_DAYS)
@staticmethod
def subscribe_pending_playbook():
records = SIRPPlaybook.get_pending_playbooks()
MODULES/EDR-Rule-11-Suspicious-Process-Spawned-by-Office.py
+6 -1
View File
@@ -4,8 +4,10 @@ from datetime import datetime
from typing import Optional, Union, Dict, Any
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.graph import END
from langgraph.graph import StateGraph
from langgraph.graph.state import CompiledStateGraph
from langgraph.types import Command
from pydantic import BaseModel, Field
from Lib.api import string_to_string_time, get_current_time_str
@@ -36,7 +38,10 @@ class Module(LanggraphModule):
# Get raw alert from stream
alert = self.read_message()
if alert is None:
return
return Command(
update={"alert": {}},
goto=END
)
# for splunk webhook
# alert = json.loads(alert["_raw"])
MODULES/ES-Rule-21-Phishing_user_report_mail.py
+6 -2
View File
@@ -4,8 +4,9 @@ from datetime import datetime
from typing import Optional, Union, Dict, Any
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.graph import StateGraph
from langgraph.graph import StateGraph, END
from langgraph.graph.state import CompiledStateGraph
from langgraph.types import Command
from pydantic import BaseModel, Field
from Lib.api import string_to_string_time, get_current_time_str
@@ -41,7 +42,10 @@ class Module(LanggraphModule):
# Read raw alert from Redis stream
alert = self.read_message()
if alert is None:
return
return Command(
update={"alert": {}},
goto=END
)
# Example: For Splunk webhooks, uncomment the following line
# alert = json.loads(alert["_raw"])
MODULES/ES-Rule-22-Phishing_user_report_mail.py
+2 -2
View File
@@ -28,7 +28,7 @@ class Module(BaseModule):
# Get raw alert from stream
alert = self.read_message()
if alert is None:
return
return None
# Parse data, this is an example of handling JSON data sent by Splunk Webhook
# alert = json.loads(alert["_raw"])
@@ -40,7 +40,7 @@ class Module(BaseModule):
alert["headers"] = headers
self.agent_state.alert_raw = alert
return
return None
def alert_analyze_node(self):
client = Dify()
MODULES/NDR-Rule-05-Suspect-C2-Communication.py
+6 -2
View File
@@ -3,8 +3,9 @@ from enum import Enum
from typing import Optional, Union, Dict, Any
from langchain_core.messages import HumanMessage
from langgraph.graph import StateGraph
from langgraph.graph import StateGraph, END
from langgraph.graph.state import CompiledStateGraph
from langgraph.types import Command
from pydantic import BaseModel, Field, ConfigDict
from Lib.api import get_current_time_str
@@ -57,7 +58,10 @@ class Module(LanggraphModule):
# Get raw alert from stream
alert = self.read_message()
if alert is None:
return
return Command(
update={"alert": {}},
goto=END
)
state.alert = alert
artifact: list = alert.get("artifact")
PLUGINS/Redis/CONFIG.example.py
+2 -1
View File
@@ -4,4 +4,5 @@
# 根据你的环境修改Redis的连接地址
# 你通过 http://192.168.241.128:8001 和 default/redis-stack-password-for-agentic-soc-platform 登录 Redis Insight
REDIS_URL = "redis://:redis-stack-password-for-agentic-soc-platform@192.168.241.128:6379/"
REDIS_STREAM_STORE_DAYS = 7 # 消息在Redis Stream中保存的天数
REDIS_STREAM_MAX_LENGTH = 10000 # Redis Stream的最大长度
REDIS_MAX_CONNECTIONS = 10 # Redis连接池大小,建议为模块数量的2倍
PLUGINS/Redis/redis_client.py
+17 -14
View File
@@ -1,22 +1,25 @@
import redis
from Lib.log import logger
from PLUGINS.Redis.CONFIG import REDIS_URL
from PLUGINS.Redis.CONFIG import REDIS_URL, REDIS_MAX_CONNECTIONS
class RedisClient(object):
# 类变量，用于存储全局唯一的连接池
_stream_pool = None
def __init__(self):
pass
@classmethod
def get_stream_connection(cls):
"""获取单例连接池中的连接"""
if cls._stream_pool is None:
logger.info("Initializing Global Redis Stream Connection Pool...")
cls._stream_pool = redis.ConnectionPool.from_url(
f"{REDIS_URL}0",
decode_responses=True,
max_connections=REDIS_MAX_CONNECTIONS, # 限制最大连接数
health_check_interval=30 # 每30秒自动检测连接状态
)
@staticmethod
def get_stream_connection():
"""用于订阅类操作,无需使用连接池"""
redis_client = redis.Redis.from_url(f"{REDIS_URL}0", decode_responses=True)
# 测试连接
try:
redis_client.ping()
return redis_client
except redis.ConnectionError as e:
logger.exception(e)
raise
# 使用同一个池创建 client 实例
redis_client = redis.Redis(connection_pool=cls._stream_pool)
return redis_client
PLUGINS/Redis/redis_stream_api.py
+61 -105
View File
@@ -1,31 +1,47 @@
import datetime
import json
import time
from typing import Dict, Any, Optional, List
import redis
from Lib.configs import REDIS_CONSUMER_GROUP, REDIS_CONSUMER_NAME
from Lib.log import logger
from PLUGINS.Redis.CONFIG import REDIS_STREAM_MAX_LENGTH
from PLUGINS.Redis.redis_client import RedisClient
class RedisStreamAPI:
class RedisStreamAPI(object):
"""
Redis Stream API封装类,提供消息发送和读取功能
"""
_instance = None
def __new__(cls, *args, **kwargs):
"""实现单例模式"""
if not cls._instance:
cls._instance = super(RedisStreamAPI, cls).__new__(cls)
cls._instance._initialized = False
return cls._instance
def __init__(self):
"""初始化RedisStreamAPI类"""
if self._initialized:
return
self.redis_client = RedisClient.get_stream_connection()
def send_message(self, stream_key: str, message: Dict[str, Any]) -> Optional[str]:
self._checked_groups = set()
self._initialized = True
logger.info("RedisStreamAPI init finished.")
def send_message(self, stream_key: str, message: Dict[str, Any], maxlen: int = REDIS_STREAM_MAX_LENGTH) -> Optional[str]:
"""
发送消息到指定stream
Args:
stream_key (str): Redis stream的key名称
message (Dict[str, Any]): 要发送的消息内容
maxlen (int): stream最大长度,超过则删除最旧的消息,默认10000
Returns:
Optional[str]: 发送成功返回消息ID,失败返回None
"""
@@ -34,7 +50,9 @@ class RedisStreamAPI:
# 发送消息到stream
message_id = self.redis_client.xadd(
stream_key,
{"data": data}
{"data": data},
maxlen=maxlen,
approximate=True
)
return message_id
@@ -43,62 +61,50 @@ class RedisStreamAPI:
return None
def read_message(self, stream_key: str, consumer_group: str = None,
consumer_name: str = None, timeout: int = 0, noack: bool = False) -> Optional[Dict[str, Any]]:
consumer_name: str = None, timeout: int = 5000) -> Optional[Dict[str, Any]]:
"""
从指定stream读取一条消息
Args:
noack:
stream_key (str): Redis stream的key名称
consumer_group (str): 消费者组名称,如果为None则使用默认配置
consumer_name (str): 消费者名称,如果为None则使用默认配置
timeout (int): 读取超时时间(毫秒),默认5000毫秒
Returns:
Optional[Dict[str, Any]]: 读取到的消息,如果没有消息或出错则返回None
timeout (int): 阻塞等待时间,单位毫秒,默认5000ms
"""
try:
if consumer_group is None:
consumer_group = REDIS_CONSUMER_GROUP
if consumer_name is None:
consumer_name = REDIS_CONSUMER_NAME
if consumer_group is None:
consumer_group = REDIS_CONSUMER_GROUP
if consumer_name is None:
consumer_name = REDIS_CONSUMER_NAME
# 确保消费者组存在
flag = self._ensure_consumer_group(stream_key, consumer_group)
if not flag:
logger.error(f"无法确保消费者组 {consumer_group} 存在.")
return None
# 从消费者组读取消息
messages = self.redis_client.xreadgroup(
consumer_group,
consumer_name,
{stream_key: '>'}, # '>' 表示只读取新消息
count=1,
block=timeout,
noack=noack,
)
# 确保消费者组存在（只在第一次调用时执行）
if stream_key not in self._checked_groups:
if self._ensure_consumer_group(stream_key, consumer_group):
self._checked_groups.add(stream_key)
if not messages or not messages[0][1]:
return None
while True:
try:
messages = self.redis_client.xreadgroup(
consumer_group,
consumer_name,
{stream_key: '>'},
count=1,
block=timeout,
noack=True,
)
# 解析消息
stream_name, stream_messages = messages[0]
if not stream_messages:
return None
if not messages or not messages[0][1]:
continue
message_id, fields = stream_messages[0]
_, stream_messages = messages[0]
message_id, fields = stream_messages[0]
data = json.loads(fields["data"])
value = fields["data"]
data = json.loads(value)
logger.info(f"Received: {stream_key} -> {message_id}")
return data
# 确认消息
flag = self.redis_client.xack(stream_key, consumer_group, message_id)
logger.info(f"{consumer_group} : {consumer_name} : {message_id}")
return data
except Exception as e:
logger.exception(e)
return None
except Exception as e:
logger.error(f"Error reading from stream {stream_key}: {e}")
time.sleep(1) # 发生异常(如网络闪断)时稍作停顿，防止死循环刷屏
def read_stream_from_start(self, stream_key, start_id='0-0'):
"""
@@ -199,43 +205,19 @@ class RedisStreamAPI:
logger.exception(e)
return []
def _ensure_consumer_group(self, stream_key: str, consumer_group: str):
def _ensure_consumer_group(self, stream_key: str, consumer_group: str) -> bool:
"""
确保消费者组存在,如果不存在则创建
Args:
stream_key (str): Redis stream的key名称
consumer_group (str): 消费者组名称
静默确保消费者组存在
"""
try:
# 检查消费者组是否存在
groups = self.redis_client.xinfo_groups(stream_key)
group_names = [group['name'] for group in groups]
if consumer_group not in group_names:
# 创建消费者组
self.redis_client.xgroup_create(stream_key, consumer_group, '$', mkstream=True)
# 直接创建，利用 mkstream=True。如果流不存在会自动创建流，如果组已存在会报错
self.redis_client.xgroup_create(stream_key, consumer_group, id='$', mkstream=True)
return True
except redis.ResponseError as e:
# 如果流不存在,xinfo_groups会报错.捕获此错误并创建流和组.
if "no such key" in str(e).lower():
try:
self.redis_client.xgroup_create(stream_key, consumer_group, '$', mkstream=True)
return True
except Exception as create_e:
if "BUSYGROUP Consumer Group name already exists" in str(e):
return True
else:
logger.exception(f"创建流 {stream_key} 和组 {consumer_group} 失败: {create_e}")
return False
elif "BUSYGROUP" in str(e):
# 消费者组已存在,这是正常情况
# 如果报错信息包含 BUSYGROUP，说明组已经存在，属于正常情况
if "BUSYGROUP" in str(e):
return True
else:
logger.exception(f"检查或创建消费者组时发生未知Redis响应错误: {e}")
return False
except Exception as e:
logger.exception(e)
# logger.error(f"创建消费者组失败: {e}")
return False
def get_stream_info(self, stream_key: str) -> Optional[Dict[str, Any]]:
@@ -281,29 +263,3 @@ class RedisStreamAPI:
self.redis_client.close()
except Exception as e:
logger.exception(e)
def clean_redis_stream(self, max_age_days=30):
"""
清理Redis Stream中超过指定天数的键值对.
"""
# 计算最老允许的时间戳,单位为毫秒
# Unix时间戳(秒)* 1000
logger.info(f"Clean Redis Stream older than {max_age_days} days.")
cutoff_timestamp_ms = int((datetime.datetime.now() - datetime.timedelta(days=max_age_days)).timestamp() * 1000)
try:
for key in self.redis_client.scan_iter(match='*'):
# 检查键的类型是否为 stream
if self.redis_client.type(key) == 'stream':
# 使用 XTRIM 命令删除早于给定ID的条目
# ID的格式是 `unix_time_ms-sequence_number`
# 我们可以使用 `unix_time_ms-0` 作为删除的上限ID
trim_id = f'{cutoff_timestamp_ms}-0'
# `XTRIM` 带有 `MINID` 选项,用于删除ID小于指定ID的所有条目
trimmed_count = self.redis_client.xtrim(key, minid=trim_id)
logger.info(f"Delete {trimmed_count} items from '{key}'")
except Exception as e:
logger.exception(e)
finally:
logger.info("clean task finished.")