stix2arango 1.1.4__tar.gz → 1.1.5__tar.gz

{stix2arango-1.1.4 → stix2arango-1.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stix2arango
-Version: 1.1.4
+Version: 1.1.5
 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 → stix2arango-1.1.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stix2arango"
-version = "1.1.4"
+version = "1.1.5"
 authors = [
   { name = "dogesec" }
 ]

{stix2arango-1.1.4 → stix2arango-1.1.5}/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,15 @@ 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):
+            if _type in ["created-by", "object-marking"]:
+                continue
+            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 +78,36 @@ 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)
+                    all_ids.setdefault(ref, []).append(obj_id)
                 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 +125,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 +143,4 @@ class BundleLoader:
 
     def __del__(self):
         with contextlib.suppress(Exception):
-            os.remove(self.db_path)
+            os.remove(self.db_path)

{stix2arango-1.1.4 → stix2arango-1.1.5}/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 → stix2arango-1.1.5}/tests/src/test_bundle_loader.py

@@ -12,8 +12,8 @@ STIX_BUNDLE = {
     "type": "bundle",
     "id": "bundle--example",
     "objects": [
-        {"id": "indicator--1", "type": "indicator"},
-        {"id": "indicator--2", "type": "indicator"},
+        {"id": "indicator--1", "type": "indicator", "bad_ref": "some-ref--7", "created_by_ref": "creator--1"},
+        {"id": "indicator--2", "type": "indicator", "object_marking_refs": ["marking--1", "marking--2"]},
         {"id": "relationship--1", "type": "relationship", "source_ref": "indicator--1", "target_ref": "indicator--2"},
         {"id": "attack-pattern--3", "type": "attack-pattern"},
     ]
@@ -50,6 +50,10 @@ def test_build_groups_creates_correct_groups(temp_json_file):
     assert "indicator--2" in flat
     assert "relationship--1" in flat
     assert "attack-pattern--3" in flat
+    assert "some-ref--7" in flat
+    assert "marking--1" not in flat
+    assert "marking--2" not in flat
+    assert "creator--1" not in flat
 
 def test_load_objects_by_ids(temp_json_file):
     loader = BundleLoader(file_path=temp_json_file)

{stix2arango-1.1.4 → stix2arango-1.1.5}/tests/src/test_utils.py

@@ -52,6 +52,22 @@ def test_get_embedded_refs():
         ("abcde-abcd-efgh", ["ref8"]),
     ]
 
+def test_get_embedded_refs_empty():
+    assert utils.get_embedded_refs(
+        {
+            "abc_ref": "ref1",
+            "empty_ref": "", #skipped entirely
+            "some_empty_refs": ["ref10", "", "ref9"], # empty ref skipped
+            "abcd_refs": ["ref1", "ref2"],
+            "abcde": [{"abcdef_ref": "ref7"}, {"abcd_efgh_ref": "ref8"}],
+        }
+    ) == [
+        ("abc", ["ref1"]),
+        ("some-empty", ["ref10", "ref9"]),
+        ("abcd", ["ref1", "ref2"]),
+        ("abcde-abcdef", ["ref7"]),
+        ("abcde-abcd-efgh", ["ref8"]),
+    ]
 
 def test_get_embedded_refs__attributes_whitelist():
     assert utils.get_embedded_refs(