stix2arango 1.1.4__py3-none-any.whl → 1.1.6__py3-none-any.whl

stix2arango/stix2arango/bundle_loader.py

@@ -10,6 +10,9 @@ import ijson
 import json
 from collections import Counter
 
+from stix2arango.utils import get_embedded_refs
+
+
 class BundleLoader:
     def __init__(self, file_path, chunk_size_min=20_000, db_path=""):
         self.file_path = Path(file_path)
@@ -19,34 +22,37 @@ class BundleLoader:
 
         self.db_path = db_path
         if not self.db_path:
-            self.temp_path = tempfile.NamedTemporaryFile(prefix='s2a_bundle_loader--', suffix='.sqlite')
+            self.temp_path = tempfile.NamedTemporaryFile(
+                prefix="s2a_bundle_loader--", suffix=".sqlite"
+            )
             self.db_path = self.temp_path.name
         self._init_db()
 
     def _init_db(self):
         """Initialize SQLite DB with objects table."""
         self.conn = sqlite3.connect(self.db_path)
-        self.conn.execute('''
+        self.conn.execute(
+            """
             CREATE TABLE IF NOT EXISTS objects (
                 id TEXT PRIMARY KEY,
                 type TEXT,
                 raw TEXT
             )
-        ''')
-        self.conn.execute('PRAGMA synchronous = OFF;')
-        self.conn.execute('PRAGMA journal_mode = MEMORY;')
-        self.conn.execute('PRAGMA temp_store = MEMORY;')
+        """
+        )
+        self.conn.execute("PRAGMA synchronous = OFF;")
+        self.conn.execute("PRAGMA journal_mode = MEMORY;")
+        self.conn.execute("PRAGMA temp_store = MEMORY;")
         self.conn.commit()
 
-
     def save_to_sqlite(self, objects):
         """Save one STIX object to the SQLite database."""
-        self.inserted = getattr(self, 'inserted', 0)
+        self.inserted = getattr(self, "inserted", 0)
 
         try:
             self.conn.executemany(
                 "INSERT OR REPLACE INTO objects (id, type, raw) VALUES (?, ?, ?)",
-                [(obj['id'], obj['type'], json.dumps(obj)) for obj in objects]
+                [(obj["id"], obj["type"], json.dumps(obj)) for obj in objects],
             )
         except sqlite3.IntegrityError as e:
             print(f"Failed to insert len({objects}) objects: {e}")
@@ -55,6 +61,13 @@ class BundleLoader:
         self.inserted += len(objects)
         # logging.info(f"inserted {self.inserted}")
 
+    @staticmethod
+    def get_refs(obj):
+        refs = []
+        for _type, targets in get_embedded_refs(obj):
+            refs.extend(targets)
+        return refs
+
     def build_groups(self):
         """
         Iterates the STIX bundle and uses union-find to group IDs such that for every
