mirror of
https://github.com/FunnyWolf/agentic-soc-platform.git
synced 2026-08-22 13:12:56 +02:00
first commit
This commit is contained in:
+171
@@ -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
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
Vendored
Vendored
+49
@@ -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
|
||||
Vendored
+24
@@ -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
|
||||
Vendored
+102
@@ -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
|
||||
Vendored
+21
@@ -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
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
@@ -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.
|
||||
|
||||
+92
@@ -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()
|
||||
@@ -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'
|
||||
@@ -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
@@ -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")
|
||||
@@ -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
|
||||
@@ -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格式
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# AI SOC Framework
|
||||
|
||||
基于 LLM 的告警分析框架,通过模块化的方式调用Langchain/Langgraph/Dify进行告警分析
|
||||
|
||||
## 功能
|
||||
|
||||
* **模块化引擎**: 动态加载和执行告警分析模块
|
||||
* **LLM集成**:
|
||||
* **LLM接口**: 包含Dify和Langgraph的集成接口及样例模块
|
||||
* **工单接口**: 包含Thehive及Nocodb的集成接口及样例模块
|
||||
* **事件驱动**: 使用 Redis Stream 作为消息总线,实现模块化告警流式处理
|
||||
|
||||
## 架构图
|
||||
|
||||

|
||||
|
||||
## 快速开始
|
||||
|
||||
### 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/) 许可证。
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 201 KiB |
@@ -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()
|
||||
@@ -0,0 +1,11 @@
|
||||
pydantic
|
||||
watchdog
|
||||
redis
|
||||
langgraph
|
||||
langgraph-checkpoint-redis
|
||||
langchain-core
|
||||
langchain-community
|
||||
langchain-openai
|
||||
thehive4py
|
||||
requests
|
||||
flask
|
||||
Reference in New Issue
Block a user