NuclearOX/ai-malware-detection-pipeline

# AI Malware Detection Pipeline An end-to-end static malware analysis platform powered by a 3-agent AI pipeline. Upload a binary, get automatic PE analysis, MITRE ATT&CK mapping, YARA/Sigma rule generation, Ghidra decompilation, and a critical evaluation — all in one interface. ## Architecture Frontend (React + Vite) │ ▼ Backend (FastAPI) │ ├── pefile → PE structure, imports, sections, entropy ├── CAPA → Capability detection ├── FLOSS → String deobfuscation ├── Ghidra → Deep decompilation (optional, via pyhidra) ├── VirusTotal API → Hash reputation (70+ AV engines) │ └── Groq API (3-agent pipeline) ├── Agent 1: Triage (llama-4-scout-17b) ├── Agent 2: Threat Intel (llama-3.3-70b) └── Agent 3: Critic (qwen3-32b) **3-Agent Pipeline (Online mode):** | Agent | Model | Role | |-------|-------|------| | Triage | `meta-llama/llama-4-scout-17b-16e-instruct` | PE classification, capabilities extraction | | Threat Intel | `llama-3.3-70b-versatile` | MITRE mapping, YARA/Sigma generation | | Critic | `qwen/qwen3-32b` | Critical evaluation, confidence scores, behavioral timeline | ## Requirements - [Docker](https://www.docker.com/) and Docker Compose - A [Groq API key](https://console.groq.com/) — free, no credit card required - A [MalwareBazaar API key](https://bazaar.abuse.ch/) — free - A [VirusTotal API key](https://www.virustotal.com) — free (4 req/min, 500/day) - Optional: [Ollama](https://ollama.com/) for local model inference ## Setup ### 1. Clone the repository git clone cd ai-malware-detection-pipeline ### 2. Configure environment variables Create a .env file in backend/, using .env.example as reference: cp backend/.env.example backend/.env Edit `backend/.env` and fill in your API keys: GROQ_API_KEY=your_groq_api_key_here MALWAREBAZAAR_API_KEY=your_malwarebazaar_api_key_here VIRUSTOTAL_API_KEY=your_virustotal_api_key_here | Key | Where to get it | Required | |-----|----------------|----------| | `GROQ_API_KEY` | [console.groq.com](https://console.groq.com/) | Yes (online mode) | | `MALWAREBAZAAR_API_KEY` | [bazaar.abuse.ch](https://bazaar.abuse.ch/) | Yes (threat feed) | | `VIRUSTOTAL_API_KEY` | [virustotal.com](https://www.virustotal.com) | Yes (VT lookup) | ### 3. Start the full stack The first time, it is necessary to build the docker image: docker compose up --build The following times, just start the container (containing both backend and frontend): docker compose up Both backend and frontend start automatically: - Frontend: `http://localhost:5173` - Backend API: `http://localhost:8000` ## Usage ### Analyzing a sample 1. Go to the **[ ANALYSIS ]** tab 2. Select **⚡ Online (Groq)** or **💻 Local (Ollama)** 3. Optionally enable **🔬 Ghidra Deep Analysis** 4. Upload a binary (`.exe`, `.dll`, `.bin`, etc.) 5. Click **Start Analysis** The pipeline runs in sequence: | Step | Tool | Output | |------|------|--------| | 1 | pefile | PE headers, sections, imports, entropy | | 2 | CAPA | Functional capabilities | | 3 | FLOSS | Deobfuscated strings | | 4 | Ghidra | Decompiled suspicious functions (if enabled) | | 5 | VirusTotal | Hash reputation across 70+ AV engines | | 6 | Agent 1 / Triage | Malware family classification | | 7 | Agent 2 / Threat Intel | MITRE ATT&CK, YARA rule, Sigma rule | | 8 | Agent 3 / Critic | Confidence scores, behavioral timeline, critical evaluation | ### Debug mode Enable **🔧 Static Analysis Only** to skip the AI pipeline and show only static results (PE, CAPA, FLOSS, VT, Ghidra). Useful for testing tools without consuming API tokens. ### Threat Stream The **[ THREAT_STREAM ]** tab shows the latest samples from MalwareBazaar in real time. Click **WEB** to open a sample on MalwareBazaar, download it manually, then upload it in the Analysis tab. ### Local mode (Ollama) Make sure Ollama is running locally, then pull a model: ollama pull gemma3:1b # lightweight, fast ollama pull llama3.2 # better quality ollama pull mistral # good balance Select **💻 Local (Ollama)**, enter the model name, and run the analysis. ## Export After analysis, use the **JSON** and **PDF** buttons to export the full report: - **JSON** — structured data with all analysis fields, suitable for further processing - **PDF** — formatted report for documentation or presentation ## Project structure . ├── backend/ │ ├── ai/ │ │ └── ai_client.py # 3-agent Groq pipeline + Ollama support │ ├── analyzers/ │ │ ├── pe_analyzer.py # PE parsing, CAPA, FLOSS, Ghidra, entropy │ │ └── extract_suspicious_functions.py # (legacy, unused) │ ├── api/ │ │ ├── routes.py # FastAPI endpoints │ │ ├── malwarebazaar.py # MalwareBazaar API client │ │ └── virustotal.py # VirusTotal API client │ ├── main.py # FastAPI app entry point │ ├── Dockerfile # Ubuntu 24.04 + Java 21 + CAPA + FLOSS + Ghidra │ ├── requirements.txt │ ├── .env # Your API keys (not committed) │ └── .env.example # Template ├── frontend/ │ ├── src/ │ │ ├── components/ │ │ │ ├── AnalysisDashboard.tsx │ │ │ └── ThreatStream.tsx │ │ ├── services/ │ │ │ └── aiService.ts │ │ ├── types.ts │ │ └── App.tsx │ ├── index.html │ ├── package.json │ └── vite.config.ts ├── docker-compose.yml └── README.md ## Notes - **It is advisable to not upload real malware on a production machine.** Use an isolated VM (e.g. Kali Linux). - All uploaded files are automatically deleted from the server after analysis. - Ghidra analysis takes ~3 minutes per binary due to full disassembly and decompilation. - Online mode makes 3 sequential API calls to Groq — total AI time is typically 30–90 seconds. - The AI pipeline may produce false positives or hallucinations on legitimate binaries. Always review the **LLM Critical Self-Evaluation** section. - Groq free tier limits: 30 RPM, 6K–30K TPM depending on model. Rate limit errors are retried automatically (up to 3 times with backoff). ## Documentation - [LLM Critical Evaluation](docs/llm_limitations.md) — limits, hallucination risks, and mitigation strategies

标签：自动化攻击