AI-driven-MITRE-Attack

POC : AI-driven MITRE Attack Detection and Alert Enrichment

This repository demonstrates a machine learning pipeline for detecting MITRE ATT&CK techniques from logs and enriching the output using a local LLM.

---

Overview

The project is divided into two main components:

1. Machine Learning Model (ML) – classifies logs to MITRE ATT&CK techniques
2. LLM Enrichment – enriches the ML prediction with analyst-friendly explanations and actionable insights

---

Machine Learning Model

Purpose

The ML model is trained to predict a MITRE technique (or BENIGN) from log events.
This allows automation of detection and categorization of potentially malicious behavior.

Approach

• Feature: commandline field from logs
• Vectorization: TF-IDF (Term Frequency – Inverse Document Frequency)
• Model: RandomForestClassifier (robust, interpretable, CPU-friendly)
• Target: mitre_label (e.g., T1059.001, T1105, BENIGN)

Training

python scripts/train_mitre_model.py

This script:

• Loads the dataset dataset_full_160k.csv

• Splits data into train/test sets

• Converts command lines to TF-IDF vectors

• Trains the Random Forest model

• Evaluates performance (precision, recall, F1-score, confusion matrix)

• Saves the trained model and vectorizer:

  • models/mitre_ml_model.pkl

  • models/tfidf_vectorizer.pkl

LLM Enrichment

Purpose

Once the ML model predicts a MITRE technique, the LLM enriches the result by providing:

• Technique explanation

• Why the command matches the technique

• Attacker intent

• Recommended investigation steps

• Suggested detection rules

This step bridges raw ML prediction and SOC analyst actionable insights.

Usage

Local LLM (e.g., Phi-3 via Ollama) is called with a prompt containing:

• ML prediction

• Raw command line

python scripts/enrich_with_llm.py

Python Dependencies

pip3 install -r requirements.txt

Librairies

• pandas
• scikit-learn
• requests

Notes

• test_model.py is used for testing your model, it is optional.
• The scripts are designed to be run locally, in a Python 3.13+ environment with the listed dependencies.
• Install Ollama on your machine and add the Phi-3 mini model.