Skip to content

Commit

Permalink
minor edits
Browse files Browse the repository at this point in the history
  • Loading branch information
rpiazza committed Jan 13, 2025
1 parent ddf36d4 commit 9653bd7
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 4 deletions.
36 changes: 36 additions & 0 deletions stix2/datastore/relational_db/demo.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
import datetime as dt

from database_backends.postgres_backend import PostgresBackend
import sys
import json

import stix2
from stix2.datastore.relational_db.relational_db import RelationalDBStore
import stix2.properties

# needed so the relational db code knows to create tables for this
from incident import incident, event, task, impact
from identity_contact_information import identity_contact_information
from observed_string import observed_string


def main():
with open(sys.argv[1], "r") as f:
bundle = stix2.parse(json.load(f), allow_custom=True)
store = RelationalDBStore(
PostgresBackend("postgresql://localhost/stix-data-sink", force_recreate=True),
True,
None,
True,
print_sql=True,
)

if store.sink.db_backend.database_exists:
for obj in bundle.objects:
store.add(obj)
else:
print("database does not exist")


if __name__ == '__main__':
main()
8 changes: 4 additions & 4 deletions stix2/datastore/relational_db/relational_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from stix2.parsing import parse


def _add(store, stix_data, allow_custom=True, version="2.1"):
def _add(sink, stix_data, allow_custom=True, version="2.1"):
"""Add STIX objects to MemoryStore/Sink.
Adds STIX objects to an in-memory dictionary for fast lookup.
Expand All @@ -32,12 +32,12 @@ def _add(store, stix_data, allow_custom=True, version="2.1"):
if isinstance(stix_data, list):
# STIX objects are in a list- recurse on each object
for stix_obj in stix_data:
_add(store, stix_obj, allow_custom, version)
_add(sink, stix_obj, allow_custom, version)

elif stix_data["type"] == "bundle":
# adding a json bundle - so just grab STIX objects
for stix_obj in stix_data.get("objects", []):
_add(store, stix_obj, allow_custom, version)
_add(sink, stix_obj, allow_custom, version)

else:
# Adding a single non-bundle object
Expand All @@ -46,7 +46,7 @@ def _add(store, stix_data, allow_custom=True, version="2.1"):
else:
stix_obj = parse(stix_data, allow_custom, version)

store.insert_object(stix_obj)
sink.insert_object(stix_obj)


class RelationalDBStore(DataStoreMixin):
Expand Down

0 comments on commit 9653bd7

Please sign in to comment.