stix2arango 1.1.2__tar.gz → 1.1.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stix2arango
-Version: 1.1.2
+Version: 1.1.4
 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
@@ -60,20 +60,6 @@ Note, the installation assumes ArangoDB is already installed locally.
 
 [You can install ArangoDB here](https://arangodb.com/download/). stix2arango is compatible with both the Enterprise and Community versions.
 
-#### A note for Mac users
-
-Fellow Mac users, ArangoDB can be installed and run using homebrew as follows;
-
-```shell
-## Install
-brew install arangodb
-## Run
-brew services start arangodb
-## will now be accessible in a browser at: http://127.0.0.1:8529 . Default username is root with no password set (leave blank) 
-## Stop
-brew services stop arangodb
-```
-
 ### Configuration options
 
 stix2arango has various settings that are defined in an `.env` file.
@@ -100,12 +86,14 @@ python3 stix2arango.py \
 Where;
 
 * `--file` (required): is the path to the valid STIX 2.1 bundle .json file
-* `--database` (required): is the name of the Arango database the objects should be stored in. If database does not exist, stix2arango will create it
+* `--database` (required): is the name of the Arango database the objects should be stored in. 
+* `--create_db` (default `true`): If database does not exist, stix2arango will create it. You can set to `false` to stop this behaviour (and avoid the risk of incorrect DBs being created). Generally setting to `false` is a good idea if you know the databases exist. This setting will only work if the Arango user being used to authenticate has permissions to create new databases.
 * `--collection` (required): is the name of the Arango collection in the database specified the objects should be stored in. If the collection does not exist, stix2arango will create it
 * `--stix2arango_note` (optional): Will be stored under the `_stix2arango_note` custom attribute in ArangoDB. Useful as can be used in AQL. `a-z` characters only. Max 24 chars.
 * `--ignore_embedded_relationships` (optional, boolean):  if `true` passed, this will stop ANY embedded relationships from being generated. This applies for all object types (SDO, SCO, SRO, SMO). If you want to target certain object types see `ignore_embedded_relationships_sro` and `ignore_embedded_relationships_sro` flags. ` Default is `false`
 * `--ignore_embedded_relationships_sro` (optional, boolean): if `true` passed, will stop any embedded relationships from being generated from SRO objects (`type` = `relationship`). Default is `false`
-* `--ignore_embedded_relationships_smo` (optional, boolean): if `true` passed, will stop any embedded relationships from being generated from SMO objects (`type` = `marking-definition`, `extension-definition`, `language-content`). Default is `false`
+* `--ignore_embedded_relationships_smo` (optional, boolean): if `true` passed, will stop any embedded relationships from being generated from SMO objects (`type` = `marking-defirnition`, `extension-definition`, `language-content`). Default is `false`
+* `--include_embedded_relationships_attributes` (optional, stix `_ref` or `_refs` attribute): if you only want to create embedded relationships from certain keys (attributes) in a STIX object you can pass a list of attributes here. e.g. `object_refs created_by_ref` . In this example, embedded relationships to all objects listed in `object_refs` and objects in `created_by_ref` will be created between source (the objects that house these attibutes) and destinations (the objects listed as values for these attributes)
 * `--is_large_file` (pass flag): Use this mode when the bundle is very large (>100mb), this will chunk the input into multiple files before loading into memory.
 
 For example, [using the MITRE ATT&CK Enterprise bundle](https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json);
@@ -132,6 +120,18 @@ python3 stix2arango.py \
 	--is_large_file
 ```
 
+If you want to include embedded relationships for `created_by_ref` and `object_marking_refs` attibutes collection, you would run;
+
+```shell
+python3 stix2arango.py \
+	--file cti_knowledge_base_store/mitre-attack-enterprise/enterprise-attack-15_1.json \
+	--database stix2arango_demo \
+	--collection demo_2 \
+	--stix2arango_note v15.1 \
+	--include_embedded_relationships_attributes object_refs created_by_ref \
+	--is_large_file
+```
+
 #### A note on embedded relationships
 
 stix2arango can handle all embedded references to other STIX objects under `_ref` and `_refs` properties in a STIX object when `--ignore_embedded_relationships` is set to false.

{stix2arango-1.1.2 → stix2arango-1.1.4}/README.md

@@ -37,20 +37,6 @@ Note, the installation assumes ArangoDB is already installed locally.
 
 [You can install ArangoDB here](https://arangodb.com/download/). stix2arango is compatible with both the Enterprise and Community versions.
 
-#### A note for Mac users
-
-Fellow Mac users, ArangoDB can be installed and run using homebrew as follows;
-
-```shell
-## Install
-brew install arangodb
-## Run
-brew services start arangodb
-## will now be accessible in a browser at: http://127.0.0.1:8529 . Default username is root with no password set (leave blank) 
-## Stop
-brew services stop arangodb
-```
-
 ### Configuration options
 
 stix2arango has various settings that are defined in an `.env` file.
@@ -77,12 +63,14 @@ python3 stix2arango.py \
 Where;
 
 * `--file` (required): is the path to the valid STIX 2.1 bundle .json file
-* `--database` (required): is the name of the Arango database the objects should be stored in. If database does not exist, stix2arango will create it
+* `--database` (required): is the name of the Arango database the objects should be stored in. 
+* `--create_db` (default `true`): If database does not exist, stix2arango will create it. You can set to `false` to stop this behaviour (and avoid the risk of incorrect DBs being created). Generally setting to `false` is a good idea if you know the databases exist. This setting will only work if the Arango user being used to authenticate has permissions to create new databases.
 * `--collection` (required): is the name of the Arango collection in the database specified the objects should be stored in. If the collection does not exist, stix2arango will create it
 * `--stix2arango_note` (optional): Will be stored under the `_stix2arango_note` custom attribute in ArangoDB. Useful as can be used in AQL. `a-z` characters only. Max 24 chars.
 * `--ignore_embedded_relationships` (optional, boolean):  if `true` passed, this will stop ANY embedded relationships from being generated. This applies for all object types (SDO, SCO, SRO, SMO). If you want to target certain object types see `ignore_embedded_relationships_sro` and `ignore_embedded_relationships_sro` flags. ` Default is `false`
 * `--ignore_embedded_relationships_sro` (optional, boolean): if `true` passed, will stop any embedded relationships from being generated from SRO objects (`type` = `relationship`). Default is `false`
-* `--ignore_embedded_relationships_smo` (optional, boolean): if `true` passed, will stop any embedded relationships from being generated from SMO objects (`type` = `marking-definition`, `extension-definition`, `language-content`). Default is `false`
+* `--ignore_embedded_relationships_smo` (optional, boolean): if `true` passed, will stop any embedded relationships from being generated from SMO objects (`type` = `marking-defirnition`, `extension-definition`, `language-content`). Default is `false`
+* `--include_embedded_relationships_attributes` (optional, stix `_ref` or `_refs` attribute): if you only want to create embedded relationships from certain keys (attributes) in a STIX object you can pass a list of attributes here. e.g. `object_refs created_by_ref` . In this example, embedded relationships to all objects listed in `object_refs` and objects in `created_by_ref` will be created between source (the objects that house these attibutes) and destinations (the objects listed as values for these attributes)
 * `--is_large_file` (pass flag): Use this mode when the bundle is very large (>100mb), this will chunk the input into multiple files before loading into memory.
 
 For example, [using the MITRE ATT&CK Enterprise bundle](https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json);
@@ -109,6 +97,18 @@ python3 stix2arango.py \
 	--is_large_file
 ```
 
+If you want to include embedded relationships for `created_by_ref` and `object_marking_refs` attibutes collection, you would run;
+
+```shell
+python3 stix2arango.py \
+	--file cti_knowledge_base_store/mitre-attack-enterprise/enterprise-attack-15_1.json \
+	--database stix2arango_demo \
+	--collection demo_2 \
+	--stix2arango_note v15.1 \
+	--include_embedded_relationships_attributes object_refs created_by_ref \
+	--is_large_file
+```
+
 #### A note on embedded relationships
 
 stix2arango can handle all embedded references to other STIX objects under `_ref` and `_refs` properties in a STIX object when `--ignore_embedded_relationships` is set to false.

{stix2arango-1.1.2 → stix2arango-1.1.4}/pyproject.toml

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

stix2arango-1.1.4/stix2arango/__main__.py

@@ -0,0 +1,81 @@
+import argparse
+from stix2arango.stix2arango import Stix2Arango
+
+
+def parse_bool(value: str):
+    value = value.lower()
+    # ["false", "no", "n"]
+    return value in ["yes", "y", "true", "1"]
+
+def parse_ref(value: str):
+    if not (value.endswith('_ref') or value.endswith('_refs')):
+        raise argparse.ArgumentTypeError('value must end with _ref or _refs')
+    return value
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Import STIX JSON into ArangoDB")
+    parser.add_argument("--file", required=True, help="Path to STIX JSON file")
+    parser.add_argument(
+        "--is_large_file",
+        action="store_true",
+        help="Use large file mode [Use this mode when the bundle is very large, this will enable you stix2arango to chunk before loading into memory]",
+    )
+    parser.add_argument("--database", required=True, help="ArangoDB database name")
+    parser.add_argument(
+        "--create_db",
+        default=True,
+        type=parse_bool,
+        help="whether or not to skip the creation of database, requires admin permission",
+    )
+    parser.add_argument("--collection", required=True, help="ArangoDB collection name")
+    parser.add_argument(
+        "--stix2arango_note", required=False, help="Note for the import", default=""
+    )
+    parser.add_argument(
+        "--ignore_embedded_relationships",
+        required=False,
+        help="Ignore Embedded Relationship for the import",
+        type=parse_bool,
+        default=False,
+    )
+    parser.add_argument(
+        "--ignore_embedded_relationships_sro",
+        required=False,
+        help="Ignore Embedded Relationship for imported SROs",
+        type=parse_bool,
+        default=False,
+    )
+    parser.add_argument(
+        "--ignore_embedded_relationships_smo",
+        required=False,
+        help="Ignore Embedded Relationship for imported SMOs",
+        type=parse_bool,
+        default=False,
+    )
+    parser.add_argument(
+        "--include_embedded_relationships_attributes",
+        required=False,
+        help="Only create embedded relationships for keys",
+        action="extend",
+        nargs="+",
+        type=parse_ref
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_arguments()
+    stix_obj = Stix2Arango(
+        database=args.database,
+        collection=args.collection,
+        file=args.file,
+        create_db=args.create_db,
+        stix2arango_note=args.stix2arango_note,
+        ignore_embedded_relationships=args.ignore_embedded_relationships,
+        ignore_embedded_relationships_sro=args.ignore_embedded_relationships_sro,
+        ignore_embedded_relationships_smo=args.ignore_embedded_relationships_smo,
+        is_large_file=args.is_large_file,
+        include_embedded_relationships_attributes=args.include_embedded_relationships_attributes,
+    )
+    stix_obj.run()

{stix2arango-1.1.2 → stix2arango-1.1.4}/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.2 → stix2arango-1.1.4}/stix2arango/stix2arango/stix2arango.py

@@ -42,6 +42,7 @@ class Stix2Arango:
         ignore_embedded_relationships=False,
         ignore_embedded_relationships_sro=True,
         ignore_embedded_relationships_smo=True,
+        include_embedded_relationships_attributes=None,
         bundle_id=None,
         username=config.ARANGODB_USERNAME,
         password=config.ARANGODB_PASSWORD,
@@ -80,15 +81,16 @@ 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
         self.ignore_embedded_relationships = ignore_embedded_relationships
         self.ignore_embedded_relationships_smo = ignore_embedded_relationships_smo
         self.ignore_embedded_relationships_sro = ignore_embedded_relationships_sro
+        self.include_embedded_relationships_attributes = include_embedded_relationships_attributes
         self.object_key_mapping = {}
         if create_collection:
             self.create_s2a_indexes()
@@ -472,14 +474,16 @@ class Stix2Arango:
         for obj in tqdm(bundle_objects, desc="upload_embedded_edges"):
             if obj["id"] not in inserted_object_ids:
                 continue
-            if (
+            if self.include_embedded_relationships_attributes:
+                pass
+            elif (
                 self.ignore_embedded_relationships_smo and obj["type"] in SMO_TYPES
             ) or (
                 self.ignore_embedded_relationships_sro and obj["type"] == "relationship"
             ):
                 continue
 
-            for ref_type, targets in utils.get_embedded_refs(obj):
+            for ref_type, targets in utils.get_embedded_refs(obj, attributes=self.include_embedded_relationships_attributes):
                 utils.create_relationship_obj(
                     obj=obj,
                     source=obj.get("id"),
@@ -578,7 +582,7 @@ class Stix2Arango:
             self.filename, all_objects
         )
 
-        if not self.ignore_embedded_relationships:
+        if (not self.ignore_embedded_relationships) or self.include_embedded_relationships_attributes:
             module_logger.info(
                 "Creating new embedded relationships using _refs and _ref"
             )

{stix2arango-1.1.2 → stix2arango-1.1.4}/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
@@ -67,10 +67,6 @@ 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"
@@ -120,7 +116,7 @@ def remove_duplicates(objects):
     return list(objects_hashmap.values())
 
 
-def get_embedded_refs(object: list | dict, xpath: list = []):
+def get_embedded_refs(object: list | dict, xpath: list = [], attributes=None):
     embedded_refs = []
     if isinstance(object, dict):
         for key, value in object.items():
@@ -129,11 +125,13 @@ def get_embedded_refs(object: list | dict, xpath: list = []):
             if match := EMBEDDED_RELATIONSHIP_RE.fullmatch(key):
                 relationship_type = "-".join(xpath + match.group(1).split("_"))
                 targets = value if isinstance(value, list) else [value]
+                if attributes and key not in attributes:
+                    continue
                 embedded_refs.append((relationship_type, targets))
             elif isinstance(value, list):
-                embedded_refs.extend(get_embedded_refs(value, xpath + [key]))
+                embedded_refs.extend(get_embedded_refs(value, xpath + [key], attributes=attributes))
     elif isinstance(object, list):
         for obj in object:
             if isinstance(obj, dict):
-                embedded_refs.extend(get_embedded_refs(obj, xpath))
+                embedded_refs.extend(get_embedded_refs(obj, xpath, attributes=attributes))
     return embedded_refs

stix2arango-1.1.4/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.4/tests/src/test_main.py

@@ -0,0 +1,131 @@
+import sys
+from unittest.mock import patch
+import pytest
+from stix2arango.__main__ import parse_arguments, main
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        [],
+        ["--database", "somedb"],
+        ["--file", "somefile"],
+        ["--file", "somefile", "--database", "somedb"],
+        ["--collection", "somefile", "--database", "somedb"],
+    ],
+)
+def test_database_and_file_required(monkeypatch, args):
+    monkeypatch.setattr(sys, "argv", ["prog", *args])
+    with pytest.raises(SystemExit) as e:
+        parse_arguments()
+
+
+@pytest.mark.parametrize(
+    "args,parsed_output",
+    [
+        (
+            ["--collection", "somefile", "--database", "somedb", "--file", "some-file"],
+            {
+                "file": "some-file",
+                "is_large_file": False,
+                "database": "somedb",
+                "create_db": True,
+                "collection": "somefile",
+                "stix2arango_note": "",
+                "ignore_embedded_relationships": False,
+                "ignore_embedded_relationships_sro": False,
+                "ignore_embedded_relationships_smo": False,
+                "include_embedded_relationships_attributes": None,
+            },
+        ),
+        (
+            [
+                "--collection",
+                "somefile",
+                "--database",
+                "somedb",
+                "--file",
+                "some-file",
+                "--ignore_embedded_relationships_smo",
+                "yes",
+            ],
+            {
+                "file": "some-file",
+                "is_large_file": False,
+                "database": "somedb",
+                "create_db": True,
+                "collection": "somefile",
+                "stix2arango_note": "",
+                "ignore_embedded_relationships": False,
+                "ignore_embedded_relationships_sro": False,
+                "ignore_embedded_relationships_smo": True,
+                "include_embedded_relationships_attributes": None,
+            },
+        ),
+        (
+            [
+                "--collection",
+                "somefile",
+                "--database",
+                "somedb",
+                "--file",
+                "some-file",
+                "--ignore_embedded_relationships_smo",
+                "yes",
+                "--include_embedded_relationships_attributes",
+                "abc_ref",
+                "abcdef_refs",
+            ],
+            {
+                "file": "some-file",
+                "is_large_file": False,
+                "database": "somedb",
+                "create_db": True,
+                "collection": "somefile",
+                "stix2arango_note": "",
+                "ignore_embedded_relationships": False,
+                "ignore_embedded_relationships_sro": False,
+                "ignore_embedded_relationships_smo": True,
+                "include_embedded_relationships_attributes": ["abc_ref", "abcdef_refs"],
+            },
+        ),
+    ],
+)
+def test_parse_args(monkeypatch, args, parsed_output):
+    args = ["prog", *args]
+    monkeypatch.setattr(sys, "argv", args)
+    assert parse_arguments().__dict__ == parsed_output
+
+
+def test_main_calls_correctly(monkeypatch):
+    args = [
+        "prog",
+        "--collection",
+        "somefile",
+        "--database",
+        "somedb",
+        "--file",
+        "some-file",
+        "--ignore_embedded_relationships_smo",
+        "yes",
+        "--include_embedded_relationships_attributes",
+        "abc_ref",
+        "abcdef_refs",
+    ]
+    monkeypatch.setattr(sys, "argv", args)
+    with patch("stix2arango.__main__.Stix2Arango") as mock_s2a_init:
+        main()
+        mock_s2a_init.assert_called_once_with(
+            **{
+                "file": "some-file",
+                "is_large_file": False,
+                "database": "somedb",
+                "create_db": True,
+                "collection": "somefile",
+                "stix2arango_note": "",
+                "ignore_embedded_relationships": False,
+                "ignore_embedded_relationships_sro": False,
+                "ignore_embedded_relationships_smo": True,
+                "include_embedded_relationships_attributes": ["abc_ref", "abcdef_refs"],
+            }
+        )

stix2arango-1.1.4/tests/src/test_utils.py

@@ -0,0 +1,172 @@
+from datetime import datetime
+from unittest.mock import MagicMock, patch
+import pytest
+from stix2arango import utils
+from stix2arango.config import DOGESEC_IDENTITY
+
+
+def test_load_file_from_url():
+    assert (
+        utils.load_file_from_url(DOGESEC_IDENTITY)["id"]
+        == "identity--9779a2db-f98c-5f4b-8d08-8ee04e02dbb5"
+    )
+
+    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_embedded_refs__attributes_whitelist():
+    assert utils.get_embedded_refs(
+        {
+            "abc_ref": "ref1",
+            "abcd_refs": ["ref1", "ref2"],
+            "abcde": [{"abcdef_ref": "ref7"}, {"abcd_efgh_ref": "ref8"}],
+        },
+        attributes=["abcd_efgh_ref", "abc_ref"],
+    ) == [
+        ("abc", ["ref1"]),
+        (
+            "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"],
+        },
+    ]

