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'}}