ソースを参照

new dataset format: multi-lists -> two arrays

René Knaebel 1年前
コミット
5743127b7f
2個のファイルの変更67行の追加46行の削除
  1. 67
    42
      dataset.py
  2. 0
    4
      main.py

+ 67
- 42
dataset.py ファイルの表示

@@ -5,10 +5,22 @@ import numpy as np
5 5
 import pandas as pd
6 6
 from tqdm import tqdm
7 7
 
8
+chars = dict((char, idx + 1) for (idx, char) in
9
+             enumerate(string.ascii_lowercase + string.punctuation + string.digits))
10
+
8 11
 
9 12
 def get_character_dict():
10
-    return dict((char, idx) for (idx, char) in
11
-                enumerate(string.ascii_lowercase + string.punctuation))
13
+    return chars
14
+
15
+
16
+def encode_char(c):
17
+    if c in chars:
18
+        return chars[c]
19
+    else:
20
+        return 0
21
+
22
+
23
+encode_char = np.vectorize(encode_char)
12 24
 
13 25
 
14 26
 def get_user_chunks(dataFrame, windowSize=10, overlapping=False,
@@ -49,16 +61,19 @@ def get_user_chunks(dataFrame, windowSize=10, overlapping=False,
49 61
                     curDomains = useData['domain']
50 62
             outDomainLists.append(list(curDomains))
51 63
             outDFFrames.append(useData)
64
+    if len(outDomainLists[-1]) != windowSize:
65
+        outDomainLists.pop(-1)
66
+        outDFFrames.pop(-1)
52 67
     return (outDomainLists, outDFFrames)
53 68
 
54 69
 
55 70
 def get_domain_features(domain, vocab, max_length=40):
56
-    curFeature = np.zeros([max_length, ])
71
+    encoding = np.zeros((max_length,))
57 72
     for j in range(np.min([len(domain), max_length])):
58 73
         curCharacter = domain[-j]
59 74
         if curCharacter in vocab:
60
-            curFeature[j] = vocab[curCharacter]
61
-    return curFeature
75
+            encoding[j] = vocab[curCharacter]
76
+    return encoding
62 77
 
63 78
 
64 79
 def get_flow_features(flow):
@@ -86,66 +101,76 @@ def get_cisco_features(curDataLine, urlSIPDict):
86 101
         return np.zeros([numCiscoFeatures, ]).ravel()
87 102
 
88 103
 
89
-def create_dataset_from_flows(user_flow_df, char_dict, maxLen, windowSize=10, use_cisco_features=False):
90
-    domainLists = []
91
-    dfLists = []
104
+def create_dataset_from_flows(user_flow_df, char_dict, max_len, window_size=10, use_cisco_features=False):
105
+    domains = []
106
+    features = []
92 107
     print("get chunks from user data frames")
93 108
     for i, user_flow in enumerate(get_flow_per_user(user_flow_df)):
94
-        (domainListsTmp, dfListsTmp) = get_user_chunks(user_flow, windowSize=windowSize,
95
-                                                       overlapping=True, maxLengthInSeconds=-1)
96
-        domainLists += domainListsTmp
97
-        dfLists += dfListsTmp
109
+        (domain_windows, feature_windows) = get_user_chunks(user_flow,
110
+                                                            windowSize=window_size,
111
+                                                            overlapping=True,
112
+                                                            maxLengthInSeconds=-1)
113
+        domains += domain_windows
114
+        features += feature_windows
98 115
         # TODO: remove later
99 116
         if i >= 10:
100 117
             break
101 118
 
102 119
     print("create training dataset")
103 120
     return create_dataset_from_lists(
104
-        domains=domainLists, dfs=dfLists, vocab=char_dict,
105
-        maxLen=maxLen,
121
+        domains=domains, features=features, vocab=char_dict,
122
+        max_len=max_len,
106 123
         use_cisco_features=use_cisco_features, urlSIPDIct=dict(),
107
-        window_size=windowSize)
124
+        window_size=window_size)
108 125
 
109 126
 
110
-def create_dataset_from_lists(domains, dfs, vocab, maxLen,
127
+def create_dataset_from_lists(domains, features, vocab, max_len,
111 128
                               use_cisco_features=False, urlSIPDIct=dict(),
112 129
                               window_size=10):
130
+    """
131
+    combines domain and feature windows to sequential training data
132
+    :param domains: list of domain windows
133
+    :param features: list of feature windows
134
+    :param vocab: 
135
+    :param max_len: 
136
+    :param use_cisco_features: idk
137
+    :param urlSIPDIct: idk
138
+    :param window_size: size of the flow window
139
+    :return: 
140
+    """
113 141
     # TODO: check for hits vs vth consistency
114
-    if 'hits' in dfs[0].keys():
115
-        hitName = 'hits'
116
-    elif 'virusTotalHits' in dfs[0].keys():
117
-        hitName = 'virusTotalHits'
142
+    # if 'hits' in dfs[0].keys():
143
+    #     hits_col = 'hits'
144
+    # elif 'virusTotalHits' in dfs[0].keys():
145
+    #     hits_col = 'virusTotalHits'
146
+    hits_col = "virusTotalHits"
147
+
118 148
     numFlowFeatures = 3
119 149
     numCiscoFeatures = 30
120 150
     numFeatures = numFlowFeatures
121 151
     if use_cisco_features:
122 152
         numFeatures += numCiscoFeatures
123
-    Xs = []
153
+    sample_size = len(domains)
124 154
     hits = []
125 155
     names = []
126 156
     servers = []
127 157
     trusted_hits = []
128
-    for i in range(window_size):
129
-        Xs.append(np.zeros([len(domains), maxLen]))
130
-        Xs.append(np.zeros([len(domains), numFeatures]))
131
-
132
-    for i in tqdm(np.arange(len(domains)), miniters=10):
133
-        ctr = 0
134
-        for j in range(np.min([window_size, len(domains[i])])):
135
-            Xs[ctr][i, :] = get_domain_features(domains[i][j], vocab, maxLen)
136
-            ctr += 1
137
-            if use_cisco_features:
138
-                Xs[ctr][i, 0:numFlowFeatures] = get_flow_features(dfs[i].iloc[j])
139
-                Xs[ctr][i, numFlowFeatures:] = get_cisco_features(dfs[i].iloc[j], urlSIPDIct)
140
-            else:
141
-                Xs[ctr][i, :] = get_flow_features(dfs[i].iloc[j])
142
-            ctr += 1
143
-
144
-        hits.append(np.max(dfs[i][hitName]))
145
-        names.append(np.unique(dfs[i]['user_hash']))
146
-        servers.append(np.max(dfs[i]['serverLabel']))
147
-        trusted_hits.append(np.max(dfs[i]['trustedHits']))
148
-    return Xs, np.array(hits), np.array(names), np.array(servers), np.array(trusted_hits)
158
+
159
+    domain_features = np.zeros((sample_size, window_size, max_len))
160
+    flow_features = np.zeros((sample_size, window_size, numFeatures))
161
+
162
+    for i in tqdm(np.arange(sample_size), miniters=10):
163
+        for j in range(window_size):
164
+            domain_features[i, j] = get_domain_features(domains[i][j], vocab, max_len)
165
+            flow_features[i, j] = get_flow_features(features[i].iloc[j])
166
+            # TODO: cisco features?
167
+
168
+        hits.append(np.max(features[i][hits_col]))
169
+        names.append(np.unique(features[i]['user_hash']))
170
+        servers.append(np.max(features[i]['serverLabel']))
171
+        trusted_hits.append(np.max(features[i]['trustedHits']))
172
+    X = [domain_features, flow_features]
173
+    return X, np.array(hits), np.array(names), np.array(servers), np.array(trusted_hits)
149 174
 
150 175
 
151 176
 def discretize_label(values, threshold):

+ 0
- 4
main.py ファイルの表示

@@ -104,10 +104,6 @@ def main():
104 104
     client_labels = client_labels[idx]
105 105
     server_labels = server_tr[idx]
106 106
 
107
-    # TODO: remove when features are flattened
108
-    for i in range(len(X_tr)):
109
-        X_tr[i] = X_tr[i][idx]
110
-
111 107
     shared_cnn = models.get_shared_cnn(len(char_dict) + 1, embeddingSize, maxLen,
112 108
                                        domainFeatures, kernel_size, domainFeatures, 0.5)
113 109
 

読み込み中…
キャンセル
保存