Threat Assessment using ChatGPT and the MITRE ATT&CK framework

In this recipe, you will learn how to leverage ChatGPT and the OpenAI API to conduct a threat assessment by providing a threat, attack, or campaign name. By combining the power of ChatGPT with the MITRE ATT&CK framework, you will be able to generate detailed threat reports, tactics, techniques, and procedures (TTPs) mappings, and associated indicators of compromise (IoCs). This information will enable cybersecurity professionals to analyze attack vectors in their environment and extend their capabilities into threat hunting.

Building upon the skills acquired in Chapter 1, this recipe will guide you through establishing the system role of a cybersecurity analyst and engineering effective prompts that generate well-formatted output, including tables. You will learn how to design prompts to obtain the desired output from ChatGPT using both the ChatGPT web UI and a Python script. Additionally, you will learn how to use the OpenAI API to generate a comprehensive threat report in a Microsoft Word file format.

Getting ready

Before diving into the recipe, you should already have your OpenAI account set up and obtained your API key. If not, revisit Chapter 1 for details. You will also need to do the following:

1. Install the python-docx library: Ensure you have the python-docx library installed in your Python environment, as it will be used to generate Microsoft Word files. You can install it using the pip install python-docx command.
2. Familiarize yourself with the MITRE ATT&CK framework: To make the most of this recipe, it’s helpful to have a basic understanding of the MITRE ATT&CK framework. Visit https://attack.mitre.org/ for more information and resources.
3. List sample threats: Prepare a list of sample threat names, attack campaigns, or adversary groups to use as examples while working through the recipe.

How to do it…

By following these steps, you can successfully utilize ChatGPT to generate a TTP-based threat report using the MITRE ATT&CK framework and proper Markdown formatting. We will be specifying the name of a threat and applying prompt engineering techniques. ChatGPT will then generate a well-formatted report with valuable insights that can assist you in threat analysis, attack vector assessment, and even in gathering IoCs for threat hunting:

1. Begin by logging in to your ChatGPT account and navigating to the ChatGPT web UI.
2. Start a new conversation with ChatGPT by clicking the New chat button.
3. Enter the following prompt to establish a system role:

   You are a professional cyber threat analyst and MITRE ATT&CK Framework expert.

4. Replace {threat_name} in the user prompt below with the threat name of your choice (in our example, we will use WannaCry). You can either combine this prompt with the system role or enter it separately:

   Provide a detailed report about {threat_name}, using the following template (and proper markdown language formatting, headings, bold keywords, tables, etc.):
   Threat Name (Heading 1)
   Summary (Heading 2)
   Short executive summary
   Details (Heading 2)
   Description and details including history/background, discovery, characteristics and TTPs, known incidents
   MITRE ATT&CK TTPs (Heading 2)
   Table containing all of the known MITRE ATT&CK TTPs that the {threat_name} attack uses. Include the following columns: Tactic, Technique ID, Technique Name, Procedure (How WannaCry uses it)
   Indicators of Compromise (Heading 2)
   Table containing all of the known indicators of compromise. Include the following columns: Type, Value, Description

Hint

Just as with the previous recipe, you can perform this in the OpenAI Playground and use Chat mode to enter the role in the System window, and the prompt in the User message window.

Figure 2.3 shows the system role and user prompt entered into the OpenAI Playground:

Figure 2.3 – OpenAI Playground method

5. After entering the appropriate system role and user prompt, press Enter.
6. ChatGPT will process the prompt and generate a formatted threat report with Markdown language formatting, headings, bold keywords, tables, and other elements specified in the prompt.

   Figure 2.4 and Figure 2.5 illustrate an example ChatGPT generation of a threat report using Markdown language formatting with a table:

Figure 2.4 – ChatGPT threat report narrative output

Figure 2.5 – ChatGPT threat report table output

7. Review the generated report to ensure it contains the desired information and formatting. If necessary, adjust your user prompt and resubmit it to improve the output.

Hint

Sometimes, ChatGPT will stop generating before it has completed the entire out. This is due to the token limit of the model being used. In such cases, you can click on the Continue Generating button.

How it works…

Just as we did in the Applying ChatGPT Roles (Application: AI CISO) recipe in Chapter 1, when you assign a role to ChatGPT, you provide a specific context or persona for the model to work with. This helps the model generate responses that are tailored to the given role, resulting in more accurate, relevant, and detailed content. The model will generate content that aligns with the expertise and perspective of the assigned role, offering better insights, opinions, or recommendations.

