POC : Detección de Tácticas MITRE y Enriquecimiento de Alertas Impulsado por IA

Este repositorio demuestra un pipeline de aprendizaje automático para detectar técnicas MITRE ATT&CK a partir de registros y enriquecer la salida usando un LLM local.

Descripción general

El proyecto está dividido en dos componentes principales:

1. Modelo de Aprendizaje Automático (ML) – clasifica registros en técnicas MITRE ATT&CK
2. Enriquecimiento con LLM – enriquece la predicción del ML con explicaciones amigables para analistas e información procesable

Modelo de Aprendizaje Automático

Propósito

El modelo ML se entrena para predecir una técnica MITRE (o BENIGN) a partir de eventos de registro.
Esto permite automatizar la detección y categorización de comportamientos potencialmente maliciosos.

Enfoque

• Característica: campo commandline de los registros
• Vectorización: TF-IDF (Frecuencia de Término – Frecuencia Inversa de Documento)
• Modelo: RandomForestClassifier (robusto, interpretable, amigable con la CPU)
• Objetivo: mitre_label (ej., T1059.001, T1105, BENIGN)

Entrenamiento

python scripts/train_mitre_model.py

Este script:

• Carga el conjunto de datos dataset_full_160k.csv
• Divide los datos en conjuntos de entrenamiento/prueba
• Convierte líneas de comandos en vectores TF-IDF
• Entrena el modelo Random Forest
• Evalúa el rendimiento (precisión, exhaustividad, puntuación F1, matriz de confusión)
• Guarda el modelo entrenado y el vectorizador:
  • models/mitre_ml_model.pkl
  • models/tfidf_vectorizer.pkl

Enriquecimiento con LLM

Propósito

Una vez que el modelo ML predice una técnica MITRE, el LLM enriquece el resultado proporcionando:

• Explicación de la técnica
• Por qué el comando coincide con la técnica
• Intención del atacante
• Pasos de investigación recomendados
• Reglas de detección sugeridas

Este paso une la predicción bruta del ML con información procesable para el analista del SOC.

Uso

Se llama a un LLM local (ej., Phi-3 vía Ollama) con un prompt que contiene:

• Predicción del ML
• Línea de comandos original

python scripts/enrich_with_llm.py

Dependencias de Python

pip3 install -r requirements.txt

Librerías

• pandas
• scikit-learn
• requests

Notas

• test_model.py se usa para probar tu modelo, es opcional.
• Los scripts están diseñados para ejecutarse localmente, en un entorno Python 3.13+ con las dependencias listadas.
• Instala Ollama en tu máquina y añade el modelo Phi-3 mini.