From 45d3020518bb97b67235d5e41a9ea5d032450b42 Mon Sep 17 00:00:00 2001
From: "Desai, Kartikey H" <khdesai@mitre.org>
Date: Fri, 17 May 2019 14:21:35 -0500
Subject: [PATCH] Fixes #257

---
 stix2/test/v21/test_bundle.py | 19 +++++++++++++++++--
 stix2/v21/bundle.py           | 14 +++++++++++++-
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/stix2/test/v21/test_bundle.py b/stix2/test/v21/test_bundle.py
index dc8ae53..3a8e996 100644
--- a/stix2/test/v21/test_bundle.py
+++ b/stix2/test/v21/test_bundle.py
@@ -183,8 +183,23 @@ 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
+    assert bundle['type'] == "bundle"
+    assert len(bundle["indicator--00000000-0000-4000-8000-000000000001"]) == 1
+
+
+def test_bundle_obj_id_not_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    assert bundle.type == "bundle"
+    assert bundle.id.startswith("bundle--")
+    assert type(bundle.objects[0]) is stix2.v21.Indicator
+    assert bundle.objects[0].type == 'indicator'
+    assert bundle.objects[1].type == 'malware'
+    assert bundle.objects[2].type == 'relationship'
+
+    with pytest.raises(KeyError) as excinfo:
+        bundle.get_obj('blah blah')
+    assert "does not match the id property of any of the bundle" in str(excinfo.value)
 
 
 def test_parse_unknown_type():
diff --git a/stix2/v21/bundle.py b/stix2/v21/bundle.py
index 9cf0557..77b78c1 100644
--- a/stix2/v21/bundle.py
+++ b/stix2/v21/bundle.py
@@ -35,4 +35,16 @@ class Bundle(_STIXBase):
         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)
+        found_objs = [elem for elem in self.objects if elem.id == obj_uuid]
+        if found_objs == []:
+            raise KeyError("'%s' does not match the id property of any of the bundle's objects" % obj_uuid)
+        return found_objs
+
+    def __getitem__(self, key):
+        try:
+            return super(Bundle, self).__getitem__(key)
+        except KeyError:
+            try:
+                return self.get_obj(key)
+            except KeyError:
+                raise KeyError("'%s' is neither a valid bundle property nor does it match the id property of any of the bundle's objects" % key)