add args to CustomExtension, implements db class

stix2/custom.py

@@ -89,7 +89,7 @@ def __init__(self, **kwargs):
     return _CustomObservable
 
 
-def _custom_extension_builder(cls, object_kind, type, properties, version, base_class):
+def _custom_extension_builder(cls, applies_to, type, properties, version, base_class):
 
     properties = _get_properties_dict(properties)
     toplevel_properties = None
@@ -98,7 +98,7 @@ def _custom_extension_builder(cls, object_kind, type, properties, version, base_
     # it exists.  How to treat the other properties which were given depends on
     # the extension type.
     extension_type = getattr(cls, "extension_type", None)
-    object_kind = object_kind
+    applies_to = applies_to
     if extension_type:
         # I suppose I could also go with a plain string property, since the
         # value is fixed... but an enum property seems more true to the
@@ -129,7 +129,7 @@ class _CustomExtension(cls, base_class):
 
         _type = type
         _properties = nested_properties
-        _object_kind = object_kind
+        _applies_to = applies_to
         if extension_type == "toplevel-property-extension":
             _toplevel_properties = toplevel_properties
 

stix2/datastore/relational_db/postgres_database_connection.py

@@ -1,4 +1,5 @@
 import postgres
+from sqlalchemy import create_engine
 
 from stix2.datastore.relational_db import DatabaseConnection
 
@@ -7,6 +8,7 @@ class PostgresDatabaseConnection(DatabaseConnection):
 
     def __init__(self, host, dbname, user):
         self.db = postgres.Postgres(url=f"host={host} dbname={dbname} user={user}")
+        self.engine = create_engine(f"postgresql://{host}/{dbname}", max_identifier_length=200)
 
     def execute(self, sql_statement, bindings):
         self.db.run(sql_statement, parameters=bindings)

stix2/datastore/relational_db/relational_db.py

@@ -1,9 +1,26 @@
-from stix2.base import _Observable, _STIXBase
+from sqlalchemy import MetaData
+from sqlalchemy.schema import CreateTable
+
+from stix2.base import _STIXBase
 from stix2.datastore import DataSink
-from stix2.datastore.relational_db.sql_bindings_creation import (
-    generate_insert_for_object,
+from stix2.datastore.relational_db.table_creation import (
+    create_core_tables, generate_object_table,
 )
 from stix2.parsing import parse
+from stix2.v21.base import _DomainObject, _Extension, _Observable
+
+
+def _get_all_subclasses(cls):
+    all_subclasses = []
+
+    for subclass in cls.__subclasses__():
+        all_subclasses.append(subclass)
+        all_subclasses.extend(_get_all_subclasses(subclass))
+    return all_subclasses
+
+
+def insert_object(store, stix_obj, is_sdo):
+    pass
 
 
 def _add(store, stix_data, allow_custom=True, version=None):
@@ -40,13 +57,7 @@ def _add(store, stix_data, allow_custom=True, version=None):
         else:
             stix_obj = parse(stix_data, allow_custom, version)
 
-        sql_binding_tuples = generate_insert_for_object(
-            store.database_connection,
-            stix_obj,
-            isinstance(stix_obj, _Observable),
-        )
-        for (sql, bindings) in sql_binding_tuples:
-            store.database_connection.execute(sql, bindings)
+        insert_object(store, stix_obj, isinstance(stix_obj, _Observable))
 
 
 class RelationalDBSink(DataSink):
@@ -73,18 +84,41 @@ class RelationalDBSink(DataSink):
             If part of a MemoryStore, the dict is shared with a MemorySource
 
     """
-    def __init__(self, database_connection, stix_data=None, allow_custom=True, version=None, _store=False):
+    def __init__(
+        self, database_connection, allow_custom=True, version=None,
+        instantiate_database=False,
+    ):
         super(RelationalDBSink, self).__init__()
         self.allow_custom = allow_custom
-
+        self.metadata = MetaData()
         self.database_connection = database_connection
 
-        if _store:
-            self._data = stix_data
-        else:
-            self._data = {}
-            if stix_data:
-                _add(self, stix_data, allow_custom, version)
+        if instantiate_database:
+            self._instantiate_database()
+
+    def _create_table_objects(self):
+        tables = create_core_tables(self.metadata)
+        for sdo_class in _get_all_subclasses(_DomainObject):
+            new_tables = generate_object_table(sdo_class, self.metadata, True)
+            tables.extend(new_tables)
+        for sdo_class in _get_all_subclasses(_Observable):
+            tables.extend(generate_object_table(sdo_class, self.metadata, False))
+        for sdo_class in _get_all_subclasses(_Extension):
+            if hasattr(sdo_class, "_applies_to"):
+                is_sdo = sdo_class._applies_to == "sdo"
+            else:
+                is_sdo = False
+            tables.extend(generate_object_table(sdo_class, self.metadata, is_sdo, is_extension=True))
+        return tables
+
+    def _instantiate_database(self):
+        self._create_table_objects()
+        self.metadata.create_all(self.database_connection.engine)
+
+    def generate_stix_schema(self):
+        tables = self._create_table_objects()
+        for t in tables:
+            print(CreateTable(t).compile(self.database_connection.engine))
 
     def add(self, stix_data, version=None):
         _add(self, stix_data, self.allow_custom, version)

stix2/datastore/relational_db/relational_db_testing.py

@@ -3,11 +3,10 @@
 import pytz
 
 import stix2
-# from stix2.datastore.relational_db.relational_db import RelationalDBSink
-from stix2.datastore.relational_db.table_creation import (
-    create_core_sdo_table, generate_object_table, get_all_subclasses,
+from stix2.datastore.relational_db.postgres_database_connection import (
+    PostgresDatabaseConnection,
 )
-from stix2.v21.base import _DomainObject, _Extension, _Observable
+from stix2.datastore.relational_db.relational_db import RelationalDBSink
 
 directory_stix_object = stix2.Directory(
     path="/foo/bar/a",
@@ -98,23 +97,8 @@ def file_example_with_PDFExt_Object():
 
 
 def main():
-    # sink = RelationalDBSink(PostgresDatabaseConnection("localhost", "stix-data-sink", "rpiazza"))
-    # sink.add(directory_stix_object)
-    # sink.add(s)
-    # reg_key = windows_registry_key_example()
-    # sink.add(reg_key)
-    # f = file_example_with_PDFExt_Object()
-    # sink.add(f)
-    # mal = malware_with_all_required_properties()
-    # sink.add(mal)
-
-    create_core_sdo_table()
-    for sdo_class in get_all_subclasses(_DomainObject):
-        generate_object_table(sdo_class)
-    for sdo_class in get_all_subclasses(_Observable):
-        generate_object_table(sdo_class)
-    for sdo_class in get_all_subclasses(_Extension):
-        generate_object_table(sdo_class, is_extension=True)
+    store = RelationalDBSink(PostgresDatabaseConnection("localhost", "stix-data-sink", "rpiazza"))
+    store.generate_stix_schema()
 
 
 if __name__ == '__main__':