IDEA_to_STIX fixed port mapping

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)