## https://sploitus.com/exploit?id=F1053B70-49D7-57BC-81FC-95131B4548AB
# π‘οΈ Security Investigation System
This is a **Proof of Concept project** for automatic analysis and tracing of security incidents, built using **Multi-Agent Collaboration** and **Neo4j Graph Database**. The system aims to verify the feasibility of simulating security expert thinking through multi-agent collaboration. It assumes that the backend has already written security alerts and related entity data (IPs, hosts, processes, etc.) into the Neo4j database. The system automatically generates query strategies, extracts relevant evidence from the graph, and ultimately produces a complete analysis report of the attack chain.
## π Key Features
- **π§ Multi-Agent Collaboration Architecture**: The system consists of four independent but closely collaborating agents: Coordinator, Planner, Retriever, and Analyzer, along with a Reporter responsible for summarizing the results.
- **πΈοΈ Real-time Attack Graph (Palantir Gotham Style)**: The frontend includes a dynamic knowledge graph driven by D3.js. The backend automatically extracts entities (IPs, domain names, hosts, processes) and relationships, rendering the attack chain during analysis.
- **β‘ Asynchronous Streaming Processing (SSE)**: The backendβs reasoning process is displayed in real time via Server-Sent Events. The βDeep Thinkingβ option is available.
- **π Autonomous Iterative Process**: The system automatically decides whether further investigation is needed based on the clues obtained from previous queries, until a complete picture of the attack is reconstructed or the maximum depth is reached.
- **βΈοΈ Instant Stoppage**: Since security investigations can be lengthy, the system allows for immediate stopping of the investigation and generation of a report at the current progress.
## ποΈ Architecture and Principle Explanation
### 1. Directory Structure
```
lab/
β βββ agents/ # Core logic of multi-agents
β β βββ coordinator.py # Coordinator: Responsible for understanding user intentions and extracting core entities
β β βββ planner.py # Planner: Determines next steps based on existing clues
β β βββ retriever.py # Retriever: Translates plans into Cypher queries for execution in Neo4j, supporting concurrency
β β βββ analyzer.py # Analyzer: Analyzes data to determine relevance of clues and decides whether to continue iterating
β β βββ reporter.py # Reporter: Summarizes all clues after the investigation is completed, generating a final Markdown report
β βββ graph/ # Definition of LangGraph state machine
β β βββ nodes.py # Encapsulation of node execution logic
β β βββ state.py # Global context state definition (TypedDict)
β β βββ workflow.py # Definition of cyclic graph workflow
β βββ tools/ # Tools/driver layer
β β βββ cypher_templates.py # Cypher query templates
β β βββ graphUtils.py # Responsible for extracting nodes/edges from Neo4j for D3.js rendering
β β βββ neo4j_tools.py # Encapsulates operations related to the Neo4j Driver
β βββ frontend/ # Front-end display layer (Native HTML/JS/CSS)
β β βββ index.html # Main page with left/right sidebar layout (timeline + attack graph)
β β βββ style.css # Gotham style theme, including graph floating window styling
β β βββ app.js # SSE event listening, DOM manipulation, state management
β β βββ attack_graph.js # Encapsulates Force-directed Graph logic using D3.js (v7)
β β βββ d3.v7.min.js # Offline D3.js library
β βββ server.py # FastAPI service for providing SSE interface (/investigate)
β βββ .env # Environment variables configuration (Neo4j, LLM Keys, etc.)
β βββ Dockerfile # Docker image build definition
β βββ docker-compose.yml # Definition of multi-container orchestration
β βββ seed_data.py # Script for importing simulated attack data
β βββ README.md # This document
```
### 2. Multi-Agent Processing Mechanism and State Transition
The system uses **LangGraph** to control the workflow. Its core logic is a **cyclic graph**:
1. **`Coordinator` (Starting Point)`: Parses the userβs initial input (e.g., βAnalyze signs of intrusion on 192.168.1.100β) and extracts initial entities (e.g., IP). These entities are added to the global state.
2. **Iterative Loop**:
- **`Planner`**: Reads the collected clues and decides what to investigate next. For example, if only an IP is known, it might decide to βquery the IPβs network connections in the past 24 hoursβ. A series of specific actions are generated.
- **`Retriever`**: Receives the actions from the Planner. It translates natural language intentions into Cypher queries to execute in Neo4j. Multiple actions can be executed concurrently. After each round, the Retriever also performs a **subgraph query** to retrieve all related edges between known entities, which are then sent to the frontend for rendering.
- **`Analyzer`**: Analyzes the data retrieved by the Retriever. Invalid data is discarded; valuable clues (e.g., downloading malicious payloads and finding their corresponding file hashes) are added to the `evidences` in the state.
At the same time, the Analyzer **is responsible for routing decisions**: If enough information has been gathered to form a complete picture, or if the set of iterations has reached its limit, then the process moves on to the Reporter. If new related clues are discovered (e.g., new internal IP addresses), the process jumps back to the Planner to start the next round. 3. **`Reporter` (Terminal)**: Based on all collected evidences, a detailed, human-readable Markdown report is generated. ### 3. How the attack graph is dynamically generated
The attack graph on the right isnβt loaded entirely from the beginning; it emerges gradually as the agent investigates. - Traditional query methods typically return discrete entities like βRETURN a, bβ. In this project, after each round of queries by the Retriever, all previously discovered node IDs are aggregated, and a subgraph query is sent to Neo4j: βMATCH (a)-[r]->(b) WHERE a.id IN $ids OR b.id IN $ids RETURN a, r, b LIMIT 200β. This ensures that whenever the agent discovers even a small piece of information, the nodes involved in that path can be automatically linked and highlighted in the graph. ---
## π Deployment and startup methods
### 1. Quick start (Recommended with Docker)
If you have already installed [Docker](https://www.docker.com/) and [Docker Compose](https://docs.docker.com/compose/), you can use the following commands for quick deployment:
1. **Configure environment variables**: Copy `.env.example` to `.env` and fill in your LLM API Key. ```bash
cp .env.example .env
# Edit the .env file and fill in OPENAI_API_KEY, etc.
```
2. **Start the system**:
```bash
docker compose up -d --build
```
This command will automatically start the Neo4j database and backend API services (including front-end static resources). 3. **Import simulated data**: Since the newly started Neo4j instance is empty, we need to import the built-in simulated attack scenario data:
```bash
# Execute directly on the host (requires neo4j dependencies) or via container:
docker exec -it security-agent-backend python seed_data.py
```
4. **Access the system**: Visit **http://localhost:8000** to start the investigation. ---
### 2. Local manual deployment (Source code execution)
#### Environment preparation
- **Python**: 3.10+
- **Database**: A running Neo4j instance (Recommended with Docker). #### Dependency installation
Itβs recommended to use a virtual environment:
```bash
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt
```
#### Start the service
```bash
source venv/bin/activate
python server.py
```
The default port is 8000 (can be changed in `server.py`). ---
## π οΈ Secondary development guidelines
If you wish to expand or customize this project, here are some common approaches:
### 1. Adapting to different graph database schemas
The systemβs Cypher generation heavily relies on the Database Schema provided by the LLM. The system automatically calls `neo4j_client.getSchema()`. If your graph contains special naming conventions (e.g., multiple labels, special attribute names), you can enhance the prompts for large models in `agents/retriever.py`, or add hardcoded few-shot examples in `tools/cypher_templates.py`. ### 2. Adding new agents
If you need to add external threat intelligence queries (e.g., VirusTotal):
1. Create a new agent named `ti_lookup.py` under `agents/`. 2. Use Langchain to call the external API to obtain hash-based reputation information. 3. Package this as a workflow node in `graph/nodes.py`. 4. Modify the routing logic in `graph/workflow.py`, for example, so that the Analyzer routes the request to the `ti_lookup` node first, before returning it to the Planner. ### 3. Customizing the front-end attack graph (D3.js)
The core logic of the front-end graph is in `frontend/attack_graph.js`. - **Customizing node colors/icons**: You can directly modify the `NODE_COLORS` and `NODE_ICONS` dictionaries in `tools/graph_utils.py`. These will be automatically applied to the front-end when the backend processes them. - **Adjusting force-directed parameters**: In the `init()` method of `attack_graph.js`, modify the collision radius, edge distance, and repulsive force to accommodate large-scale nodes. - **Expanding hover tooltips**: Currently, hover events capture up to 8 attributes. This can be controlled in the `showTooltipNode()` function using `slice(0, 8)`. If your business nodes contain very long texts, you may need to exclude certain fields from the hover tooltip. ### 4. Enabling/configuring `Deep Thinking` parameters
By default, the front-end UI provides a toggle switch for [Deep Thinking]. When enabled, requests sent to the underlying large model (e.g., DeepSeek r1) include a flag indicating whether deep thinking is enabled. This parameter is passed into the `InvestigationState` through the request body in `server.py`. All agent nodes can check whether this feature is enabled by accessing `state.get("deep_thinking")`. ---
> *βSecurity is a process, not a product.β β Bruce Schneier*