stix2arango 1.0.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

stix2arango/services/arangodb_service.py

@@ -220,7 +220,8 @@ class ArangoDBService:
         )
         out = [dict(zip(('id', '_key', 'modified', '_record_modified', '_is_latest', '_id'), obj_tuple)) for obj_tuple in out]
         annotated, deprecated = annotate_versions(out)
-        self.db.collection(collection_name).update_many(annotated, sync=True, keep_none=False)
+        for chunk in utils.chunked(annotated, 10_000):
+            self.db.collection(collection_name).update_many(chunk, sync=True, keep_none=False)
         return deprecated
 
 
@@ -297,8 +298,9 @@ class ArangoDBService:
     @contextlib.contextmanager
     def transactional(self, write=None, exclusive=None, sync=True):
         original_db = self.db
-        transactional_db = self.db.begin_transaction(allow_implicit=True, write=write, exclusive=exclusive, sync=sync)
+        transactional_db = self.db.begin_transaction(allow_implicit=True, write=write, exclusive=exclusive, sync=sync, lock_timeout=300)
         try:
+            logging.info(f"entering transaction: {transactional_db.transaction_status()}")
             self.db = transactional_db
             yield self
             transactional_db.commit_transaction()
@@ -306,4 +308,5 @@ class ArangoDBService:
             transactional_db.abort_transaction()
             raise
         finally:
+            logging.info(f"exiting transaction: {transactional_db.transaction_status()}")
             self.db = original_db

stix2arango/stix2arango/stix2arango.py

@@ -16,12 +16,14 @@ from tqdm import tqdm
 from ..services.arangodb_service import ArangoDBService
 from jsonschema import validate
 from arango.collection import StandardCollection
+import arango.exceptions
 
 
 from .. import utils
 
 module_logger = logging.getLogger("data_ingestion_service")
 SMO_TYPES = ["marking-definition", "extension-definition", "language-content"]
+LARGE_FILE_SIZE = 80 * 1024 * 1024
 
 
 class Stix2Arango:
@@ -52,9 +54,9 @@ class Stix2Arango:
         """
         `modify_fn` should modify in-place, returned value is discarded
         """
-        
+
         self.alter_functions = []
-        
+
         self.core_collection_vertex, self.core_collection_edge = (
             utils.get_vertex_and_edge_collection_names(collection)
         )
@@ -76,7 +78,7 @@ class Stix2Arango:
         self.arangodb_extra_data = {}
 
         self.file = file
-        self.is_large_file = is_large_file
+        self._is_large_file = is_large_file
         self.note = stix2arango_note or ""
         self.identity_ref = utils.load_file_from_url(config.STIX2ARANGO_IDENTITY)
         self.default_ref_objects = [
@@ -98,6 +100,10 @@ class Stix2Arango:
         if self.file:
             self.filename = Path(self.file).name
 
+    @property
+    def is_large_file(self):
+        return self._is_large_file or os.path.getsize(self.file) > LARGE_FILE_SIZE
+
     def alter_objects(self, objects: list[dict]):
         for obj in objects:
             obj.update(self.arangodb_extra_data)
@@ -106,7 +112,9 @@ class Stix2Arango:
                     fn(obj)
                 except Exception as e:
                     logging.warning(f"alter function {fn} failed on {obj}")
-                    logging.warning(f"alter function {fn} failed on {obj}", exc_info=True)
+                    logging.warning(
+                        f"alter function {fn} failed on {obj}", exc_info=True
+                    )
 
     def add_object_alter_fn(self, modify_fn):
         if not callable(modify_fn):
@@ -157,16 +165,37 @@ class Stix2Arango:
                     )
                 )
 
+    def create_analyzer(self, *args, **kwargs):
+        try:
+            return self.arango.db.create_analyzer(*args, **kwargs)
+        except arango.exceptions.AnalyzerCreateError as e:
+            if e.error_code != 10:
+                raise
+    
     def create_taxii_views(self):
         views = set()
+        self.create_analyzer(
+            name="date_transform",
+            analyzer_type="aql",
+            features=[],
+            properties={
+                "queryString": "RETURN DATE_TIMESTAMP(@param)*1000 + TO_NUMBER(LAST(REGEX_MATCHES(@param, '.+\\\\.(\\\\d{6}).*')))%1000",
+                "collapsePositions": False,
+                "keepNull": True,
+                "batchSize": 1000,
+                "memoryLimit": 10485760,
+                "returnType": "number",
+            },
+        )
         for name, collection in self.arango.collections.items():
