diff --git a/IDEA_to_STIX/IdeaToStix.py b/IDEA_to_STIX/IdeaToStix.py index 58070a1ecfe8c5e3b8ff95c328107feaa6ca3af8..f11baefa3b010f65294e545ee1852cf6fb54effc 100644 --- a/IDEA_to_STIX/IdeaToStix.py +++ b/IDEA_to_STIX/IdeaToStix.py @@ -74,7 +74,7 @@ class StixGenerator(object): object_counter = [object_counter[-1]] return network_values, object_counter, objects - def one_network_traffic_object(self, src_network_references=None, dst_network_references=None): + def one_network_traffic_object(self, src_network_references=None, dst_network_references=None, port=None): network_traffic = { 'type': "network-traffic" } @@ -83,27 +83,39 @@ class StixGenerator(object): network_traffic['src_ref'] = [str(ip_key) for ip_key in src_network_references['Ip_addr_references']] if src_network_references.get('Proto'): network_traffic['protocols'] = src_network_references['Proto'] - if src_network_references.get('Port'): - network_traffic['src_port'] = src_network_references['Port'][0] + if port: + network_traffic['src_port'] = port if dst_network_references: if dst_network_references.get('Ip_addr_references'): network_traffic['dst_ref'] = [str(ip_key) for ip_key in dst_network_references['Ip_addr_references']] if dst_network_references.get('Proto'): network_traffic['protocols'] = dst_network_references['Proto'] - if dst_network_references.get('Port'): - network_traffic['dst_port'] = dst_network_references['Port'][0] + if port: + network_traffic['dst_port'] = port return network_traffic def all_network_traffic_objects(self, src_network_references, dst_network_references, object_counter): objects = {} if src_network_references: for network_record in src_network_references: - objects[str(object_counter)] = self.one_network_traffic_object(network_record) - object_counter += 1 + if network_record.get('Port'): + for port in network_record['Port']: + objects[str(object_counter)] = self.one_network_traffic_object( + src_network_references=network_record, port=port) + object_counter += 1 + else: + objects[str(object_counter)] = self.one_network_traffic_object(network_record) + object_counter += 1 if dst_network_references: for network_record in dst_network_references: - objects[str(object_counter)] = self.one_network_traffic_object(None, network_record) - object_counter += 1 + if network_record.get('Port'): + for port in network_record['Port']: + objects[str(object_counter)] = self.one_network_traffic_object( + dst_network_references=network_record, port=port) + object_counter += 1 + else: + objects[str(object_counter)] = self.one_network_traffic_object(None, network_record) + object_counter += 1 return objects, object_counter def external_references(self, refs): @@ -116,7 +128,7 @@ class StixGenerator(object): 'external_id': record.split(":")[1]}) return ext_references - def observed_data_object(self, identity, data, file, labels=False): + def observed_data_object(self, identity, data, labels=False): observed_data = { 'type': "observed-data", 'id': "observed-data--" + str(uuid4()), @@ -127,7 +139,6 @@ class StixGenerator(object): 'number-observed': data['ConnCount'] if data.get('ConnCount') else 1, 'x_idea_original_data': data } - print(file) if data.get('Ref'): observed_data['external_references'] = self.external_references(data['Ref']) if labels: @@ -185,10 +196,10 @@ def get_args(): help="Path to directory of IDEA files you want to convert.") return parser -def generate_sighting_message(data, category, file): +def generate_sighting_message(data, category): stix_gen = StixGenerator() identity = stix_gen.identity_object(data.get('Node')) - observed_data = stix_gen.observed_data_object(identity['id'], data, file) + observed_data = stix_gen.observed_data_object(identity['id'], data) alert_object = stix_gen.alert_object(category, data.get('Ref')) sighting_object = stix_gen.sighting_object(identity['id'], observed_data['id'], alert_object['id'], data['DetectTime'], data.get('ConnCount'), data.get('EventTime'), @@ -196,10 +207,10 @@ def generate_sighting_message(data, category, file): return [json.dumps(sighting_object), json.dumps(identity), json.dumps(alert_object), json.dumps(observed_data)] -def generate_observable_message(data, file): +def generate_observable_message(data): stix_gen = StixGenerator() identity = stix_gen.identity_object(data.get('Node')) - observed_data = stix_gen.observed_data_object(identity['id'], data, file, True) + observed_data = stix_gen.observed_data_object(identity['id'], data, True) return [json.dumps(identity), json.dumps(observed_data)] @@ -222,9 +233,9 @@ def main(): if type in data['Category'][0]: sighting_message = type if sighting_message: - output = generate_sighting_message(data, sighting_message, file) + output = generate_sighting_message(data, sighting_message) else: - output = generate_observable_message(data, file) + output = generate_observable_message(data) output_file = open(os.path.join(os.getcwd(), "STIX_converted_messages", "STIX_converted_"+file), 'w') for object in output: json.dump(json.JSONDecoder().decode(object), output_file)