handle hex_property in list

Author: rpiazza

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__)
 

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

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
 

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,20 +399,23 @@ 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(
             generate_insert_for_array_in_table(
                 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:

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:

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,71 +215,76 @@ 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,
     )
 
     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()