cti-python-stix2/stix2/sources/memory.py

203 lines
6.1 KiB
Python
Raw Normal View History

2017-07-12 16:58:31 +02:00
"""
Python STIX 2.0 Memory Source/Sink
Classes:
MemoryStore
MemorySink
MemorySource
TODO: Test everything.
TODO: Use deduplicate() calls only when memory corpus is dirty (been added to)
can save a lot of time for successive queries
Notes:
Not worrying about STIX versioning. The in memory STIX data at anytime
will only hold one version of a STIX object. As such, when save() is called,
the single versions of all the STIX objects are what is written to file.
2017-07-12 16:58:31 +02:00
"""
import collections
2017-07-12 16:58:31 +02:00
import json
import os
from stix2validator import validate_instance
2017-07-12 16:58:31 +02:00
from stix2 import Bundle
2017-08-31 20:03:12 +02:00
from stix2.sources import DataSink, DataSource, DataStore
from stix2.sources.filters import Filter
2017-07-12 16:58:31 +02:00
def _add(store, stix_data):
"""Adds stix objects to MemoryStore/Source/Sink."""
if isinstance(stix_data, collections.Mapping):
# stix objects are in a bundle
# verify STIX json data
r = validate_instance(stix_data)
# make dictionary of the objects for easy lookup
if r.is_valid:
for stix_obj in stix_data["objects"]:
store.data[stix_obj["id"]] = stix_obj
else:
raise ValueError("Error: data passed was found to not be valid by the STIX 2 Validator: \n%s", r.as_dict())
elif isinstance(stix_data, list):
# stix objects are in a list
for stix_obj in stix_data:
r = validate_instance(stix_obj)
if r.is_valid:
store.data[stix_obj["id"]] = stix_obj
else:
raise ValueError("Error: STIX object %s is not valid under STIX 2 validator.\n%s", stix_obj["id"], r)
else:
raise ValueError("stix_data must be in bundle format or raw list")
2017-07-12 16:58:31 +02:00
class MemoryStore(DataStore):
"""
"""
def __init__(self, stix_data):
"""
Notes:
It doesn't make sense to create a MemoryStore by passing
in existing MemorySource and MemorySink because there could
be data concurrency issues. Just as easy to create new MemoryStore.
"""
super(MemoryStore, self).__init__()
self.data = {}
2017-07-12 16:58:31 +02:00
if stix_data:
_add(self, stix_data)
self.source = MemorySource(stix_data=self.data, _store=True)
self.sink = MemorySink(stix_data=self.data, _store=True)
2017-07-12 16:58:31 +02:00
def save_to_file(self, file_path):
2017-07-20 21:34:09 +02:00
return self.sink.save_to_file(file_path=file_path)
2017-07-12 16:58:31 +02:00
def load_from_file(self, file_path):
return self.source.load_from_file(file_path=file_path)
2017-07-12 16:58:31 +02:00
class MemorySink(DataSink):
"""
"""
def __init__(self, stix_data, _store=False):
2017-07-12 16:58:31 +02:00
"""
Args:
stix_data (dictionary OR list): valid STIX 2.0 content in
bundle or a list.
_store (bool): if the MemorySink is a part of a DataStore,
in which case "stix_data" is a direct reference to
shared memory with DataSource.
2017-07-12 16:58:31 +02:00
"""
super(MemorySink, self).__init__()
self.data = {}
2017-07-12 16:58:31 +02:00
if _store:
self.data = stix_data
elif stix_data:
self.add(stix_data)
2017-07-12 16:58:31 +02:00
def add(self, stix_data):
"""
"""
_add(self, stix_data)
2017-07-12 16:58:31 +02:00
def save_to_file(self, file_path):
2017-07-12 16:58:31 +02:00
"""
"""
json.dump(Bundle(self.data.values()), file_path, indent=4)
class MemorySource(DataSource):
def __init__(self, stix_data, _store=False):
2017-07-12 16:58:31 +02:00
"""
Args:
stix_data (dictionary OR list): valid STIX 2.0 content in
bundle or list.
_store (bool): if the MemorySource is a part of a DataStore,
in which case "stix_data" is a direct reference to shared
memory with DataSink.
2017-07-12 16:58:31 +02:00
"""
super(MemorySource, self).__init__()
self.data = {}
2017-07-12 16:58:31 +02:00
if _store:
self.data = stix_data
elif stix_data:
_add(self, stix_data)
2017-07-12 16:58:31 +02:00
def get(self, stix_id, _composite_filters=None):
"""
"""
if _composite_filters is None:
# if get call is only based on 'id', no need to search, just retrieve from dict
try:
stix_obj = self.data[stix_id]
except KeyError:
stix_obj = None
return stix_obj
# if there are filters from the composite level, process full query
query = [Filter("id", "=", stix_id)]
2017-07-12 16:58:31 +02:00
all_data = self.query(query=query, _composite_filters=_composite_filters)
# reduce to most recent version
stix_obj = sorted(all_data, key=lambda k: k['modified'])[0]
return stix_obj
def all_versions(self, stix_id, _composite_filters=None):
"""
Notes:
Since Memory sources/sinks don't handle multiple versions of a
2017-08-30 17:18:11 +02:00
STIX object, this operation is unnecessary. Translate call to get().
Args:
stix_id (str): The id of the STIX 2.0 object to retrieve. Should
return a list of objects, all the versions of the object
specified by the "id".
2017-07-12 16:58:31 +02:00
Returns:
(list): STIX object that matched ``stix_id``.
2017-07-12 16:58:31 +02:00
"""
2017-07-12 16:58:31 +02:00
return [self.get(stix_id=stix_id, _composite_filters=_composite_filters)]
def query(self, query=None, _composite_filters=None):
"""
"""
if query is None:
query = []
# combine all query filters
if self.filters:
query.extend(list(self.filters))
2017-07-12 16:58:31 +02:00
if _composite_filters:
query.extend(_composite_filters)
# Apply STIX common property filters.
2017-07-12 16:58:31 +02:00
all_data = self.apply_common_filters(self.data.values(), query)
return all_data
def load_from_file(self, file_path):
"""
"""
file_path = os.path.abspath(file_path)
stix_data = json.load(open(file_path, "r"))
r = validate_instance(stix_data)
if r.is_valid:
for stix_obj in stix_data["objects"]:
self.data[stix_obj["id"]] = stix_obj
raise ValueError("Error: STIX data loaded from file (%s) was found to not be validated by STIX 2 Validator.\n%s", file_path, r)