+            logging.info(f'creating taxii index for {name}')
             collection.add_index(
                 dict(
                     type="inverted",
                     name="taxii_search",
                     sparse=True,
                     fields=[
-                        "_record_created",
+                        dict(name="_record_created", analyzer="date_transform"),
                         "modified",
                         "id",
                         "_taxii.visible",
@@ -182,7 +211,12 @@ class Stix2Arango:
                     },
                 )
             )
-            views.add('ats__' + name.removesuffix('_vertex_collection').removesuffix('_edge_collection'))
+            views.add(
+                "ats__"
+                + name.removesuffix("_vertex_collection").removesuffix(
+                    "_edge_collection"
+                )
+            )
 
     def create_default_indexes(self):
         for name, collection in self.arango.collections.items():
@@ -354,17 +388,23 @@ class Stix2Arango:
         module_logger.info(
             f"Inserting objects into database. Total objects: {len(objects)}"
         )
-        with self.arango.transactional(exclusive=[self.core_collection_edge, self.core_collection_vertex]):
+        with self.arango.transactional(
+            exclusive=[self.core_collection_edge, self.core_collection_vertex]
+        ):
             inserted_object_ids, existing_objects = (
                 self.arango.insert_several_objects_chunked(
                     objects, self.core_collection_vertex
                 )
             )
             deprecated_key_ids = self.arango.update_is_latest_several_chunked(
-                inserted_object_ids, self.core_collection_vertex, self.core_collection_edge
+                inserted_object_ids,
+                self.core_collection_vertex,
+                self.core_collection_edge,
             )
 
-        self.update_object_key_mapping(self.core_collection_vertex, objects, existing_objects)
+        self.update_object_key_mapping(
+            self.core_collection_vertex, objects, existing_objects
+        )
         return inserted_object_ids, existing_objects, deprecated_key_ids
 
     def update_object_key_mapping(self, collection, objects, existing_objects={}):
@@ -409,16 +449,22 @@ class Stix2Arango:
         module_logger.info(
             f"Inserting relationship into database. Total objects: {len(objects)}"
         )
-        with self.arango.transactional(exclusive=[self.core_collection_edge, self.core_collection_vertex]):
+        with self.arango.transactional(
+            exclusive=[self.core_collection_edge, self.core_collection_vertex]
+        ):
             inserted_object_ids, existing_objects = (
                 self.arango.insert_relationships_chunked(
                     objects, self.object_key_mapping, self.core_collection_edge
                 )
             )
             deprecated_key_ids = self.arango.update_is_latest_several_chunked(
-                inserted_object_ids, self.core_collection_edge, self.core_collection_edge
+                inserted_object_ids,
+                self.core_collection_edge,
+                self.core_collection_edge,
             )
-        self.update_object_key_mapping(self.core_collection_edge, objects, existing_objects)
+        self.update_object_key_mapping(
+            self.core_collection_edge, objects, existing_objects
+        )
         return inserted_object_ids, deprecated_key_ids
 
     def map_embedded_relationships(self, bundle_objects, inserted_object_ids):
@@ -441,7 +487,7 @@ class Stix2Arango:
                     targets=targets,
                     relationship=ref_type,
                     arango_obj=self,
-                    bundle_id=self.bundle_id or '',
+                    bundle_id=self.bundle_id or "",
                     insert_statement=objects,
                     extra_data=self.arangodb_extra_data,
                 )
@@ -454,11 +500,11 @@ class Stix2Arango:
         inserted_object_ids = []
         existing_objects = {}
         for chunk in utils.chunked(objects, 20_000):
-            with self.arango.transactional(exclusive=[self.core_collection_edge, self.core_collection_vertex]):
-                inserted, existing = (
-                    self.arango.insert_relationships_chunked(
-                        chunk, self.object_key_mapping, self.core_collection_edge
-                    )
+            with self.arango.transactional(
+                exclusive=[self.core_collection_edge, self.core_collection_vertex]
+            ):
+                inserted, existing = self.arango.insert_relationships_chunked(
+                    chunk, self.object_key_mapping, self.core_collection_edge
                 )
                 inserted_object_ids.extend(inserted)
                 existing_objects.update(existing)
@@ -541,8 +587,9 @@ class Stix2Arango:
                 all_objects, inserted_object_ids + inserted_relationship_ids
             )
 
