stix2arango 0.0.5__py3-none-any.whl

stix2arango/__main__.py

@@ -0,0 +1,27 @@
+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()

stix2arango/config.py

@@ -0,0 +1,41 @@
+import os
+import logging
+from dotenv import load_dotenv
+from uuid import UUID
+
+load_dotenv()
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="[%(asctime)s] %(levelname)s - %(message)s",  # noqa D100 E501
+    datefmt="%Y-%m-%d - %H:%M:%S",
+)
+ARANGODB_HOST = os.getenv("ARANGODB_HOST")
+ARANGODB_PORT = os.getenv("ARANGODB_PORT")
+ARANGODB_USERNAME = os.getenv("ARANGODB_USERNAME")
+ARANGODB_PASSWORD = os.getenv("ARANGODB_PASSWORD")
+
+json_schema = {
+    "type": "object",
+    "properties": {
+        "type": {"type": "string", "const": "bundle"},
+        "id": {"type": "string"},
+        "objects": {"type": "array", "items": {"type": "object"}}
+    },
+    "required": ["type", "id", "objects"]
+}
+STIX2ARANGO_IDENTITY = "https://github.com/muchdogesec/stix4doge/raw/main/objects/identity/stix2arango.json" # this is stix2arango identity
+DOGESEC_IDENTITY = "https://github.com/muchdogesec/stix4doge/raw/main/objects/identity/dogesec.json" # this is stix2arango identity
+
+STIX2ARANGO_MARKING_DEFINITION = "https://raw.githubusercontent.com/muchdogesec/stix4doge/main/objects/marking-definition/stix2arango.json" # this is stix2arango marking-definition
+
+IDENTITY_REFS = [
+    STIX2ARANGO_IDENTITY,
+    DOGESEC_IDENTITY
+]
+MARKING_DEFINITION_REFS = [
+   STIX2ARANGO_MARKING_DEFINITION
+]
+DEFAULT_OBJECT_URL = MARKING_DEFINITION_REFS + IDENTITY_REFS
+
+namespace = UUID("72e906ce-ca1b-5d73-adcd-9ea9eb66a1b4")

stix2arango/services/__init__.py

@@ -0,0 +1 @@
+from .arangodb_service import ArangoDBService

stix2arango/services/arangodb_service.py

