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

stix2arango/__main__.py

@@ -1,27 +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(
+        "--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(
+        "--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(
+        "--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(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()
+    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/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/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/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.2.dist-info → stix2arango-1.1.4.dist-info}/METADATA

@@ -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.dist-info → stix2arango-1.1.4.dist-info}/RECORD

@@ -1,16 +1,16 @@
 stix2arango/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-stix2arango/__main__.py,sha256=wbR_iO70Vld2NYiml6Kz4rH396uOiNwTtjNBl4AHZEg,1987
+stix2arango/__main__.py,sha256=zsCi_bfDULLDkqlRwXyGhFuLvSRcvESEc4MMN7h1lbQ,2835
 stix2arango/config.py,sha256=NZFrcnEfz-0QBrut2Rh7xMF78v0bk6U6y2TY_7mHxSs,1407
-stix2arango/utils.py,sha256=CmyRH1gmymu1vA3KI4vijek0oS4ODoMBIXmMuoIcFM4,4690
+stix2arango/utils.py,sha256=bUKJBQ2owbCQKWs_m-VYjYCHuQLykizabE4D3LPspW8,4636
 stix2arango/services/__init__.py,sha256=E87fB-dxI4mPxMVs00jdLhjp9jFhkVfjhMKIqGLRJlY,45
-stix2arango/services/arangodb_service.py,sha256=qgE7yb-ysqyJ1oUjqAV6eGa53WvWHAFAd6I7aKZ6KyU,11960
+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/stix2arango.py,sha256=t5uW6-Zshk0GvEDxwZBn7mzg0goaBCIm_7hhruCimKc,22046
+stix2arango/stix2arango/stix2arango.py,sha256=HJXDqA9NWxXVQSHPmbpkEKurpWEbZmy5bng5SQ1OsjE,22412
 stix2arango/templates/marking-definition.json,sha256=0q9y35mUmiF6xIWSLpkATL4JTHGSCNyLbejqZiQ0AuE,3113
-stix2arango-1.1.2.dist-info/METADATA,sha256=p1UJ9P5Lq4CKUiFhMiHrq8FFvuwCA79lJHpv9j50TSU,6873
-stix2arango-1.1.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-stix2arango-1.1.2.dist-info/entry_points.txt,sha256=k2WnxMsHFLoyC6rqfvjhIMS1zwtWin51-MbNCGmSMYE,58
-stix2arango-1.1.2.dist-info/licenses/LICENSE,sha256=BK8Ppqlc4pdgnNzIxnxde0taoQ1BgicdyqmBvMiNYgY,11364
-stix2arango-1.1.2.dist-info/RECORD,,
+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,,