MITRE ATT&CK AI Agent

Warning

USE WITH CAUTION: This tool utilizes a Python REPL agent that has the capability to execute arbitrary code on your machine. It is highly recommended to run this tool within a sandboxed environment (like a Docker container) or a virtual environment to prevent unintended system modifications.

This is a CLI tool powered by LangChain and Pandas that allows you to interact with MITRE ATT&CK data using natural language queries. It uses an OpenAI-compatible API (like LM Studio or OpenAI) to process your requests and analyze the data.

Prerequisites

• Python 3.10+
• An OpenAI-compatible API server running (e.g., LM Studio, LocalAI, or OpenAI API)

Installation

1. Clone the repository or navigate to the project directory.

2. It is highly recommended to use a virtual environment. We recommend using uv for fast dependency management.

   Using uv (Recommended):

   # Install uv if you haven't already
   pip install uv
   
   # Create a virtual environment
   uv venv
   
   # Activate the virtual environment
   # On macOS/Linux:
   source .venv/bin/activate
   # On Windows:
   # .venv\Scripts\activate
   
   # Install dependencies
   uv pip install -r requirements.txt

   Using standard pip:

   python -m venv venv
   source venv/bin/activate  # or venv\Scripts\activate on Windows
   pip install -r requirements.txt

Configuration

Create a .env file in the root directory with your API configuration. You can copy the example below:

OPENAI_API_BASE=http://localhost:1234/v1
OPENAI_API_KEY=lm-studio

Alternatively, you can pass these as command-line arguments.

Usage

Run the agent interactively:

python langchain_mitre_agent.py

Run a specific query immediately:

python langchain_mitre_agent.py "How many techniques are there in the dataset?"

Command Line Arguments

• query: Optional query to run immediately.
• --urls: List of data URLs to load (defaults to MITRE Enterprise ATT&CK JSON).
• --base-url: OpenAI API Base URL (default: env var OPENAI_API_BASE or http://localhost:1234/v1).
• --api-key: OpenAI API Key (default: env var OPENAI_API_KEY or lm-studio).
• --model: Model name to use (optional, will prompt if not set).

Features

• Interactive Mode: Chat with the agent about the data.
• Data Loading: Automatically downloads and processes MITRE ATT&CK data.
• Python REPL: The agent can execute Python code to perform complex data analysis on the loaded Pandas DataFrame.
• Rich Output: Uses the rich library for beautiful terminal output.