@@ -0,0 +1,303 @@
+import contextlib
+import os
+import json
+import logging
+import re
+import time
+from typing import Any
+import arango.database
+from arango.collection import StandardCollection
+from arango import ArangoClient
+from arango.exceptions import ArangoServerError
+
+from datetime import datetime, timezone
+from tqdm import tqdm
+
+from stix2arango.services.version_annotator import annotate_versions
+
+from .. import config
+from .. import utils
+from pprint import pprint
+
+module_logger = logging.getLogger("data_ingestion_service")
+
+
+class ArangoDBService:
+
+    def __init__(
+        self,
+        db,
+        vertex_collections,
+        edge_collections,
+        relationship=None,
+        create_db=False,
+        create=False,
+        username=None,
+        password=None,
+        host_url=None,
+        **kwargs,
+    ):
+        self.ARANGO_DB = self.get_db_name(db)
+        self.ARANGO_GRAPH = f"{self.ARANGO_DB.split('_database')[0]}_graph"
+        self.COLLECTIONS_VERTEX = vertex_collections
+        self.COLLECTIONS_EDGE = edge_collections
+        self.FORCE_RELATIONSHIP = [relationship] if relationship else None
+        self.missing_collection = True
+
+        module_logger.info("Establishing connection...")
+        client = ArangoClient(hosts=host_url)
+        self._client = client
+
+        if create_db:
+            module_logger.info(f"create db `{self.ARANGO_DB}` if not exist")
+            self.sys_db = client.db("_system", username=username, password=password)
+
+            module_logger.info("_system database - OK")
+
+            if not self.sys_db.has_database(self.ARANGO_DB):
+                self.create_database(self.ARANGO_DB)
+
+        self.db = client.db(
+            self.ARANGO_DB, username=username, password=password, verify=True
+        )
+
+        if self.db.has_graph(self.ARANGO_GRAPH):
+            self.cti2stix_graph = self.db.graph(self.ARANGO_GRAPH)
+        elif create:
+            self.cti2stix_graph = self.db.create_graph(self.ARANGO_GRAPH)
+
+        self.collections: dict[str, StandardCollection] = {}
+        for collection in self.COLLECTIONS_VERTEX:
+            if create:
+                self.collections[collection] = self.create_collection(collection)
+
+            self.collections[collection] = self.db.collection(collection)
+
+        for collection in self.COLLECTIONS_EDGE:
+
+            if create:
+                try:
+                    self.cti2stix_objects_relationship = (
+                        self.cti2stix_graph.create_edge_definition(
+                            edge_collection=collection,
+                            from_vertex_collections=self.COLLECTIONS_VERTEX,
+                            to_vertex_collections=self.COLLECTIONS_VERTEX,
+                        )
+                    )
+                except Exception as e:
+                    module_logger.debug(
+                        f"create edge collection {collection} failed with {e}"
+                    )
+
+            self.cti2stix_objects_relationship = self.cti2stix_graph.edge_collection(
+                collection
+            )
+            self.collections[collection] = self.cti2stix_objects_relationship
+
+        module_logger.info("ArangoDB Connected now!")
+
+    def create_database(self, db_name):
+        try:
+            self.sys_db.create_database(db_name)
+        except arango.exceptions.DatabaseCreateError as e:
+            module_logger.debug(f"create database {db_name} failed with {e}")
+
+    def create_collection(self, collection_name):
+        try:
+            return self.db.create_collection(collection_name)
+        except arango.exceptions.CollectionCreateError as e:
+            module_logger.warning(
+                f"create collection {collection_name} failed with {e}"
+            )
+            return self.db.collection(collection_name)
+
+    def execute_raw_query(self, query: str, bind_vars=None, **kwargs) -> list:
+        try:
+            cursor = self.db.aql.execute(query, bind_vars=bind_vars, **kwargs)
+            result = [doc for doc in cursor]
+            return result
+        except arango.exceptions.AQLQueryExecuteError:
+            module_logger.error(f"AQL exception in the query: {query}")
+            raise
+
+    def insert_several_objects(self, objects: list[dict], collection_name: str) -> None:
+        if not collection_name:
+            module_logger.info(f"Object has unknown type: {objects}")
+            return
+
+        for _, obj in enumerate(objects):
+            now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+            obj["_is_latest"] = False
+            obj["_record_created"] = obj.get("_record_created", now)
+            obj["_record_modified"] = now
+            obj["_key"] = obj.get("_key", f'{obj["id"]}+{now}')
+
+            if obj["type"] == "relationship":
+                obj.update(
+                    _target_type=obj["target_ref"].split("--")[0],
+                    _source_type=obj["source_ref"].split("--")[0],
+                )
+        new_insertions = objects #[obj for obj in objects if f'{obj["id"]};{obj["_record_md5_hash"]}' not in existing_objects]
+        existing_objects = {}
+
+        d = self.db.collection(collection_name).insert_many(new_insertions, overwrite_mode="ignore", sync=True)
+        for i, ret in enumerate(d):
+            obj = objects[i]
+            if isinstance(ret, arango.exceptions.DocumentInsertError):
+                if ret.error_code == 1210:
+                    existing_objects[f'{obj["id"]};{obj["_record_md5_hash"]}'] = collection_name + '/' + re.search(r'conflicting key: (.*)', ret.message).group(1)
+                    if 'relationship--7a1682a3-fec3-53ed-8dd5-68f54b1d6f7a' in ret.message:
+                        print(1)
+                else:
+                    raise ret
+        return [obj["id"] for obj in new_insertions], existing_objects
+
+    def insert_several_objects_chunked(
+        self, objects, collection_name, chunk_size=1000, remove_duplicates=True
+    ):
+        if remove_duplicates:
+            original_length = len(objects)
+            objects = utils.remove_duplicates(objects)
+            logging.info(
+                "removed {count} duplicates from imported objects.".format(
+                    count=original_length - len(objects)
+                )
+            )
+
+        progress_bar = tqdm(
+            utils.chunked(objects, chunk_size),
+            total=len(objects),
+            desc="insert_several_objects_chunked",
+        )
+        inserted_objects = []
+        existing_objects = {}
+        for chunk in progress_bar:
+            inserted, existing = self.insert_several_objects(chunk, collection_name)
+            inserted_objects.extend(inserted)
+            existing_objects.update(existing)
+            progress_bar.update(len(chunk))
+        return inserted_objects, existing_objects
+
+    def insert_relationships_chunked(
+        self,
+        relationships: list[dict[str, Any]],
+        id_to_key_map: dict[str, str],
+        collection_name: str,
+        chunk_size=1200,
+    ):
+        for relationship in relationships:
+            source_key = id_to_key_map.get(relationship["source_ref"])
+            target_key = id_to_key_map.get(relationship["target_ref"])
+
+            relationship["_stix2arango_ref_err"] = not (target_key and source_key)
+            relationship["_from"] = self.fix_edge_ref(source_key or relationship["_from"])
+            relationship["_to"] = self.fix_edge_ref(target_key or relationship["_to"])
+            relationship["_record_md5_hash"] = relationship.get(
+                "_record_md5_hash", utils.generate_md5(relationship)
+            )
+        return self.insert_several_objects_chunked(
+            relationships, collection_name, chunk_size=chunk_size
+        )
+    
+    @staticmethod
+    def fix_edge_ref(_id):
+        c, _, _key = _id.partition('/')
+        if not c:
+            c = "missing_collection"
+        return f"{c}/{_key}"
+
+    def update_is_latest_several(self, object_ids, collection_name):
+        # returns newly deprecated _ids
+        query = """
+            FOR doc IN @@collection OPTIONS {indexHint: "s2a_search", forceIndexHint: true}
+            FILTER doc.id IN @object_ids
+            RETURN [doc.id, doc._key, doc.modified, doc._record_modified, doc._is_latest, doc._id]
+        """
+        out = self.execute_raw_query(
+            query,
+            bind_vars={
+                "@collection": collection_name,
+                "object_ids": object_ids,
+            },
+        )
+        out = [dict(zip(('id', '_key', 'modified', '_record_modified', '_is_latest', '_id'), obj_tuple)) for obj_tuple in out]
+        annotated, deprecated = annotate_versions(out)
+        self.db.collection(collection_name).update_many(annotated, sync=True, keep_none=False)
+        return deprecated
+
+
+
+    def update_is_latest_several_chunked(
+        self, object_ids, collection_name, edge_collection=None, chunk_size=5000
+    ):
+        logging.info(f"Updating _is_latest for {len(object_ids)} newly inserted items")
+        progress_bar = tqdm(
+            utils.chunked(object_ids, chunk_size),
+            total=len(object_ids),
+            desc="update_is_latest_several_chunked",
+        )
+        deprecated_key_ids = []  # contains newly deprecated _ids
+        for chunk in progress_bar:
+            deprecated_key_ids.extend(
+                self.update_is_latest_several(chunk, collection_name)
+            )
+            progress_bar.update(len(chunk))
+
+        logging.info(
+            f"Updating relationship's _is_latest for {len(deprecated_key_ids)} items"
+        )
+        self.deprecate_relationships(deprecated_key_ids, edge_collection)
+        return deprecated_key_ids
+
+    def deprecate_relationships(
+        self, deprecated_key_ids: list, edge_collection: str, chunk_size=5000
+    ):
+        keys = self.get_relationships_to_deprecate(deprecated_key_ids, edge_collection)
+        self.db.collection(edge_collection).update_many(
+            tuple(dict(_key=_key, _is_latest=False) for _key in keys),
+            silent=True,
+            raise_on_document_error=True,
+        )
+        return len(keys)
+
+    def get_relationships_to_deprecate(
+        self, deprecated_key_ids: list, edge_collection: str
+    ):
+        query = """
+        FOR doc IN @@collection OPTIONS {indexHint: "s2a_search_edge", forceIndexHint: true}
+        FILTER doc._from IN @deprecated_key_ids AND doc._is_latest == TRUE
+        RETURN doc._id
+        """
+        items_to_deprecate_full: set[str] = {*deprecated_key_ids}
+
+        while deprecated_key_ids:
+            deprecated_key_ids = self.execute_raw_query(
+                query,
+                bind_vars={
+                    "@collection": edge_collection,
+                    "deprecated_key_ids": deprecated_key_ids,
+                },
+            )
+            items_to_deprecate_full.update(deprecated_key_ids)
+        return [_id.split("/", 1)[1] for _id in items_to_deprecate_full]
+
+    @staticmethod
+    def get_db_name(name):
+        ENDING = "_database"
+        if name.endswith(ENDING):
+            return name
+        return name + ENDING
+    
+    @contextlib.contextmanager
+    def transactional(self, write=None, exclusive=None, sync=True):
+        original_db = self.db
+        transactional_db = self.db.begin_transaction(allow_implicit=True, write=write, exclusive=exclusive, sync=sync)
+        try:
+            self.db = transactional_db
+            yield self
+            transactional_db.commit_transaction()
+        except:
+            transactional_db.abort_transaction()
+            raise
+        finally:
+            self.db = original_db

