diff --git a/dataset.py b/dataset.py index 031a9b2..9c8364f 100644 --- a/dataset.py +++ b/dataset.py @@ -13,7 +13,6 @@ def get_character_dict(): def get_user_chunks(dataFrame, windowSize=10, overlapping=False, maxLengthInSeconds=300): - # print('maxLength: ' + str(maxLengthInSeconds)) maxMilliSeconds = maxLengthInSeconds * 1000 outDomainLists = [] outDFFrames = [] @@ -39,7 +38,6 @@ def get_user_chunks(dataFrame, windowSize=10, overlapping=False, userIDs = np.arange(len(dataFrame)) for blockID in np.arange(numBlocks): curIDs = userIDs[blockID:blockID + windowSize] - # print(curIDs) useData = dataFrame.iloc[curIDs] curDomains = useData['domain'] if maxLengthInSeconds != -1: @@ -64,17 +62,20 @@ def get_domain_features(domain, vocab, max_length=40): def get_flow_features(flow): - useKeys = ['duration', 'bytes_down', 'bytes_up'] - curFeature = np.zeros([len(useKeys), ]) - for i, curKey in enumerate(useKeys): + keys = ['duration', 'bytes_down', 'bytes_up'] + features = np.zeros([len(keys), ]) + for i, key in enumerate(keys): + # TODO: does it still works after exceptions occur -- default: zero! + # i wonder whether something brokes + # if there are exceptions regarding to inconsistent feature length try: - curFeature[i] = np.log1p(flow[curKey]).astype(float) + features[i] = np.log1p(flow[key]).astype(float) except: pass - return curFeature + return features -def getCiscoFeatures(curDataLine, urlSIPDict): +def get_cisco_features(curDataLine, urlSIPDict): numCiscoFeatures = 30 try: ciscoFeatures = urlSIPDict[str(curDataLine['domain']) + str(curDataLine['server_ip'])] @@ -94,20 +95,21 @@ def create_dataset_from_flows(user_flow_df, char_dict, maxLen, threshold=3, wind overlapping=True, maxLengthInSeconds=-1) domainLists += domainListsTmp dfLists += dfListsTmp + # TODO: remove later if i >= 10: break print("create training dataset") return create_dataset_from_lists( - domains=domainLists, dfs=dfLists, charachterDict=char_dict, + domains=domainLists, dfs=dfLists, vocab=char_dict, maxLen=maxLen, threshold=threshold, use_cisco_features=use_cisco_features, urlSIPDIct=dict(), - windowSize=windowSize) + window_size=windowSize) -def create_dataset_from_lists(domains, dfs, charachterDict, maxLen, threshold=3, +def create_dataset_from_lists(domains, dfs, vocab, maxLen, threshold=3, use_cisco_features=False, urlSIPDIct=dict(), - windowSize=10): + window_size=10): if 'hits' in dfs[0].keys(): hitName = 'hits' elif 'virusTotalHits' in dfs[0].keys(): @@ -117,38 +119,46 @@ def create_dataset_from_lists(domains, dfs, charachterDict, maxLen, threshold=3, numFeatures = numFlowFeatures if use_cisco_features: numFeatures += numCiscoFeatures - outputFeatures = [] - label = [] + Xs = [] + ys = [] hits = [] - trainNames = [] - for i in range(windowSize): - outputFeatures.append(np.zeros([len(domains), maxLen])) - outputFeatures.append(np.zeros([len(domains), numFeatures])) + names = [] + servers = [] + trusted_hits = [] + for i in range(window_size): + Xs.append(np.zeros([len(domains), maxLen])) + Xs.append(np.zeros([len(domains), numFeatures])) for i in tqdm(np.arange(len(domains)), miniters=10): - curCounter = 0 - # print('len domainList: ' + str(len(domainLists[i]))) - # print('len df: ' + str(len(dfLists[i]))) - for j in range(np.min([windowSize, len(domains[i])])): - outputFeatures[curCounter][i, :] = get_domain_features(domains[i][j], charachterDict, maxLen) - curCounter += 1 + ctr = 0 + for j in range(np.min([window_size, len(domains[i])])): + Xs[ctr][i, :] = get_domain_features(domains[i][j], vocab, maxLen) + ctr += 1 if use_cisco_features: - outputFeatures[curCounter][i, 0:numFlowFeatures] = get_flow_features(dfs[i].iloc[j]) - outputFeatures[curCounter][i, numFlowFeatures:] = get_cisco_features(dfs[i].iloc[j], urlSIPDIct) + Xs[ctr][i, 0:numFlowFeatures] = get_flow_features(dfs[i].iloc[j]) + Xs[ctr][i, numFlowFeatures:] = get_cisco_features(dfs[i].iloc[j], urlSIPDIct) else: - outputFeatures[curCounter][i, :] = get_flow_features(dfs[i].iloc[j]) - curCounter += 1 - curLabel = 0.0 - if np.max(dfs[i][hitName]) >= threshold: - curLabel = 1.0 - elif np.max(dfs[i][hitName]) == -1: - curLabel = -1.0 - elif np.max(dfs[i][hitName]) > 0 and np.max(dfs[i][hitName]) < threshold: - curLabel = -2.0 - label.append(curLabel) + Xs[ctr][i, :] = get_flow_features(dfs[i].iloc[j]) + ctr += 1 + + ys.append(discretize_label(dfs[i][hitName], threshold)) hits.append(np.max(dfs[i][hitName])) - trainNames.append(np.unique(dfs[i]['user_hash'])) - return (outputFeatures, np.array(label), np.array(hits), np.array(trainNames)) + names.append(np.unique(dfs[i]['user_hash'])) + servers.append(np.max(dfs[i]['serverLabel'])) + trusted_hits.append(np.max(dfs[i]['trustedHits'])) + return Xs, np.array(ys), np.array(hits), np.array(names), np.array(servers), np.array(trusted_hits) + + +def discretize_label(values, threshold): + maxVal = np.max(values) + if maxVal >= threshold: + return 1.0 + elif maxVal == -1: + return -1.0 + elif 0 < maxVal < threshold: + return -2.0 + else: + return 0.0 def get_user_flow_data(): diff --git a/main.py b/main.py index 703af91..fc3aa4d 100644 --- a/main.py +++ b/main.py @@ -37,20 +37,21 @@ def main(): user_flow_df = dataset.get_user_flow_data() print("create training dataset") - (X_tr, y_tr, hits_tr, names_tr) = dataset.create_dataset_from_flows( + (X_tr, y_tr, hits_tr, names_tr, server_tr, trusted_hits_tr) = dataset.create_dataset_from_flows( user_flow_df, char_dict, maxLen=maxLen, threshold=threshold, windowSize=windowSize) pos_idx = np.where(y_tr == 1.0)[0] neg_idx = np.where(y_tr == 0.0)[0] + idx = np.concatenate((pos_idx, neg_idx)) - use_idx = np.concatenate((pos_idx, neg_idx)) - - y_tr = y_tr[use_idx] - # hits_tr = hits_tr[use_idx] - # names_tr = names_tr[use_idx] + y_tr = y_tr[idx] + hits_tr = hits_tr[idx] + names_tr = names_tr[idx] + server_tr = server_tr[idx] + trusted_hits_tr = trusted_hits_tr[idx] for i in range(len(X_tr)): - X_tr[i] = X_tr[i][use_idx] + X_tr[i] = X_tr[i][idx] # TODO: WTF? I don't get it... shared_cnn = models.get_shared_cnn(len(char_dict) + 1, embeddingSize, maxLen,