Extract Procedural Details of the TTP based on the actor mentioned in OSINT reportings
import pickle
import requests
import keyring
from bs4 import BeautifulSoup
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain import PromptTemplate
from langchain.chains import RetrievalQAWithSourcesChain
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
import os
import json
import openai


# Keyring ChatGPT API token storage
service_id = 'Working'
user_name = 'ChatGPTAI'

#OpenAI API key
retrieved_key = keyring.get_password(service_id, user_name)

#langchain set up requirement to set up key like this
os.environ[f"OPENAI_API_KEY"] = retrieved_key

#calling LLM model
llm=ChatOpenAI(temperature=0.05,model_name='gpt-3.5-turbo-16k')
#==================================================Load the mitre Techniques JSON data from the file===============================================
input_filename = "mitre_techniques.json"
with open(input_filename, "r") as json_file:
    techniques_data = json.load(json_file)


#this is to query to T code entered from the Streamlit application
def queryTcode(tcode):
    # Input a Mitre T code
    input_t_code = tcode
    # Search for the provided T code in the data
    matching_technique = None
    for technique in techniques_data.get("techniques", []):
        if technique["Tcode"] == input_t_code:
            matching_technique = technique
            break


    description = matching_technique['description']
    Tname = matching_technique['name']
    Tcode = matching_technique['Tcode']
    # Define the prompt -> set to singular question for testing
    prompt_text = f'''generate a detailed question based on description below that can answer the Procedural level of the TTPs. 
                    output must:contain only the question and nothing else. don't contain T code or mitre. must include the name of the attack: {Tname}. No mentions of TTPs or procedural leve.
                    output must be: "how did the threat actor/malware specifically use" 
                    Mitre: {Tcode}
                    description:{description}'''

    output = llm.predict(prompt_text)
    print(output)
    return output


# Getting Dection rule being written
def DetectRuleCreate(llm,RuleType,tcode,procedure):
    # Hv to specificy Detection rule type such as KQL,Splunk, Elastic Search, YARA, SIGMA
    # Ex Detction Ruling=KQL
    DetectionRuling = str(RuleType)
    # Search for the provided T code in the data
    matching_technique = None
    for technique in techniques_data.get("techniques", []):
        if technique["Tcode"] == tcode:
            matching_technique = technique
            break
    if matching_technique is None:
        return "Technique not found"
    mitigation=matching_technique['mitigation']
    description = matching_technique['description']
    Tname = matching_technique['name']
    Tcode = matching_technique['Tcode']
    ptext=f'''write me {DetectionRuling} detection rule based on the procedure. Here all the necessary data. Output must be accurate to the MITRE Technique presented. only give me the rule in the detection rule format and nothing else included.
    Mitre T code:{Tcode}
    T code name: {Tname}
    Description: {description}
    Mitigations:{mitigation}
    Procedure of the actor: {procedure}
'''
    output = llm.predict(ptext)
    #print(output)
    return output



#==================================================================Set Attributes======================================

class Document:
    def __init__(self, page_content, metadata):
        self.page_content = page_content
        self.metadata = metadata

    def __repr__(self):
        return f"Document(page_content='{self.page_content}', metadata={self.metadata})"


#--------------------------------------------- LANGCHAIN HTML PARSING LOGIC -----------------------------------------
#call langchain variables
#what langchain.document_loaders do -> Document Loaders are responsible for loading documents into the LangChain system. They handle various types of documents, including PDFs, and convert them into a format that can be processed by the LangChain system.

#text splitter splits the text to smaller bit -> will need to tune seperator in future for errors
def get_text_chunks_langchain(text,link,title):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=13000, chunk_overlap=4400, separators=["\\n"])
    texts = text_splitter.split_text(text)
    docs = [Document(page_content=t, metadata={'source': link , 'title': title}) for t in texts]
    return docs

#PFIx this later
#getting the mitre TTPs when u put in the T code. use emdedding to extract the question via another QA retrieval
def QARetrieval(llm, VectorStore,tcode):
    for technique in techniques_data.get("techniques", []):
        if technique["Tcode"] == tcode:
            matching_technique = technique
            break
    if matching_technique is None:
        return "Technique not found"
    Tname = matching_technique['name']
    # chatGPT API Prompt template
    template=f'''{queryTcode(tcode)} output requirement: 1.put into a paragraph which bullet point specific technical commands,codes,or indicators that is related to the question and remove anything not relating to {Tname} and must only start with *threat group name or malware name only *used/exploited/orestablished*'''
    answer = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=VectorStore.as_retriever(search_type="mmr"))
    chain = answer.run(template)
    return chain


