LightML Module Development

Lightml Module Development

The lightml module is a subset of the OffStrike EDR platform, it uses a Python-based ML model to classify security log messages into MITRE ATT&CK framework categories.

Overview

The script processes security logs data, performs text pre-processing, trains a machine learning model (Multinomial Naive Bayes), and predicts the corresponding MITRE ATT&CK techniques for each log entry. It also leverages the mitreattack-python library to access and display information from the MITRE ATT&CK framework.

Features

• Log Data Processing: Reads and processes security log messages.
• Text Preprocessing:
  • Tokenization (word splitting)
  • Stop word removal
  • Punctuation removal
  • Lowercasing
• Machine Learning Classification: Uses TF-IDF vectorization and Multinomial Naive Bayes to classify log messages.
• MITRE ATT&CK Integration: Uses the mitreattack-python library to retrieve and display MITRE ATT&CK technique information.
• Clear Ouput: Provides a classification report and a DataFrame with predicted MITRE ATT&CK techniques.

Project

---

lightml/
│
├── lightml.py             # Main script to train the model and make predictions
├── requirements.txt       # List of dependencies required for the project
└── README.md              # Project documentation

Requirements

• Python 3.x
• pandas
• scikit-learn
• nltk
• mitreattack-python
• numpy

Installation

To set up the project, follow these steps:

Clone the

---

git clone https://gitlab.offstrike.com/defender/lightml.git
cd lightml

pip install -r requirements.txt

Usage

1. Training the Model: The model is trained using a sample dataset of security logs and corresponding MITRE ATT&CK techniques.
2. Making Predictions: The model will predict the corresponding MITRE ATT&CK techniques for each security log entry.
3. Run the Script:

   python3 lightml.py
   

Model Training

The model is built using a Naive Bayes classifier within a scikit-learn pipeline. It leverages the following components:

• TF-IDF Vectorization: Converts text data into numerical format.
• Multinomial Naive Bayes: A classification algorithm suitable for text classification tasks.

The training data consists of log entries paired with their corresponding MITRE ATT&CK techniques. The model is trained using cross-validation and hyperparameter tuning to optimize performance.

Predicting MITRE Techniques

The function predict_mitre_technique(logs) takes a list of plain security logs, preprocesses them, and returns the predicted MITRE techniques. This allows users to quickly analyze new log entries without needing manual labeling.

Evaluation

During training, the model’s performance is evaluated using a classification report, which provides metrics such as precision, recall, and F1-score. This helps assess how well the model generalizes to unseen data.

Current Performance

Below is the output of how the script currently performs on a very small training dataset. The current performance is very promising, with potential to get better and handle different techniques within the MITRE ATT&CK framework than what is currently in the test sample.

ML Classification Report:
               precision    recall  f1-score   support

   T1059.004       1.00      0.00      0.00         1
       T1078       0.50      1.00      0.67         1

    accuracy                           0.50         2
   macro avg       0.75      0.50      0.33         2
weighted avg       0.75      0.50      0.33         2

Log: Failed login attempt for user guest from 192.168.1.101 
Predicted MITRE Technique: T1078

Log: User admin logged in successfully 
Predicted MITRE Technique: T1078

Log: powershell command executed to list processes 
Predicted MITRE Technique: T1059.001

Log: New user created with administrative privileges 
Predicted MITRE Technique: T1078

Log: File downloaded from suspicious external IP 
Predicted MITRE Technique: T1071.001