AntiDriftDeepMalware/dataset.py at master · nirosen/AntiDriftDeepMalware · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import os
from torch.utils.data import Dataset
import pandas as pd
import torch
import numpy as np
import pickle


# TODO: merge dataset files - tensors  syscalls_tensors_df + scores balanced_order_df
# and load little amount instead of all!!

class TensorPickleSCNNCSVDataset(Dataset):
    def __init__(self, syscalls_tensors_df, balanced_order_df, reports_count, batch_size, max_sequence_length, time_weight):

        if balanced_order_df is None:
            return None

        self.balanced_order_df = balanced_order_df
        self.syscalls_tensors_df = syscalls_tensors_df
        self.max_length = max_sequence_length

        self.sequences = [self.syscalls_tensors_df.loc[row['hash']]['tensor'] for _, row in self.balanced_order_df.iterrows()
                          if row['hash'] in self.syscalls_tensors_df.index]

        self.scores = [self.syscalls_tensors_df.loc[row['hash']]['score'] for _, row in self.balanced_order_df.iterrows()
                       if row['hash'] in self.syscalls_tensors_df.index]

        # self.sequences_count = 0
        if reports_count == 0 or reports_count > len(self.sequences): #  check that reports_count is not larger then  len(self.sequences)
            self.sequences_count = len(self.sequences)
        else:
            # calculate the dataset portion size:
            # self.portion = arguments.train_portion if ds_type is "train" else arguments.valid_portion if ds_type is "validation" else arguments.test_portion
            # self.sequences_count = int(reports_count*self.portion)
            # check larger then BS:
            self.sequences_count = int(reports_count)  # * self.portion)  #self.portion = arguments.train_portion if ds_type is "train" else arguments.valid_portion if ds_type is "validation" else arguments.test_portion
            if self.sequences_count < batch_size:
                raise SystemError(
                    f"reports_count={reports_count} is too small because"
                    f"sequences_count < batch_size.\n"
                    f"please increase \"reports_count\" argument.")

        # get time distance for sample weight
        self.balanced_order_df.first_seen_obj = pd.to_datetime(self.balanced_order_df.first_seen_obj)
        self.times = [row['first_seen_obj'] for _, row in self.balanced_order_df.iterrows()
                      if row['hash'] in self.syscalls_tensors_df.index]
        self.hashes = [row['hash'] for _, row in self.balanced_order_df.iterrows()
                      if row['hash'] in self.syscalls_tensors_df.index]
        self.max_first_seen_obj = max(self.times)
        self.min_first_seen_obj = min(self.times)
        self.time_range_total_seconds = (self.max_first_seen_obj-self.min_first_seen_obj).seconds
        self.time_weight = time_weight

        print(f"Initialized TensorPickleCSV Dataset instance object of size: {self.sequences_count}.")

    def __len__(self):
        return self.sequences_count

    def __getitem__(self, item):
        return self.__load_tensor_syscalls_sequence(item, self.sequences[item], self.max_length, self.scores[item])

    def __load_tensor_syscalls_sequence(self, item, sequence_tensor, max_length, score):
        syscalls_label = 1 if score > 5 else 0

        # TODO: check this:
        # Load all tensors onto GPU 1
        # torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
        # sequence_tensor = torch.load(sequence_file_path)[:max_length]  # TODO: check max_length !
        sequence_tensor = torch.Tensor(sequence_tensor[:max_length])
        seq_length = len(sequence_tensor)

        # padding each seq with zeros until max_seq_len = 10,000
        padded_sequence_tensor = torch.zeros(self.max_length, dtype=torch.uint8)
        padded_sequence_tensor[0:sequence_tensor.shape[0]] = sequence_tensor

        # # adding samples_weights to each batch  # TODO: not good! this can only return 1 sample weight
        # current_batch_size = padded_sequence_tensor.shape[1]
        # x = torch.randn(current_batch_size, 10)
        # samples_weights = torch.empty(current_batch_size).uniform_(0,1)

        # divide by 86400 to change days to seconds
        time_to_max = (self.max_first_seen_obj - self.times[item]).seconds
        a = self.time_weight
        sample_weight = np.exp(-a*time_to_max/self.time_range_total_seconds)  # TODO: normalize the time_to_max:  time_to_max/(latest-first)
        #sample_weight = np.exp(-a*time_to_max)

        # time_to_max = (self.max_first_seen_obj - self.balanced_order_df.iloc[item].first_seen_obj).days
        # time_from_min = (self.balanced_order_df.iloc[item].first_seen_obj - self.max_first_seen_obj).days
        # time_from_min = (self.times[item] - self.max_first_seen_obj).days

        return padded_sequence_tensor, syscalls_label, seq_length, sample_weight, self.times[item], self.hashes[item]


#
# class RandomSCNNDataset(Dataset):
#     def __init__(self, arguments, ds_type, on_memory=False):
#         # self.log = log
#         self.sequences_count = arguments.reports_count
#         self.min_length = arguments.min_sequence_length
#         self.max_length = arguments.max_sequence_length
#         self.on_memory = on_memory
#         self.max_syscalls_types = arguments.categories_input_layer
#         self.portion = arguments.train_portion if ds_type is TRAIN else arguments.valid_portion if ds_type is VALIDATION else arguments.test_portion
#
#         self.sequences = self.__generate_random_sequences()
#         self.sequences_count = len(self.sequences)
#
#         # self.log.info(
#         #     f"Initialized RandomSCNNDataset instance object for dataset type: {ds_type} of size: {self.sequences_count}.")
#
#     def __len__(self):
#         return self.sequences_count
#
#     def __getitem__(self, item):
#         return self.sequences[item]
#
#     def __generate_random_sequences(self):
#         # calculate the dataset portion size
#         sequences_count = int(self.sequences_count * self.portion)
#
#         # create a list of lengths for each sequence (let them have different lengths)
#         sequence_lengths = np.random.randint(low=self.min_length, high=self.max_length + 1, size=sequences_count)
#
#         # create the list of sequences according to the length
#         syscalls_sequences = [np.random.randint(low=0, high=self.max_syscalls_types, size=seq_length) for seq_length in
#                               sequence_lengths]
#         syscalls_sequences = [np.pad(sequence, (0, self.max_length - len(sequence)), mode='constant', constant_values=0)
#                               for sequence in syscalls_sequences]
#
#         # create the score of the sequences
#         syscalls_class = np.random.choice([0, 1], size=sequences_count, p=[0.8, 0.2]).astype(np.int64)
#
#         # define and set the result dictionary
#         sequences = {}
#         for index in range(sequences_count):
#             sequences[index] = {
#                 "index": index,
#                 "sequence": syscalls_sequences[index].astype(np.int64),
#                 "label": syscalls_class[index]
#             }
#
#         return sequences
#