stix2arango 1.1.4__tar.gz → 1.1.6__tar.gz

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

@@ -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 → stix2arango-1.1.6}/pyproject.toml

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

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

@@ -5,31 +5,51 @@ import tempfile
 from pathlib import Path
 from unittest.mock import patch
 
-from stix2arango.stix2arango.bundle_loader import BundleLoader  # adjust this import path
+from stix2arango.stix2arango.bundle_loader import (
+    BundleLoader,
+)  # adjust this import path
 
 # Sample bundle with related and unrelated objects
 STIX_BUNDLE = {
     "type": "bundle",
     "id": "bundle--example",
     "objects": [
-        {"id": "indicator--1", "type": "indicator"},
-        {"id": "indicator--2", "type": "indicator"},
-        {"id": "relationship--1", "type": "relationship", "source_ref": "indicator--1", "target_ref": "indicator--2"},
+        {
+            "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"},
-    ]
+        {"id": "marking--1", "type": "marking-definition"},
+    ],
 }
 
+
 @pytest.fixture
 def temp_json_file():
-    with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.json') as f:
+    with tempfile.NamedTemporaryFile(delete=False, mode="w", suffix=".json") as f:
         json.dump(STIX_BUNDLE, f)
         return f.name
 
+
 def test_init_db_and_tempfile_creation():
     loader = BundleLoader(file_path=tempfile.mkstemp()[1])
     assert isinstance(loader.db_path, str)
     assert Path(loader.db_path).exists()
 
+
 def test_save_to_sqlite_inserts_objects(temp_json_file):
     loader = BundleLoader(file_path=temp_json_file)
     objects = STIX_BUNDLE["objects"]
@@ -39,17 +59,24 @@ def test_save_to_sqlite_inserts_objects(temp_json_file):
     count = cursor.fetchone()[0]
     assert count == len(objects)
 
+
 def test_build_groups_creates_correct_groups(temp_json_file):
-    loader = BundleLoader(file_path=temp_json_file, chunk_size_min=1)
+    loader = BundleLoader(file_path=temp_json_file, chunk_size_min=3)
     groups = loader.build_groups()
 
     # There should be at least 2 groups: one for related indicators, one for attack-pattern
     assert isinstance(groups, list)
     flat = [id_ for group in groups for id_ in group]
-    assert "indicator--1" in flat
-    assert "indicator--2" in flat
-    assert "relationship--1" in flat
-    assert "attack-pattern--3" in flat
+    assert set(flat) == {
+        "some-ref--7",
+        "indicator--2",
+        "indicator--1",
+        "relationship--1",
+        "creator--1",
+        "attack-pattern--3",
+        "marking--1",
+    }
+
 
 def test_load_objects_by_ids(temp_json_file):
     loader = BundleLoader(file_path=temp_json_file)
@@ -59,6 +86,7 @@ def test_load_objects_by_ids(temp_json_file):
     assert isinstance(result, list)
     assert {obj["id"] for obj in result} == {"indicator--1", "attack-pattern--3"}
 
+
 def test_get_objects_returns_objects(temp_json_file):
     loader = BundleLoader(file_path=temp_json_file)
     loader.save_to_sqlite(STIX_BUNDLE["objects"])
@@ -67,6 +95,7 @@ def test_get_objects_returns_objects(temp_json_file):
     assert isinstance(result, list)
     assert result[0]["id"] == "indicator--1"
 
+
 def test_chunks_generator(temp_json_file):
     loader = BundleLoader(file_path=temp_json_file, chunk_size_min=1)
     chunks = list(loader.chunks)
@@ -76,18 +105,20 @@ def test_chunks_generator(temp_json_file):
         assert isinstance(chunk, list)
         assert all("id" in obj for obj in chunk)
 
+
 def test_chunks_with_no_existing_groups_calls_build_groups(temp_json_file):
     loader = BundleLoader(file_path=temp_json_file, chunk_size_min=1)
-    with patch.object(loader, 'build_groups', wraps=loader.build_groups) as mocked:
+    with patch.object(loader, "build_groups", wraps=loader.build_groups) as mocked:
         chunks = list(loader.chunks)
         mocked.assert_called_once()
         assert len(chunks) > 0
 
+
 def test_no_crash_on_missing_relationship_fields(temp_json_file):
     # Remove `source_ref` and `target_ref` to simulate bad data
     bad_bundle = {
         "type": "bundle",
-        "objects": [{"id": "bad--1", "type": "relationship"}]
+        "objects": [{"id": "bad--1", "type": "relationship"}],
     }
     with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
         json.dump(bad_bundle, f)
@@ -104,4 +135,3 @@ def test_sqlite_file_removed_after_gc(temp_json_file):
     assert db_path.exists()
     del loader
     assert not db_path.exists(), "should already be deleted"
-

{stix2arango-1.1.4 → stix2arango-1.1.6}/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(