Backend.py
hokman0414/MITREATTACKProceduralGPT
import requests
from bs4 import BeautifulSoup

def ScrapeSite(link):
    payload = {'api_key': 'ENTER SCRAPERAPI KEY', 'url': link}
    r = requests.get('http://api.scraperapi.com', params=payload)
    if r.status_code == 200:
        ScrappedData = r.text
        return ScrappedData
    else:
        ScrappedData = 'Was Not Sucessful'
        return ScrappedData

def extract_text_from(html):
    soup = BeautifulSoup(html, features="html.parser")
    text = soup.get_text()
    title_tag = soup.find('title').text
    lines = (line.strip() for line in text.splitlines())
    pagecontent = '\\n'.join(line for line in lines if line)

    return {'title': title_tag,'pagecontent':pagecontent}



LinkScrapeData.py
hokman0414/MITREATTACKProceduralGPT
import streamlit as st
import Backend
from LinkScrapeData import ScrapeSite,extract_text_from
from langchain.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings
import os
import keyring
import json
import pickle
from Backend import llm
from langchain import PromptTemplate

#vector store is contained within Main.py and not backend.py

#------------------------------------------STREAMLIT APPLICATION------------------------------------------
# Set page title and icon
st.set_page_config(page_title="Mitre-Tcode ProcedureGPT App", page_icon="🛡️")

# Add a cool title header
st.title("🛡️ Welcome to the Mitre-Tcode ProcedureGPT App 🛡️")

# Add some description text
st.write("This app allows you to explore Mitre-Tcodes procedures using the power of GPT-3.5.")

# Input for article links as a text box
st.write("Enter article link")
article_links_text = st.text_area("Article Link")

# Input for Mitre T-code varaible
t_code = st.text_input("Enter a Mitre T-code:")

# Initialize a list to store entered article links
article_links = []

# Check if both article links and T-code are empty
if not article_links_text.strip() and not t_code.strip():
    st.warning("Please enter at least one article link and Mitre T-code.")
else:
    # Submit button to save the entered data into variables
    if st.button("Submit"):
        # Split the article links by line breaks and save them into the list
        article_links = article_links_text.split("\n")

        # Print entered article links
        st.subheader("Procedure TTP")
        #Article link Dictionary and it's HTML codes
        LinkDataSet = {link: extract_text_from(ScrapeSite(link)) for link in article_links}

        #format dataset to fit lanchain and chunkdata for embedding
        Total_chuncked_data = []
        for link, info in LinkDataSet.items():
            title = info['title']
            pagecontent = info['pagecontent']
            #calling function to format Langchain Data
            formatted_langchainChunkingData=Backend.get_text_chunks_langchain(pagecontent,link,title)
            Total_chuncked_data.append(formatted_langchainChunkingData)
        #stat the embedding process by ingesting all the data
        # --------------------storing embeddings in vectorStore--------------------------------
        embeddings =OpenAIEmbeddings()
        vectorStore_openAI = FAISS.from_documents(Total_chuncked_data[0], embeddings)
        with open("faiss_store_openai.pkl", 'wb') as f:
            pickle.dump(vectorStore_openAI, f)
        with open("faiss_store_openai.pkl", 'rb') as f:
            VectorStore = pickle.load(f)
        #questions asking
        # Prompt_template.format( Tcode="T1566")


        #question='What are the various methods and tactics used by adversaries in phishing attacks, including targeting individuals or organizations, using malicious attachments or links, and employing social engineering techniques?'
        answer = Backend.QARetrieval(llm,VectorStore,t_code.upper())
        st.write(answer)

#---------------------You can add more cool features and components here!---------------------------------------------------

        #st.subheader("Detection Rule")
        #Write Detection rule -> enalbe optional by uncommenting
        #RuleType = 'YARA'
        #DetectRule = Backend.DetectRuleCreate(llm, RuleType, t_code.upper(),answer)
        #st.code(DetectRule,language='python')

main.py
hokman0414/MITREATTACKProceduralGPT