From 55cf00d7f0f5e0703c80a88844dddcfb4c1b4fdd Mon Sep 17 00:00:00 2001 From: Chris Lenk Date: Wed, 15 Nov 2017 10:37:17 -0500 Subject: [PATCH] Move relationships() to DataSources --- stix2/environment.py | 41 ++--------------- stix2/sources/__init__.py | 88 +++++++++++++++++++++++++++++++++++++ stix2/sources/filesystem.py | 39 ++++++++++++++++ stix2/sources/memory.py | 39 ++++++++++++++++ stix2/sources/taxii.py | 39 ++++++++++++++++ 5 files changed, 209 insertions(+), 37 deletions(-) diff --git a/stix2/environment.py b/stix2/environment.py index 6a8250f..b018c01 100644 --- a/stix2/environment.py +++ b/stix2/environment.py @@ -2,7 +2,6 @@ import copy from .core import parse as _parse from .sources import CompositeDataSource, DataStore -from .sources.filters import Filter class ObjectFactory(object): @@ -173,41 +172,9 @@ class Environment(object): else: return None - 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 of Relationship objects involving the given STIX object. - - """ - results = [] - filters = [Filter('type', '=', 'relationship')] - + def relationships(self, *args, **kwargs): try: - obj_id = obj.get('id', '') + return self.source.relationships(*args, **kwargs) 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 + raise AttributeError('Environment has no data source') + relationships.__doc__ = DataStore.relationships.__doc__ diff --git a/stix2/sources/__init__.py b/stix2/sources/__init__.py index b3e8a29..22f4027 100644 --- a/stix2/sources/__init__.py +++ b/stix2/sources/__init__.py @@ -16,6 +16,7 @@ import uuid from six import with_metaclass +from stix2.sources.filters import Filter from stix2.utils import deduplicate @@ -89,6 +90,28 @@ class DataStore(object): """ return self.source.query(*args, **kwargs) + def relationships(self, *args, **kwargs): + """Retrieve Relationships involving the given STIX object. + + Translate relationships() call to the appropriate DataSource call. + + 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. + + """ + return self.source.relationships(*args, **kwargs) + def add(self, *args, **kwargs): """Method for storing STIX objects. @@ -191,6 +214,28 @@ 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. + + 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. + + """ + class CompositeDataSource(DataSource): """Controller for all the attached DataSources. @@ -354,6 +399,49 @@ class CompositeDataSource(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`. + + Federated relationships retrieve method - iterates through all + DataSources defined in "data_sources". + + 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") + + for ds in self.data_sources: + if not target_only: + results.extend(ds.query(filters + [Filter('source_ref', '=', obj_id)])) + if not source_only: + results.extend(ds.query(filters + [Filter('target_ref', '=', obj_id)])) + + return results + def add_data_source(self, data_source): """Attach a DataSource to CompositeDataSource instance diff --git a/stix2/sources/filesystem.py b/stix2/sources/filesystem.py index e92c525..db22faa 100644 --- a/stix2/sources/filesystem.py +++ b/stix2/sources/filesystem.py @@ -308,6 +308,45 @@ 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. diff --git a/stix2/sources/memory.py b/stix2/sources/memory.py index 308d0d0..af0dd02 100644 --- a/stix2/sources/memory.py +++ b/stix2/sources/memory.py @@ -301,6 +301,45 @@ 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")) diff --git a/stix2/sources/taxii.py b/stix2/sources/taxii.py index 8eb5069..257bbd5 100644 --- a/stix2/sources/taxii.py +++ b/stix2/sources/taxii.py @@ -222,6 +222,45 @@ 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.