MalwareBazaar_AMP_Transfer.py

from file_locations import all_simple_custom_detections_url, amp_main_file, logger, \
    main_mb_df_file, main_simple_custom_detections_url, mb_url, sha_output_files
import gc
import os
import pandas as pd
import requests
from requests.auth import HTTPBasicAuth
import time

# Create start time
start_time = time.time()

# Make df more reader friendly in 'Run' windows
pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)
pd.set_option('expand_frame_repr', False)

# Add MalwareBazaar API key below
api_key = ''

# Add AMP API Username & Password below
amp_api_username = ''
amp_api_pw = ''

tags = [
    'AgentTesla',
    'AZORult',
    'Babuk',
    'BatLoader',
    'BlackBasta',
    'BlackCat',
    'BruteRatel',
    'BumbleBee',
    'CobaltStrike',
    'Emotet',
    'FormBook',
    'Hive',
    'IceID',
    'Lockbit',
    'Lokibot',
    'Luna',
    'Magniber',
    'Metasploit',
    'Meterpreter',
    'Mimikatz',
    'Nighthawk',
    'NJRat',
    'Qbot',
    'Raccoon',
    'RedLine',
    'Remcos',
    'Sliver',
    'SOCGholish',
    'Trickbot',
    'Ursnif',
    'Venus',
    'WannaCry'
]


def pull_from_malware_bazaar_api():
    # Loop through MalwareBazaar tags and normalize data
    for tag in tags:
        try:
            headers = {'API-KEY': api_key}
            data = {'query': 'get_taginfo',
                    'tag': tag,
                    'limit': '1000'}

            response = requests.post(mb_url, headers=headers, data=data)
            # Transform get request into json format.
            data = response.json()['data']

        except Exception as e:
            logger.error(f'Failed to pull data from Malware Bazaar API for {tag}. Most likely no results found.')
            logger.error(f'Exception is {e} for {tag}')

        try:
            # Create dataframe after normalizing
            df = pd.json_normalize(data)

            # If empty ignore it
            if df.shape[0] == 1:
                logger.warning(f'No matches found for {tag}!')

            # If not empty
            else:
                # Keep only SHA-256 hashes
                df = df['sha256_hash']

                # Create name of the SHA-256 output file on the fly
                sha_output_file = os.path.join(os.environ['USERPROFILE'], 'Desktop',
                                               'Python_Manipulated_Files', 'MalwareBazaar', tag + '.txt')
                # Create txt file on the fly
                df.to_csv(sha_output_file, index=False, header=0)
                logger.info(f'Successfully sent text file to folder. Total count is {df.shape[0]} for {tag}.')

        except Exception as e:
            logger.error(f'Failed to send text file to folder for {tag}')
            logger.error(f'Exception is {e} for {tag}')
            continue

    # Calculate how long script has taken until this point.
    logger.info(f'Code took {((time.time() - start_time) / 60):.3f} minutes to execute so far.')
    logger.info(f'------------------------------------------------------------------------------------------')


def create_master_df_from_malware_bazaar():
    # Delete main_df file if it exists
    try:
        os.remove(main_mb_df_file)
        logger.info(f'Successfully removed main_df_file!')

    except Exception as e:
        logger.warning(f'main_df file not found! Exception is {e}.')

    # Create main df
    main_df = pd.DataFrame()

    # Start loop to aggregate MalwareBazaar hashes
    for filename in os.listdir(sha_output_files):
        f = os.path.join(sha_output_files, filename)
        # Check if file
        if os.path.isfile(f):
            # This file should not exist yet. Thus ignore it if it does
            if filename != 'AA_Quarantine_&_Prevent_Install.txt':
                df = pd.read_csv(f, header=None)

                # Adding how many rows
                logger.info(f'{filename} is adding another {df.shape[0]} rows to main df.')

                # Append df to main_df
                main_df = main_df.append(df)

                # Delete df reference then free up memory via garbage collection
                del df
                gc.collect()

                # Main df has this many rows
                logger.info(f'{main_df.shape[0]} after appending {filename}.')

    # Send aggregated dataframe to directory
    main_df.to_csv(main_mb_df_file, index=False, header=0)
    logger.info(f'Successfully created AA_Quarantine_&_Prevent_Install in directory!')
    logger.info(f'Code took {((time.time() - start_time) / 60):.3f} minutes to execute so far.')
    logger.info(f'------------------------------------------------------------------')


