rene/ma_cisco_malware
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

WIP

Browse Source
This commit is contained in:
René Knaebel 2017-06-30 10:42:21 +02:00
parent bbd63fd1da
commit 7ae68cc30e
2 changed files with 41 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

83
dataset.py
View File

@ -6,12 +6,6 @@ import pandas as pd
from tqdm import tqdm 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(): def get_character_dict():
return dict((char, idx) for (idx, char) in return dict((char, idx) for (idx, char) in
enumerate(string.ascii_lowercase + string.punctuation)) enumerate(string.ascii_lowercase + string.punctuation))
@ -60,23 +54,21 @@ def get_user_chunks(dataFrame, windowSize=10, overlapping=False,
return (outDomainLists, outDFFrames) return (outDomainLists, outDFFrames)
def getFeatureVecForDomain(domain, characterDict, maxLen=40): def get_domain_features(domain, vocab, max_length=40):
curFeature = np.zeros([maxLen, ]) curFeature = np.zeros([max_length, ])
for j in range(np.min([len(domain), maxLen])): for j in range(np.min([len(domain), max_length])):
# print(j)
curCharacter = domain[-j] curCharacter = domain[-j]
if curCharacter in characterDict: if curCharacter in vocab:
curFeature[j] = characterDict[curCharacter] curFeature[j] = vocab[curCharacter]
return curFeature return curFeature
def getFlowFeatures(curDataLine): def get_flow_features(flow):
useKeys = ['duration', 'bytes_down', 'bytes_up'] useKeys = ['duration', 'bytes_down', 'bytes_up']
curFeature = np.zeros([len(useKeys), ]) curFeature = np.zeros([len(useKeys), ])
for i in range(len(useKeys)): for i, curKey in enumerate(useKeys):
curKey = useKeys[i]
try: try:
curFeature[i] = np.log1p(curDataLine[curKey]).astype(float) curFeature[i] = np.log1p(flow[curKey]).astype(float)
except: except:
pass pass
return curFeature return curFeature
@ -93,13 +85,13 @@ def getCiscoFeatures(curDataLine, urlSIPDict):
return np.zeros([numCiscoFeatures, ]).ravel() 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 = [] domainLists = []
dfLists = [] dfLists = []
print("get chunks from user data frames") print("get chunks from user data frames")
for i, user_flow in enumerate(get_flow_per_user(user_flow_df)): for i, user_flow in enumerate(get_flow_per_user(user_flow_df)):
(domainListsTmp, dfListsTmp) = get_user_chunks(user_flow, windowSize=windowSize, (domainListsTmp, dfListsTmp) = get_user_chunks(user_flow, windowSize=windowSize,
overlapping=False, maxLengthInSeconds=-1) overlapping=True, maxLengthInSeconds=-1)
domainLists += domainListsTmp domainLists += domainListsTmp
dfLists += dfListsTmp dfLists += dfListsTmp
if i >= 10: 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") print("create training dataset")
return create_dataset_from_lists( return create_dataset_from_lists(
domainLists=domainLists, dfLists=dfLists, charachterDict=char_dict, domains=domainLists, dfs=dfLists, charachterDict=char_dict,
maxLen=maxLen, threshold=threshold, maxLen=maxLen, threshold=threshold,
flagUseCiscoFeatures=False, urlSIPDIct=dict(), use_cisco_features=use_cisco_features, urlSIPDIct=dict(),
windowSize=windowSize) windowSize=windowSize)
def create_dataset_from_lists(domainLists, dfLists, charachterDict, maxLen, threshold=3, def create_dataset_from_lists(domains, dfs, charachterDict, maxLen, threshold=3,
flagUseCiscoFeatures=False, urlSIPDIct=dict(), use_cisco_features=False, urlSIPDIct=dict(),
windowSize=10): windowSize=10):
if 'hits' in dfLists[0].keys(): if 'hits' in dfs[0].keys():
hitName = 'hits' hitName = 'hits'
elif 'virusTotalHits' in dfLists[0].keys(): elif 'virusTotalHits' in dfs[0].keys():
hitName = 'virusTotalHits' hitName = 'virusTotalHits'
numFlowFeatures = 3 numFlowFeatures = 3
numCiscoFeatures = 30 numCiscoFeatures = 30
numFeatures = numFlowFeatures numFeatures = numFlowFeatures
if flagUseCiscoFeatures: if use_cisco_features:
numFeatures += numCiscoFeatures numFeatures += numCiscoFeatures
outputFeatures = [] outputFeatures = []
label = [] label = []
hits = [] hits = []
trainNames = [] trainNames = []
for i in range(windowSize): for i in range(windowSize):
outputFeatures.append(np.zeros([len(domainLists), maxLen])) outputFeatures.append(np.zeros([len(domains), maxLen]))
outputFeatures.append(np.zeros([len(domainLists), numFeatures])) 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 curCounter = 0
# print('len domainList: ' + str(len(domainLists[i]))) # print('len domainList: ' + str(len(domainLists[i])))
# print('len df: ' + str(len(dfLists[i]))) # print('len df: ' + str(len(dfLists[i])))
for j in range(np.min([windowSize, len(domainLists[i])])): for j in range(np.min([windowSize, len(domains[i])])):
outputFeatures[curCounter][i, :] = getFeatureVecForDomain(domainLists[i][j], charachterDict, maxLen) outputFeatures[curCounter][i, :] = get_domain_features(domains[i][j], charachterDict, maxLen)
curCounter += 1 curCounter += 1
if flagUseCiscoFeatures: if use_cisco_features:
outputFeatures[curCounter][i, 0:numFlowFeatures] = getFlowFeatures(dfLists[i].iloc[j]) outputFeatures[curCounter][i, 0:numFlowFeatures] = get_flow_features(dfs[i].iloc[j])
outputFeatures[curCounter][i, numFlowFeatures:] = getCiscoFeatures(dfLists[i].iloc[j], urlSIPDIct) outputFeatures[curCounter][i, numFlowFeatures:] = get_cisco_features(dfs[i].iloc[j], urlSIPDIct)
else: else:
outputFeatures[curCounter][i, :] = getFlowFeatures(dfLists[i].iloc[j]) outputFeatures[curCounter][i, :] = get_flow_features(dfs[i].iloc[j])
curCounter += 1 curCounter += 1
curLabel = 0.0 curLabel = 0.0
if np.max(dfLists[i][hitName]) >= threshold: if np.max(dfs[i][hitName]) >= threshold:
curLabel = 1.0 curLabel = 1.0
elif np.max(dfLists[i][hitName]) == -1: elif np.max(dfs[i][hitName]) == -1:
curLabel = -1.0 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 curLabel = -2.0
label.append(curLabel) label.append(curLabel)
hits.append(np.max(dfLists[i][hitName])) hits.append(np.max(dfs[i][hitName]))
trainNames.append(np.unique(dfLists[i]['user_hash'])) trainNames.append(np.unique(dfs[i]['user_hash']))
return (outputFeatures, np.array(label), np.array(hits), np.array(trainNames)) return (outputFeatures, np.array(label), np.array(hits), np.array(trainNames))
def get_user_flow_data(): def get_user_flow_data():
# load train and test data from joblib df = pd.read_csv("data/rk_data.csv.gz")
# created with createTrainDataMultipleTaskLearning.py df.drop("Unnamed: 0", 1, inplace=True)
# rk: changed to csv file df.set_index(keys=['user_hash'], drop=False, inplace=True)
trainDFs = pd.read_csv("data/rk_data.csv.gz") return df
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
def get_flow_per_user(df): def get_flow_per_user(df):

6
main.py
View File

@ -5,6 +5,12 @@ import dataset
import models 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(): def main():
# parameter # parameter
innerCNNFilters = 512 innerCNNFilters = 512