MitreMap - Infer MITRE technique from Threat Intel Data.ipynb

Kernel: Python 3.8.0 ('venv_conda': venv)

MitreMap - Infer MITRE technique from Threat Intel Data

Notebook Version: 1.0
Notebook Author: Vani Asawa

Python Version: >=Python 3.8
Platforms Supported: Azure Machine Learning Notebooks

Data Source Required: None

GPU Compute Required: No
GPU Compute Recommended: Yes

Requirements Path: ../mitremap-notebook/requirements.txt
Packages Downloaded:

• ipywidgets==7.5.1

• transformers==4.5.1

• torch==1.10.2

• msticpy==2.1.2

• nltk==3.6.2

• iocextract==1.13.1

• shap==0.41.0

Overview

This notebook allows a user to map descriptive text of an incident on to relevant MITRE ATT&CK Enterprise techniques. It uses a GPT2 language model to associate terms in the description with similar descriptions in past incidents. It also extracts relevant Indicators of Compromise from the text.

You can use the notebook with one of several pre-trained models or train your own model using your own threat reports or public sources.

Motivation

Please refer to Motivation and Goals to learn more.

Prerequisites

Please do not run the notebook cells all at once. The cells need to be run sequentially and successfully executed before proceeding with the remainder of the notebook.

Table of Contents

0. Installations [One-Time Setup]

1. Imports

2. Configure Input Data and Model Parameters

3. Run

4. Results

0. Installations [One-Time Setup]

Please refer to One-Time Setup to configure the virtual environment, install the required packages, and download the model artifacts.

Use the Powershell or BASH script below to download the model artifacts

Estimated Time to download the model artifacts - 5-10 minutes

Option 1: Powershell

!PowerShell ./model.ps1

Option 2: BASH

%%bash
./model.sh

Download the utils whl to use the inference packages

%pip install utils-1.0-py3-none-any.whl

Re-start the kernel and run the Notebook from 1. Imports.

1. Imports

The modules used to run this notebook can be found under mitremap-notebook/utils/*

import os
import sys
sys.path.append(os.getcwd())

import utils
from utils import main, inference, configs

2. Configure Input Data and Model Parameters,

Please refer to Input Parameters to learn more about setting the input parameter configurations.

Start using the notebook with one of the threat intel examples in the markdown script 😊

config_widgets = configs.configure_model_parameters()
for k in config_widgets.keys():
       display(config_widgets[k])

3. Run

Time to run the main.go function depends on the -

1. Length of the Threat Intel Report, and

2. If Model Explainability is set to True

For our sample threat reports in the markdown script, you can expect -

• < 1 minute without model explainability, and

• 1-2 minutes with model explainability.

configs, inference_df, iocs_df = main.go(
    config_widgets
)

4. Results

• configs: Stores the input configurations set by the customer

• inference_df: Stores the inference results for the threat intel data

• iocs_df: Stores the IOCs extracted from the threat intel data.

Use the inference.print_detailed_report(inference_df, configs) to obtain a printed summary of the MITRE technique predictions.

inference.print_detailed_report(
    inference_df,
    configs
)

print('Summary Statistics for Inference Dataframe: ')
print('Shape of Inference Dataframe: ', inference_df.shape)
if not inference_df.empty:
    print('Sample rows: ')
    display(inference_df.head(5))
else:
    print('No results obtained.')

print('Summary Statistics for IOCs Dataframe: ')
print('Shape of IOCs Dataframe: ', iocs_df.shape)
if not iocs_df.empty:
    print('Distinct counts for each category of IOCs: ')
    display(iocs_df.groupby('IOC_Type').count().rename(columns={'IOC_Value': 'Count'}))
    print('Sample rows: ')
    display(iocs_df.head(5))
else:
    print('No IOCs obtained.')