Add check report before inserting by name and date, fix depedency to STIX2

Author: Samuel Hassine

examples/stix2/import.py

@@ -9,7 +9,7 @@
 config = yaml.load(open(os.path.dirname(__file__) + '/../config.yml'))
 
 # File to import
-file_to_import = config['mitre']['repository_path_cti'] + '/enterprise-attack/enterprise-attack.json'
+file_to_import = config['mitre']['repository_path_cti'] + '/apt1.json'
 
 # OpenCTI initialization
 opencti = OpenCTI(config['opencti']['api_url'], config['opencti']['api_key'], config['opencti']['log_file'], config['opencti']['verbose'])

pycti/opencti.py

@@ -91,6 +91,14 @@ def check_existing_stix_domain_entity(self, stix_id=None, name=None, type=None):
             object_result = self.search_stix_domain_entity_by_name(name, type)
         return object_result
 
+    def check_existing_report(self, stix_id=None, name=None, published=None):
+        object_result = None
+        if stix_id is not None:
+            object_result = self.get_stix_domain_entity_by_stix_id(stix_id)
+        if object_result is None and name is not None and published is not None:
+            object_result = self.search_report_by_name_and_date(name, published)
+        return object_result
+
     def update_settings_field(self, id, key, value):
         self.log('Updating settings field ' + key + ' of ' + id + '...')
         query = """
@@ -259,6 +267,29 @@ def search_stix_domain_entity_by_name(self, name_or_alias, type='Stix-Domain-Ent
         else:
             return None
 
+    def search_reports_by_name_and_date(self, name, published):
+        query = """
+               query Reports($name: String, $published: DateTime) {
+                   reports(name: $name, published: $published) {
+                       edges {
+                           node {
+                               id
+                               entity_type
+                           }
+                       }
+                   }
+               }
+           """
+        result = self.query(query, {'name': name, 'published': published})
+        return self.parse_multiple(result['data']['reports'])
+
+    def search_report_by_name_and_date(self, name, published):
+        result = self.search_reports_by_name_and_date(name, published)
+        if len(result) > 0:
+            return result[0]
+        else:
+            return None
+
     def update_stix_domain_entity_field(self, id, key, value):
         self.log('Updating field ' + key + ' of ' + id + '...')
         query = """
@@ -2731,7 +2762,7 @@ def create_report_if_not_exists(self,
                                     modified=None
                                     ):
         if stix_id is not None:
-            object_result = self.get_stix_domain_entity_by_stix_id(stix_id)
+            object_result = self.check_existing_report(stix_id, name, published)
         else:
             object_result = None
         if object_result is not None:

pycti/opencti_stix2.py

@@ -949,7 +949,7 @@ def import_relationship(self, stix_relation, update=False):
         if date is None:
             date = datetime.datetime.utcnow().replace(microsecond=0, tzinfo=datetime.timezone.utc).isoformat()
 
-        stix_relation = self.opencti.create_relation_if_not_exists(
+        stix_relation_result = self.opencti.create_relation_if_not_exists(
             source_id,
             source_type,
             target_id,
@@ -967,8 +967,9 @@ def import_relationship(self, stix_relation, update=False):
             stix_relation['created'] if 'created' in stix_relation else None,
             stix_relation['modified'] if 'modified' in stix_relation else None,
         )
-        if stix_relation is not None:
-            stix_relation_id = stix_relation['id']
+        if stix_relation_result is not None:
+            stix_relation_result_id = stix_relation_result['id']
+            self.mapping_cache[stix_relation['id']] = stix_relation_result_id
         else:
             return None
 
@@ -1040,7 +1041,7 @@ def import_relationship(self, stix_relation, update=False):
                 # Add refs to report
                 self.opencti.add_object_ref_to_report_if_not_exists(report_id, source_id)
                 self.opencti.add_object_ref_to_report_if_not_exists(report_id, target_id)
-                self.opencti.add_object_ref_to_report_if_not_exists(report_id, stix_relation_id)
+                self.opencti.add_object_ref_to_report_if_not_exists(report_id, stix_relation_result_id)
 
     def resolve_author(self, title):
         if 'fireeye' in title.lower() or 'mandiant' in title.lower():

setup.py

@@ -12,7 +12,7 @@
     print("warning: pypandoc module not found, could not convert Markdown to RST")
     read_md = lambda f: open(f, 'r').read()
 
-VERSION = "1.2.4"
+VERSION = "1.2.5"
 
 class VerifyVersionCommand(install):
     description = 'verify that the git tag matches our version'
@@ -50,7 +50,7 @@ def run(self):
         'Topic :: Software Development :: Libraries :: Python Modules'
     ],
     include_package_data=True,
-    install_requires=['requests', 'PyYAML', 'python-dateutil', 'datefinder'],
+    install_requires=['requests', 'PyYAML', 'python-dateutil', 'datefinder', 'stix2', 'pytz'],
     cmdclass={
         'verify': VerifyVersionCommand,
     }