stix2arango-1.1.2/stix2arango/__main__.py

@@ -1,27 +0,0 @@
-import argparse
-from stix2arango.stix2arango import Stix2Arango
-
-def parse_bool(value: str):
-    value = value.lower()
-    # ["false", "no", "n"]
-    return value in ["yes", "y", "true", "1"]
-
-def parse_arguments():
-    parser = argparse.ArgumentParser(description="Import STIX JSON into ArangoDB")
-    parser.add_argument("--file", required=True, help="Path to STIX JSON file")
-    parser.add_argument("--is_large_file", action="store_true", help="Use large file mode [Use this mode when the bundle is very large, this will enable you stix2arango to chunk before loading into memory]")
-    parser.add_argument("--database", required=True, help="ArangoDB database name")
-    parser.add_argument("--create_db", default=True, type=parse_bool, help="whether or not to skip the creation of database, requires admin permission")
-    parser.add_argument("--collection", required=True, help="ArangoDB collection name")
-    parser.add_argument("--stix2arango_note", required=False, help="Note for the import", default="")
-    parser.add_argument("--ignore_embedded_relationships", required=False, help="Ignore Embedded Relationship for the import", type=parse_bool, default=False)
-    parser.add_argument("--ignore_embedded_relationships_sro", required=False, help="Ignore Embedded Relationship for imported SROs", type=parse_bool, default=False)
-    parser.add_argument("--ignore_embedded_relationships_smo", required=False, help="Ignore Embedded Relationship for imported SMOs", type=parse_bool, default=False)
-
-    return parser.parse_args()
-
-
-def main():
-    args = parse_arguments()
-    stix_obj = Stix2Arango(args.database, args.collection, file=args.file, create_db=args.create_db, stix2arango_note=args.stix2arango_note, ignore_embedded_relationships=args.ignore_embedded_relationships, ignore_embedded_relationships_sro=args.ignore_embedded_relationships_sro, ignore_embedded_relationships_smo=args.ignore_embedded_relationships_smo, is_large_file=args.is_large_file)
-    stix_obj.run()