Skip to content

Commit

Permalink
all dbs run
Browse files Browse the repository at this point in the history
  • Loading branch information
rpiazza committed Dec 12, 2024
1 parent adf41ab commit fab7199
Show file tree
Hide file tree
Showing 8 changed files with 59 additions and 50 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

from stix2.properties import HexProperty


class DatabaseBackend:
def __init__(self, database_connection_url, force_recreate=False, **kwargs: Any):
self.database_connection_url = database_connection_url
Expand Down Expand Up @@ -40,11 +41,11 @@ def determine_schema_name(stix_object):

@staticmethod
def schema_for(stix_class):
return ""
return None

@staticmethod
def schema_for_core():
return ""
return None

# =========================================================================
# sql type methods
Expand Down Expand Up @@ -133,3 +134,7 @@ def process_value_for_insert(self, stix_type, value):
return bytes.fromhex(value)
else:
return value

def next_id(self, data_sink):
with self.database_connection.begin() as trans:
return trans.execute(data_sink.sequence)
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import os
from typing import Any

from sqlalchemy import Text, VARCHAR
from sqlalchemy import Text, VARCHAR, CheckConstraint
from sqlalchemy.schema import CreateSchema

from stix2.base import (
Expand Down Expand Up @@ -80,6 +80,6 @@ def array_allowed():

@staticmethod
def create_regex_constraint_expression(column_name, pattern):
return f"{column_name} REGEXP {pattern}"
return CheckConstraint(f"{column_name} REGEXP {pattern}")


Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import os
from typing import Any

from sqlalchemy import TIMESTAMP, LargeBinary, Text
from sqlalchemy import TIMESTAMP, CheckConstraint, Text
from sqlalchemy.schema import CreateSchema

from stix2.base import (
Expand Down Expand Up @@ -88,4 +88,4 @@ def array_allowed():

@staticmethod
def create_regex_constraint_expression(column_name, pattern):
return f"{column_name} ~ {pattern}"
return CheckConstraint(f"{column_name} ~ {pattern}")
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
class SQLiteBackend(DatabaseBackend):
default_database_connection_url = f"sqlite:///stix-data-sink.db"

temp_sequence_count = 0

def __init__(self, database_connection_url=default_database_connection_url, force_recreate=False, **kwargs: Any):
super().__init__(database_connection_url, force_recreate=force_recreate, **kwargs)

Expand Down Expand Up @@ -64,4 +66,9 @@ def array_allowed():

@staticmethod
def create_regex_constraint_expression(column_name, pattern):
return f"{column_name} ~ {pattern}"
return None

@staticmethod
def next_id(data_sink):
SQLiteBackend.temp_sequence_count += 1
return SQLiteBackend.temp_sequence_count
22 changes: 14 additions & 8 deletions stix2/datastore/relational_db/input_creation.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ def generate_insert_information(self, dictionary_name, stix_object, **kwargs):
if is_valid_type(ListProperty, valid_types):
value_binding = "values"
if not data_sink.db_backend.array_allowed():
next_id = data_sink.next_id()
next_id = data_sink.db_backend.next_id(data_sink)
table_child = data_sink.tables_dictionary[
canonicalize_table_name(table_name + "_" + dictionary_name + "_" + "values", schema_name)]
child_table_inserts = generate_insert_for_dictionary_list(table_child, next_id, value)
Expand Down Expand Up @@ -261,7 +261,7 @@ def generate_insert_information( # noqa: F811
elif isinstance(self.contained, EmbeddedObjectProperty):
insert_statements = list()
for value in stix_object[name]:
next_id = data_sink.next_id()
next_id = db_backend.next_id(data_sink)
table = data_sink.tables_dictionary[canonicalize_table_name(table_name + "_" + name, schema_name)]
bindings = {
"id": foreign_key_value,
Expand Down Expand Up @@ -346,14 +346,15 @@ def generate_insert_for_array_in_table(table, values, foreign_key_value, column_
def generate_insert_for_external_references(data_sink, stix_object):
insert_statements = list()
next_id = None
object_table = data_sink.tables_dictionary["common.external_references"]
schema_name = data_sink.db_backend.schema_for_core()
object_table = data_sink.tables_dictionary[canonicalize_table_name("external_references", schema_name)]
for er in stix_object["external_references"]:
bindings = {"id": stix_object["id"]}
for prop in ["source_name", "description", "url", "external_id"]:
if prop in er:
bindings[prop] = er[prop]
if "hashes" in er:
next_id = data_sink.next_id()
next_id = data_sink.db_backend.next_id(data_sink)
bindings["hash_ref_id"] = next_id
# else:
# # hash_ref_id is non-NULL, so -1 means there are no hashes
Expand All @@ -363,7 +364,12 @@ def generate_insert_for_external_references(data_sink, stix_object):

if "hashes" in er:
insert_statements.extend(
generate_insert_for_hashes(data_sink, "hashes", er, "external_references", "common", foreign_key_value=next_id),
generate_insert_for_hashes(data_sink,
"hashes",
er,
"external_references",
schema_name,
foreign_key_value=next_id),
)

return insert_statements
Expand All @@ -385,7 +391,7 @@ def generate_insert_for_granular_markings(data_sink, granular_markings_table, st
bindings["selectors"] = granular_marking.get("selectors")
insert_statements.append(insert(granular_markings_table).values(bindings))
else:
next_id = data_sink.next_id()
next_id = db_backend.next_id(data_sink)
bindings["selectors"] = next_id
insert_statements.append(insert(granular_markings_table).values(bindings))
table = data_sink.tables_dictionary[
Expand Down Expand Up @@ -419,9 +425,9 @@ def generate_insert_for_granular_markings(data_sink, granular_markings_table, st
def generate_insert_for_core(data_sink, stix_object, core_properties, stix_type_name, schema_name):
db_backend = data_sink.db_backend
if stix_type_name in ["sdo", "sro", "common"]:
core_table = data_sink.tables_dictionary[db_backend.schema_for_core() + "." + "core_sdo"]
core_table = data_sink.tables_dictionary[canonicalize_table_name("core_sdo", db_backend.schema_for_core())]
else:
core_table = data_sink.tables_dictionary[db_backend.schema_for_core() + "." + "core_sco"]
core_table = data_sink.tables_dictionary[canonicalize_table_name("core_sco", db_backend.schema_for_core())]
insert_statements = list()
core_bindings = {}

Expand Down
7 changes: 1 addition & 6 deletions stix2/datastore/relational_db/relational_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
from stix2.datastore.relational_db.utils import canonicalize_table_name
from stix2.parsing import parse


def _add(store, stix_data, allow_custom=True, version="2.1"):
"""Add STIX objects to MemoryStore/Sink.
Expand Down Expand Up @@ -138,7 +137,7 @@ def __init__(
create_table_objects(
self.metadata, stix_object_classes,
)
self.sequence = Sequence("my_general_seq", metadata=self.metadata, start=1, schema="common")
self.sequence = Sequence("my_general_seq", metadata=self.metadata, start=1, schema=db_backend.schema_for_core())

self.allow_custom = allow_custom

Expand Down Expand Up @@ -182,10 +181,6 @@ def clear_tables(self):
print(f'delete_stmt: {delete_stmt}')
trans.execute(delete_stmt)

def next_id(self):
with self.db_backend.database_connection.begin() as trans:
return trans.execute(self.sequence)


class RelationalDBSource(DataSource):
def __init__(
Expand Down
8 changes: 4 additions & 4 deletions stix2/datastore/relational_db/relational_db_testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -290,9 +290,9 @@ def test_dictionary():

def main():
store = RelationalDBStore(
MariaDBBackend(f"mariadb+pymysql://admin:admin@127.0.0.1:3306/rdb", force_recreate=True),
#MariaDBBackend(f"mariadb+pymysql://admin:admin@127.0.0.1:3306/rdb", force_recreate=True),
#PostgresBackend("postgresql://localhost/stix-data-sink", force_recreate=True),
#SQLiteBackend("sqlite:///stix-data-sink.db", force_recreate=True),
SQLiteBackend("sqlite:///stix-data-sink.db", force_recreate=True),

True,
None,
Expand All @@ -305,9 +305,9 @@ def main():
ap = kill_chain_test()
store.add(ap)

x=email_message
# x=email_message

store.add(x)
# store.add(x)

td = test_dictionary()

Expand Down
46 changes: 21 additions & 25 deletions stix2/datastore/relational_db/table_creation.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
# from collections import OrderedDict

from sqlalchemy import ( # create_engine,; insert,
ARRAY, CheckConstraint, Column, ForeignKey, Integer, Table, Text,
UniqueConstraint,
Expand All @@ -18,7 +16,7 @@
ObjectReferenceProperty, Property, ReferenceProperty, StringProperty,
TimestampProperty, TypeProperty,
)
from stix2.v21.base import _Extension
from stix2.v21.base import _Extension, _Observable
from stix2.v21.common import KillChainPhase


Expand Down Expand Up @@ -61,13 +59,14 @@ def derive_column_name(prop):


def create_object_markings_refs_table(metadata, db_backend, sco_or_sdo):
schema_name = db_backend.schema_for_core()
return create_ref_table(
metadata,
db_backend,
{"marking-definition"},
"object_marking_refs_" + sco_or_sdo,
"common.core_" + sco_or_sdo + ".id",
"common",
canonicalize_table_name("core_" + sco_or_sdo, schema_name) + ".id",
schema_name,
0,
)

Expand Down Expand Up @@ -174,7 +173,7 @@ def create_granular_markings_table(metadata, db_backend, sco_or_sdo):
Column(
"marking_ref",
db_backend.determine_sql_type_for_reference_property(),
CheckConstraint(db_backend.create_regex_constraint_expression("marking_ref", reg_ex)),
db_backend.create_regex_constraint_expression("marking_ref", reg_ex),
),
]
if db_backend.array_allowed():
Expand Down Expand Up @@ -229,12 +228,13 @@ def create_granular_markings_table(metadata, db_backend, sco_or_sdo):

def create_external_references_tables(metadata, db_backend):
reg_ex = "'^[a-z][a-z0-9-]+[a-z0-9]--[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-5][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$'" # noqa: E131
schema_name = db_backend.schema_for_core()
columns = [
Column(
"id",
db_backend.determine_sql_type_for_key_as_id(),
ForeignKey("common.core_sdo" + ".id", ondelete="CASCADE"),
CheckConstraint(db_backend.create_regex_constraint_expression("id", reg_ex))),
ForeignKey(canonicalize_table_name("core_sdo", schema_name) + ".id", ondelete="CASCADE"),
db_backend.create_regex_constraint_expression("id", reg_ex)),
Column("source_name", db_backend.determine_sql_type_for_string_property()),
Column("description", db_backend.determine_sql_type_for_string_property()),
Column("url", db_backend.determine_sql_type_for_string_property()),
Expand All @@ -243,8 +243,8 @@ def create_external_references_tables(metadata, db_backend):
Column("hash_ref_id", db_backend.determine_sql_type_for_key_as_int(), autoincrement=False),
]
return [
Table("external_references", metadata, *columns, schema="common"),
create_hashes_table("hashes", metadata, db_backend, "common", "external_references", Integer),
Table("external_references", metadata, *columns, schema=schema_name),
create_hashes_table("hashes", metadata, db_backend, schema_name, "external_references", Integer),
]


Expand All @@ -256,7 +256,7 @@ def create_core_table(metadata, db_backend, stix_type_name):
Column(
"id",
db_backend.determine_sql_type_for_key_as_id(),
CheckConstraint(db_backend.create_regex_constraint_expression("id", reg_ex)),
db_backend.create_regex_constraint_expression("id", reg_ex),
primary_key=True,
),
Column("spec_version", db_backend.determine_sql_type_for_string_property(), default="2.1"),
Expand All @@ -267,7 +267,7 @@ def create_core_table(metadata, db_backend, stix_type_name):
Column(
"created_by_ref",
db_backend.determine_sql_type_for_reference_property(),
CheckConstraint(db_backend.create_regex_constraint_expression("created_by_ref", reg_ex))),
db_backend.create_regex_constraint_expression("created_by_ref", reg_ex)),
Column("created", db_backend.determine_sql_type_for_timestamp_property()),
Column("modified", db_backend.determine_sql_type_for_timestamp_property()),
Column("revoked", db_backend.determine_sql_type_for_boolean_property()),
Expand Down Expand Up @@ -399,8 +399,7 @@ def generate_table_information(self, name, db_backend, **kwargs): # noqa: F811
name,
self.determine_sql_type(db_backend),
# this regular expression might accept or reject some legal base64 strings
CheckConstraint(db_backend.create_regex_constraint_expression(name, "'^[-A-Za-z0-9+/]*={0,3}$'")
),
db_backend.create_regex_constraint_expression(name, "'^[-A-Za-z0-9+/]*={0,3}$'"),
nullable=not self.required,
)

Expand Down Expand Up @@ -527,7 +526,7 @@ def generate_table_information(self, name, db_backend, **kwargs): # noqa: F811
return Column(
name,
self.determine_sql_type(db_backend),
CheckConstraint(db_backend.create_regex_constraint_expression(name, f"'^{enum_re}$'")),
db_backend.create_regex_constraint_expression(name, f"'^{enum_re}$'"),
nullable=not self.required,
)

Expand Down Expand Up @@ -609,9 +608,7 @@ def generate_table_information(self, name, db_backend, **kwargs): # noqa: F811
name,
db_backend.determine_sql_type_for_key_as_id(),
ForeignKey(foreign_key_column, ondelete="CASCADE"),
CheckConstraint(
db_backend.create_regex_constraint_expression(name, id_req_exp)
),
db_backend.create_regex_constraint_expression(name, id_req_exp),
primary_key=True,
nullable=not (self.required),
)
Expand Down Expand Up @@ -723,14 +720,13 @@ def ref_column(name, specifics, db_backend, auth_type=0):
types = "|".join(specifics)
if auth_type == 0:
reg_ex = f"'^({types})" + "--[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-5][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$'" # noqa: F811
constraint = \
CheckConstraint(db_backend.create_regex_constraint_expression(name, reg_ex))
constraint = db_backend.create_regex_constraint_expression(name, reg_ex)
else:
reg_ex = "'--[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-5][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$')"
constraint = \
CheckConstraint(db_backend.create_regex_constraint_expression(f"NOT({name}", f"'^({types})'") + " AND " +
db_backend.create_regex_constraint_expression(name, reg_ex))
return Column(name, db_backend.determine_sql_type_for_reference_property(), constraint)
# constraint = \
# CheckConstraint(db_backend.create_regex_constraint_expression(f"NOT({name}", f"'^({types})'") + " AND " +
# db_backend.create_regex_constraint_expression(name, reg_ex))
return Column(name, db_backend.determine_sql_type_for_reference_property()) # , constraint)
else:
return Column(
name,
Expand Down Expand Up @@ -805,7 +801,7 @@ def generate_object_table(
core_properties = list()
columns = list()
tables = list()
if schema_name == "sco":
if issubclass(stix_object_class, _Observable):
core_table = "sco"
else:
# sro, smo common properties are the same as sdo's
Expand Down

0 comments on commit fab7199

Please sign in to comment.