When we provide a threat name and direct ChatGPT to reference the MITRE ATT&CK framework, we are able to leverage its massive dataset, which includes detailed information about threats and the MITRE ATT&CK framework. As a result, it is able to correlate the two and quickly give us the relevant threat information as it pertains to the TTPs identified in the framework.

Important note

When using the current version of ChatGPT and the OpenAI API as of the time of this writing, the dataset is only trained up through September 2021. Therefore, it will not have knowledge of any threat data after that. However, we will cover techniques later in this book on how to use the API and Python to feed recent data into the request.

By providing a clear template for the output in your prompt, you guide ChatGPT to generate responses that adhere to the specified structure and formatting. This helps ensure that the generated content is consistent, well organized, and suitable for use in reports, presentations, or other formal documents. The model will focus on generating content that matches the formatting and structure you’ve provided while still delivering the information you requested. See the Enhancing Output with Templates (Application: Threat Report) and Formatting Output as a Table (Application: Security Controls Table) recipes in Chapter 1 for further details.

There’s more…

You can extend the power and flexibility of this recipe by using the OpenAI API with a Python script to generate a threat report, similar to the one created in the ChatGPT web UI. Here’s how you do it:

1. Start by importing the necessary libraries:

   import openai
   from openai import OpenAI
   import os
   from docx import Document
   from tqdm import tqdm
   import threading
   import time

2. Set up the OpenAI API the same as we did in the Setting the OpenAI API key as an Environment Variable recipe in Chapter 1:

   openai.api_key = os.getenv("OPENAI_API_KEY")

3. Create a function to generate a report using the OpenAI API:

   def generate_report(threat_name: str) -> str:
       ...
       return response['choices'][0]['message']['content'].strip()

   This function takes a threat name as input and sends it as part of a prompt to the OpenAI API. It returns the generated text from the API response.

4. Create a function to convert the generated text, which is in Markdown format, to a Microsoft Word document:

   def markdown_to_docx(markdown_text: str, output_file: str):
       ...
       document.save(output_file)

   This function takes the generated text in Markdown format and an output filename. It parses the Markdown text and creates a Word document with the appropriate formatting.

5. Create a function to extract tables from the Markdown text:

   def extract_tables(markdown_text: str):
       ...
       return tables

   This function iterates through the Markdown text and extracts any tables it finds.

6. Create a function to display the elapsed time while waiting for the API call:

   def display_elapsed_time():
       ...

   This function shows the elapsed time in seconds while waiting for the API call to complete.

7. Get the threat name from user input:

   threat_name = input("Enter the name of a cyber threat: ")

8. Start a separate thread to display the elapsed time while making the API call:

   api_call_completed = False
   elapsed_time_thread = threading.Thread(target=display_elapsed_time)
   elapsed_time_thread.start()

9. Make the API call and handle exceptions:

   try:
       report = generate_report(threat_name)
       api_call_completed = True
       elapsed_time_thread.join()
   except Exception as e:
       ...

10. Save the generated report as a Word document:

    docx_output_file = f"{threat_name}_report.docx"

11. Generate the report and handle exceptions:

    try:
        with tqdm(total=1, desc="Generating report and files") as pbar:
            markdown_to_docx(report, docx_output_file)
        print("\nReport and tables generated successfully!")
    except Exception as e:
        ...

Here is how the completed script should look:

import openai
from openai import OpenAI
import os
from docx import Document
from tqdm import tqdm
import threading
import time
# Set up the OpenAI API
openai.api_key = os.getenv("OPENAI_API_KEY")
# Function to generate a report using the OpenAI API
def generate_report(threat_name: str) -> str:
    # Define the conversation messages
    messages = [
        {"role": "system", "content": "You are a professional cyber threat analyst and MITRE ATT&CK Framework expert."},
        {"role": "user", "content": f'Provide a detailed report about {threat_name}, using the following template (and proper markdown language formatting, headings, bold keywords, tables, etc.):\n\n\
        Threat Name (Heading 1)\n\n\
        Summary (Heading 2)\n\
        Short executive summary\n\n\
        Details (Heading 2)\n\
        Description and details including history/background, discovery, characteristics and TTPs, known incidents\n\n\
        MITRE ATT&CK TTPs (Heading 2)\n\
        Table containing all of the known MITRE ATT&CK TTPs that the {threat_name} attack uses. Include the following columns: Tactic, Technique ID, Technique Name, Procedure (How {threat_name} uses it)\n\n\
        Indicators of Compromise (Heading 2)\n\
        Table containing all of the known indicators of compromise. Include the following collumns: Type, Value, Description\n\n\  '}
    ]
    client = OpenAI()
    # Call the OpenAI API
    response = client.chat.completions.create
        model="gpt-3.5-turbo",
        messages=messages,
        max_tokens=2048,
        n=1,
        stop=None,
        temperature=0.7,
    )
    # Return the generated text
    return response.choices[0].message.content.strip()
