diff --git a/stix2/datastore/__init__.py b/stix2/datastore/__init__.py index 39229805..715c6e6b 100644 --- a/stix2/datastore/__init__.py +++ b/stix2/datastore/__init__.py @@ -212,7 +212,7 @@ def add(self, *args, **kwargs): """ try: return self.sink.add(*args, **kwargs) - except AttributeError as ex: + except AttributeError: msg = "%s has no data sink to put objects in" raise AttributeError(msg % self.__class__.__name__) diff --git a/stix2/datastore/relational_db/database_backends/database_backend_base.py b/stix2/datastore/relational_db/database_backends/database_backend_base.py index 118db629..bc675b3e 100644 --- a/stix2/datastore/relational_db/database_backends/database_backend_base.py +++ b/stix2/datastore/relational_db/database_backends/database_backend_base.py @@ -1,15 +1,13 @@ from typing import Any -from sqlalchemy import ( - create_engine, Boolean, Float, Integer, LargeBinary, Text, TIMESTAMP, -) +from sqlalchemy import Boolean, Float, Integer, Text, create_engine from sqlalchemy_utils import create_database, database_exists, drop_database from stix2.base import ( - _DomainObject, _Extension, _MetaObject, _Observable, _RelationshipObject, - _STIXBase, + _DomainObject, _MetaObject, _Observable, _RelationshipObject, ) + class DatabaseBackend: def __init__(self, database_connection_url, force_recreate=False, **kwargs: Any): self.database_connection_url = database_connection_url diff --git a/stix2/datastore/relational_db/database_backends/postgres_backend.py b/stix2/datastore/relational_db/database_backends/postgres_backend.py index e80fd1bd..0d24056f 100644 --- a/stix2/datastore/relational_db/database_backends/postgres_backend.py +++ b/stix2/datastore/relational_db/database_backends/postgres_backend.py @@ -1,23 +1,13 @@ import os from typing import Any -from sqlalchemy import ( - TIMESTAMP, LargeBinary, Text, -) +from sqlalchemy import TIMESTAMP, LargeBinary, Text from sqlalchemy.schema import CreateSchema from stix2.base import ( - _DomainObject, _Extension, _MetaObject, _Observable, _RelationshipObject, - _STIXBase, + _DomainObject, _MetaObject, _Observable, _RelationshipObject, ) from stix2.datastore.relational_db.utils import schema_for -from stix2.properties import ( - BinaryProperty, BooleanProperty, DictionaryProperty, - EmbeddedObjectProperty, EnumProperty, ExtensionsProperty, FloatProperty, - HashesProperty, HexProperty, IDProperty, IntegerProperty, ListProperty, - ObjectReferenceProperty, Property, ReferenceProperty, StringProperty, - TimestampProperty, TypeProperty, -) from .database_backend_base import DatabaseBackend diff --git a/stix2/datastore/relational_db/input_creation.py b/stix2/datastore/relational_db/input_creation.py index ff308517..a850755b 100644 --- a/stix2/datastore/relational_db/input_creation.py +++ b/stix2/datastore/relational_db/input_creation.py @@ -248,7 +248,11 @@ def generate_insert_information( # noqa: F811 return insert_statements else: if db_backend.array_allowed(): - return {name: stix_object[name]} + if isinstance(self.contained, HexProperty): + return {name: [bytes.fromhex(x) for x in stix_object[name]]} + else: + return {name: stix_object[name]} + else: insert_statements = list() table = data_sink.tables_dictionary[ @@ -258,13 +262,14 @@ def generate_insert_information( # noqa: F811 ) ] for elem in stix_object[name]: - bindings = {"id": stix_object["id"], name: elem} + bindings = { + "id": stix_object["id"], + name: bytes.fromhex(elem) if isinstance(self.contained, HexProperty) else elem, + } insert_statements.append(insert(table).values(bindings)) return insert_statements - - @add_method(ReferenceProperty) def generate_insert_information(self, name, stix_object, **kwargs): # noqa: F811 return {name: stix_object[name]} @@ -300,7 +305,6 @@ def generate_insert_for_array_in_table(table, values, foreign_key_value, column_ def generate_insert_for_external_references(data_sink, stix_object): - db_backend = data_sink.db_backend insert_statements = list() next_id = None object_table = data_sink.tables_dictionary["common.external_references"] @@ -395,7 +399,7 @@ def generate_insert_for_core(data_sink, stix_object, core_properties, stix_type_ core_insert_statement = insert(core_table).values(core_bindings) insert_statements.append(core_insert_statement) - if "labels" in stix_object: + if "labels" in stix_object and "labels" in child_table_properties: label_table_name = canonicalize_table_name(core_table.name + "_labels", data_sink.db_backend.schema_for_core()) labels_table = data_sink.tables_dictionary[label_table_name] insert_statements.extend( @@ -403,12 +407,15 @@ def generate_insert_for_core(data_sink, stix_object, core_properties, stix_type_ labels_table, stix_object["labels"], stix_object["id"], - column_name="label" - )) + column_name="label", + ), + ) if "object_marking_refs" in stix_object: - object_marking_table_name = canonicalize_table_name("object_marking_refs", - data_sink.db_backend.schema_for_core()) + object_marking_table_name = canonicalize_table_name( + "object_marking_refs", + data_sink.db_backend.schema_for_core(), + ) if stix_type_name != "sco": object_markings_ref_table = data_sink.tables_dictionary[object_marking_table_name + "_sdo"] else: diff --git a/stix2/datastore/relational_db/relational_db.py b/stix2/datastore/relational_db/relational_db.py index 56cd32d0..2dc6009e 100644 --- a/stix2/datastore/relational_db/relational_db.py +++ b/stix2/datastore/relational_db/relational_db.py @@ -1,19 +1,14 @@ -from sqlalchemy import MetaData, create_engine, delete, select -from sqlalchemy.schema import CreateSchema, CreateTable, Sequence -from sqlalchemy_utils import create_database, database_exists, drop_database +from sqlalchemy import MetaData, delete +from sqlalchemy.schema import CreateTable, Sequence -from stix2.base import ( - _DomainObject, _MetaObject, _Observable, _RelationshipObject, _STIXBase, -) +from stix2.base import _STIXBase from stix2.datastore import DataSink, DataSource, DataStoreMixin from stix2.datastore.relational_db.input_creation import ( generate_insert_for_object, ) from stix2.datastore.relational_db.query import read_object from stix2.datastore.relational_db.table_creation import create_table_objects -from stix2.datastore.relational_db.utils import ( - canonicalize_table_name, schema_for, table_name_for, -) +from stix2.datastore.relational_db.utils import canonicalize_table_name from stix2.parsing import parse @@ -88,6 +83,7 @@ def __init__( source=RelationalDBSource( db_backend, metadata=self.metadata, + allow_custom=allow_custom, ), sink=RelationalDBSink( db_backend, @@ -190,9 +186,10 @@ def next_id(self): with self.db_backend.database_connection.begin() as trans: return trans.execute(self.sequence) + class RelationalDBSource(DataSource): def __init__( - self, db_backend, *stix_object_classes, metadata=None, + self, db_backend, allow_custom, *stix_object_classes, metadata=None, ): """ Initialize this source. Only one of stix_object_classes and metadata @@ -217,6 +214,8 @@ def __init__( self.db_backend = db_backend + self.allow_custom = allow_custom + if metadata: self.metadata = metadata else: diff --git a/stix2/datastore/relational_db/relational_db_testing.py b/stix2/datastore/relational_db/relational_db_testing.py index 926803a2..8f200b36 100644 --- a/stix2/datastore/relational_db/relational_db_testing.py +++ b/stix2/datastore/relational_db/relational_db_testing.py @@ -64,7 +64,7 @@ def malware_with_all_required_properties(): ref2 = stix2.v21.ExternalReference( source_name="ACME Threat Intel", description="Threat report", - url="http://www.example.com/threat-report.pdf" + url="http://www.example.com/threat-report.pdf", ) now = dt.datetime(2016, 5, 12, 8, 17, 27, tzinfo=pytz.utc) @@ -203,10 +203,10 @@ def custom_obj(): @stix2.CustomObject( "test-object", [ - ("prop_name", stix2.properties.ListProperty(stix2.properties.BinaryProperty())) + ("prop_name", stix2.properties.ListProperty(stix2.properties.BinaryProperty())), ], "extension-definition--15de9cdb-3515-4271-8479-8141154c5647", - is_sdo=True + is_sdo=True, ) class TestClass: pass @@ -215,53 +215,58 @@ class TestClass: def test_binary_list(): return TestClass(prop_name=["AREi", "7t3M"]) + @stix2.CustomObject( "test2-object", [ - ("prop_name", stix2.properties.ListProperty( - stix2.properties.HexProperty() - )) + ( + "prop_name", stix2.properties.ListProperty( + stix2.properties.HexProperty(), + ), + ), ], "extension-definition--15de9cdb-4567-4271-8479-8141154c5647", - is_sdo=True - ) - + is_sdo=True, +) class Test2Class: - pass + pass + def test_hex_list(): return Test2Class( - prop_name=["1122", "fedc"] + prop_name=["1122", "fedc"], ) + @stix2.CustomObject( "test3-object", [ - ("prop_name", - stix2.properties.DictionaryProperty( - valid_types=[ - stix2.properties.IntegerProperty, - stix2.properties.FloatProperty, - stix2.properties.StringProperty - ] - ) - ) + ( + "prop_name", + stix2.properties.DictionaryProperty( + valid_types=[ + stix2.properties.IntegerProperty, + stix2.properties.FloatProperty, + stix2.properties.StringProperty, + ], + ), + ), ], "extension-definition--15de9cdb-1234-4271-8479-8141154c5647", - is_sdo=True - ) + is_sdo=True, +) class Test3Class: pass def test_dictionary(): return Test3Class( - prop_name={"a": 1, "b": 2.3, "c": "foo"} + prop_name={"a": 1, "b": 2.3, "c": "foo"}, ) def main(): store = RelationalDBStore( PostgresBackend("postgresql://localhost/stix-data-sink", force_recreate=True), - False, + True, None, True, print_sql=True, @@ -269,17 +274,17 @@ def main(): if store.sink.db_backend.database_exists: - # td = test_dictionary() - # - # store.add(td) - # - # th = test_hex_list() - # - # store.add(th) - # - # tb = test_binary_list() - # - # store.add(tb) + td = test_dictionary() + + store.add(td) + + th = test_hex_list() + + store.add(th) + + tb = test_binary_list() + + store.add(tb) co = custom_obj() diff --git a/stix2/datastore/relational_db/table_creation.py b/stix2/datastore/relational_db/table_creation.py index d1feeddb..d101db48 100644 --- a/stix2/datastore/relational_db/table_creation.py +++ b/stix2/datastore/relational_db/table_creation.py @@ -1,8 +1,8 @@ # from collections import OrderedDict from sqlalchemy import ( # create_engine,; insert, - ARRAY, TIMESTAMP, Boolean, CheckConstraint, Column, ForeignKey, - Integer, Table, Text, UniqueConstraint + ARRAY, CheckConstraint, Column, ForeignKey, Integer, Table, Text, + UniqueConstraint, ) from stix2.datastore.relational_db.add_method import add_method @@ -27,7 +27,7 @@ def create_array_column(property_name, contained_sql_type, optional): property_name, ARRAY(contained_sql_type), CheckConstraint(f"{property_name} IS NULL or array_length({property_name}, 1) IS NOT NULL"), - nullable=optional + nullable=optional, ) @@ -183,11 +183,12 @@ def create_granular_markings_table(metadata, db_backend, sco_or_sdo): columns.append(create_array_column("selectors", db_backend.determine_sql_type_for_string_property(), False)) else: - columns.append(Column( - "selectors", - db_backend.determine_sql_type_for_key_as_int(), - unique=True - ) + columns.append( + Column( + "selectors", + db_backend.determine_sql_type_for_key_as_int(), + unique=True, + ), ) child_columns = [ @@ -206,8 +207,12 @@ def create_granular_markings_table(metadata, db_backend, sco_or_sdo): nullable=False, ), ] - tables.append(Table(canonicalize_table_name("granular_marking_" + sco_or_sdo + "_" + "selector"), - metadata, *child_columns, schema=schema_name)) + tables.append( + Table( + canonicalize_table_name("granular_marking_" + sco_or_sdo + "_" + "selector"), + metadata, *child_columns, schema=schema_name, + ), + ) tables.append( Table( "granular_marking_" + sco_or_sdo, @@ -218,8 +223,9 @@ def create_granular_markings_table(metadata, db_backend, sco_or_sdo): OR (lang IS NOT NULL AND marking_ref IS NULL)""", ), - schema=schema_name - )) + schema=schema_name, + ), + ) return tables @@ -279,12 +285,16 @@ def create_core_table(metadata, db_backend, stix_type_name): if db_backend.array_allowed(): columns.append(create_array_column("labels", db_backend.determine_sql_type_for_string_property(), True)) else: - tables.append(create_array_child_table(metadata, - db_backend, - table_name, - "_labels", - "label", - db_backend.determine_sql_type_for_string_property())) + tables.append( + create_array_child_table( + metadata, + db_backend, + table_name, + "_labels", + "label", + db_backend.determine_sql_type_for_string_property(), + ), + ) else: columns.append(Column("defanged", db_backend.determine_sql_type_for_boolean_property(), default=False)) @@ -294,7 +304,8 @@ def create_core_table(metadata, db_backend, stix_type_name): metadata, *columns, schema=db_backend.schema_for_core(), - )) + ), + ) return tables @@ -422,9 +433,11 @@ def generate_table_information(self, name, db_backend, metadata, schema_name, ta else: contained_class = self.valid_types[0].contained columns.append( - create_array_column("value", - contained_class.determine_sql_type(db_backend), - False) + create_array_column( + "value", + contained_class.determine_sql_type(db_backend), + False, + ), ) else: for column_type in self.valid_types: @@ -586,6 +599,7 @@ def generate_table_information(self, name, db_backend, **kwargs): # noqa: F811 @add_method(ListProperty) def generate_table_information(self, name, db_backend, metadata, schema_name, table_name, **kwargs): # noqa: F811 is_extension = kwargs.get('is_extension') + is_embedded_object = kwargs.get('is_embedded_object') tables = list() # handle more complex embedded object before deciding if the ARRAY type is usable if isinstance(self.contained, EmbeddedObjectProperty): @@ -637,13 +651,23 @@ def generate_table_information(self, name, db_backend, metadata, schema_name, ta schema_name, ), ] - elif ((isinstance(self.contained, (StringProperty, IntegerProperty, FloatProperty)) and not db_backend.array_allowed()) or + elif (( + isinstance( + self.contained, + (BinaryProperty, BooleanProperty, StringProperty, IntegerProperty, FloatProperty, HexProperty), + ) and + not db_backend.array_allowed() + ) or isinstance(self.contained, EnumProperty)): columns = list() + if is_embedded_object: + id_type = db_backend.determine_sql_type_for_key_as_int() + else: + id_type = db_backend.determine_sql_type_for_key_as_id() columns.append( Column( "id", - self.contained.determine_sql_type(db_backend), + id_type, ForeignKey( canonicalize_table_name(table_name, schema_name) + ".id", ondelete="CASCADE", diff --git a/stix2/datastore/relational_db/utils.py b/stix2/datastore/relational_db/utils.py index 92806b38..ca387e46 100644 --- a/stix2/datastore/relational_db/utils.py +++ b/stix2/datastore/relational_db/utils.py @@ -1,9 +1,6 @@ from collections.abc import Iterable, Mapping import inflection -from sqlalchemy import ( # create_engine,; insert, - TIMESTAMP, Boolean, Float, Integer, LargeBinary, Text, -) from stix2.properties import ( BinaryProperty, BooleanProperty, FloatProperty, HexProperty,