generate_vectors.py

from datasets.datasets import load_data
import os
from utils.utils import load_parameters, set_parameter
import time
import pickle


if __name__ == "__main__":
    saved_vectors_directory = "sample_dataset_saved_feature_vectors"

    if not os.path.exists(saved_vectors_directory):
        os.mkdir(saved_vectors_directory)

    malicious_vector_filepath = os.path.join(saved_vectors_directory, "malicious")
    benign_vector_filepath = os.path.join(saved_vectors_directory, "benign")

    if not os.path.exists(malicious_vector_filepath):
        os.mkdir(malicious_vector_filepath)

    if not os.path.exists(benign_vector_filepath):
        os.mkdir(benign_vector_filepath)

    parameters = load_parameters("parameters.ini")

    # This flag must be on to generate, changes return of PortableExecutableDataset to return filepath as well
    set_parameter("parameters.ini", "dataset", "generate_feature_vector_files", "True")
    set_parameter("parameters.ini", "hyperparam", "training_batch_size", "1")
    set_parameter("parameters.ini", "hyperparam", "test_batch_size", "1")

    train_dataloader_dict, valid_dataloader_dict, test_dataloader_dict, num_features = load_data(parameters)

    print(len(train_dataloader_dict['malicious'].dataset) +
          len(test_dataloader_dict['malicious'].dataset) + len(valid_dataloader_dict['malicious'].dataset))
    print(len(train_dataloader_dict['benign'].dataset) + len(test_dataloader_dict['benign'].dataset) +
          len(valid_dataloader_dict['benign'].dataset))

    for i, data_dict in enumerate([train_dataloader_dict, valid_dataloader_dict, test_dataloader_dict]):
        for filetype in data_dict:
            dataloader = data_dict[filetype]

            for index, data in enumerate(dataloader):
                print(index, filetype)

                vector, label, filepath = data

                filename = filepath[0].split("/")[-1]

                if filetype == 'malicious':
                    pickle.dump(vector, open(os.path.join(malicious_vector_filepath, filename), 'wb'))
                else:
                    pickle.dump(vector, open(os.path.join(benign_vector_filepath, filename), 'wb'))