first commit

This commit is contained in:
FunnyWolf
2025-09-07 20:02:11 +08:00
commit 3c003016c2
28 changed files with 1511 additions and 0 deletions
+171
View File
@@ -0,0 +1,171 @@
### Django template
*.log
*.pot
*.pyc
__pycache__/
local_settings.py
db.sqlite3
db.sqlite3-journal
media
# If your build process includes running collectstatic, then you probably don't need or want to include staticfiles/
# in your Git repository. Update and uncomment the following line accordingly.
# <django-project-name>/staticfiles/
### Python template
# Byte-compiled / optimized / DLL files
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
# Django stuff:
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/#use-with-ide
.pdm.toml
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
.idea/
Test/*
CONFIG.py
Test
+26
View File
@@ -0,0 +1,26 @@
FLASK_LISTEN_PORT = 7000
FLASK_LISTEN_HOST = "0.0.0.0"
REDIS_HOST = 'localhost'
REDIS_PORT = 6379
REDIS_DB = 0
REDIS_PASSWORD = "redis-stack-password-for-ai-soc-framework"
OPENAI_API_KEY = "sk-xxx"
OPENAI_BASE_URL = "https://api.openai.com/v1"
OPENAI_MODEL = ""
OPENAI_PROXY = ""
DIFY_BASE_URL = "https://api.dify.ai/v1"
DIFY_PROXY = None
DIFY_API_KEY = {
"Phishing_User_Report_Splunk_Dify_Nocodb": "app-xxx"
}
THEHIVE_URL = "https://192.168.1.114:443"
THEHIVE_API_KEY = "xxx"
NOCODB_URL = "http://192.168.1.114:8080"
NOCODB_TOKEN = "xxx"
NOCODB_ALERT_TABLE_ID = "xxx"
View File
+92
View File
@@ -0,0 +1,92 @@
import importlib
import os
import threading
import time
from Core.module_loader import start_watching
from Lib.log import logger
class Engine:
def __init__(self):
self.modules = {}
self.modules_dir = "MODULES"
self.observer = None
def start(self):
if not os.path.isdir(self.modules_dir):
os.makedirs(self.modules_dir)
self._load_initial_modules()
self.observer = start_watching(self, self.modules_dir)
logger.info("Engine started successfully, beginning module monitoring")
def stop(self):
if self.observer:
self.observer.stop()
self.observer.join()
for module_name in list(self.modules.keys()):
self.unload_module(module_name)
logger.info("All modules have been stopped")
def _load_initial_modules(self):
for filename in os.listdir(self.modules_dir):
if filename.endswith(".py") and not filename.startswith("_"):
module_name = filename.replace(".py", "")
file_path = os.path.join(self.modules_dir, filename)
self.load_module(module_name, file_path)
def run_loop(self, module_name, instance):
while module_name in self.modules:
logger.debug(f"Start Running module: {module_name}")
try:
instance.run()
except Exception as e:
logger.exception(e)
logger.debug(f"Finish Running module: {module_name}")
time.sleep(0.1)
def load_module(self, module_name: str, file_path: str):
if module_name in self.modules:
return
logger.info(f"Loading module: {module_name}")
try:
spec = importlib.util.spec_from_file_location(module_name, file_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
module_class = getattr(module, "Module")
threads = []
for i in range(module_class.thread_num):
thread_name = f"{module_name}_thread_{i}"
instance = module_class()
instance._thread_name = thread_name
thread = threading.Thread(target=self.run_loop, args=(module_name, instance), name=thread_name)
thread.daemon = True
threads.append(thread)
self.modules[module_name] = threads
for thread in self.modules[module_name]:
thread.start()
logger.info(f"Module '{module_name}' started successfully")
except Exception as e:
logger.error(f"Failed to load module: {e}")
def unload_module(self, module_name: str):
if module_name not in self.modules:
return
logger.info(f"Unloading module: {module_name}")
del self.modules[module_name]
logger.info(f"Module '{module_name}' unloaded successfully")
def reload_module(self, module_name: str, file_path: str):
logger.info(f"Reloading module: {module_name}")
self.unload_module(module_name)
self.load_module(module_name, file_path)
+42
View File
@@ -0,0 +1,42 @@
import os
from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer
from Lib.log import logger
class ModuleChangeHandler(FileSystemEventHandler):
def __init__(self, engine):
self.engine = engine
def _is_valid_module(self, path: str) -> bool:
filename = os.path.basename(path)
return filename.endswith(".py") and not filename.startswith(("_", "."))
def on_created(self, event):
if not event.is_directory and self._is_valid_module(event.src_path):
logger.info(f"New module file detected: {event.src_path}")
module_name = os.path.basename(event.src_path).replace(".py", "")
self.engine.load_module(module_name, event.src_path)
def on_deleted(self, event):
if not event.is_directory and self._is_valid_module(event.src_path):
logger.info(f"Module file deleted: {event.src_path}")
module_name = os.path.basename(event.src_path).replace(".py", "")
self.engine.unload_module(module_name)
def on_modified(self, event):
if not event.is_directory and self._is_valid_module(event.src_path):
logger.info(f"Module file modified: {event.src_path}")
module_name = os.path.basename(event.src_path).replace(".py", "")
self.engine.reload_module(module_name, event.src_path)
def start_watching(engine, path: str) -> Observer:
event_handler = ModuleChangeHandler(engine)
observer = Observer()
observer.schedule(event_handler, path, recursive=False)
observer.start()
logger.info(f"Starting to monitor directory: '{path}'")
return observer
+12
View File
@@ -0,0 +1,12 @@
services:
redis-stack:
image: redis/redis-stack:latest
container_name: redis-stack
restart: always
ports:
- "6379:6379"
- "8001:8001"
environment:
- REDIS_ARGS=--requirepass redis-stack-password-for-ai-soc-framework
volumes:
- ./redis-data:/data
View File
+49
View File
@@ -0,0 +1,49 @@
from typing import Dict, Any
import requests
from CONFIG import DIFY_BASE_URL, DIFY_PROXY
from Lib.log import logger
requests.packages.urllib3.disable_warnings()
class DifyClient(object):
def __init__(self):
self.base_url = DIFY_BASE_URL
def run_workflow(self, api_key: str, inputs: Dict[str, Any], user: str = "default_user") -> Dict[str, Any]:
url = f"{self.base_url}/workflows/run"
headers = {
"Authorization": f"Bearer {api_key}",
}
payload = {
"inputs": inputs,
"response_mode": "blocking",
"user": user
}
proxies = None
if DIFY_PROXY:
proxies = {
"http": DIFY_PROXY,
"https": DIFY_PROXY,
}
try:
response = requests.post(url,
headers=headers,
json=payload,
proxies=proxies,
)
response.raise_for_status()
response_data = response.json()
logger.debug(f"Dify API response: {response_data}")
data = response_data.get("data", {})
if data and data.get("status") == "succeeded":
outputs = data.get("outputs")
return outputs
else:
return {}
except Exception as e:
raise
+24
View File
@@ -0,0 +1,24 @@
import requests
from CONFIG import NOCODB_URL, NOCODB_TOKEN, NOCODB_ALERT_TABLE_ID
class NocodbClient(object):
def __init__(self):
pass
@staticmethod
def create_alert(record: dict):
headers = {"xc-token": NOCODB_TOKEN}
url = f"{NOCODB_URL}/api/v2/tables/{NOCODB_ALERT_TABLE_ID}/records"
try:
response = requests.post(url,
headers=headers,
json=record)
response.raise_for_status()
response_data = response.json()
return response_data
except Exception as e:
raise
+102
View File
@@ -0,0 +1,102 @@
import httpx
import urllib3
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
from CONFIG import OPENAI_MODEL, OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_PROXY
class OpenAIAPI(object):
def __init__(self):
self.api_key = None
self.base_url = None
self.model = None
self.temperature = 0
self.alive = False
def set_api_key(self, api_key):
self.api_key = api_key
def set_base_url(self, base_url: str):
self.base_url = base_url.rstrip('/')
def set_temperature(self, temperature: float):
self.temperature = temperature
def set_model(self, model: str):
self.model = model
def get_model(self, model_kwargs=None):
if model_kwargs is None:
model_kwargs = {}
self.set_api_key(OPENAI_API_KEY)
self.set_base_url(OPENAI_BASE_URL)
self.set_model(OPENAI_MODEL)
http_client = None
if OPENAI_PROXY:
http_client = httpx.Client(proxy=OPENAI_PROXY)
return ChatOpenAI(
base_url=self.base_url,
api_key=self.api_key,
model=self.model,
temperature=self.temperature,
http_client=http_client,
model_kwargs=model_kwargs,
)
def is_alive(self):
model = self.get_model()
# 基础连通性测试
parser = StrOutputParser()
chain = model | parser
messages = [
("system", "give you `ping` reply `pong`."),
("human", "ping"),
]
try:
ai_msg = chain.invoke(messages)
self.alive = False
return True
except Exception as e:
self.alive = False
return False
def is_support_function_calling(self):
# Function calling 能力测试
def test_func(x: str) -> str:
"""A test function that returns the input string."""
return x
model = self.get_model()
try:
model_with_tools = model.bind_tools([test_func])
test_messages = [
("system", "When user says test, call test_func with 'hello' as argument."),
("human", "test"),
]
response = model_with_tools.invoke(test_messages)
if not response.tool_calls:
return False
except Exception as e:
return False
return True
@staticmethod
def is_model_alive(model: ChatOpenAI):
parser = StrOutputParser()
chain = model | parser
messages = [
("system", "give you `ping` reply `pong`."),
("human", "ping"),
]
try:
ai_msg = chain.invoke(messages)
return True
except Exception as e:
return False
+21
View File
@@ -0,0 +1,21 @@
from thehive4py import TheHiveApi
from thehive4py.types.alert import InputAlert, OutputAlert
from CONFIG import THEHIVE_URL, THEHIVE_API_KEY
class TheHiveClient(object):
def __init__(self):
self.hive = TheHiveApi(
url=THEHIVE_URL,
apikey=THEHIVE_API_KEY,
verify=False
)
def alert_create(self, alert_data: InputAlert):
try:
ouput_alert: OutputAlert = self.hive.alert.create(alert=alert_data)
return ouput_alert
except Exception as e:
print(f"Error creating alert: {e}")
return None
+58
View File
@@ -0,0 +1,58 @@
from flask import Flask, request, jsonify
from CONFIG import FLASK_LISTEN_PORT, FLASK_LISTEN_HOST
from Forwarder.log import logger
from Lib.redis_stream_api import RedisStreamAPI
redis_stream_api = RedisStreamAPI()
app = Flask(__name__)
@app.route("/api/v1/webhook/splunk", methods=["POST"])
def receive_splunk_webhook():
if not request.is_json:
return jsonify({"status": "error", "message": "Request must be JSON"}), 400
data = request.get_json()
result = data.pop('result', {})
search_name = data.get('search_name')
sid = data.get('sid')
host = data.get('app')
owner = data.get('owner')
results_link = data.get('results_link')
logger.info(f"Splunk webhook: {data}")
redis_stream_api.send_message(search_name, result)
logger.info("Message sent to Redis stream")
# 5. 返回成功响应
return jsonify({"status": "success", "message": "Webhook received"}), 200
@app.route("/api/v1/webhook/kibana", methods=["POST"])
def receive_kibana_webhook():
if not request.is_json:
return jsonify({"status": "error", "message": "Request must be JSON"}), 400
data = request.get_json()
rule_name = data.get('rule').get("name")
hits = data.get('context').get("hits")
for hit in hits:
_source = hit.pop('_source', {})
logger.info(f"elasticsearch webhook: {hit}")
redis_stream_api.send_message(rule_name, _source)
logger.info("Message sent to Redis stream")
# 5. 返回成功响应
return jsonify({"status": "success", "message": "Webhook received"}), 200
@app.route("/")
def index():
return "Forwarder is running"
if __name__ == '__main__':
logger.info(f"Starting Flask server on http://{FLASK_LISTEN_HOST}:{FLASK_LISTEN_PORT}")
logger.info(f"Splunk Webhook URL : http://{FLASK_LISTEN_HOST}:{FLASK_LISTEN_PORT}/api/v1/webhook/splunk")
logger.info(f"Kibana Webhook URL : http://{FLASK_LISTEN_HOST}:{FLASK_LISTEN_PORT}/api/v1/webhook/kibana")
app.run(host=FLASK_LISTEN_HOST, port=FLASK_LISTEN_PORT)
+23
View File
@@ -0,0 +1,23 @@
import logging
import sys
def setup_logging(log_file='main.log'):
"""
Configures the root logger for the entire application.
- Logs to both a file and the console.
- Sets a standardized format for log messages.
"""
logging.basicConfig(
level=logging.INFO,
format='%(levelname)s - %(asctime)s - [%(name)s] - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S',
handlers=[
logging.FileHandler(log_file),
logging.StreamHandler(sys.stdout)
]
)
setup_logging()
logger = logging.getLogger("forwarder")
+22
View File
@@ -0,0 +1,22 @@
MIT License
Copyright (c) 2025 funnywolf
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
View File
+92
View File
@@ -0,0 +1,92 @@
import os
from langchain_core.prompts import SystemMessagePromptTemplate
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph.state import CompiledStateGraph
from CONFIG import DIFY_API_KEY
from Lib.configs import MODULE_DATA_DIR, REDIS_CONSUMER_GROUP
from Lib.llmapi import AgentState
from Lib.log import logger
from Lib.redis_stream_api import RedisStreamAPI
class BaseModule(object):
thread_num = 1
def __init__(self):
self._thread_name = None
self.logger = logger
self.agent_state: AgentState = AgentState(messages=[], alert_raw={}, temp_data={}, analyze_result={})
# debug
self.debug_alert_name = None
self.debug_message_id = None # 设置为非None以启用Debug模式
@property
def module_name(self):
"""获取模块加载路径"""
if self.debug_alert_name is None:
return self.__module__.split(".")[-1]
else:
return self.debug_alert_name
def read_message(self) -> dict:
"""读取消息"""
redis_stream_api = RedisStreamAPI()
if self.debug_message_id is not None:
message = redis_stream_api.read_stream_from_start(self.module_name, start_id=self.debug_message_id)
else:
message = redis_stream_api.read_message(stream_key=self.module_name, consumer_group=REDIS_CONSUMER_GROUP, consumer_name=self._thread_name)
return message
def get_dify_api_key(self, app_name=None):
if app_name is None:
app_name = self.module_name
return DIFY_API_KEY.get(app_name)
def run(self):
raise NotImplementedError
class LanggraphModule(BaseModule):
def __init__(self):
super().__init__()
self.graph: CompiledStateGraph = None
## LLM PART
@staticmethod
def get_checkpointer():
checkpointer = MemorySaver()
return checkpointer
def load_system_prompt_template(self, filename):
"""从模块对应的 MODULES_DATA 目录加载 md 文件内容
Args:
filename: md 文件名称,不需要包含 .md 后缀
Returns:
str: md 文件的内容
"""
template_path = os.path.join(MODULE_DATA_DIR, self.module_name, f"{filename}.md")
try:
with open(template_path, 'r', encoding='utf-8') as f:
system_prompt_template: SystemMessagePromptTemplate = SystemMessagePromptTemplate.from_template(f.read())
return system_prompt_template
except Exception as e:
logger.warning(f"Failed to load prompt template {template_path}: {str(e)}")
raise e
def run_graph(self):
self.graph.checkpointer.delete_thread(self.module_name)
config = RunnableConfig()
config["configurable"] = {"thread_id": self.module_name}
self.agent_state = AgentState(messages=[], alert_raw={}, temp_data={}, analyze_result={})
for event in self.graph.stream(self.agent_state, config, stream_mode="values"):
self.logger.debug(event)
def run(self):
self.run_graph()
+6
View File
@@ -0,0 +1,6 @@
import os
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MODULE_DATA_DIR = os.path.join(BASE_DIR, 'MODULES_DATA')
REDIS_CONSUMER_GROUP = 'AI_SOC_FRAMEWORK_GROUP'
REDIS_CONSUMER_NAME = 'AI_SOC_FRAMEWORK_CONSUMER_0'
+11
View File
@@ -0,0 +1,11 @@
from typing import Annotated, Any, Dict, List
from langgraph.graph import add_messages
from pydantic import BaseModel
class AgentState(BaseModel):
messages: Annotated[List[Any], add_messages]
alert_raw: Dict[str, Any]
temp_data: Dict[str, Any]
analyze_result: Dict[str, Any]
+23
View File
@@ -0,0 +1,23 @@
import logging
import sys
def setup_logging(log_file='main.log'):
"""
Configures the root logger for the entire application.
- Logs to both a file and the console.
- Sets a standardized format for log messages.
"""
logging.basicConfig(
level=logging.INFO,
format='%(levelname)s - %(asctime)s - [%(name)s] - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S',
handlers=[
logging.FileHandler(log_file),
logging.StreamHandler(sys.stdout)
]
)
setup_logging()
logger = logging.getLogger("ai-soc-framework")
+32
View File
@@ -0,0 +1,32 @@
import redis
from Lib.log import logger
from CONFIG import (
REDIS_HOST, REDIS_PORT, REDIS_DB, REDIS_PASSWORD
)
class RedisClient(object):
def __init__(self):
pass
@staticmethod
def get_stream_connection():
"""用于订阅类操作,无需使用连接池"""
redis_client = redis.Redis(
host=REDIS_HOST,
port=REDIS_PORT,
db=REDIS_DB,
password=REDIS_PASSWORD,
decode_responses=True,
)
# 测试连接
try:
redis_client.ping()
return redis_client
except redis.ConnectionError as e:
logger.exception(e)
raise
+265
View File
@@ -0,0 +1,265 @@
import json
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 Lib.redis_client import RedisClient
class RedisStreamAPI:
"""
Redis Stream API封装类,提供消息发送和读取功能
"""
def __init__(self):
"""初始化RedisStreamAPI类"""
self.redis_client = RedisClient.get_stream_connection()
def send_message(self, stream_key: str, message: Dict[str, Any]) -> Optional[str]:
"""
发送消息到指定stream
Args:
stream_key (str): Redis stream的key名称
message (Dict[str, Any]): 要发送的消息内容
Returns:
Optional[str]: 发送成功返回消息ID,失败返回None
"""
try:
data = json.dumps(message, ensure_ascii=False)
# 发送消息到stream
message_id = self.redis_client.xadd(
stream_key,
{"data": data}
)
return message_id
except Exception as e:
logger.exception(e)
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]]:
"""
从指定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
"""
try:
if consumer_group is None:
consumer_group = REDIS_CONSUMER_GROUP
if consumer_name is None:
consumer_name = REDIS_CONSUMER_NAME
# 确保消费者组存在
self._ensure_consumer_group(stream_key, consumer_group)
# 从消费者组读取消息
messages = self.redis_client.xreadgroup(
consumer_group,
consumer_name,
{stream_key: '>'}, # '>' 表示只读取新消息
count=1,
block=timeout,
noack=noack,
)
if not messages or not messages[0][1]:
return None
# 解析消息
stream_name, stream_messages = messages[0]
if not stream_messages:
return None
message_id, fields = stream_messages[0]
value = fields["data"]
data = json.loads(value)
# 确认消息
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
def read_stream_from_start(self, stream_key, start_id='0-0'):
"""
从指定 Stream 的开头重复读取消息。
:param topic: Stream 的名称。
:param count: 要读取的消息数量。
"""
try:
messages = self.redis_client.xread(
count=1,
block=0,
streams={stream_key: start_id}
)
if not messages or not messages[0][1]:
return None
# 解析消息
stream_name, stream_messages = messages[0]
if not stream_messages:
return None
message_id, fields = stream_messages[0]
value = fields["data"]
data = json.loads(value)
return data
except Exception as e:
logger.exception(e)
return None
def acknowledge_message(self, stream_key: str, message_id: str,
consumer_group: str = None) -> bool:
"""
确认消息已被处理
Args:
stream_key (str): Redis stream的key名称
message_id (str): 消息ID
consumer_group (str): 消费者组名称,如果为None则使用默认配置
Returns:
bool: 确认成功返回True,失败返回False
"""
try:
if consumer_group is None:
consumer_group = REDIS_CONSUMER_GROUP
# 确认消息
result = self.redis_client.xack(stream_key, consumer_group, message_id)
if result:
return True
else:
return False
except Exception as e:
logger.exception(e)
return False
def get_pending_messages(self, stream_key: str, consumer_group: str = None,
consumer_name: str = None) -> List[Dict[str, Any]]:
"""
获取待处理的消息
Args:
stream_key (str): Redis stream的key名称
consumer_group (str): 消费者组名称,如果为None则使用默认配置
consumer_name (str): 消费者名称,如果为None则使用默认配置
Returns:
List[Dict[str, Any]]: 待处理的消息列表
"""
try:
if consumer_group is None:
consumer_group = REDIS_CONSUMER_GROUP
if consumer_name is None:
consumer_name = REDIS_CONSUMER_NAME
# 获取待处理消息
pending_messages = self.redis_client.xpending(
stream_key, consumer_group, '-', '+', 100, consumer_name
)
messages = []
for message_id, consumer, idle_time, delivery_count in pending_messages:
messages.append({
'message_id': message_id,
'consumer': consumer,
'idle_time': idle_time,
'delivery_count': delivery_count
})
return messages
except Exception as e:
logger.exception(e)
return []
def _ensure_consumer_group(self, stream_key: str, consumer_group: str):
"""
确保消费者组存在,如果不存在则创建
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)
except redis.ResponseError as e:
if "BUSYGROUP" in str(e):
pass
else:
logger.exception(e)
except Exception as e:
logger.exception(e)
def get_stream_info(self, stream_key: str) -> Optional[Dict[str, Any]]:
"""
获取stream信息
Args:
stream_key (str): Redis stream的key名称
Returns:
Optional[Dict[str, Any]]: stream信息,失败返回None
"""
try:
info = self.redis_client.xinfo_stream(stream_key)
return info
except Exception as e:
logger.exception(e)
return None
def delete_stream(self, stream_key: str) -> bool:
"""
删除stream
Args:
stream_key (str): Redis stream的key名称
Returns:
bool: 删除成功返回True,失败返回False
"""
try:
result = self.redis_client.delete(stream_key)
if result:
return True
else:
return False
except Exception as e:
logger.exception(e)
return False
def close(self):
"""关闭Redis连接"""
try:
self.redis_client.close()
except Exception as e:
logger.exception(e)
@@ -0,0 +1,101 @@
import json
import uuid
from typing import Optional, Union, Dict, Any
from pydantic import BaseModel, Field
from thehive4py.types.alert import InputAlert, OutputAlert
from External.difyclient import DifyClient
from External.thehiveclient import TheHiveClient
from Lib.base import BaseModule
class AnalyzeResult(BaseModel):
"""用于从文本中提取用户信息的结构"""
is_phishing: bool = Field(description="是否为钓鱼邮件,True或False")
confidence: float = Field(description="信心指数,范围0到1之间")
reasoning: Optional[Union[str, Dict[str, Any]]] = Field(description="推理过程", default=None)
class Module(BaseModule):
def __init__(self):
super().__init__()
self.thehive_client = TheHiveClient()
def alert_preprocess_node(self):
"""预处理告警数据"""
# 获取stream中的原始告警
self.agent_state.alert_raw = self.read_message()
# 解析数据,此处是获取Elasticsearch Webhook发送的JSON数据的处理样例
alert = self.agent_state.alert_raw
headers = alert["headers"]
headers = {"From": headers["From"], "To": headers["To"], "Subject": headers["Subject"], "Date": headers["Date"],
"Return-Path": headers["Return-Path"],
"Authentication-Results": headers["Authentication-Results"]}
alert["headers"] = headers
self.agent_state.alert_raw = alert
return
def alert_analyze_node(self):
api_key = self.get_dify_api_key()
client = DifyClient()
inputs = {
"alert_raw": json.dumps(self.agent_state.alert_raw)
}
result = client.run_workflow(
api_key=api_key,
inputs=inputs,
user=self.module_name
)
self.agent_state.analyze_result = result.get("analyze_result")
return
def alert_output_node(self):
"""处理分析结果"""
analyze_result: AnalyzeResult = AnalyzeResult(**self.agent_state.analyze_result)
alert_raw = self.agent_state.alert_raw
to = alert_raw["headers"]["To"]
subject = alert_raw["headers"]["Subject"]
if analyze_result.is_phishing and analyze_result.confidence > 0.8:
severity = 2
else:
severity = 0
# 发送到thehive
input_alert: InputAlert = {
"type": "phishing",
"source": "user_report",
"sourceRef": str(uuid.uuid4()),
"title": self.module_name,
"description": f"```json{alert_raw}```",
"tags": ["phishing", "user_report"],
"severity": severity,
"summary": f"{analyze_result.model_dump()}",
"observables": [
{"dataType": "mail", "data": to},
{"dataType": "mail-subject", "data": subject},
],
}
output_alert: OutputAlert = self.thehive_client.alert_create(input_alert)
self.logger.debug(output_alert)
return
def run(self):
self.alert_preprocess_node()
self.alert_analyze_node()
self.alert_output_node()
if __name__ == "__main__":
module = Module()
module.debug_alert_name = "Phishing_User_Report_V2"
module.debug_message_id = "0-0"
module.run()
@@ -0,0 +1,187 @@
import json
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.state import CompiledStateGraph
from pydantic import BaseModel, Field
from External.nocodbclient import NocodbClient
from External.opanaiapi import OpenAIAPI
from External.thehiveclient import TheHiveClient
from Lib.base import LanggraphModule
from Lib.llmapi import AgentState
class AnalyzeResult(BaseModel):
"""用于从文本中提取用户信息的结构"""
is_phishing: bool = Field(description="是否为钓鱼邮件,True或False")
confidence: float = Field(description="信心指数,范围0到1之间")
reasoning: Optional[Union[str, Dict[str, Any]]] = Field(description="推理过程", default=None)
class Module(LanggraphModule):
thread_num = 2
def __init__(self):
super().__init__()
self.thehive_client = TheHiveClient()
self.init()
def init(self):
def alert_preprocess_node(state: AgentState):
"""预处理告警数据"""
# 获取stream中的原始告警
alert_raw = self.read_message()
if alert_raw is None:
return
# 解析数据,此处是获取Splunk Webhook发送的JSON数据的处理样例
alert_raw = json.loads(alert_raw["_raw"])
headers = alert_raw["headers"]
headers = {"From": headers["From"], "To": headers["To"], "Subject": headers["Subject"], "Date": headers["Date"],
"Return-Path": headers["Return-Path"],
"Authentication-Results": headers["Authentication-Results"]}
alert_raw["headers"] = headers
state.alert_raw = alert_raw
return state
# 定义node
def alert_analyze_node(state: AgentState):
"""AI分析告警数据"""
# 加载system prompt
system_prompt_template = self.load_system_prompt_template(f"senior_phishing_expert")
# 演示如何生成动态提示词
current_date = datetime.now().strftime("%Y-%m-%d")
system_message = system_prompt_template.format(current_date=current_date)
# 构建few-shot示例
few_shot_examples = [
HumanMessage(
content=json.dumps({
"headers": {
"From": "\"Wang Lei, Project Manager\" <lei.wang@example-corp.com>",
"To": "\"Li Na, Marketing Department\" <na.li@example-corp.com>",
"Subject": "Project Alpha Weekly Status Report",
"Date": "Tue, 2 Sep 2025 10:15:00 +0800",
"Return-Path": "lei.wang@example-corp.com",
"Authentication-Results": "mx.example-corp.com; spf=pass smtp.mail=lei.wang@example-corp.com;"
},
"body": {
"plain_text": "Hi Li Na,\n\nPlease find attached the weekly status report for Project Alpha.\n\nThis week, we have completed the initial design phase and are on track to begin development next Monday as planned. Please review the attached document and let me know if you have any feedback before our sync-up meeting on Wednesday.\n\nThanks,\n\nBest Regards\nWang Lei / 王雷\nProject Manager / 项目经理\nTechnology Department / 技术部\nExample Corporation / 示例公司\nMobile: +86 13800138000\nEmail / 邮箱: lei.wang@example-corp.com\n",
"html": ""
},
"attachments": [
{
"filename": "Project_Alpha_Weekly_Report_W35.pdf",
"filepath": "attachments/Project_Alpha_Weekly_Report_W35.pdf",
"content_type": "application/pdf"
}
]
})
),
AIMessage(
content=str(AnalyzeResult(is_phishing=False, confidence=0.95,
reasoning="The email is from a known colleague within the same organization, discussing a legitimate project.").model_dump())
),
HumanMessage(
content=json.dumps({
"headers": {
"From": "\"Microsoft Support\" <support-noreply@microsft.com>",
"To": "\"Valued Customer\" <user@example.com>",
"Subject": "紧急:您的账户已被暂停,需要立即验证 Urgent: Your Account is Suspended, Immediate Verification Required",
"Date": "Tue, 2 Sep 2025 14:30:10 +0800",
"Return-Path": "<bounce-scam@phish-delivery.net>",
"Authentication-Results": "mx.example.com; spf=fail smtp.mail=support-noreply@microsft.com; dkim=fail header.d=microsft.com; dmarc=fail (p=REJECT sp=REJECT) header.from=microsft.com",
"X-Coremail-Antispam": "1Uf129KBjvdXoW7GF18tw4xZF4xWF4rtw4kCrg_yoWfZFg_GF4DC348Wrnxtr15J398ZwnFy3ZFgrZ8CF9a9r4DZrZ8X3WkXa4kJr98K3y8C3WfJw1fXFW3ArnrZa93tF15tjkaLaAFLSUrUUUUUb8apTn2vfkv8UJUUUU8Yxn0WfASr-VFAUDa7-sFnT9fnUUvcSsGvfC2KfnxnUUI43ZEXa7IU04v35UUUUU=="
},
"body": {
"plain_text": "尊敬的用户,\n\n我们的系统检测到您的帐户存在异常登录活动。为了保护您的安全,我们已临时暂停您的帐户。\n\n请立即点击以下链接以验证您的身份并恢复您的帐户访问权限:\n\nhttps://login.microsoftonline.com/common/oauth2/v2.0/authorize?client_id=... (请注意,这只是显示文本,实际链接是恶意的)\n\n如果您不在24小时内完成验证,您的帐户将被永久锁定。\n\n感谢您的合作。\n\n微软安全团队\n\n---\n\nDear User,\n\nOur system has detected unusual sign-in activity on your account. For your security, we have temporarily suspended your account.\n\nPlease click the link below immediately to verify your identity and restore access:\n\nhttp://secure-login-update-required.com/reset-password?user=user@example.com\n\nIf you do not verify within 24 hours, your account will be permanently locked.\n\nThank you for your cooperation.\n\nThe Microsoft Security Team",
"html": "<html><head></head><body><p>尊敬的用户,</p><p>我们的系统检测到您的帐户存在异常登录活动。为了保护您的安全,我们已临时暂停您的帐户。</p><p>请立即点击以下链接以验证您的身份并恢复您的帐户访问权限:</p><p><a href='http://secure-login-update-required.com/reset-password?user=user@example.com'>https://login.microsoftonline.com/common/oauth2/v2.0/authorize?client_id=...</a></p><p>如果您不在24小时内完成验证,您的帐户将被永久锁定。</p><p>感谢您的合作。</p><p><b>微软安全团队</b></p></body></html>"
},
"attachments": [
{
"filename": "Account_Verification_Form.html",
"filepath": "attachments/Account_Verification_Form.html",
"content_type": "text/html"
}
]
})
),
AIMessage(
content=str(AnalyzeResult(is_phishing=True, confidence=0.92,
reasoning="The email contains several red flags: the sender's domain is misspelled, the Return-Path is from a suspicious domain, SPF and DKIM checks fail, and the email urges immediate action with threatening language. Additionally, the provided links do not match official Microsoft URLs.").model_dump())
),
]
# 构建消息列表
messages = [
system_message,
*few_shot_examples,
HumanMessage(content=json.dumps(state.alert_raw)),
]
# 运行
openai_api = OpenAIAPI()
model_kwargs = {
"extra_body": {
"enable_thinking": False
}
}
llm = openai_api.get_model(model_kwargs)
llm = llm.with_structured_output(AnalyzeResult)
response: AnalyzeResult = llm.invoke(messages)
state.analyze_result = response.model_dump()
return state
def alert_output_node(state: AgentState):
"""处理分析结果"""
analyze_result: AnalyzeResult = AnalyzeResult(**state.analyze_result)
if analyze_result.is_phishing and analyze_result.confidence > 0.8:
severity = "HIGH"
else:
severity = "INFO"
# 发送到nocodb
payload = {
"Title": self.module_name,
"SEVERITY": severity,
"RAW": json.dumps(state.alert_raw),
"Description": json.dumps(state.alert_raw),
"Summary": json.dumps(state.analyze_result),
"Status": "NEW",
"Source": "Splunk"
}
result = NocodbClient.create_alert(payload)
return state
# 编译graph
workflow = StateGraph(AgentState)
workflow.add_node("alert_preprocess_node", alert_preprocess_node)
workflow.add_node("alert_analyze_node", alert_analyze_node)
workflow.add_node("alert_output_node", alert_output_node)
workflow.set_entry_point("alert_preprocess_node")
workflow.add_edge("alert_preprocess_node", "alert_analyze_node")
workflow.add_edge("alert_analyze_node", "alert_output_node")
workflow.set_finish_point("alert_output_node")
self.graph: CompiledStateGraph = workflow.compile(checkpointer=self.get_checkpointer())
return True
if __name__ == "__main__":
module = Module()
module.debug_alert_name = "Phishing_User_Report_V1"
module.debug_message_id = "0-0"
module.run()
@@ -0,0 +1,42 @@
# 高级钓鱼邮件分析专家
## 角色定位
你是一位专业的钓鱼邮件分析专家,由先进的AI技术驱动。你擅长分析一封邮件是否为钓鱼邮件。
## 目标
通过分析用户提供的邮件信息来分析邮件是否为钓鱼邮件。
## 所需输入
JSON格式的邮件信息,包含以下字段:
- 发件人邮箱地址
- 收件人邮箱地址
- 邮件主题
- 邮件正文内容
- 附件信息(如果有)
- 链接信息(如果有)
## 任务描述
- 多维度分析用户提供的邮件信息
- 确认邮件是否为钓鱼邮件
- 提供详细的分析报告,说明判断依据
## 思考步骤
1. 分析发件人是否为可信来源
2. 分析From和Reply-To地址是否匹配
3. 检查邮件主题是否包含可疑关键词
4. 分析邮件正文内容,寻找可疑链接或附件
5. 检查邮件中的链接是否指向可疑网站
6. 检查邮件中的附件是否包含恶意软件
7. 综合以上分析,判断邮件是否为钓鱼邮件
8. 提供详细的分析报告,说明判断依据
## 注意事项
- 今天是:{current_date}
- 回复必须使用JSON格式
+75
View File
@@ -0,0 +1,75 @@
# AI SOC Framework
基于 LLM 的告警分析框架,通过模块化的方式调用Langchain/Langgraph/Dify进行告警分析
## 功能
* **模块化引擎**: 动态加载和执行告警分析模块
* **LLM集成**:
* **LLM接口**: 包含Dify和Langgraph的集成接口及样例模块
* **工单接口**: 包含Thehive及Nocodb的集成接口及样例模块
* **事件驱动**: 使用 Redis Stream 作为消息总线,实现模块化告警流式处理
## 架构图
![img.png](Static/img.png)
## 快速开始
### 1. 环境准备
* Python 3.12+
* Docker 和 Docker Compose (用于运行 Redis Stack)
### 2. 开发环境
1. **克隆项目**
```bash
git clone https://github.com/FunnyWolf/ai-soc-framework
cd ai-soc-framework
```
2. **安装依赖**
建议在虚拟环境中使用 pip 安装依赖:
```bash
python -m venv venv
source venv/bin/activate
pip install -r requirements.txt
```
3. **启动依赖服务**
项目使用 Redis 作为消息队列。我们提供了 Docker Compose 文件来快速启动一个 Redis Stack 实例。
```bash
cd Docker/redis_stack
docker-compose up -d
```
4. **配置**
复制配置文件模板,并根据您的环境修改 `CONFIG.py` 文件,填入各个服务的 API 密钥、URL 和其他必要参数。
```bash
cp CONFIG.example.py CONFIG.py
```
5. **启动 Forwarder**
Forwarder 负责将SIEM的 Webhook 推送的告警数据写入 Redis Stream (当前适配Splunk和ELK)
```bash
cd Forwarder
python app.py
```
5. **运行**
完成上述步骤后,在项目根目录运行主程序:
```bash
python main.py
```
程序将启动核心引擎,加载 `MODULES` 目录下的所有模块,并开始监听和处理事件。
## 模块开发
您可以参考 `MODULES` 目录下的现有模块,开发自己的自动化流程。每个模块都是一个独立的 Python 文件,框架会自动加载并运行它。
## 许可证
该项目采用 [MIT](https://choosealicense.com/licenses/mit/) 许可证。
BIN
View File
Binary file not shown.

After

Width:  |  Height:  |  Size: 201 KiB

+24
View File
@@ -0,0 +1,24 @@
import time
from Core.engine import Engine
from Lib.log import logger
def main():
try:
engine = Engine()
engine.start()
while True:
time.sleep(1)
except KeyboardInterrupt:
logger.info("Program interrupted by user")
except Exception as e:
logger.error(f"Program error occurred: {e}")
finally:
if 'engine' in locals() and engine:
engine.stop()
logger.info("Program has been closed")
if __name__ == '__main__':
main()
+11
View File
@@ -0,0 +1,11 @@
pydantic
watchdog
redis
langgraph
langgraph-checkpoint-redis
langchain-core
langchain-community
langchain-openai
thehive4py
requests
flask