From a61344a8aafeaca7831a36f4ec78095d9af5f634 Mon Sep 17 00:00:00 2001 From: "Desai, Kartikey H" Date: Tue, 14 May 2019 13:48:54 -0400 Subject: [PATCH] Add get_obj function to bundle.py to make accessing bundles easier --- stix2/test/v21/test_bundle.py | 2 ++ stix2/v21/bundle.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/stix2/test/v21/test_bundle.py b/stix2/test/v21/test_bundle.py index 7adea92..dc8ae53 100644 --- a/stix2/test/v21/test_bundle.py +++ b/stix2/test/v21/test_bundle.py @@ -183,6 +183,8 @@ def test_parse_bundle(version): assert bundle.objects[0].type == 'indicator' assert bundle.objects[1].type == 'malware' assert bundle.objects[2].type == 'relationship' + assert bundle.get_obj('malware--00000000-0000-4000-8000-000000000003').type == 'malware' + assert bundle.get_obj('blah blah') is None def test_parse_unknown_type(): diff --git a/stix2/v21/bundle.py b/stix2/v21/bundle.py index c9e083a..9cf0557 100644 --- a/stix2/v21/bundle.py +++ b/stix2/v21/bundle.py @@ -33,3 +33,6 @@ class Bundle(_STIXBase): self._properties['objects'].contained.allow_custom = kwargs.get('allow_custom', False) super(Bundle, self).__init__(**kwargs) + + def get_obj(self, obj_uuid): + return next((elem for elem in self.objects if elem['id'] == obj_uuid), None)