def main(self):
# path to the pcap to get the update from
if len(sys.argv) < 2:
pcap_path = '/pcaps/eval.pcap'
else:
pcap_path = sys.argv[1]
source_mac = None
key = None
split_path = 'None'
try:
split_path = os.path.split(pcap_path)[-1]
split_path = split_path.split('.')
split_path = split_path[0].split('-')
key = split_path[0].split('_')[1]
except Exception as e:
self.logger.debug('Could not get key because %s', str(e))
# ignore misc files
if (split_path[-1] != 'miscellaneous'):
# Initialize and load the model
if len(sys.argv) > 2:
load_path = sys.argv[2]
else:
load_path = '/models/RandomForestModel.pkl'
# Compute model hash
with open(load_path, 'rb') as handle:
model_hash = hashlib.md5(handle.read()).hexdigest()
model = Model(duration=None,
hidden_size=None,
model_type='RandomForest')
model.load(load_path)
self.logger.debug('Loaded model from %s', load_path)
# Get representations from the model
reps, source_mac, timestamps, preds, others = model.get_representation(
pcap_path, source_ip=source_mac, mean=False)
if preds is not None:
self.logger.debug('Generating predictions')
last_update, prev_rep = self.common.get_previous_state(
source_mac, timestamps[0])
_, mean_rep = self.common.average_representation(
reps,
timestamps,
prev_representation=prev_rep,
last_update=last_update)
mean_preds = model.classify_representation(mean_rep)
if len(sys.argv) > 2:
for p in mean_preds:
self.logger.debug(p)
# Update the stored representation
if reps is not None:
self.logger.debug('Updating stored data')
r_key = self.common.update_data(source_mac, reps,
timestamps, preds, others,
model_hash)
# Get the sessions that the model looked at
sessions = model.sessions
# Clean the sessions
clean_sessions = []
inferred_mac = None
for session_dict in sessions:
cleaned_sessions, inferred_mac = \
clean_session_dict(
session_dict,
source_address=source_mac
)
clean_sessions.append(cleaned_sessions)
if source_mac is None:
source_mac = inferred_mac
# Make simple decisions based on vector differences and update times
timestamp = timestamps[0].timestamp()
labels, confs = zip(*preds)
abnormality = 0
# abnormality = eval_pcap(
# pcap_path, self.conf_labels, self.time_const, label=labels[0], rnn_size=self.rnn_size, model_path='/models/RandomForestModel.pkl', model_type='RandomForest')
prev_s = self.common.get_address_info(source_mac, timestamp)
decision = self.common.basic_decision(key, source_mac, prev_s,
timestamp, labels, confs,
abnormality)
self.logger.debug('Created message')
for i in range(3):
self.logger.info(labels[i] + ' : ' +
str(round(confs[i], 3)))
# update Redis with decision
try:
self.r.hmset(r_key, decision)
except Exception as e:
self.logger.error(
'Failed to update keys in Redis because: {0}'.format(
str(e)))
# Get json message
message = json.dumps(decision)
self.logger.info('Message: ' + message)
if not self.skip_rabbit:
self.common.connect_rabbit()
self.common.channel.basic_publish(
exchange=self.common.exchange,
routing_key=self.common.routing_key,
body=message)
else:
message = {}
message[key] = {'valid': False}
message = json.dumps(message)
self.logger.info('Not enough sessions in pcap')
if not self.skip_rabbit:
self.common.connect_rabbit()
self.common.channel.basic_publish(
exchange=self.common.exchange,
routing_key=self.common.routing_key,
body=message)
if not self.skip_rabbit:
try:
self.common.connection.close()
except Exception as e:
self.logger.error(
'Unable to close rabbit connection because: {0}'.format(
str(e)))
return
config = json.load(config_file)
duration = config['duration']
hidden_size = config['state size']
labels = config['labels']
# Get the data directory
if len(sys.argv) < 2:
data_dir = '/pcaps'
else:
data_dir = sys.argv[1]
m = MLPClassifier((hidden_size),
alpha=0.1,
activation='relu',
max_iter=1000)
# Initialize the model
model = Model(duration=duration,
hidden_size=hidden_size,
labels=labels,
model=m,
model_type='OneLayer')
# Train the model
model.train(data_dir)
# Save the model to the specified path
if len(sys.argv) == 3:
save_path = sys.argv[2]
else:
save_path = 'models/OneLayerModel.pkl'
model.save(save_path)
if __name__ == '__main__':
# Load model params from config
with open('opts/config.json') as config_file:
config = json.load(config_file)
duration = config['duration']
labels = config['labels']
# Get the data directory
if len(sys.argv) < 2:
data_dir = '/pcaps'
else:
data_dir = sys.argv[1]
m = RandomForestClassifier(n_estimators=100,
min_samples_split=5,
class_weight='balanced')
# Initialize the model
model = Model(duration=duration,
labels=labels,
model=m,
model_type='RandomForest')
# Train the model
model.train(data_dir)
# Save the model to the specified path
if len(sys.argv) == 3:
save_path = sys.argv[2]
else:
save_path = 'models/RandomForestModel.pkl'
model.save(save_path)
format(str(e)))
if len(sys.argv) < 2:
data_dir = '/pcaps'
else:
data_dir = sys.argv[1]
# Load model from specified path
if len(sys.argv) > 2:
load_path = sys.argv[2]
else:
load_path = '/models/OneLayerModel.pkl'
if len(sys.argv) >= 4:
save_path = sys.argv[3]
else:
save_path = 'models/OneLayerModel.pkl'
model = Model(duration=None, hidden_size=None, model_type='OneLayer')
logger.info('Loading model from %s', load_path)
model.load(load_path)
# Initialize results dictionary
results = {}
results['labels'] = model.labels
# Get the true label assignments
logger.info('Getting label assignments')
with open('opts/label_assignments.json') as handle:
label_assignments = json.load(handle)
# Walk through testing directory and get all the pcaps
logger.info('Getting pcaps')
pcaps = []