label_malware.py

import multiprocessing
import tqdm
import os
import json
import make_entropy_vis
import pandas as pd
import numpy as np
sigresults = {
    "Signed": "Valid",
    "A required certificate is not within its validity period when verifying against the current system clock or the timestamp in the signed file.": "Expired",
    "A certificate was explicitly revoked by its issuer.": "Revoked"
}


class ParseFailure(Exception):
    pass


def parse_VT_file(filename):
    with open(filename) as json_file:
        json_text = json.load(json_file)
        name = json_text['submission']['filename']
        sha256 = json_text['sha256']

        try:
            detected = json_text['scans']['Microsoft']['detected']
            detectedname = json_text['scans']['Microsoft']['result']
        except KeyError:
            detected = False
        if detected is False:
            try:
                detected = json_text['scans']['BitDefender']['detected']
                detectedname = json_text['scans']['BitDefender']['result']
            except KeyError:
                detected = False
                detectedname = None

        try:
            sigstatus = json_text['additional_info']['sigcheck']['verified']
            try:
                signature = sigresults[sigstatus]
            except ValueError:
                signature = sigstatus
            signers = json_text['additional_info']['sigcheck']['signers']

        except KeyError:
            signature = 'unsigned'
            signers = 'none'
        link = json_text['permalink']
        if os.path.exists(filename[:-5]):
            image = make_entropy_vis.worker(filename[:-5])['entropy_norm'].tolist()
            exists = True
        else:
            image = np.zeros((8, 8)).tolist()
            exists = False
    return [sha256, name, signature, signers, detected, detectedname, link, image, exists]


if __name__ == '__main__':
    path_to_json = './malware'
    json_files = [os.path.join(path_to_json, pos_json)for pos_json in os.listdir(path_to_json) if pos_json.endswith('.json')]

    pool = multiprocessing.Pool(processes=8)

    results = list(tqdm.tqdm(pool.imap_unordered(parse_VT_file, json_files), total=len(json_files)))

    mal_data = pd.DataFrame(results, columns=['sha256', 'filename', 'signed', 'signers', 'detected', 'detectedname', 'link', 'image', 'exists'])
    mal_data['class'] = mal_data['detectedname'].str.split(':').str[0]
    mal_data[mal_data.exists].drop('exists', axis=1).to_csv("./data/labelledVTEntropy2.csv")