VT_Classify.py

#!/usr/bin/env python3

import os

def generate_hash_files():
    
    base_path = "/media/cape/Malware/Nico/Dataset_Resized_Malicious_Binaries"
    
    dirs = os.listdir(base_path)
    dirs.remove("dir_tree.txt")
    dirs.remove("total_samples.txt")

    for dir in dirs:
        with open(os.path.join(base_path,dir+".txt"), "w") as file:
            for hash in os.listdir(os.path.join(base_path,dir)):
                file.writelines(os.path.splitext(hash)[0]+"\n")
        
import requests
import json
import shutil

def check_family():
    
    base_path = "/media/cape/Malware/Nico/Dataset_Resized_Malicious_Binaries"
    FAMILIES = ["babuk", "blackmatter", "cerber", "chaos", "wastedlocker", "conti", "cryptbot", "dharma", "formbook", "gandcrab", "hive", "lockbit", "medusalocker", "nitro", "phobos", "sodinokibi", "virlock", "loki", "stop", "razy", "marte", "msil", "snatch"]

    txt_files = [file for file in os.listdir(base_path) if file.endswith(".txt")]
    txt_files.remove("dir_tree.txt")
    txt_files.remove("total_samples.txt")

    
    for txt_file in txt_files[:1]:
        # Change the name of the .txt file
        txt_file = "medusalocker"
        with open(os.path.join(base_path,txt_file+".txt"), "r") as file:
            hashes = file.readlines()

        print(len(hashes))
        i=0

        for hash in hashes:
            
            families = {
                "conti": 0,
                "cryptbot": 0,
                "dharma": 0,
                "formbook": 0,
                "gandcrab": 0,
                "hive": 0,
                "lockbit": 0,
                "nitro": 0,
                "phobos": 0,
                "sodinokibi": 0,
                "virlock": 0,
                "loki": 0,
                "stop": 0,
                "stopcrypt": 0,
                "mint": 0,
                "razy": 0,
                "marte": 0,
                "crysis": 0,
                "msil": 0,
                "snatch": 0,
                "medusalocker": 0,
                "babuk": 0,
                "blackmatter": 0, 
                "cerber": 0, 
                "chaos": 0,
                "wastedlocker": 0
            }

            url = "https://www.virustotal.com/api/v3/files/"
            url = url+hash.strip("\n")
            headers = {
                "accept": "application/json",
                "x-apikey": "PUT_HERE_THE_API_KEY"
            }

            print(i, "- ", hash.strip("\n"))

            json_report = json.loads(requests.get(url, headers=headers).text)

            if 'popular_threat_classification' in json_report.get('data', {}).get('attributes', {}):
                if 'popular_threat_name' in json_report.get('data', {}).get('attributes', {}).get('popular_threat_classification', {}):
                    stats = json_report['data']['attributes']['popular_threat_classification']['popular_threat_name']
                    stat = []

                    for s in stats:
                        stat.append(s['value'])

                    for s in stat:
                        if s in FAMILIES:
                            families[s] = families[s] + 1

                    to_remove = True
                    for key, value in families.items():
                        if value > 0:
                            to_remove = False

                    for key, value in families.items():
                        if key == txt_file and value > 0:
                            print("The file correctly belongs to its family")
                            break
                        
                        elif key == "stopcrypt" and value > 0:
                            print("The file belongs to a different family")
                            shutil.move(os.path.join(base_path,txt_file,hash.strip("\n")+".exe"), os.path.join(base_path,"stop",hash.strip("\n")+".exe"))
                            print("Moved ", hash.strip("\n"), " to the correct family, stop")
                            break

                        elif key == "crysis" and value > 0:
                            print("The file belongs to a different family")
                            shutil.move(os.path.join(base_path,txt_file,hash.strip("\n")+".exe"), os.path.join(base_path,"dharma",hash.strip("\n")+".exe"))
                            print("Moved ", hash.strip("\n"), " to the correct family, crysis")
                            break

                        elif key in FAMILIES and value > 0:
                            print("The file belongs to a different family")
                            shutil.move(os.path.join(base_path,txt_file,hash.strip("\n")+".exe"), os.path.join(base_path,key,hash.strip("\n")+".exe"))
                            print("Moved ", hash.strip("\n"), " to the correct family, ", key)
                            break

                    if to_remove:
                        os.remove(os.path.join(base_path,txt_file,hash.strip("\n")+".exe"))
                        print("Removed ", hash.strip("\n"), " because it doesn't belongs to any family.")

                    i = i+1
                else:
                    os.remove(os.path.join(base_path,txt_file,hash.strip("\n")+".exe"))
                    print("Removing because is a general ransomware")
            else:
                os.remove(os.path.join(base_path,txt_file,hash.strip("\n")+".exe"))
                print("Removing because is a general ransomware")






def exit_program():
    print("Exiting the program.")
    exit()

if __name__ == "__main__":
    switch_dict = {
        '1': generate_hash_files,
        '2': check_family,
        '3': exit_program,
    }

    while(True):
        choose = input("1: Generate the file containing all the hash\n2: Send the samples to VirusTotal\n3: Exit\n")

        selected_function = switch_dict.get(choose, lambda: print("Invalid option"))

        selected_function()