Reorganize DataStore code for dereferencing
- Move `relationships()` to DataStore like `related_to()` is. If a DataStore implementation needs a different way to handle relationship dereferencing (e.g. TAXII in the future, or CompositeDataSource), it can overwrite these functions. - Reduce code duplication. - Check for presence of Data Source/Sink in all DataStores, not just in Environment.stix2.0
Chris Lenk
parent
29dec997a0
commit
f0331f8b9b
|
|
@ -105,30 +105,12 @@ class Environment(object):
|
|||
return self.factory.create(*args, **kwargs)
|
||||
create.__doc__ = ObjectFactory.create.__doc__
|
||||
|
||||
def get(self, *args, **kwargs):
|
||||
try:
|
||||
return self.source.get(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('Environment has no data source to query')
|
||||
get.__doc__ = DataStore.get.__doc__
|
||||
|
||||
def all_versions(self, *args, **kwargs):
|
||||
"""Retrieve all versions of a single STIX object by ID.
|
||||
"""
|
||||
try:
|
||||
return self.source.all_versions(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('Environment has no data source to query')
|
||||
all_versions.__doc__ = DataStore.all_versions.__doc__
|
||||
|
||||
def query(self, *args, **kwargs):
|
||||
"""Retrieve STIX objects matching a set of filters.
|
||||
"""
|
||||
try:
|
||||
return self.source.query(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('Environment has no data source to query')
|
||||
query.__doc__ = DataStore.query.__doc__
|
||||
get = DataStore.__dict__['get']
|
||||
all_versions = DataStore.__dict__['all_versions']
|
||||
query = DataStore.__dict__['query']
|
||||
relationships = DataStore.__dict__['relationships']
|
||||
related_to = DataStore.__dict__['related_to']
|
||||
add = DataStore.__dict__['add']
|
||||
|
||||
def add_filters(self, *args, **kwargs):
|
||||
try:
|
||||
|
|
@ -142,13 +124,6 @@ class Environment(object):
|
|||
except AttributeError:
|
||||
raise AttributeError('Environment has no data source')
|
||||
|
||||
def add(self, *args, **kwargs):
|
||||
try:
|
||||
return self.sink.add(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('Environment has no data sink to put objects in')
|
||||
add.__doc__ = DataStore.add.__doc__
|
||||
|
||||
def parse(self, *args, **kwargs):
|
||||
return _parse(*args, **kwargs)
|
||||
parse.__doc__ = _parse.__doc__
|
||||
|
|
@ -171,17 +146,3 @@ class Environment(object):
|
|||
return self.get(creator_id)
|
||||
else:
|
||||
return None
|
||||
|
||||
def relationships(self, *args, **kwargs):
|
||||
try:
|
||||
return self.source.relationships(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('Environment has no data source')
|
||||
relationships.__doc__ = DataStore.relationships.__doc__
|
||||
|
||||
def related_to(self, *args, **kwargs):
|
||||
try:
|
||||
return self.source.related_to(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('Environment has no data source')
|
||||
related_to.__doc__ = DataStore.related_to.__doc__
|
||||
|
|
|
|||
|
|
@ -59,7 +59,10 @@ class DataStore(object):
|
|||
object specified by the "id".
|
||||
|
||||
"""
|
||||
return self.source.get(*args, **kwargs)
|
||||
try:
|
||||
return self.source.get(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('%s has no data source to query' % self.__class__.__name__)
|
||||
|
||||
def all_versions(self, *args, **kwargs):
|
||||
"""Retrieve all versions of a single STIX object by ID.
|
||||
|
|
@ -73,7 +76,10 @@ class DataStore(object):
|
|||
stix_objs (list): a list of STIX objects
|
||||
|
||||
"""
|
||||
return self.source.all_versions(*args, **kwargs)
|
||||
try:
|
||||
return self.source.all_versions(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('%s has no data source to query' % self.__class__.__name__)
|
||||
|
||||
def query(self, *args, **kwargs):
|
||||
"""Retrieve STIX objects matching a set of filters.
|
||||
|
|
@ -88,7 +94,10 @@ class DataStore(object):
|
|||
stix_objs (list): a list of STIX objects
|
||||
|
||||
"""
|
||||
return self.source.query(*args, **kwargs)
|
||||
try:
|
||||
return self.source.query(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('%s has no data source to query' % self.__class__.__name__)
|
||||
|
||||
def relationships(self, *args, **kwargs):
|
||||
"""Retrieve Relationships involving the given STIX object.
|
||||
|
|
@ -110,7 +119,10 @@ class DataStore(object):
|
|||
(list): List of Relationship objects involving the given STIX object.
|
||||
|
||||
"""
|
||||
return self.source.relationships(*args, **kwargs)
|
||||
try:
|
||||
return self.source.relationships(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('%s has no data source to query' % self.__class__.__name__)
|
||||
|
||||
def related_to(self, *args, **kwargs):
|
||||
"""Retrieve STIX Objects that have a Relationship involving the given
|
||||
|
|
@ -134,7 +146,10 @@ class DataStore(object):
|
|||
(list): List of STIX objects related to the given STIX object.
|
||||
|
||||
"""
|
||||
return self.source.related_to(*args, **kwargs)
|
||||
try:
|
||||
return self.source.related_to(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('%s has no data source to query' % self.__class__.__name__)
|
||||
|
||||
def add(self, *args, **kwargs):
|
||||
"""Method for storing STIX objects.
|
||||
|
|
@ -146,7 +161,10 @@ class DataStore(object):
|
|||
stix_objs (list): a list of STIX objects
|
||||
|
||||
"""
|
||||
return self.sink.add(*args, **kwargs)
|
||||
try:
|
||||
return self.sink.add(*args, **kwargs)
|
||||
except AttributeError:
|
||||
raise AttributeError('%s has no data sink to put objects in' % self.__class__.__name__)
|
||||
|
||||
|
||||
class DataSink(with_metaclass(ABCMeta)):
|
||||
|
|
@ -238,11 +256,8 @@ class DataSource(with_metaclass(ABCMeta)):
|
|||
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def relationships(self, obj, relationship_type=None, source_only=False, target_only=False):
|
||||
"""
|
||||
Implement: The specific data source API calls, processing,
|
||||
functionality required for dereferencing relationships.
|
||||
"""Retrieve Relationships involving the given STIX object.
|
||||
|
||||
Only one of `source_only` and `target_only` may be `True`.
|
||||
|
||||
|
|
@ -259,6 +274,26 @@ class DataSource(with_metaclass(ABCMeta)):
|
|||
(list): List of Relationship objects involving the given STIX object.
|
||||
|
||||
"""
|
||||
results = []
|
||||
filters = [Filter('type', '=', 'relationship')]
|
||||
|
||||
try:
|
||||
obj_id = obj.get('id', '')
|
||||
except AttributeError:
|
||||
obj_id = obj
|
||||
|
||||
if relationship_type:
|
||||
filters.append(Filter('relationship_type', '=', relationship_type))
|
||||
|
||||
if source_only and target_only:
|
||||
raise ValueError("Search either source only or target only, but not both")
|
||||
|
||||
if not target_only:
|
||||
results.extend(self.query(filters + [Filter('source_ref', '=', obj_id)]))
|
||||
if not source_only:
|
||||
results.extend(self.query(filters + [Filter('target_ref', '=', obj_id)]))
|
||||
|
||||
return results
|
||||
|
||||
def related_to(self, obj, relationship_type=None, source_only=False, target_only=False):
|
||||
"""Retrieve STIX Objects that have a Relationship involving the given
|
||||
|
|
@ -486,6 +521,9 @@ class CompositeDataSource(DataSource):
|
|||
(list): List of Relationship objects involving the given STIX object.
|
||||
|
||||
"""
|
||||
if not self.has_data_sources():
|
||||
raise AttributeError('CompositeDataSource has no data sources')
|
||||
|
||||
results = []
|
||||
filters = [Filter('type', '=', 'relationship')]
|
||||
|
||||
|
|
@ -508,6 +546,56 @@ class CompositeDataSource(DataSource):
|
|||
|
||||
return results
|
||||
|
||||
def related_to(self, obj, relationship_type=None, source_only=False, target_only=False):
|
||||
"""Retrieve STIX Objects that have a Relationship involving the given
|
||||
STIX object.
|
||||
|
||||
Only one of `source_only` and `target_only` may be `True`.
|
||||
|
||||
Federated related objects method - iterates through all
|
||||
DataSources defined in "data_sources".
|
||||
|
||||
Args:
|
||||
obj (STIX object OR dict OR str): The STIX object (or its ID) whose
|
||||
related objects will be looked up.
|
||||
relationship_type (str): Only retrieve objects related by this
|
||||
Relationships type.
|
||||
source_only (bool): Only examine Relationships for which this
|
||||
object is the source_ref. Default: False.
|
||||
target_only (bool): Only examine Relationships for which this
|
||||
object is the target_ref. Default: False.
|
||||
|
||||
Returns:
|
||||
(list): List of STIX objects related to the given STIX object.
|
||||
|
||||
"""
|
||||
if not self.has_data_sources():
|
||||
raise AttributeError('CompositeDataSource has no data sources')
|
||||
|
||||
results = []
|
||||
for ds in self.data_sources:
|
||||
rels = ds.relationships(obj, relationship_type, source_only, target_only)
|
||||
|
||||
try:
|
||||
obj_id = obj.get('id', '')
|
||||
except AttributeError:
|
||||
obj_id = obj
|
||||
|
||||
for ds in self.data_sources:
|
||||
for r in rels:
|
||||
if not source_only:
|
||||
# relationships() found relationships where target_ref is obj_id
|
||||
source_id = r.source_ref
|
||||
if source_id != obj_id: # needed if target_only is also false
|
||||
results.append(ds.get(source_id))
|
||||
if not target_only:
|
||||
# relationships() found relationships where source_ref is obj_id
|
||||
target_id = r.target_ref
|
||||
if target_id != obj_id: # needed if source_only is also false
|
||||
results.append(ds.get(target_id))
|
||||
|
||||
return results
|
||||
|
||||
def add_data_source(self, data_source):
|
||||
"""Attach a DataSource to CompositeDataSource instance
|
||||
|
||||
|
|
|
|||
|
|
@ -308,45 +308,6 @@ class FileSystemSource(DataSource):
|
|||
|
||||
return stix_objs
|
||||
|
||||
def relationships(self, obj, relationship_type=None, source_only=False, target_only=False):
|
||||
"""Retrieve Relationships involving the given STIX object.
|
||||
|
||||
Only one of `source_only` and `target_only` may be `True`.
|
||||
|
||||
Args:
|
||||
obj (STIX object OR dict OR str): The STIX object (or its ID) whose
|
||||
relationships will be looked up.
|
||||
relationship_type (str): Only retrieve Relationships of this type.
|
||||
source_only (bool): Only retrieve Relationships for which this
|
||||
object is the source_ref. Default: False.
|
||||
target_only (bool): Only retrieve Relationships for which this
|
||||
object is the target_ref. Default: False.
|
||||
|
||||
Returns:
|
||||
(list): List of Relationship objects involving the given STIX object.
|
||||
|
||||
"""
|
||||
results = []
|
||||
filters = [Filter('type', '=', 'relationship')]
|
||||
|
||||
try:
|
||||
obj_id = obj.get('id', '')
|
||||
except AttributeError:
|
||||
obj_id = obj
|
||||
|
||||
if relationship_type:
|
||||
filters.append(Filter('relationship_type', '=', relationship_type))
|
||||
|
||||
if source_only and target_only:
|
||||
raise ValueError("Search either source only or target only, but not both")
|
||||
|
||||
if not target_only:
|
||||
results.extend(self.query(filters + [Filter('source_ref', '=', obj_id)]))
|
||||
if not source_only:
|
||||
results.extend(self.query(filters + [Filter('target_ref', '=', obj_id)]))
|
||||
|
||||
return results
|
||||
|
||||
def _parse_file_filters(self, query):
|
||||
"""Extract STIX common filters.
|
||||
|
||||
|
|
|
|||
|
|
@ -301,45 +301,6 @@ class MemorySource(DataSource):
|
|||
|
||||
return all_data
|
||||
|
||||
def relationships(self, obj, relationship_type=None, source_only=False, target_only=False):
|
||||
"""Retrieve Relationships involving the given STIX object.
|
||||
|
||||
Only one of `source_only` and `target_only` may be `True`.
|
||||
|
||||
Args:
|
||||
obj (STIX object OR dict OR str): The STIX object (or its ID) whose
|
||||
relationships will be looked up.
|
||||
relationship_type (str): Only retrieve Relationships of this type.
|
||||
source_only (bool): Only retrieve Relationships for which this
|
||||
object is the source_ref. Default: False.
|
||||
target_only (bool): Only retrieve Relationships for which this
|
||||
object is the target_ref. Default: False.
|
||||
|
||||
Returns:
|
||||
(list): List of Relationship objects involving the given STIX object.
|
||||
|
||||
"""
|
||||
results = []
|
||||
filters = [Filter('type', '=', 'relationship')]
|
||||
|
||||
try:
|
||||
obj_id = obj.get('id', '')
|
||||
except AttributeError:
|
||||
obj_id = obj
|
||||
|
||||
if relationship_type:
|
||||
filters.append(Filter('relationship_type', '=', relationship_type))
|
||||
|
||||
if source_only and target_only:
|
||||
raise ValueError("Search either source only or target only, but not both")
|
||||
|
||||
if not target_only:
|
||||
results.extend(self.query(filters + [Filter('source_ref', '=', obj_id)]))
|
||||
if not source_only:
|
||||
results.extend(self.query(filters + [Filter('target_ref', '=', obj_id)]))
|
||||
|
||||
return results
|
||||
|
||||
def load_from_file(self, file_path, allow_custom=False, version=None):
|
||||
file_path = os.path.abspath(file_path)
|
||||
stix_data = json.load(open(file_path, "r"))
|
||||
|
|
|
|||
|
|
@ -222,45 +222,6 @@ class TAXIICollectionSource(DataSource):
|
|||
|
||||
return stix_objs
|
||||
|
||||
def relationships(self, obj, relationship_type=None, source_only=False, target_only=False):
|
||||
"""Retrieve Relationships involving the given STIX object.
|
||||
|
||||
Only one of `source_only` and `target_only` may be `True`.
|
||||
|
||||
Args:
|
||||
obj (STIX object OR dict OR str): The STIX object (or its ID) whose
|
||||
relationships will be looked up.
|
||||
relationship_type (str): Only retrieve Relationships of this type.
|
||||
source_only (bool): Only retrieve Relationships for which this
|
||||
object is the source_ref. Default: False.
|
||||
target_only (bool): Only retrieve Relationships for which this
|
||||
object is the target_ref. Default: False.
|
||||
|
||||
Returns:
|
||||
(list): List of Relationship objects involving the given STIX object.
|
||||
|
||||
"""
|
||||
results = []
|
||||
filters = [Filter('type', '=', 'relationship')]
|
||||
|
||||
try:
|
||||
obj_id = obj.get('id', '')
|
||||
except AttributeError:
|
||||
obj_id = obj
|
||||
|
||||
if relationship_type:
|
||||
filters.append(Filter('relationship_type', '=', relationship_type))
|
||||
|
||||
if source_only and target_only:
|
||||
raise ValueError("Search either source only or target only, but not both")
|
||||
|
||||
if not target_only:
|
||||
results.extend(self.query(filters + [Filter('source_ref', '=', obj_id)]))
|
||||
if not source_only:
|
||||
results.extend(self.query(filters + [Filter('target_ref', '=', obj_id)]))
|
||||
|
||||
return results
|
||||
|
||||
def _parse_taxii_filters(self, query):
|
||||
"""Parse out TAXII filters that the TAXII server can filter on.
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue