Skip to content

Commit 7df7b9e

Browse files
rpiazzarpiazza
committed
more on ARRAY
1 parent a84867c commit 7df7b9e

File tree

7 files changed

+144
-113
lines changed

7 files changed

+144
-113
lines changed

‎stix2/datastore/relational_db/database_backends/database_backend_base.py‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,11 @@
1-
21
from typing import Any
3-
import os
42

5-
from sqlalchemy import create_engine
6-
from sqlalchemy_utils import create_database, database_exists, drop_database
73
from sqlalchemy import ( # create_engine,; insert,
84
ARRAY, TIMESTAMP, Boolean, CheckConstraint, Column, Float, ForeignKey,
9-
Integer, LargeBinary, Table, Text, UniqueConstraint,
5+
Integer, LargeBinary, Table, Text, UniqueConstraint, create_engine,
106
)
7+
from sqlalchemy_utils import create_database, database_exists, drop_database
8+
119

1210
class DatabaseBackend:
1311
def __init__(self, database_connection_url, force_recreate=False, **kwargs: Any):
@@ -26,7 +24,7 @@ def _create_schemas(self):
2624
pass
2725

2826
@staticmethod
29-
def _determine_schema_name(stix_object):
27+
def determine_schema_name(stix_object):
3028
return ""
3129

3230
def _create_database(self):
@@ -81,6 +79,3 @@ def determine_sql_type_for_key_as_id(): # noqa: F811
8179
@staticmethod
8280
def array_allowed():
8381
return False
84-
85-
86-

‎stix2/datastore/relational_db/database_backends/postgres_backend.py‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
import os
22
from typing import Any
3-
from sqlalchemy.schema import CreateSchema
3+
44
from sqlalchemy import ( # create_engine,; insert,
55
ARRAY, TIMESTAMP, Boolean, CheckConstraint, Column, Float, ForeignKey,
66
Integer, LargeBinary, Table, Text, UniqueConstraint,
77
)
8-
9-
from stix2.datastore.relational_db.utils import schema_for
8+
from sqlalchemy.schema import CreateSchema
109

