import os 
import sys
import argparse
import numpy
import scipy
import array 
import imageio
from PIL import Image



def convert_file(source_path,destination_directory):
    f = open(source_path,'rb')
    ln = os.path.getsize(source_path);
    width = 256;
    rem = ln%width;
    lag = width-rem
    lag_arr=numpy.zeros((lag,),dtype=int)
    a = array.array("B"); #Read the file content as unsigned char (1 byte minimum)    
    a.fromfile(f,ln); #Discard the last rem bytes to align on width boundary
    arr=numpy.append(a,lag_arr)
    f.close();
    g = numpy.reshape(arr,(int(len(arr)/width),width)); #reshape to 2D array
    g = numpy.uint8(g); #Do not normalize the pixel values by dividing by 255
    (source_directory, filename) = os.path.split(source_path);
    saved_path = os.path.join(destination_directory, filename+ '.png');
    imageio.imwrite(saved_path,g);
    print ("Convert: " + source_path + " \tTo: " + saved_path + "\n")


def get_file_paths(directory):
    targets = []
    for path in os.listdir(directory):
        targets.append(os.path.join(directory,path))
    return targets

def convert_to_images(benign_path,malware_path,image_path):
    #Convert all files to image representation
    malware_paths = get_file_paths(malware_path)
    #Set the directory path for converted image files of malware classification
    malware_image_directory = os.path.join(image_path,"malware")
    print("Convert exe files in: " + malware_path + " to images in: " + malware_image_directory)
    print("-----The number of files: %s------\n" % len(malware_paths))

    #convert each malware file
    for malware in malware_paths:
        convert_file(malware, malware_image_directory)
    #Get full paths for all malware image files
    malware_image_files = get_file_paths(malware_image_directory);

    #Get full paths for all benign files
    benign_paths = get_file_paths(benign_path);
    benign_image_directory = os.path.join(image_path,"benignware");
    print("Convert exe files in: " + benign_path + " to images in: " + benign_image_directory);
    print("\n-----\n");

    #COnvert each benign file
    for path in benign_paths:
        convert_file(path, benign_image_directory)
    #Get full paths for all benign image files
    benign_image_files = get_file_paths(benign_image_directory);
    #print out totals
    print("\nTotal number of malicious exe files: ", len(malware_paths));
    print("\nTotal number of malicious image files: ", len(malware_image_files));
    print("\nTotal number of benign exe files: ", len(benign_paths));
    print("\nTotal number of benign image files: ", len(benign_image_files));

parser = argparse.ArgumentParser("get object vectors for files")
parser.add_argument("--malware_paths",default=None,help="/home/sowmya/malware_classification/malware")
parser.add_argument("--benignware_paths",default=None,help="/home/sowmya/malware_classification/benignware")
parser.add_argument("--image_paths",default=None,help="/home/sowmya/malware_classification/images")

args = parser.parse_args()

if args.malware_paths and args.benignware_paths and args.image_paths:
    convert_to_images(args.benignware_paths,args.malware_paths,args.image_paths)
else:
    print ("You did not specify all the three paths")
    parser.print_help()
