'''
Read data
Create Pandas DataFrame with index: Archive Name ; columns: words, label
words: ByteArray consisting of bytes from archive member; Label: BENIGN, MALWARE
Pickle Pandas DataFrames and save
'''

import zipfile as zf
import pandas as pd
import os

BENIGN_DATA = os.path.join("Data/benign")
MALWARE_DATA = os.path.join("Data/malware-nodup")


def get_filelist(dirname):
    '''
    Get a list of filenames in a given directory.

    :param dirname: Directory path
    :return: list containing file names
    '''
    filelist = []
    for (_, _, filenames) in os.walk(dirname):
        for f in filenames:
            if ".zip" in f:
                filelist.append(f)

    return filelist


def read_data(dirname, filelist, label):
    '''
    Extract data from all zipfiles in a directory and return as dataframe.

    :param dirname: Directory path
    :param filelist: List containing filenames
    :param label: String -- BENIGN/MALWARE
    :return: pandas dataframe -- index is filename, columns are words and label
    '''
    print("Current Directory: " + dirname)
    out_array = bytearray()

    index = []

    for f in filelist:
        index.append(f[:-4])

    columns = ['words', 'label']

    df = pd.DataFrame(index=index, columns=columns)

    count = 1

    for f in filelist:
        df.loc[f[:-4]]['words'] = []
        with zf.ZipFile(dirname + "/" + f, "r") as currzip:
            for file in currzip.namelist():
                print(str(count) + " / " + str(len(index)) +
                      " Reading - Archive: " + f[:-4] + " File: " + file)
                out_array.extend(currzip.read(file))
        for w in out_array:
            df.loc[f[:-4]]['words'].append(str(w))
        df.loc[f[:-4]]['label'] = label
        out_array = bytearray()
        count += 1

    print("Read Done")

    return df


def main():
    benign_list = get_filelist(BENIGN_DATA)
    malware_list = get_filelist(MALWARE_DATA)

    benign_data = read_data(BENIGN_DATA, benign_list, "BENIGN")
    benign_data.to_pickle('Data/benign_data.pickle', compression='gzip', protocol=4)

    del benign_data

    malware_data = read_data(MALWARE_DATA, malware_list, "MALWARE")
    malware_data.to_pickle('Data/malware_data.pickle', compression='gzip', protocol=4)

    return 0


if __name__ == '__main__':
    import time
    start_time = time.time()
    main()
    print('Time taken: ' + str(time.time() - start_time) + ' seconds')
