diff --git a/dataset.py b/dataset.py index 6bcd700..031a9b2 100644 --- a/dataset.py +++ b/dataset.py @@ -6,12 +6,6 @@ import pandas as pd from tqdm import tqdm -# config = tf.ConfigProto(log_device_placement=True) -# config.gpu_options.per_process_gpu_memory_fraction = 0.5 -# config.gpu_options.allow_growth = True -# session = tf.Session(config=config) - - def get_character_dict(): return dict((char, idx) for (idx, char) in enumerate(string.ascii_lowercase + string.punctuation)) @@ -60,23 +54,21 @@ def get_user_chunks(dataFrame, windowSize=10, overlapping=False, return (outDomainLists, outDFFrames) -def getFeatureVecForDomain(domain, characterDict, maxLen=40): - curFeature = np.zeros([maxLen, ]) - for j in range(np.min([len(domain), maxLen])): - # print(j) +def get_domain_features(domain, vocab, max_length=40): + curFeature = np.zeros([max_length, ]) + for j in range(np.min([len(domain), max_length])): curCharacter = domain[-j] - if curCharacter in characterDict: - curFeature[j] = characterDict[curCharacter] + if curCharacter in vocab: + curFeature[j] = vocab[curCharacter] return curFeature -def getFlowFeatures(curDataLine): +def get_flow_features(flow): useKeys = ['duration', 'bytes_down', 'bytes_up'] curFeature = np.zeros([len(useKeys), ]) - for i in range(len(useKeys)): - curKey = useKeys[i] + for i, curKey in enumerate(useKeys): try: - curFeature[i] = np.log1p(curDataLine[curKey]).astype(float) + curFeature[i] = np.log1p(flow[curKey]).astype(float) except: pass return curFeature @@ -93,13 +85,13 @@ def getCiscoFeatures(curDataLine, urlSIPDict): return np.zeros([numCiscoFeatures, ]).ravel() -def create_dataset_from_flows(user_flow_df, char_dict, maxLen, threshold=3, windowSize=10): +def create_dataset_from_flows(user_flow_df, char_dict, maxLen, threshold=3, windowSize=10, use_cisco_features=False): domainLists = [] dfLists = [] print("get chunks from user data frames") for i, user_flow in enumerate(get_flow_per_user(user_flow_df)): (domainListsTmp, dfListsTmp) = get_user_chunks(user_flow, windowSize=windowSize, - overlapping=False, maxLengthInSeconds=-1) + overlapping=True, maxLengthInSeconds=-1) domainLists += domainListsTmp dfLists += dfListsTmp if i >= 10: @@ -107,68 +99,63 @@ def create_dataset_from_flows(user_flow_df, char_dict, maxLen, threshold=3, wind print("create training dataset") return create_dataset_from_lists( - domainLists=domainLists, dfLists=dfLists, charachterDict=char_dict, + domains=domainLists, dfs=dfLists, charachterDict=char_dict, maxLen=maxLen, threshold=threshold, - flagUseCiscoFeatures=False, urlSIPDIct=dict(), + use_cisco_features=use_cisco_features, urlSIPDIct=dict(), windowSize=windowSize) -def create_dataset_from_lists(domainLists, dfLists, charachterDict, maxLen, threshold=3, - flagUseCiscoFeatures=False, urlSIPDIct=dict(), +def create_dataset_from_lists(domains, dfs, charachterDict, maxLen, threshold=3, + use_cisco_features=False, urlSIPDIct=dict(), windowSize=10): - if 'hits' in dfLists[0].keys(): + if 'hits' in dfs[0].keys(): hitName = 'hits' - elif 'virusTotalHits' in dfLists[0].keys(): + elif 'virusTotalHits' in dfs[0].keys(): hitName = 'virusTotalHits' numFlowFeatures = 3 numCiscoFeatures = 30 numFeatures = numFlowFeatures - if flagUseCiscoFeatures: + if use_cisco_features: numFeatures += numCiscoFeatures outputFeatures = [] label = [] hits = [] trainNames = [] for i in range(windowSize): - outputFeatures.append(np.zeros([len(domainLists), maxLen])) - outputFeatures.append(np.zeros([len(domainLists), numFeatures])) + outputFeatures.append(np.zeros([len(domains), maxLen])) + outputFeatures.append(np.zeros([len(domains), numFeatures])) - for i in tqdm(np.arange(len(domainLists)), miniters=10): + 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(domainLists[i])])): - outputFeatures[curCounter][i, :] = getFeatureVecForDomain(domainLists[i][j], charachterDict, maxLen) + for j in range(np.min([windowSize, len(domains[i])])): + outputFeatures[curCounter][i, :] = get_domain_features(domains[i][j], charachterDict, maxLen) curCounter += 1 - if flagUseCiscoFeatures: - outputFeatures[curCounter][i, 0:numFlowFeatures] = getFlowFeatures(dfLists[i].iloc[j]) - outputFeatures[curCounter][i, numFlowFeatures:] = getCiscoFeatures(dfLists[i].iloc[j], urlSIPDIct) + 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) else: - outputFeatures[curCounter][i, :] = getFlowFeatures(dfLists[i].iloc[j]) + outputFeatures[curCounter][i, :] = get_flow_features(dfs[i].iloc[j]) curCounter += 1 curLabel = 0.0 - if np.max(dfLists[i][hitName]) >= threshold: + if np.max(dfs[i][hitName]) >= threshold: curLabel = 1.0 - elif np.max(dfLists[i][hitName]) == -1: + elif np.max(dfs[i][hitName]) == -1: curLabel = -1.0 - elif np.max(dfLists[i][hitName]) > 0 and np.max(dfLists[i][hitName]) < threshold: + elif np.max(dfs[i][hitName]) > 0 and np.max(dfs[i][hitName]) < threshold: curLabel = -2.0 label.append(curLabel) - hits.append(np.max(dfLists[i][hitName])) - trainNames.append(np.unique(dfLists[i]['user_hash'])) + 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)) def get_user_flow_data(): - # load train and test data from joblib - # created with createTrainDataMultipleTaskLearning.py - # rk: changed to csv file - trainDFs = pd.read_csv("data/rk_data.csv.gz") - trainDFs.drop("Unnamed: 0", 1, inplace=True) - trainDFs.set_index(keys=['user_hash'], drop=False, inplace=True) - users = trainDFs['user_hash'].unique().tolist() - u0 = trainDFs.loc[trainDFs.user_hash == users[0]] - return trainDFs + df = pd.read_csv("data/rk_data.csv.gz") + df.drop("Unnamed: 0", 1, inplace=True) + df.set_index(keys=['user_hash'], drop=False, inplace=True) + return df def get_flow_per_user(df): diff --git a/main.py b/main.py index b4f0ada..703af91 100644 --- a/main.py +++ b/main.py @@ -5,6 +5,12 @@ import dataset import models +# config = tf.ConfigProto(log_device_placement=True) +# config.gpu_options.per_process_gpu_memory_fraction = 0.5 +# config.gpu_options.allow_growth = True +# session = tf.Session(config=config) + + def main(): # parameter innerCNNFilters = 512