def test_results_template():
logger = logging.getLogger(__name__)
instance = ResultsOutput(logger, 'testver', True)
assert instance.results_template(
'/some/dir/trace_ab34_2001-01-01_02_03-client-ip-1-2-3-4.pcap', False, {}) == {'ab34': {'pcap_labels': 'ip-1-2-3-4', 'valid': False}, 'pcap': 'trace_ab34_2001-01-01_02_03-client-ip-1-2-3-4.pcap'}
assert instance.valid_template(
99.0, '192.168.1.1', '0e:00:00:00:00:01', 'normal', False, ['arole'], ['1.0']) == {'decisions': {'behavior': 'normal', 'investigate': False}, 'classification': {'labels': ['arole'], 'confidences': ['1.0']}, 'timestamp': 99.0, 'source_ip': '192.168.1.1', 'source_mac': '0e:00:00:00:00:01'}
def test_results_output():
logger = logging.getLogger(__name__)
instance = ResultsOutput(logger, 'testver', True)
instance.rabbit_host = '127.0.0.1'
results_json = json.dumps({
"filename1": [{"top_role": "role1", "role_list": [["role1", 0.9], ["role2", 0.8], ["role3", 0.7]]}],
"filename2": [{"top_role": "Unknown", "role_list": [["role1", 0.2], ["role2", 0.1], ["role3", 0.01]]}]})
instance.output_from_result_json('1', '/some/file.pcap', results_json)
def output_results(self, result_json_str, run_complete):
if run_complete:
if self.final_stage == 'algorithm' and self.operation == 'predict':
if self.output and os.path.isdir(self.output):
uid = os.getenv('id', 'None')
file_path = os.getenv('file_path', self.in_path)
results_outputter = ResultsOutput(self.logger, uid, file_path)
result_json_file_name = os.path.join(self.output, 'predict.json')
results_outputter.output_from_result_json(result_json_str, result_json_file_name)
def test_rabbit_smoke_good():
logger = logging.getLogger(__name__)
instance = ResultsOutput(logger, 'testver', True)
instance.rabbit_host = '127.0.0.1'
instance.output_msg('x', 'y', 'z')
instance.output_invalid('1', '/some/file.pcap', 'file.pcap')
instance.output_valid('1', '/some/file.pcap', 'file.pcap', 99.0, '1.2.3.4', '0e:00:00:00:00:01', ['arole'], ['1.0'])
def test_rabbit_smoke_bad():
logger = logging.getLogger(__name__)
instance = ResultsOutput(logger, 'testver', True)
for badhost, badport in (
('nosuchthing', 9999),
('127.0.0.1', 65537)):
instance = ResultsOutput(logger, 'testver', True)
instance.rabbit_host = badhost
instance.rabbit_port = badport
instance.output_msg('x', 'y', 'z')
def test_output_from_result_json():
logger = logging.getLogger(__name__)
instance = ResultsOutput(logger, 'testver', 'path/')
result_json = {
'/dir/trace_ab12_2001-01-01_02_03-client-ip-1-2-3-4.pcap': [{
'top_role':
'foo',
'source_ip':
'1.2.3.4',
'source_mac':
'01:02:03:04:05:06',
'timestamp':
999,
'role_list': [('bsomething', 0.7), ('asomething', 0.6),
('csomething', 0.5)]
}],
}
reformatted_result_json_file = os.devnull
reformatted_json = instance.output_from_result_json(
json.dumps(result_json), reformatted_result_json_file)
assert reformatted_json == {
'tool': 'networkml',
'data': {
'mac_addresses': {
'01:02:03:04:05:06': {
'uid': 'testver',
'file_path': 'path/',
'pcap': '',
'pcap_key': '',
'pcap_labels': None,
'timestamp': 999,
'source_ip': '1.2.3.4',
'decisions': {
'investigate': False
},
'classification': {
'labels': ['bsomething', 'asomething', 'csomething'],
'confidences': (0.7, 0.6, 0.5)
}
}
}
}
} # nosec - fine in a test.
def run_stages(self):
stages = ('parser', 'featurizer', 'algorithm')
stage_runners = {
'parser': self.run_parser_stage,
'featurizer': self.run_featurizer_stage,
'algorithm': self.run_algorithm_stage}
try:
first_stage_index = stages.index(self.first_stage)
final_stage_index = stages.index(self.final_stage)
except ValueError:
self.logger.error('Unknown first/final stage name')
return
if first_stage_index > final_stage_index:
self.logger.error('Invalid first and final stage combination')
return
run_schedule = stages[first_stage_index:(final_stage_index+1)]
result = self.in_path
self.logger.info(f'running stages: {run_schedule}')
run_complete = False
try:
for stage in run_schedule:
runner = stage_runners[stage]
result = runner(result)
run_complete = True
except Exception as err:
self.logger.error(f'Could not run stage: {err}')
uid = os.getenv('id', 'None')
file_path = os.getenv('file_path', self.in_path)
results_outputter = ResultsOutput(
self.logger, __version__, self.rabbit)
if run_complete:
if self.final_stage == 'algorithm' and self.operation == 'predict':
if self.output:
if os.path.isdir(self.output):
result_json_file = os.path.join(self.output, 'predict.json')
else:
result_json_file = self.output
with open(result_json_file, 'w') as result_json:
result_json.write(result)
results_outputter.output_from_result_json(uid, file_path, result)
else:
results_outputter.output_invalid(uid, file_path, file_path)
else:
results_outputter.output_invalid(uid, file_path, file_path)
def test_parse_pcap_name():
logger = logging.getLogger(__name__)
instance = ResultsOutput(logger, 'testver', True)
assert instance.parse_pcap_name('notaposeidontracefile.pcap') == (
'notaposeidontracefile', None)
assert instance.parse_pcap_name('trace_but_invalid') == (
None, None)
assert instance.parse_pcap_name('trace_ab12_2001-01-01_02_03-client-ip-1-2-3-4.pcap') == (
'ab12', 'ip-1-2-3-4')
assert instance.parse_pcap_name('trace_8adfcc152604e75d37a1a2ac62124ae859105239_2020-01-21_21_31_44-client-ip-17-253-66-125-17-253-66-125-192-168-3-2-udp-frame-ntp-wsshort-ip-eth-port-123.pcap') == (
'8adfcc152604e75d37a1a2ac62124ae859105239', 'ip-17-253-66-125-17-253-66-125-192-168-3-2-udp-frame-ntp-wsshort-ip-eth-port-123')
assert instance.parse_pcap_name('trace_8198b3326dcb032a2bfbb8030339ff2159b9993d_2020-02-19_03_16_21.pcap') == (
'8198b3326dcb032a2bfbb8030339ff2159b9993d', None)
assert instance.parse_pcap_name('trace_ab12_2001-01-01_02_03-miscellaneous-stuff.pcap') == (
None, None)
def test_rabbit_msg_template():
logger = logging.getLogger(__name__)
instance = ResultsOutput(logger, 'testver', True)
assert instance.rabbit_msg_template('x', 'y', 'z') == {
'id': 'x', 'type': 'metadata', 'file_path': 'y', 'data': 'z', 'results': {'tool': 'networkml', 'version': 'testver'}}