-        
-        with self.arango.transactional(exclusive=[self.core_collection_edge, self.core_collection_vertex]):
+        with self.arango.transactional(
+            exclusive=[self.core_collection_edge, self.core_collection_vertex]
+        ):
             self.arango.deprecate_relationships(
                 deprecated_key_ids1, self.core_collection_edge
             )

{stix2arango-1.0.2.dist-info → stix2arango-1.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stix2arango
-Version: 1.0.2
+Version: 1.1.0
 Summary: stix2arango is a command line tool that takes a group of STIX 2.1 objects in a bundle and inserts them into ArangoDB. It can also handle updates to existing objects in ArangoDB imported in a bundle.
 Project-URL: Homepage, https://github.com/muchdogesec/stix2arango
 Project-URL: Issues, https://github.com/muchdogesec/stix2arango/issues

{stix2arango-1.0.2.dist-info → stix2arango-1.1.0.dist-info}/RECORD

@@ -3,14 +3,14 @@ stix2arango/__main__.py,sha256=wbR_iO70Vld2NYiml6Kz4rH396uOiNwTtjNBl4AHZEg,1987
 stix2arango/config.py,sha256=NZFrcnEfz-0QBrut2Rh7xMF78v0bk6U6y2TY_7mHxSs,1407
 stix2arango/utils.py,sha256=heln_kXBSTLo0R4n1IrfCHbRnCHTQf9Y2URhvwQ0Smc,4845
 stix2arango/services/__init__.py,sha256=E87fB-dxI4mPxMVs00jdLhjp9jFhkVfjhMKIqGLRJlY,45
-stix2arango/services/arangodb_service.py,sha256=mXBQUeKR9IHErmvuZ73AiASZKawZOpuRiJQXWaN8bx8,11706
+stix2arango/services/arangodb_service.py,sha256=qgE7yb-ysqyJ1oUjqAV6eGa53WvWHAFAd6I7aKZ6KyU,11960
 stix2arango/services/version_annotator.py,sha256=Sd1MIaXzK0fpNopNxRoB_3etodzAjX5D_p3uGQSWzOI,2946
 stix2arango/stix2arango/__init__.py,sha256=OqxWEEsHqR1QQpznM5DbFJ5bO5numKYtoYhjXYJMEyg,36
 stix2arango/stix2arango/bundle_loader.py,sha256=qi-0E_bMIMPZXzISvjhrWX8K-f7iFv9vOekldOGVczU,4603
-stix2arango/stix2arango/stix2arango.py,sha256=UVv1OVy_BzqXYN9cWzjKK3QUSRuS-7jf3n_IgGQwx_8,20769
+stix2arango/stix2arango/stix2arango.py,sha256=PyrDzW8YuKzibpMA0UImwpDBf1BeiP5pUz5tE_nVwG4,22073
 stix2arango/templates/marking-definition.json,sha256=0q9y35mUmiF6xIWSLpkATL4JTHGSCNyLbejqZiQ0AuE,3113
-stix2arango-1.0.2.dist-info/METADATA,sha256=uuH8tlqansoJvxNkqCxf4h_h1MteaRUQIlSQQG2DNOA,6873
-stix2arango-1.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-stix2arango-1.0.2.dist-info/entry_points.txt,sha256=k2WnxMsHFLoyC6rqfvjhIMS1zwtWin51-MbNCGmSMYE,58
-stix2arango-1.0.2.dist-info/licenses/LICENSE,sha256=BK8Ppqlc4pdgnNzIxnxde0taoQ1BgicdyqmBvMiNYgY,11364
-stix2arango-1.0.2.dist-info/RECORD,,
+stix2arango-1.1.0.dist-info/METADATA,sha256=DHq9KWi0BV3juGOP-JLh5y9rxB-htxKu_l4vrw19fjg,6873
+stix2arango-1.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+stix2arango-1.1.0.dist-info/entry_points.txt,sha256=k2WnxMsHFLoyC6rqfvjhIMS1zwtWin51-MbNCGmSMYE,58
+stix2arango-1.1.0.dist-info/licenses/LICENSE,sha256=BK8Ppqlc4pdgnNzIxnxde0taoQ1BgicdyqmBvMiNYgY,11364
+stix2arango-1.1.0.dist-info/RECORD,,