import os
import numpy as np
import matplotlib.pyplot as plt


def file_size(file):
    import os
    file_size = os.stat(file)
    file_size_in_kb = file_size.st_size / 1000
    return file_size_in_kb


def get_width(file_size):
    if file_size < 10:
        width = 32
    elif 10 <= file_size < 30:
        width = 64
    elif 30 <= file_size < 60:
        width = 128
    elif 60 <= file_size < 100:
        width = 256
    elif 100 <= file_size < 200:
        width = 384
    elif 200 <= file_size < 500:
        width = 512
    elif 500 <= file_size < 1000:
        width = 768
    else:
        width = 1024
    return width


def convert_exe_to_image(file, save_path):
    """

    :param file: exe file
    :param save_path: folder where to save images, eg:C:/images/
    :return: image in grayscale
    """
    names = file.split('/')
    filename = names[-1]
    filesize = file_size(file)
    width = get_width(filesize)
    f = open(file, 'r+b')
    alist = []
    for each in f:
        for item in each:
            alist.append(item)
    myarray = np.asarray(alist)
    length = myarray.shape
    height = length[0] // width
    myarray = myarray[:height * width]
    img = np.reshape(myarray, (height, width))
    plt.gray()
    plt.imshow(img)
    plt.imsave(save_path + filename + '.png', img)


def generate_benign_samples(folder_path, save_path):
    """
    :param folder_path: folder path containing benign .exe files
    :param save_path: folder where to save images, eg:C:/images/
    :return: folder with images under name 'folder_path_images'
    """
    files = os.listdir(folder_path)
    for file in files:
        convert_exe_to_image(folder_path + '/' + file, save_path)