stix2arango/services/version_annotator.py

@@ -0,0 +1,76 @@
+from collections import defaultdict
+import os
+import time
+from typing import List, Dict
+import copy
+
+from datetime import datetime
+import os
+import typing
+from arango.client import ArangoClient
+from arango.database import StandardDatabase
+
+
+def annotate_versions(objects: List[Dict]):
+    grouped = defaultdict(list)
+
+    # Group by 'id'
+    for obj in objects:
+        grouped[obj['id']].append(obj)
+
+    result: list[dict] = []
+    deprecated = []
+
+    for obj_id, items in grouped.items():
+        # items = [copy.deepcopy(item) for item in items]
+
+        # Separate items with non-None modified
+        valid_modified = [item for item in items if item.get('modified') is not None]
+
+        # _is_latest: max(modified) -> max(_record_modified)
+        if valid_modified:
+            max_modified = max(item.get('modified') for item in valid_modified)
+            latest_candidates = [item for item in valid_modified if item.get('modified') == max_modified]
+            max_record_modified_latest = max(item['_record_modified'] for item in latest_candidates)
+        else:
+            max_modified = None
+            max_record_modified_latest = max(item['_record_modified'] for item in items)
+        # _is_earliest: min(modified) -> max(_record_modified)
+        if valid_modified:
+            min_modified = min(item.get('modified') for item in valid_modified)
+            earliest_candidates = [item for item in valid_modified if item.get('modified') == min_modified]
+            max_record_modified_earliest = max(item['_record_modified'] for item in earliest_candidates)
+        else:
+            min_modified = None
+            max_record_modified_earliest = min(item['_record_modified'] for item in items)
+
+        # _taxii_visible: for each modified (including None), select highest _record_modified
+        taxii_visible_keys = set()
+        modified_groups = defaultdict(list)
+        for item in items:
+            modified_groups[item.get('modified')].append(item)
+
+        for mod_val, group in modified_groups.items():
+            max_rec_mod = max(i['_record_modified'] for i in group)
+            for item in group:
+                if item['_record_modified'] == max_rec_mod:
+                    taxii_visible_keys.add(item['_key'])
+
+        for item in items:
+            is_latest = (
+                item.get('modified') == max_modified
+                and item['_record_modified'] == max_record_modified_latest
+            )
+
+            is_earliest = (
+                item.get('modified') == min_modified
+                and item['_record_modified'] == max_record_modified_earliest
+            )
+
+            if item.get('_is_latest') and not is_latest:
+                deprecated.append(item['_id'])
+            item['_is_latest'] = is_latest
+            item['_taxii'] = dict(visible=item['_key'] in taxii_visible_keys, first=is_earliest, last=is_latest)
+            result.append(item)
+
+    return result, deprecated

