Skip to content

Commit

Permalink
handle hex_property in list
Browse files Browse the repository at this point in the history
  • Loading branch information
rpiazza committed Nov 17, 2024
1 parent 35ce57a commit 3276661
Show file tree
Hide file tree
Showing 8 changed files with 120 additions and 100 deletions.
2 changes: 1 addition & 1 deletion stix2/datastore/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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

Expand Down
27 changes: 17 additions & 10 deletions stix2/datastore/relational_db/input_creation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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[
Expand All @@ -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]}
Expand Down Expand Up @@ -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"]
Expand Down Expand Up @@ -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:
Expand Down
19 changes: 9 additions & 10 deletions stix2/datastore/relational_db/relational_db.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down Expand Up @@ -88,6 +83,7 @@ def __init__(
source=RelationalDBSource(
db_backend,
metadata=self.metadata,
allow_custom=allow_custom,
),
sink=RelationalDBSink(
db_backend,
Expand Down Expand Up @@ -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
Expand All @@ -217,6 +214,8 @@ def __init__(

self.db_backend = db_backend

self.allow_custom = allow_custom

if metadata:
self.metadata = metadata
else:
Expand Down
75 changes: 40 additions & 35 deletions stix2/datastore/relational_db/relational_db_testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand All @@ -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()

Expand Down
Loading

0 comments on commit 3276661

Please sign in to comment.