SSamson/malware_detection_Project

1"""Binary to Image Converter."""2import os3from argparse import ArgumentParser4from PIL import Image5from tqdm import tqdm6 7 8def get_binary_data(filename):9    """Extract byte values from binary executable file and store them in list.10 11    :param file: executable file12    :return: byte value list13    """14    binary_values = []15    with open(filename, "rb") as file:16        while byte := file.read(1):17            binary_values.append(ord(byte))18 19    return binary_values20 21 22def get_size(data_length):23    """Obtain image size.24 25    Source: Malware images: visualization and automatic classification.26    :param data_length: Number of bytes in file27    :return: size as integer tuple28    """29    size = data_length30    kib = 2**1031 32    if size < 10 * kib:33        width = 3234    elif size < 30 * kib:35        width = 6436    elif size < 60 * kib:37        width = 12838    elif size < 100 * kib:39        width = 25640    elif size < 200 * kib:41        width = 38442    elif size < 500 * kib:43        width = 51244    elif size < 1000 * kib:45        width = 76846    else:47        width = 102448 49    height = size // width + 150 51    return (width, height)52 53 54def save_file(folder, filename, data, size):55    """Save PIL image to disk.56 57    :param folder: folder where images will be saved58    :param filename: binary filename59    :param data: grayscale image60    :param size: image size61    """62    image = Image.new("L", size)63    image.putdata(data)64 65    name, _ = os.path.splitext(filename)66    name = os.path.basename(name)67    imagename = os.path.join(folder, name + ".png")68 69    image.save(imagename)70 71 72if __name__ == "__main__":73 74    parser = ArgumentParser(75        description="Transform all files in a folder into PNG images"76    )77    parser.add_argument("input_folder", help="Folder with the original files")78    parser.add_argument(79        "output_folder", help="Folder where the images will be saved"80    )81    args = parser.parse_args()82 83    files = [84        os.path.join(args.input_folder, bin_file)85        for bin_file in os.listdir(args.input_folder)86    ]87    os.makedirs(args.output_folder, exist_ok=True)88    for bin_file in tqdm(files):89        greyscale_data = get_binary_data(bin_file)90        bin_size = get_size(len(greyscale_data))91        save_file(args.output_folder, bin_file, greyscale_data, bin_size)92