def pull_main_hash_list():

    # Use the request library to pull data. Pass in id and key or you will receive 401 error
    request = requests.get(main_simple_custom_detections_url,
                           auth=HTTPBasicAuth(username=amp_api_username, password=amp_api_pw))

    # Pull data in json format. Slice data based off 'data' element
    json_response = request.json()
    data = request.json()['data']['items']

    # Create dataframe
    all_hashes_in_amp_df = pd.json_normalize(data)

    # Just pull the SHA256 column
    all_hashes_in_amp_df = all_hashes_in_amp_df['sha256']

    # Build base dataframe before breaking out data in items column
    while 'next' in json_response['metadata']['links']:
        next_url = json_response['metadata']['links']['next']
        request = requests.get(next_url, auth=HTTPBasicAuth(username=amp_api_username, password=amp_api_pw))
        json_response = request.json()
        data = request.json()['data']['items']
        df_append = pd.json_normalize(data)
        df_append = df_append['sha256']
        all_hashes_in_amp_df = all_hashes_in_amp_df.append(df_append, ignore_index=True)
        logger.info(f'The shape of the SHA-256 AMP df is {all_hashes_in_amp_df.shape[0]}')

    # Send amp df to txt file
    all_hashes_in_amp_df.to_csv(amp_main_file, index=False, header=False)

    logger.info(f'Code took {((time.time() - start_time) / 60):.3f} minutes to execute so far.')
    logger.info(f'------------------------------------------------------------------')

    # Return all_hashes_in_amp_df
    return all_hashes_in_amp_df


def compare_differences(all_hashes_in_amp_df):
    # Create AMP set off main df
    amp_sha_set = set(all_hashes_in_amp_df.unique())

    # Loop through directory to start pruning unique hashes
    for filename in os.listdir(sha_output_files):
        f = os.path.join(sha_output_files, filename)
        # Check if file
        if os.path.isfile(f):

            # Create df per file
            df = pd.read_csv(f, header=None)

            # Add column name to pull set
            df.columns = ['sha256']

            # Create Malware Bazaar Set
            malware_bazaar_set = set(df['sha256'])

            # Find IPs missing from AMP by using set logic to compare MalwareBazaar with AMP.
            missing_sha_from_amp = malware_bazaar_set - amp_sha_set

            # Convert set to list
            missing_sha_list = list(missing_sha_from_amp)
            missing_sha_df = pd.DataFrame(missing_sha_list)
            # missing_sha_df.rename(columns={0: 'sha256'}, inplace=True)

            # If set is not empty
            if missing_sha_from_amp:
                missing_sha_df.to_csv(f, index=False, header=False)
                logger.info(f'The following hashes from {filename} will be added: {missing_sha_from_amp}')
                logger.info(f'\n')

            else:
                # If set is empty remove file
                os.remove(f)

    logger.info(f'Code took {((time.time() - start_time) / 60):.3f} minutes to execute so far.')
    logger.info(f'------------------------------------------------------------------')


def insert_amp_sha256():
    """
    Imports SHA_256 hashes into their respective buckets in AMP's Simple Custom Detection Lists.
    """
    # Create requests for all Simple Custom Detection GUIDs for loop
    request = requests.get(all_simple_custom_detections_url,
                           auth=HTTPBasicAuth(username=amp_api_username, password=amp_api_pw))

    # Pull data in json format. Slice data based off 'data' element
    # json_response = request.json()

    data = request.json()['data']

    # Create dataframe
    amp_df = pd.json_normalize(data)

    counter = 0

    # Iterate over MalwareBazaar directory
    for filename in os.listdir(sha_output_files):
        f = os.path.join(sha_output_files, filename)
        # Check if file
        if os.path.isfile(f):
            # Remove '.txt' so file so files names match in AMP and directory
            name_of_file = filename.replace('.txt', '')
            logger.info(name_of_file)

            # Create df from file
            df = pd.read_csv(f, header=None)

            # Add header
            df.columns = ['sha256']

            # Query AMP df
            new_amp_df = amp_df.loc[amp_df['name'] == name_of_file]

            # Reset index to pull the correct GUID
            new_amp_df = new_amp_df.reset_index(drop=True)

            # Pull GUID to loop through
            try:
                guid = new_amp_df.guid[0]
                logger.info(f'Guid is {guid}')
                del new_amp_df
            except Exception as e:
                logger.warning(f'Exception is {e}. GUID is {guid}')

            # This will hit the production environment!
            for index, row in enumerate(df.itertuples(index=False)):

                # Add hash to Simple Custom Detections
                temporary_url = f'https://api.amp.cisco.com/v1/file_lists/{guid}/files/{row.sha256}'

                # Send post request
                requests.post(temporary_url, auth=(amp_api_username, amp_api_pw))
                logger.info(f'Successfully updated {guid} with hash {row.sha256}')

                counter += 1

            # # Delete df reference then free up memory via garbage collection
            del df
            gc.collect()

    logger.info(f'Total hashes added = {counter}')
    logger.info(f'Code took {((time.time() - start_time) / 60):.3f} minutes to execute in its entirety.')
    logger.info(f'------------------------------------------------------------------')


if __name__ == '__main__':
    pull_from_malware_bazaar_api()
    create_master_df_from_malware_bazaar()
    all_hashes_in_amp_df = pull_main_hash_list()
    compare_differences(all_hashes_in_amp_df)
    insert_amp_sha256()