From e1e368c0d2d4c3916d05bbe213ac0294ee885e01 Mon Sep 17 00:00:00 2001 From: Chris Lenk Date: Fri, 20 Oct 2017 09:13:04 -0400 Subject: [PATCH] Add created_by() function to Environment --- stix2/environment.py | 14 ++++++++++++++ stix2/test/test_environment.py | 22 ++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/stix2/environment.py b/stix2/environment.py index c4816ee..bd50b7f 100644 --- a/stix2/environment.py +++ b/stix2/environment.py @@ -152,3 +152,17 @@ class Environment(object): def parse(self, *args, **kwargs): return _parse(*args, **kwargs) parse.__doc__ = _parse.__doc__ + + def created_by(self, obj): + """Retrieve the Identity refered to by the object's `created_by_ref`. + + Args: + obj: The STIX object whose `created_by_ref` property will be looked + up. + + Returns: + The STIX object's creator. + """ + + creator_id = obj.get('created_by_ref', '') + return self.get(creator_id) diff --git a/stix2/test/test_environment.py b/stix2/test/test_environment.py index 81f2cda..5758732 100644 --- a/stix2/test/test_environment.py +++ b/stix2/test/test_environment.py @@ -184,3 +184,25 @@ def test_parse_malware(): assert mal.modified == FAKE_TIME assert mal.labels == ['ransomware'] assert mal.name == "Cryptolocker" + + +def test_created_by(): + identity = stix2.Identity(**IDENTITY_KWARGS) + factory = stix2.ObjectFactory(created_by_ref=identity.id) + env = stix2.Environment(store=stix2.MemoryStore(), factory=factory) + env.add(identity) + + ind = env.create(stix2.Indicator, **INDICATOR_KWARGS) + creator = env.created_by(ind) + assert creator.id == identity.id + + +def test_created_by_no_datasource(): + identity = stix2.Identity(**IDENTITY_KWARGS) + factory = stix2.ObjectFactory(created_by_ref=identity.id) + env = stix2.Environment(factory=factory) + + ind = env.create(stix2.Indicator, **INDICATOR_KWARGS) + with pytest.raises(AttributeError) as excinfo: + env.created_by(ind) + assert 'Environment has no data source' in str(excinfo.value)