From 29dec997a06f000622d9957e4a7a7b15494b7e8a Mon Sep 17 00:00:00 2001 From: Chris Lenk Date: Thu, 16 Nov 2017 14:58:59 -0500 Subject: [PATCH] Add related_to() Function for calling relationships() but instead of just returning the Relationship objects, returns the STIX objects being refered to in those Relationships. --- stix2/environment.py | 7 ++++ stix2/sources/__init__.py | 66 ++++++++++++++++++++++++++++++++++ stix2/test/test_environment.py | 28 +++++++++++++++ stix2/test/test_filesystem.py | 25 +++++++++++++ stix2/test/test_memory.py | 26 ++++++++++++++ 5 files changed, 152 insertions(+) diff --git a/stix2/environment.py b/stix2/environment.py index b018c01..3f0742b 100644 --- a/stix2/environment.py +++ b/stix2/environment.py @@ -178,3 +178,10 @@ class Environment(object): 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__ diff --git a/stix2/sources/__init__.py b/stix2/sources/__init__.py index 22f4027..7afe974 100644 --- a/stix2/sources/__init__.py +++ b/stix2/sources/__init__.py @@ -112,6 +112,30 @@ class DataStore(object): """ return self.source.relationships(*args, **kwargs) + def related_to(self, *args, **kwargs): + """Retrieve STIX Objects that have a Relationship involving the given + STIX object. + + Translate related_to() 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 + 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. + + """ + return self.source.related_to(*args, **kwargs) + def add(self, *args, **kwargs): """Method for storing STIX objects. @@ -236,6 +260,48 @@ class DataSource(with_metaclass(ABCMeta)): """ + 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`. + + 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. + + """ + results = [] + rels = self.relationships(obj, relationship_type, source_only, target_only) + + try: + obj_id = obj.get('id', '') + except AttributeError: + obj_id = obj + + 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(self.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(self.get(target_id)) + + return results + class CompositeDataSource(DataSource): """Controller for all the attached DataSources. diff --git a/stix2/test/test_environment.py b/stix2/test/test_environment.py index e595eb9..46b49c6 100644 --- a/stix2/test/test_environment.py +++ b/stix2/test/test_environment.py @@ -283,3 +283,31 @@ def test_relationships_by_target_and_source(ds): env.relationships(MALWARE_ID, target_only=True, source_only=True) assert 'not both' in str(excinfo.value) + + +def test_related_to(ds): + env = stix2.Environment(store=ds) + mal = env.get(MALWARE_ID) + resp = env.related_to(mal) + + assert len(resp) == 3 + assert any(x['id'] == CAMPAIGN_ID for x in resp) + assert any(x['id'] == INDICATOR_ID for x in resp) + assert any(x['id'] == IDENTITY_ID for x in resp) + + +def test_related_to_by_source(ds): + env = stix2.Environment(store=ds) + resp = env.related_to(MALWARE_ID, source_only=True) + + assert len(resp) == 1 + assert resp[0]['id'] == IDENTITY_ID + + +def test_related_to_by_target(ds): + env = stix2.Environment(store=ds) + resp = env.related_to(MALWARE_ID, target_only=True) + + assert len(resp) == 2 + assert any(x['id'] == CAMPAIGN_ID for x in resp) + assert any(x['id'] == INDICATOR_ID for x in resp) diff --git a/stix2/test/test_filesystem.py b/stix2/test/test_filesystem.py index 66ca47f..455132f 100644 --- a/stix2/test/test_filesystem.py +++ b/stix2/test/test_filesystem.py @@ -446,3 +446,28 @@ def test_relationships_by_target_and_source(rel_fs_store): rel_fs_store.relationships(MALWARE_ID, target_only=True, source_only=True) assert 'not both' in str(excinfo.value) + + +def test_related_to(rel_fs_store): + mal = rel_fs_store.get(MALWARE_ID) + resp = rel_fs_store.related_to(mal) + + assert len(resp) == 3 + assert any(x['id'] == CAMPAIGN_ID for x in resp) + assert any(x['id'] == INDICATOR_ID for x in resp) + assert any(x['id'] == IDENTITY_ID for x in resp) + + +def test_related_to_by_source(rel_fs_store): + resp = rel_fs_store.related_to(MALWARE_ID, source_only=True) + + assert len(resp) == 1 + assert any(x['id'] == IDENTITY_ID for x in resp) + + +def test_related_to_by_target(rel_fs_store): + resp = rel_fs_store.related_to(MALWARE_ID, target_only=True) + + assert len(resp) == 2 + assert any(x['id'] == CAMPAIGN_ID for x in resp) + assert any(x['id'] == INDICATOR_ID for x in resp) diff --git a/stix2/test/test_memory.py b/stix2/test/test_memory.py index d26b7d3..0096916 100644 --- a/stix2/test/test_memory.py +++ b/stix2/test/test_memory.py @@ -352,3 +352,29 @@ def test_relationships_by_target_and_source(rel_mem_store): rel_mem_store.relationships(MALWARE_ID, target_only=True, source_only=True) assert 'not both' in str(excinfo.value) + + +def test_related_to(rel_mem_store): + mal = rel_mem_store.get(MALWARE_ID) + resp = rel_mem_store.related_to(mal) + + assert len(resp) == 3 + assert any(x['id'] == CAMPAIGN_ID for x in resp) + assert any(x['id'] == INDICATOR_ID for x in resp) + assert any(x['id'] == IDENTITY_ID for x in resp) + + +def test_related_to_by_source(rel_mem_store): + resp = rel_mem_store.related_to(MALWARE_ID, source_only=True) + + assert len(resp) == 1 + assert any(x['id'] == IDENTITY_ID for x in resp) + + +def test_related_to_by_target(rel_mem_store): + resp = rel_mem_store.related_to(MALWARE_ID, target_only=True) + print(resp) + + assert len(resp) == 2 + assert any(x['id'] == CAMPAIGN_ID for x in resp) + assert any(x['id'] == INDICATOR_ID for x in resp)