seprate_stage.py

from f1_calculation import calculate_score, read_unstructured_csv, check_filepath
import pandas as pd


def get_subtechniques_with_apt(mapping_dictionary):
    """
    mapping_dictionary: dictionary contian the mapping of sub-techniques, stages and apt list
    {
        sub-techniques: {
            stages: apt_list
        }
    }
    
    return: dictionary that contian mapping of stages and apt list
    sub-techniques: {
            stages: apt_list
        }
        
    """
    subtechniques_apt_dictionary_list = []
    for idx, dixt in mapping_dictionary.items():
        subtechniques_apt_dictionary_list.append(dixt)
    return subtechniques_apt_dictionary_list


def get_stage_against_apt(df, sub_technique):
    """
    df:  stages file dataframe
    
    subtechnique: name the of the subtechnique that needs to be found in the stages file
    
    return: list of all stages against the APTS
    """
    # stages will store all stages found 
    stages = []

    try:
        columns = df.columns.tolist()
        for column in columns:
            columns_input = df[column].tolist()
            if sub_technique in columns_input:
                stages.append(column)
    except Exception as e:
        print("sub technique " + sub_technique + " doesn't exist")
    return stages


def get_stage_for_top_technique(top_techniques, stage_with_apt_dictionary):
    current_attack_stage = ""
    compared_apt = list(top_techniques.keys())[0]
    for stage, apts in stage_with_apt_dictionary.items():
        if compared_apt in apts:
            current_attack_stage = stage
    return current_attack_stage

def find_stage(input_filename, stages_file, database_filename):
    """
    input_filename:  is the filename that generated by generate_report.py
    database_filename: is the file name that contain database 
    return: top_techniques, top_stage
    """

    file_input = read_unstructured_csv(input_filename)
    subtechniques_apt_list = get_subtechniques_with_apt(file_input)
    df = pd.read_excel(stages_file)
    stage_with_apt_dictionary = {}

    # iterate over the subtechniques_apt_list
    for subtechnique_apt in subtechniques_apt_list:
        for subtechnique, apts in subtechnique_apt.items():
            all_stages = get_stage_against_apt(df, subtechnique)
            if len(all_stages) > 0:
                print(subtechnique, all_stages)
                stage_with_apt_dictionary[all_stages[-1]] = apts
            else:
                print("No stage found for subtechnique: " + subtechnique)

    # top_techniques contain top 3 techniques
    top_techinques = calculate_score(database_filename=database_filename, sheet_name=input_filename)
    top_stage = get_stage_for_top_technique(top_techniques=top_techinques,
                                            stage_with_apt_dictionary=stage_with_apt_dictionary)

    return (top_techinques, top_stage)

def find_all_stages(stage_filename):
    df = pd.read_excel(stage_filename)
    return df.columns.to_list()