# Function to convert markdown text to a Word document
def markdown_to_docx(markdown_text: str, output_file: str):
    document = Document()
    # Variables to keep track of the current table
    table = None
    in_table = False
    # Iterate through the lines of the markdown text
    for line in markdown_text.split('\n'):
        # Add headings based on the markdown heading levels
        if line.startswith('# '):
            document.add_heading(line[2:], level=1)
        elif line.startswith('## '):
            document.add_heading(line[3:], level=2)
        elif line.startswith('### '):
            document.add_heading(line[4:], level=3)
        elif line.startswith('#### '):
            document.add_heading(line[5:], level=4)
        # Handle tables in the markdown text
        elif line.startswith('|'):
            row = [cell.strip() for cell in line.split('|')[1:-1]]
            if not in_table:
                in_table = True
                table = document.add_table(rows=1, cols=len(row), style='Table Grid')
                for i, cell in enumerate(row):
                    table.cell(0, i).text = cell
            else:
                if len(row) != len(table.columns):  # If row length doesn't match table, it's a separator
                    continue
                new_row = table.add_row()
                for i, cell in enumerate(row):
                    new_row.cells[i].text = cell
        # Add paragraphs for other text
        else:
            if in_table:
                in_table = False
                table = None
            document.add_paragraph(line)
    # Save the Word document
    document.save(output_file)
# Function to extract tables from the markdown text
def extract_tables(markdown_text: str):
    tables = []
    current_table = []
    # Iterate through the lines of the markdown text
    for line in markdown_text.split('\n'):
        # Check if the line is part of a table
        if line.startswith('|'):
            current_table.append(line)
        # If the table ends, save it to the tables list
        elif current_table:
            tables.append('\n'.join(current_table))
            current_table = []
    return tables
# Function to display elapsed time while waiting for the API call
def display_elapsed_time():
    start_time = time.time()
    while not api_call_completed:
        elapsed_time = time.time() - start_time
        print(f"\rCommunicating with the API - Elapsed time: {elapsed_time:.2f} seconds", end="")
        time.sleep(1)
# Get user input
threat_name = input("Enter the name of a cyber threat: ")
api_call_completed = False
elapsed_time_thread = threading.Thread(target=display_elapsed_time)
elapsed_time_thread.start()
# Handle exceptions during the API call
try:
    # Generate the report using the OpenAI API
    report = generate_report(threat_name)
    api_call_completed = True
    elapsed_time_thread.join()
except Exception as e:
    api_call_completed = True
    elapsed_time_thread.join()
    print(f"\nAn error occurred during the API call: {e}")
    exit()
# Save the report as a Word document
docx_output_file = f"{threat_name}_report.docx"
# Handle exceptions during the report generation
try:
    with tqdm(total=1, desc="Generating report and files") as pbar:
        markdown_to_docx(report, docx_output_file)
    print("\nReport and tables generated successfully!")
except Exception as e:
    print(f"\nAn error occurred during the report generation: {e}")

This script uses the OpenAI API to generate a cyber threat report as a Microsoft Word document.

The crux of this script lies in several key functions. The first function, generate_report(), takes in a cyber threat name and uses it as a prompt for the OpenAI API. It returns the generated text from the API response. This text is in Markdown format and is subsequently transformed into a Microsoft Word document by the markdown_to_docx() function.

This function parses through the Markdown text line by line, creating tables and headings as required, and finally saves it as a Word document. In parallel, there is an extract_tables() function that is designed to locate and extract any tables present within the Markdown text.

To enhance the user experience, the display_elapsed_time() function is incorporated. This function tracks and displays the time taken for the API call to complete. It runs in a separate thread, initiated before making the API call:

Figure 2.6 – Example output of the display_elapsed_time function

The API call itself, as well as the report generation, are wrapped in try-except blocks to handle any potential exceptions. Once the report is generated, it is saved as a Word document, with the filename based on the user-inputted cyber threat name.

Upon successful execution of this script, a detailed threat report in Word document format is produced, mimicking the output generated by the ChatGPT web UI. This recipe demonstrates how the OpenAI API can be adapted within a Python script to automate the generation of comprehensive reports.

Hint

You can swap out the chat-3.5-turbo model with the GPT-4 model, if you are a ChatGPT Plus subscriber, for often improved results. Just keep in mind that the GPT-4 model is a bit more expensive than the chat-3.5-turbo model.

You can also improve accuracy and get a more consistent output by lowering the temperature value.