1110
from stix2.base import (
12-
_DomainObject, _Extension, _MetaObject, _Observable, _RelationshipObject, _STIXBase,
11+
_DomainObject, _Extension, _MetaObject, _Observable, _RelationshipObject,
12+
_STIXBase,
1313
)
14-
14+
from stix2.datastore.relational_db.utils import schema_for
1515
from stix2.properties import (
1616
BinaryProperty, BooleanProperty, DictionaryProperty,
1717
EmbeddedObjectProperty, EnumProperty, ExtensionsProperty, FloatProperty,
@@ -22,6 +22,7 @@
2222

2323
from .database_backend_base import DatabaseBackend
2424

25+
2526
class PostgresBackend(DatabaseBackend):
2627
default_database_connection_url = \
2728
f"postgresql://{os.getenv('POSTGRES_USER', 'postgres')}:" + \
@@ -40,7 +41,7 @@ def _create_schemas(self):
4041
trans.execute(CreateSchema("sro", if_not_exists=True))
4142

4243
@staticmethod
43-
def _determine_schema_name(stix_object):
44+
def determine_schema_name(stix_object):
4445
if isinstance(stix_object, _DomainObject):
4546
return "sdo"
4647
elif isinstance(stix_object, _Observable):
@@ -73,6 +74,3 @@ def determine_sql_type_for_timestamp_property(): # noqa: F811
7374
@staticmethod
7475
def array_allowed():
7576
return True
76-
77-
78-

‎stix2/datastore/relational_db/input_creation.py‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,16 +30,17 @@ def generate_insert_information(self, name, stix_object, **kwargs): # noqa: F81
3030
return {name: stix_object[name]}
3131

3232

33-
3433
def instance_in_valid_types(cls, valid_types):
3534
for v in valid_types:
3635
if isinstance(v, cls):
3736
return True
3837
return False
3938

39+
4040
def is_valid_type(cls, valid_types):
4141
return cls in valid_types or instance_in_valid_types(cls, valid_types)
4242

43+
4344
@add_method(DictionaryProperty)
4445
def generate_insert_information(self, dictionary_name, stix_object, **kwargs): # noqa: F811
4546
bindings = dict()
@@ -311,7 +312,8 @@ def generate_insert_for_external_references(data_sink, stix_object):
311312
return insert_statements
312313

313314

314-
def generate_insert_for_granular_markings(granular_markings_table, stix_object):
315+
def generate_insert_for_granular_markings(data_sink, granular_markings_table, stix_object):
316+
db_backend = data_sink.db_backend
315317
insert_statements = list()
316318
granular_markings = stix_object["granular_markings"]
317319
for idx, granular_marking in enumerate(granular_markings):
@@ -322,7 +324,18 @@ def generate_insert_for_granular_markings(granular_markings_table, stix_object):
322324
marking_ref_value = granular_marking.get("marking_ref")
323325
if marking_ref_value:
324326
bindings["marking_ref"] = marking_ref_value
325-
bindings["selectors"] = granular_marking.get("selectors")
327+
if db_backend.array_allowed():
328+
bindings["selectors"] = granular_marking.get("selectors")
329+
else:
330+
table = data_sink.tables_dictionary[
331+
canonicalize_table_name(
332+
granular_markings_table + "_selector",
333+
db_backend.schema_for_core(),
334+
)
335+
]
336+
for sel in granular_marking.get("selectors"):
337+
selector_bindings = {"id": stix_object["id"], "selector": sel}
338+
insert_statements.append(insert(table).values(selector_bindings))
326339
insert_statements.append(insert(granular_markings_table).values(bindings))
327340
return insert_statements
328341

@@ -359,12 +372,13 @@ def generate_insert_for_core(data_sink, stix_object, core_properties, schema_nam
359372

360373
core_insert_statement = insert(core_table).values(core_bindings)
361374
insert_statements.append(core_insert_statement)
375+
object_marking_table_name = canonicalize_table_name("object_marking_refs", data_sink.db_backend.schema_for_core())
362376

363377
if "object_marking_refs" in stix_object:
364378
if schema_name != "sco":
365-
object_markings_ref_table = data_sink.tables_dictionary["common.object_marking_refs_sdo"]
379+
object_markings_ref_table = data_sink.tables_dictionary[object_marking_table_name + "_sdo"]
366380
else:
367-
object_markings_ref_table = data_sink.tables_dictionary["common.object_marking_refs_sco"]
381+
object_markings_ref_table = data_sink.tables_dictionary[object_marking_table_name + "_sco"]
368382
insert_statements.extend(
369383
generate_insert_for_array_in_table(
370384
object_markings_ref_table,
@@ -375,11 +389,16 @@ def generate_insert_for_core(data_sink, stix_object, core_properties, schema_nam
375389

376390
# Granular markings
377391
if "granular_markings" in stix_object:
392+
granular_marking_table_name = canonicalize_table_name(
393+
"granular_marking",
394+
data_sink.db_backend.schema_for_core(),
395+
)
378396
if schema_name != "sco":
379-
granular_marking_table = data_sink.tables_dictionary["common.granular_marking_sdo"]
397+
granular_marking_table = data_sink.tables_dictionary[granular_marking_table_name + "_sdo"]
380398
else:
381-
granular_marking_table = data_sink.tables_dictionary["common.granular_marking_sco"]
399+
granular_marking_table = data_sink.tables_dictionary[granular_marking_table_name + "_sco"]
382400
granular_input_statements = generate_insert_for_granular_markings(
401+
data_sink,
383402
granular_marking_table,
384403
stix_object,
385404
)

‎stix2/datastore/relational_db/relational_db.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,6 @@
1515
canonicalize_table_name, schema_for, table_name_for,
1616
)
1717
from stix2.parsing import parse
18-
import stix2.registry
19-
import stix2.utils
2018

2119

2220
def _add(store, stix_data, allow_custom=True, version="2.1"):
@@ -202,7 +200,7 @@ def add(self, stix_data, version=None):
202200
# return "common"
203201

204202
def insert_object(self, stix_object):
205-
schema_name = self.db_backend._determine_schema_name(stix_object)
203+
schema_name = self.db_backend.determine_schema_name(stix_object)
206204
with self.db_backend.database_connection.begin() as trans:
207205
statements = generate_insert_for_object(self, stix_object, schema_name)
208206
for stmt in statements:
@@ -259,7 +257,7 @@ def get(self, stix_id, version=None, _composite_filters=None):
259257
stix_obj = read_object(
260258
stix_id,
261259
self.metadata,
262-
conn
260+
conn,
263261
)
264262

265263
return stix_obj

‎stix2/datastore/relational_db/relational_db_testing.py‎

Lines changed: 29 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -128,14 +128,14 @@ def dictionary_test():
128128
defanged=True,
129129
ctime="1980-02-23T05:43:28.2678Z",
130130
extensions={
131-
"raster-image-ext": {
132-
"exif_tags": {
133-
"Make": "Nikon",
134-
"Model": "D7000",
135-
"XResolution": 4928,
136-
"YResolution": 3264,
137-
},
138-
},
131+
"raster-image-ext": {
132+
"exif_tags": {
133+
"Make": "Nikon",
134+
"Model": "D7000",
135+
"XResolution": 4928,
136+
"YResolution": 3264,
137+
},
138+
},
139139
},
140140
)
141141

@@ -148,34 +148,35 @@ def kill_chain_test():
148148
modified="2016-05-12T08:17:27.000Z",
149149
name="Spear Phishing",
150150
kill_chain_phases=[
151-
{
152-
"kill_chain_name": "lockheed-martin-cyber-kill-chain",
153-
"phase_name": "reconnaissance",
154-
},
151+
{
152+
"kill_chain_name": "lockheed-martin-cyber-kill-chain",
153+
"phase_name": "reconnaissance",
154+
},
155155
],
156156
external_references=[
157-
{
158-
"source_name": "capec",
159-
"external_id": "CAPEC-163",
160-
},
157+
{
158+
"source_name": "capec",
159+
"external_id": "CAPEC-163",
160+
},
161161
],
162162
granular_markings=[
163-
{
164-
"lang": "en_US",
165-
"selectors": ["kill_chain_phases"],
166-
},
167-
{
168-
"marking_ref": "marking-definition--50902d70-37ae-4f85-af68-3f4095493b42",
169-
"selectors": ["external_references"],
170-
},
163+
{
164+
"lang": "en_US",
165+
"selectors": ["kill_chain_phases"],
166+
},
167+
{
168+
"marking_ref": "marking-definition--50902d70-37ae-4f85-af68-3f4095493b42",
169+
"selectors": ["external_references"],
170+
},
171171
], )
172172

173+
173174
@stix2.CustomObject(
174175
'x-custom-type',
175-
properties=[
176-
("phases", stix2.properties.ListProperty(stix2.KillChainPhase)),
177-
("something_else", stix2.properties.IntegerProperty()),
178-
],
176+
properties=[
177+
("phases", stix2.properties.ListProperty(stix2.KillChainPhase)),
178+
("something_else", stix2.properties.IntegerProperty()),
179+
],
179180
)
180181
class CustomClass:
181182
pass
@@ -275,8 +276,6 @@ def main():
275276

276277
# store.add(tb)
277278

278-
279-
280279
co = custom_obj()
281280

282281
store.add(co)

0 commit comments

Comments
 (0)