stix2arango 1.1.3__py3-none-any.whl → 1.1.5__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/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/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,
@@ -89,6 +90,7 @@ class Stix2Arango:
         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

@@ -116,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():
@@ -125,11 +125,16 @@ 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]
+                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))
             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.3.dist-info → stix2arango-1.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stix2arango
-Version: 1.1.3
+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
@@ -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.3.dist-info → stix2arango-1.1.5.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=eVAMvXZVylM2RXzi2ph0RVW__eoSjAHWSWSG3900yjk,4487
+stix2arango/utils.py,sha256=C-VDwsCABFU2hrv1EwF7oaQUYOE2MWG40_7WGzHu0A4,4796
 stix2arango/services/__init__.py,sha256=E87fB-dxI4mPxMVs00jdLhjp9jFhkVfjhMKIqGLRJlY,45
 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=sC-br0nptUtZMzNza6v3s6rjdgJk-EG0_KErN9JN9qQ,22060
+stix2arango/stix2arango/bundle_loader.py,sha256=lqhW4RwELRdRax7erSuYv7h02Nata7SXNWRxaA97r_w,5128
+stix2arango/stix2arango/stix2arango.py,sha256=HJXDqA9NWxXVQSHPmbpkEKurpWEbZmy5bng5SQ1OsjE,22412
 stix2arango/templates/marking-definition.json,sha256=0q9y35mUmiF6xIWSLpkATL4JTHGSCNyLbejqZiQ0AuE,3113
-stix2arango-1.1.3.dist-info/METADATA,sha256=z9lCPIr6WmDFUpxzg4CQhEt_hSlBpGFjqdignPw0mSw,6873
-stix2arango-1.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-stix2arango-1.1.3.dist-info/entry_points.txt,sha256=k2WnxMsHFLoyC6rqfvjhIMS1zwtWin51-MbNCGmSMYE,58
-stix2arango-1.1.3.dist-info/licenses/LICENSE,sha256=BK8Ppqlc4pdgnNzIxnxde0taoQ1BgicdyqmBvMiNYgY,11364
-stix2arango-1.1.3.dist-info/RECORD,,
+stix2arango-1.1.5.dist-info/METADATA,sha256=Jxs4Bp4z67-Bzy3AF3fYZlX4v6f_ToemEy7SKLN4oUY,7797
+stix2arango-1.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+stix2arango-1.1.5.dist-info/entry_points.txt,sha256=k2WnxMsHFLoyC6rqfvjhIMS1zwtWin51-MbNCGmSMYE,58
+stix2arango-1.1.5.dist-info/licenses/LICENSE,sha256=BK8Ppqlc4pdgnNzIxnxde0taoQ1BgicdyqmBvMiNYgY,11364
+stix2arango-1.1.5.dist-info/RECORD,,