stix2arango/stix2arango/__init__.py

@@ -0,0 +1 @@
+from .stix2arango import Stix2Arango

stix2arango/stix2arango/bundle_loader.py

@@ -0,0 +1,126 @@
+import contextlib
+from datetime import datetime
+import logging
+import os
+from pathlib import Path
+import sqlite3
+import tempfile
+import uuid
+import ijson
+import json
+from collections import Counter
+
+class BundleLoader:
+    def __init__(self, file_path, chunk_size_min=20_000, db_path=""):
+        self.file_path = Path(file_path)
+        self.chunk_size_min = chunk_size_min
+        self.groups = None
+        self.bundle_id = "bundle--" + str(uuid.uuid4())
+
+        self.db_path = db_path
+        if not self.db_path:
+            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('''
+            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.commit()
+
+
+    def save_to_sqlite(self, objects):
+        """Save one STIX object to the SQLite database."""
+        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]
+            )
+        except sqlite3.IntegrityError as e:
+            print(f"Failed to insert len({objects}) objects: {e}")
+        else:
+            self.conn.commit()
+        self.inserted += len(objects)
+        # logging.info(f"inserted {self.inserted}")
+
+    def build_groups(self):
+        """
+        Iterates the STIX bundle and uses union-find to group IDs such that for every
+        relationship (object of type "relationship"), its own id and its source_ref
+        and target_ref end up in the same group.
+        """
+        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)
+            to_insert = []
+            for obj in objects:
+                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']]
+                    all_ids[obj_id].extend([sr, tr])
+                    all_ids.setdefault(sr, []).extend([tr, obj_id])
+                    all_ids.setdefault(tr, []).extend([sr, 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:
+                    continue
+                group_objs = {obj_id, *all_ids[obj_id]}
+                handled.update(group_objs)
+                new_group = group.union(group_objs)
+                if len(new_group) >= self.chunk_size_min:
+                    group.clear()
+                    self.groups.append(tuple(new_group))
+                else:
+                    group.update(group_objs)
+
+        from_ids(all_ids)
+        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))
+        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():
+            yield self.get_objects(group)
+
+    def __del__(self):
+        with contextlib.suppress(Exception):
+            os.remove(self.db_path)