__author__ = 'kewang'

import os
import subprocess
import multiprocessing

datapath = '/run/media/kewang/externalHD/Projects/kaggle/malware/test'


def get_hex_histogram(filename):
    if filename.endswith('.bytes') is False:
        return

    filename = os.path.join(datapath, filename)
    print(filename)

    # remove addresses in ascii hex dump files
    file_path, basename = os.path.split(filename)
    filename_base, filename_ext = os.path.splitext(basename)
    raw_bytes_filename = os.path.join(file_path, filename_base + '.raw_bytes')
    with open(raw_bytes_filename, 'w') as r_fid:
        subprocess.call([
            'cut',
            '-c',
            '9-',
            filename
        ], stdout=r_fid)

    # build hex value histogram
    hist = []
    for i in range(257):
        hist.append(0)
    with open(raw_bytes_filename, 'r') as f:
        for line in f:
            tokens = line.strip().split()
            for token in tokens:
                try:
                    value = int(token, 16)
                    hist[value] += 1
                except:
                    hist[256] += 1

    hist_filename = os.path.join(file_path, filename_base + '.hist')
    with open(hist_filename, 'w') as f_out:
        for hist_token in hist:
            f_out.write(str(hist_token) + ' ')

    return


# filename = 'values.bytes'

if __name__ == '__main__':
    files = os.listdir(datapath)

    p = multiprocessing.Pool(processes=8)
    p.map(get_hex_histogram, files)

