AI Malware Detector

Live demo: malware-detector-production.up.railway.app

AI-powered PE malware detection — upload a file and get instant classification.

Created by Jainam Maru

Table of Contents

• Features
• Quick Start
• API
• Running Tests
• Architecture
• Contributing
• License

Features

• PE Analysis — parses PE headers, sections, imports, DLLs, hashes
• ML Classification — LightGBM model trained on synthetic PE features with SHAP explanations
• FastAPI Backend — REST API with /api/scan, /api/history, /api/dashboard, /api/report/{id}
• Web UI — dark-themed drag-and-drop React frontend
• SQLite Persistence — scan history survives restarts (capped at 1000 records)

Quick Start

pip install -r requirements.txt
python main.py

Open http://localhost:8000 export BENIGN_DIRS="/path/to/pe/files,/another/dir" python model/train.py

Without `BENIGN_DIRS`, training exits with a clear error message.

### Docker

```bash
docker build -t malware-detector .
docker run -p 8000:8000 malware-detector

API

Endpoint	Method	Description
/api/health	GET	Service status
/api/scan	POST	Upload and scan a PE file
/api/history	GET	Paginated scan history
/api/report/{id}	GET	Detailed scan report
/api/dashboard	GET	Aggregate scan statistics

Example: Scan a file

curl -X POST https://malware-detector-production.up.railway.app/api/scan \
  -F "file=@sample.exe"

Response Format

Successful scan returns:

{
  "filename": "sample.exe",
  "malicious": true,
  "confidence": 0.97,
  "threat_score": 87,
  "explanation": {
    "top_features": [
      {"name": "suspicious_section_count", "value": 8, "contribution": 0.32},
      {"name": "high_entropy_sections", "value": 5, "contribution": 0.28}
    ]
  },
  "sha256": "a1b2c3d4...",
  "timestamp": "2026-07-29T12:00:00Z"
}

API Reference

Endpoint	Method	Description	Query Params
/api/health	GET	Service health check	—
/api/scan	POST	Upload and scan a PE file	file (multipart)
/api/history	GET	Paginated scan history	page, limit
/api/report/{id}	GET	Detailed scan report by ID	—
/api/dashboard	GET	Aggregate statistics	—

Authentication: The API is publicly accessible. Rate limiting is applied per IP (100 req/min).

Errors: All endpoints return {"detail": "error message"} with appropriate HTTP status codes (400, 404, 422, 500).

Running Tests

pip install -r requirements.txt
python -m pytest tests/ -v

Architecture

malware-detector/
├── api/              # FastAPI REST endpoints
│   └── routes.py
├── ml/               # ML model training + prediction
│   ├── predict.py
│   ├── train.py
│   └── synthetic_data.py
├── extractor/        # PE feature extraction
│   └── pe_extractor.py
├── models/           # Trained model storage
├── frontend/         # React CRA web UI
├── tests/            # pytest tests with generated PE fixtures
└── Dockerfile        # Docker build for Railway deployment