stix2arango 1.1.1__tar.gz → 1.1.3__tar.gz

{stix2arango-1.1.1 → stix2arango-1.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stix2arango
-Version: 1.1.1
+Version: 1.1.3
 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.1 → stix2arango-1.1.3}/docs/README.md

@@ -422,6 +422,12 @@ This is done using the STIX2 Python library, and such relationships are generate
     "target_ref": "<REF TARGET PROPERTY>",
     "object_marking_refs": [
         "<object_marking_refs SHOULD MATCH SOURCE_REF OBJECT>"
+    ],
+    "external_references": [
+        {
+            "source_name": "stix2arango",
+            "description": "embedded-relationship"
+        }
     ]
 }
 ```

{stix2arango-1.1.1 → stix2arango-1.1.3}/pyproject.toml

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

{stix2arango-1.1.1 → stix2arango-1.1.3}/stix2arango/services/arangodb_service.py

@@ -63,7 +63,7 @@ class ArangoDBService:
 
         if self.db.has_graph(self.ARANGO_GRAPH):
             self.cti2stix_graph = self.db.graph(self.ARANGO_GRAPH)
-        elif create:
+        elif create_db:
             self.cti2stix_graph = self.db.create_graph(self.ARANGO_GRAPH)
 
         self.collections: dict[str, StandardCollection] = {}
@@ -199,7 +199,7 @@ class ArangoDBService:
     
     @staticmethod
     def fix_edge_ref(_id):
-        c, _, _key = _id.partition('/')
+        c, _, _key = _id.rpartition('/')
         if not c:
             c = "missing_collection"
         return f"{c}/{_key}"

{stix2arango-1.1.1 → stix2arango-1.1.3}/stix2arango/stix2arango/stix2arango.py

@@ -80,9 +80,9 @@ class Stix2Arango:
         self.file = 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.identity_ref = utils.load_file_from_url(config.STIX2ARANGO_IDENTITY).copy()
         self.default_ref_objects = [
-            utils.load_file_from_url(link)
+            utils.load_file_from_url(link).copy()
             for link in config.MARKING_DEFINITION_REFS + config.IDENTITY_REFS
         ]
         self.bundle_id = bundle_id
@@ -469,7 +469,6 @@ class Stix2Arango:
 
     def map_embedded_relationships(self, bundle_objects, inserted_object_ids):
         objects = []
-        inserted_data = []
         for obj in tqdm(bundle_objects, desc="upload_embedded_edges"):
             if obj["id"] not in inserted_object_ids:
                 continue

{stix2arango-1.1.1 → stix2arango-1.1.3}/stix2arango/utils.py

@@ -8,7 +8,7 @@ import json
 import hashlib
 import os
 from . import config
-from datetime import datetime
+from datetime import UTC, datetime
 
 module_logger = logging.getLogger("data_ingestion_service")
 from arango.database import StandardDatabase
@@ -37,7 +37,6 @@ def create_relationship_obj(
     arango_obj=None,
     extra_data=None,
 ):
-    insert_data = []
     if not isinstance(targets, list):
         return []
     for target in targets:
@@ -68,26 +67,18 @@ def create_relationship_obj(
         relationship_object["_bundle_id"] = bundle_id
         relationship_object["_file_name"] = os.path.basename(arango_obj.file or "")
         relationship_object["_stix2arango_note"] = arango_obj.note
-        relationship_object["_record_created"] = datetime.now().strftime(
-            "%Y-%m-%dT%H:%M:%S.%f"
-        )
-        relationship_object["_record_modified"] = relationship_object["_record_created"]
         relationship_object["_is_ref"] = True
         relationship_object["type"] = "relationship"
         relationship_object["spec_version"] = "2.1"
-        # relationship_object['_record_md5_hash'] = generate_md5(relationship_object)
+        relationship_object["external_references"] = {
+            "source_name": "stix2arango",
+            "description": "embedded-relationship",
+        }
         if extra_data:
             relationship_object.update(extra_data)
 
         insert_statement.append(relationship_object)
-        insert_data.append(
-            [
-                "relationship",
-                relationship_object["id"],
-                True if "modified" in obj else False,
-            ]
-        )
-    return insert_data
+    return
 
 
 def generate_md5(obj: dict):

stix2arango-1.1.3/tests/src/test_arango_service.py

@@ -0,0 +1,61 @@
+
+
+import pytest
+from stix2arango import config
+from stix2arango.services.arangodb_service import ArangoDBService
+from stix2arango.stix2arango.stix2arango import Stix2Arango
+import arango.exceptions
+from unittest.mock import patch, MagicMock
+
+
+def test_get_db_name():
+    assert ArangoDBService.get_db_name("test") == 'test_database'
+    assert ArangoDBService.get_db_name("test_database") == 'test_database'
+
+@pytest.fixture(scope='module')
+def service():
+    s = ArangoDBService('test_s2a_arango_service', [], [], host_url=Stix2Arango.ARANGODB_URL, username=config.ARANGODB_USERNAME, password=config.ARANGODB_PASSWORD, create_db=True)
+    yield s
+    s.sys_db.delete_database('test_s2a_arango_service_database')
+
+
+def test_init_creates_graph(service):
+    assert service.db.has_graph('test_s2a_arango_service_graph')
+
+
+def test_create_db(service):
+    service.create_database('test_s2a_arango_service2')
+    assert service.sys_db.has_database('test_s2a_arango_service2')
+    service.create_database('test_s2a_arango_service2')
+    service.sys_db.delete_database('test_s2a_arango_service2')
+
+def test_create_collection(service):
+    assert service.create_collection('ss_vertex_collection').info()['name'] == 'ss_vertex_collection', "created collection must have the passed name"
+    assert service.create_collection('ss_vertex_collection').info()['name'] == 'ss_vertex_collection', "recreating should not throw error"
+
+def test_execute_raw_query(service):
+    c = service.create_collection('execute_raw_query')
+    c.insert_many([{"a": 1}, {"a": 2}, {"a": 3}])
+    assert service.execute_raw_query("FOR d IN @@collection FILTER d.a >= 2 RETURN d.a", bind_vars={'@collection': c.name}) == [2, 3]
+    with pytest.raises(arango.exceptions.AQLQueryExecuteError):
+        service.execute_raw_query("FOR d IN @@collection FILTER d.a >= 2 RETURN d.a")
+
+def test_fix_edge_ref():
+    assert ArangoDBService.fix_edge_ref('nothing') == "missing_collection/nothing"
+    assert ArangoDBService.fix_edge_ref('abcd/nothing') == "abcd/nothing"
+
+def test_transactional_resets_db(service):
+    db = service.db
+    with patch('arango.database.TransactionDatabase.commit_transaction') as mock_commit, service.transactional():
+        pass
+    assert service.db == db, "db not reset to original after context exit"
+    mock_commit.assert_called_once()
+
+
+def test_transactional_resets_db__on_failure(service):
+    db = service.db
+    with patch('arango.database.TransactionDatabase.abort_transaction') as mock_abort, pytest.raises(Exception), service.transactional():
+        raise Exception("must raise")
+    assert service.db == db, "db not reset to original after context exit"
+    mock_abort.assert_called_once()
+

stix2arango-1.1.3/tests/src/test_utils.py

@@ -0,0 +1,168 @@
+from datetime import datetime
+from unittest.mock import MagicMock, patch
+import pytest
+from stix2arango import utils
+from stix2arango.config import STIX2ARANGO_IDENTITY
+
+
+def test_load_file_from_url():
+    assert utils.load_file_from_url(STIX2ARANGO_IDENTITY) == {
+        "type": "identity",
+        "spec_version": "2.1",
+        "id": "identity--72e906ce-ca1b-5d73-adcd-9ea9eb66a1b4",
+        "created_by_ref": "identity--9779a2db-f98c-5f4b-8d08-8ee04e02dbb5",
+        "created": "2020-01-01T00:00:00.000Z",
+        "modified": "2020-01-01T00:00:00.000Z",
+        "name": "stix2arango",
+        "description": "https://github.com/muchdogesec/stix2arango",
+        "identity_class": "system",
+        "sectors": ["technology"],
+        "contact_information": "https://www.dogesec.com/contact/",
+        "object_marking_refs": [
+            "marking-definition--94868c89-83c2-464b-929b-a1a8aa3c8487",
+            "marking-definition--97ba4e8b-04f6-57e8-8f6e-3a0f0a7dc0fb",
+        ],
+    }
+
+    with pytest.raises(Exception):
+        utils.load_file_from_url("https://gogle.cm/ojkskja")
+
+
+def test_remove_duplicates():
+    assert (
+        len(
+            utils.remove_duplicates(
+                [
+                    {
+                        "id": "all-the-same",
+                    },
+                    {
+                        "id": "all-the-same",
+                    },
+                    {"id": "all-the-same", "a": 1},
+                    {"id": "all-the-same", "a": 1, "_is_latest": True},
+                    {"id": "all-the-same", "a": 2, "b": 1},
+                    {"id": "all-the-same", "a": 2, "b": 1, "c": 1},
+                    {"id": "all-the-same", "a": 2, "b": 1, "c": 1, "_is": 8},
+                ]
+            )
+        )
+        == 4
+    )
+
+
+def test_get_embedded_refs():
+    assert utils.get_embedded_refs(
+        {
+            "abc_ref": "ref1",
+            "abcd_refs": ["ref1", "ref2"],
+            "abcde": [{"abcdef_ref": "ref7"}, {"abcd_efgh_ref": "ref8"}],
+        }
+    ) == [
+        ("abc", ["ref1"]),
+        ("abcd", ["ref1", "ref2"]),
+        ("abcde-abcdef", ["ref7"]),
+        ("abcde-abcd-efgh", ["ref8"]),
+    ]
+
+
+def test_get_vertex_and_edge_collection_names():
+    assert utils.get_vertex_and_edge_collection_names("ade") == (
+        "ade_vertex_collection",
+        "ade_edge_collection",
+    )
+    assert utils.get_vertex_and_edge_collection_names("ade_edge_collection") == (
+        "ade_vertex_collection",
+        "ade_edge_collection",
+    )
+    assert utils.get_vertex_and_edge_collection_names("ade_vertex_collection") == (
+        "ade_vertex_collection",
+        "ade_edge_collection",
+    )
+
+
+def test_create_relationship__no_targets():
+    assert utils.create_relationship_obj(None, None, None, None, None, None) == []
+
+
+def test_create_relationships():
+    obj = {
+        "created": "2024-01-01T00:00:00.000000",
+        "modified": "2024-01-02T00:00:00.000000",
+        "object_marking_refs": ["marking--xyz"],
+    }
+
+    source = "indicator--source"
+    targets = ["malware--target1", "software--target--2"]
+    relationship = "indicates"
+    insert_statement = []
+    bundle_id = "bundle--xyz"
+    extra_data = {"confidence": 80, "labels": ["example"]}
+    arango_obj = MagicMock()
+    arango_obj.identity_ref = {"id": "identity--abcd"}
+    arango_obj.core_collection_vertex = "vertices"
+    arango_obj.file = "/some/path/file.txt"
+    arango_obj.note = "This is a note"
+
+    utils.create_relationship_obj(
+        obj=obj,
+        source=source,
+        targets=targets,
+        relationship=relationship,
+        insert_statement=insert_statement,
+        bundle_id=bundle_id,
+        arango_obj=arango_obj,
+        extra_data=extra_data,
+    )
+
+    assert len(insert_statement) == 2
+    assert insert_statement == [
+        {
+            "relationship_type": "indicates",
+            "created": "2024-01-01T00:00:00.000000",
+            "modified": "2024-01-02T00:00:00.000000",
+            "object_marking_refs": ["marking--xyz"],
+            "id": "relationship--8a86537f-d4bb-55bc-b06d-ddeabbd9899d",
+            "created_by_ref": "identity--abcd",
+            "source_ref": "indicator--source",
+            "target_ref": "malware--target1",
+            "_from": "vertices/indicator--source",
+            "_to": "vertices/malware--target1",
+            "_bundle_id": "bundle--xyz",
+            "_file_name": "file.txt",
+            "_stix2arango_note": "This is a note",
+            "_is_ref": True,
+            "type": "relationship",
+            "spec_version": "2.1",
+            "external_references": {
+                "source_name": "stix2arango",
+                "description": "embedded-relationship",
+            },
+            "confidence": 80,
+            "labels": ["example"],
+        },
+        {
+            "relationship_type": "indicates",
+            "created": "2024-01-01T00:00:00.000000",
+            "modified": "2024-01-02T00:00:00.000000",
+            "object_marking_refs": ["marking--xyz"],
+            "id": "relationship--fda071b0-e1e8-5831-a707-850f1b1a5a0e",
+            "created_by_ref": "identity--abcd",
+            "source_ref": "indicator--source",
+            "target_ref": "software--target--2",
+            "_from": "vertices/indicator--source",
+            "_to": "vertices/software--target--2",
+            "_bundle_id": "bundle--xyz",
+            "_file_name": "file.txt",
+            "_stix2arango_note": "This is a note",
+            "_is_ref": True,
+            "type": "relationship",
+            "spec_version": "2.1",
+            "external_references": {
+                "source_name": "stix2arango",
+                "description": "embedded-relationship",
+            },
+            "confidence": 80,
+            "labels": ["example"],
+        },
+    ]