@@ -63,30 +76,35 @@ class BundleLoader:
         """
         all_ids: dict[str, list[str]] = dict()  # All object IDs in the file
         logging.info(f"loading into {self.db_path}")
-        
-        with open(self.file_path, 'rb') as f:
-            objects = ijson.items(f, 'objects.item', use_float=True)
+
+        with open(self.file_path, "rb") as f:
+            objects = ijson.items(f, "objects.item", use_float=True)
             to_insert = []
             for obj in objects:
-                obj_id = obj.get('id')
+                obj_id = obj.get("id")
                 to_insert.append(obj)
                 all_ids.setdefault(obj_id, [])
-                if obj['type'] == 'relationship' and all(x in obj for x in ['target_ref', 'source_ref']):
-                    sr, tr = [obj['source_ref'], obj['target_ref']]
+                if obj["type"] == "relationship" and all(
+                    x in obj for x in ["target_ref", "source_ref"]
+                ):
+                    sr, tr = [obj["source_ref"], obj["target_ref"]]
                     all_ids[obj_id].extend([sr, tr])
                     all_ids.setdefault(sr, []).extend([tr, obj_id])
                     all_ids.setdefault(tr, []).extend([sr, obj_id])
+                for ref in self.get_refs(obj):
+                    all_ids[obj_id].append(ref)
                 if len(to_insert) >= self.chunk_size_min:
                     self.save_to_sqlite(to_insert)
                     to_insert.clear()
             if to_insert:
                 self.save_to_sqlite(to_insert)
-        
+
         logging.info(f"loaded {self.inserted} into {self.db_path}")
         handled = set()
 
         self.groups = []
         group = set()
+
         def from_ids(all_ids):
             for obj_id in all_ids:
                 if obj_id in handled:
@@ -104,18 +122,17 @@ class BundleLoader:
         if group:
             self.groups.append(tuple(group))
         return self.groups
-    
+
     def load_objects_by_ids(self, ids):
         """Retrieve a list of STIX objects by their IDs from the SQLite database."""
-        placeholders = ','.join(['?'] * len(ids))
+        placeholders = ",".join(["?"] * len(ids))
         query = f"SELECT raw FROM objects WHERE id IN ({placeholders})"
         cursor = self.conn.execute(query, list(ids))
         return [json.loads(row[0]) for row in cursor.fetchall()]
 
-
     def get_objects(self, group):
         return list(self.load_objects_by_ids(group))
-    
+
     @property
     def chunks(self):
         for group in self.groups or self.build_groups():
@@ -123,4 +140,4 @@ class BundleLoader:
 
     def __del__(self):
         with contextlib.suppress(Exception):
-            os.remove(self.db_path)
+            os.remove(self.db_path)

stix2arango/utils.py

@@ -125,6 +125,9 @@ def get_embedded_refs(object: list | dict, xpath: list = [], attributes=None):
             if match := EMBEDDED_RELATIONSHIP_RE.fullmatch(key):
                 relationship_type = "-".join(xpath + match.group(1).split("_"))
                 targets = value if isinstance(value, list) else [value]
+                targets = [_target for _target in targets if _target and isinstance(_target, str)]
+                if not targets:
+                    continue
                 if attributes and key not in attributes:
                     continue
                 embedded_refs.append((relationship_type, targets))

{stix2arango-1.1.4.dist-info → stix2arango-1.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stix2arango
-Version: 1.1.4
+Version: 1.1.6
 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.1.4.dist-info → stix2arango-1.1.6.dist-info}/RECORD

@@ -1,16 +1,16 @@
 stix2arango/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 stix2arango/__main__.py,sha256=zsCi_bfDULLDkqlRwXyGhFuLvSRcvESEc4MMN7h1lbQ,2835
 stix2arango/config.py,sha256=NZFrcnEfz-0QBrut2Rh7xMF78v0bk6U6y2TY_7mHxSs,1407
-stix2arango/utils.py,sha256=bUKJBQ2owbCQKWs_m-VYjYCHuQLykizabE4D3LPspW8,4636
+stix2arango/utils.py,sha256=C-VDwsCABFU2hrv1EwF7oaQUYOE2MWG40_7WGzHu0A4,4796
 stix2arango/services/__init__.py,sha256=E87fB-dxI4mPxMVs00jdLhjp9jFhkVfjhMKIqGLRJlY,45
 stix2arango/services/arangodb_service.py,sha256=jr6zXFueluCU60WOJy7XuA9Ty0zW5FzGNBJGtJzq0PY,11964
 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/bundle_loader.py,sha256=YphKnJIiHjg_nuQUf59OUthDOVIQYoOIoOqqQUlU1II,4982
 stix2arango/stix2arango/stix2arango.py,sha256=HJXDqA9NWxXVQSHPmbpkEKurpWEbZmy5bng5SQ1OsjE,22412
 stix2arango/templates/marking-definition.json,sha256=0q9y35mUmiF6xIWSLpkATL4JTHGSCNyLbejqZiQ0AuE,3113
-stix2arango-1.1.4.dist-info/METADATA,sha256=Hre8CoZ_6ic_jNFQ1ONrA8jTr3wLL54wFT-_nLYGsmY,7797
-stix2arango-1.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-stix2arango-1.1.4.dist-info/entry_points.txt,sha256=k2WnxMsHFLoyC6rqfvjhIMS1zwtWin51-MbNCGmSMYE,58
-stix2arango-1.1.4.dist-info/licenses/LICENSE,sha256=BK8Ppqlc4pdgnNzIxnxde0taoQ1BgicdyqmBvMiNYgY,11364
-stix2arango-1.1.4.dist-info/RECORD,,
+stix2arango-1.1.6.dist-info/METADATA,sha256=ga8FQKsBxCFYGOyV6YAV9gZOaMv8-0nkLiyTpTsxWKA,7797
+stix2arango-1.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+stix2arango-1.1.6.dist-info/entry_points.txt,sha256=k2WnxMsHFLoyC6rqfvjhIMS1zwtWin51-MbNCGmSMYE,58
+stix2arango-1.1.6.dist-info/licenses/LICENSE,sha256=BK8Ppqlc4pdgnNzIxnxde0taoQ1BgicdyqmBvMiNYgY,11364
+stix2arango-1.1.6.dist-info/RECORD,,