Merge pull request #264 from khdesai/stix2_issue_257

Stix2 issue 257: make accessing bundles easier

stix2/test/v20/test_bundle.py

@@ -236,3 +236,100 @@ def test_bundle_with_different_spec_objects():
         stix2.v20.Bundle(objects=data)
 
     assert "Spec version 2.0 bundles don't yet support containing objects of a different spec version." in str(excinfo.value)
+
+
+def test_bundle_obj_id_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    mal_list = bundle.get_obj("malware--00000000-0000-4000-8000-000000000003")
+    assert bundle.objects[1] == mal_list[0]
+    assert len(mal_list) == 1
+
+
+@pytest.mark.parametrize(
+    "bundle_data", [{
+        "type": "bundle",
+        "id": "bundle--00000000-0000-4000-8000-000000000007",
+        "spec_version": "2.0",
+        "objects": [
+            {
+                "type": "indicator",
+                "id": "indicator--00000000-0000-4000-8000-000000000001",
+                "created": "2017-01-01T12:34:56.000Z",
+                "modified": "2017-01-01T12:34:56.000Z",
+                "pattern": "[file:hashes.MD5 = 'd41d8cd98f00b204e9800998ecf8427e']",
+                "valid_from": "2017-01-01T12:34:56Z",
+                "labels": [
+                    "malicious-activity",
+                ],
+            },
+            {
+                "type": "malware",
+                "id": "malware--00000000-0000-4000-8000-000000000003",
+                "created": "2017-01-01T12:34:56.000Z",
+                "modified": "2017-01-01T12:34:56.000Z",
+                "name": "Cryptolocker1",
+                "labels": [
+                    "ransomware",
+                ],
+            },
+            {
+                "type": "malware",
+                "id": "malware--00000000-0000-4000-8000-000000000003",
+                "created": "2017-01-01T12:34:56.000Z",
+                "modified": "2017-12-21T12:34:56.000Z",
+                "name": "CryptolockerOne",
+                "labels": [
+                    "ransomware",
+                ],
+            },
+            {
+                "type": "relationship",
+                "id": "relationship--00000000-0000-4000-8000-000000000005",
+                "created": "2017-01-01T12:34:56.000Z",
+                "modified": "2017-01-01T12:34:56.000Z",
+                "relationship_type": "indicates",
+                "source_ref": "indicator--a740531e-63ff-4e49-a9e1-a0a3eed0e3e7",
+                "target_ref": "malware--9c4638ec-f1de-4ddb-abf4-1b760417654e",
+            },
+        ],
+    }],
+)
+def test_bundle_objs_ids_found(bundle_data):
+    bundle = stix2.parse(bundle_data)
+
+    mal_list = bundle.get_obj("malware--00000000-0000-4000-8000-000000000003")
+    assert bundle.objects[1] == mal_list[0]
+    assert bundle.objects[2] == mal_list[1]
+    assert len(mal_list) == 2
+
+
+def test_bundle_getitem_overload_property_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    assert bundle.type == "bundle"
+    assert bundle['type'] == "bundle"
+
+
+def test_bundle_getitem_overload_obj_id_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    mal_list = bundle["malware--00000000-0000-4000-8000-000000000003"]
+    assert bundle.objects[1] == mal_list[0]
+    assert len(mal_list) == 1
+
+
+def test_bundle_obj_id_not_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    with pytest.raises(KeyError) as excinfo:
+        bundle.get_obj('non existent')
+    assert "does not match the id property of any of the bundle" in str(excinfo.value)
+
+
+def test_bundle_getitem_overload_obj_id_not_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    with pytest.raises(KeyError) as excinfo:
+        bundle['non existent']
+    assert "neither a valid bundle property nor does it match the id property" in str(excinfo.value)

stix2/test/v21/test_bundle.py

@@ -207,3 +207,103 @@ def test_stix_object_property():
 
     identity = stix2.v21.Identity(name="test", identity_class="individual")
     assert prop.clean(identity) is identity
+
+
+def test_bundle_obj_id_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    mal_list = bundle.get_obj("malware--00000000-0000-4000-8000-000000000003")
+    assert bundle.objects[1] == mal_list[0]
+    assert len(mal_list) == 1
+
+
+@pytest.mark.parametrize(
+    "bundle_data", [{
+        "type": "bundle",
+        "id": "bundle--00000000-0000-4000-8000-000000000007",
+        "objects": [
+            {
+                "type": "indicator",
+                "spec_version": "2.1",
+                "id": "indicator--00000000-0000-4000-8000-000000000001",
+                "created": "2017-01-01T12:34:56.000Z",
+                "modified": "2017-01-01T12:34:56.000Z",
+                "indicator_types": [
+                    "malicious-activity",
+                ],
+                "pattern": "[file:hashes.MD5 = 'd41d8cd98f00b204e9800998ecf8427e']",
+                "valid_from": "2017-01-01T12:34:56Z",
+            },
+            {
+                "type": "malware",
+                "spec_version": "2.1",
+                "id": "malware--00000000-0000-4000-8000-000000000003",
+                "created": "2017-01-01T12:34:56.000Z",
+                "modified": "2017-01-01T12:34:56.000Z",
+                "name": "Cryptolocker1",
+                "malware_types": [
+                    "ransomware",
+                ],
+            },
+            {
+                "type": "malware",
+                "spec_version": "2.1",
+                "id": "malware--00000000-0000-4000-8000-000000000003",
+                "created": "2017-01-01T12:34:56.000Z",
+                "modified": "2017-12-21T12:34:56.000Z",
+                "name": "CryptolockerOne",
+                "malware_types": [
+                    "ransomware",
+                ],
+            },
+            {
+                "type": "relationship",
+                "spec_version": "2.1",
+                "id": "relationship--00000000-0000-4000-8000-000000000005",
+                "created": "2017-01-01T12:34:56.000Z",
+                "modified": "2017-01-01T12:34:56.000Z",
+                "relationship_type": "indicates",
+                "source_ref": "indicator--a740531e-63ff-4e49-a9e1-a0a3eed0e3e7",
+                "target_ref": "malware--9c4638ec-f1de-4ddb-abf4-1b760417654e",
+            },
+        ],
+    }],
+)
+def test_bundle_objs_ids_found(bundle_data):
+    bundle = stix2.parse(bundle_data)
+
+    mal_list = bundle.get_obj("malware--00000000-0000-4000-8000-000000000003")
+    assert bundle.objects[1] == mal_list[0]
+    assert bundle.objects[2] == mal_list[1]
+    assert len(mal_list) == 2
+
+
+def test_bundle_getitem_overload_property_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    assert bundle.type == "bundle"
+    assert bundle['type'] == "bundle"
+
+
+def test_bundle_getitem_overload_obj_id_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    mal_list = bundle["malware--00000000-0000-4000-8000-000000000003"]
+    assert bundle.objects[1] == mal_list[0]
+    assert len(mal_list) == 1
+
+
+def test_bundle_obj_id_not_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    with pytest.raises(KeyError) as excinfo:
+        bundle.get_obj('non existent')
+    assert "does not match the id property of any of the bundle" in str(excinfo.value)
+
+
+def test_bundle_getitem_overload_obj_id_not_found():
+    bundle = stix2.parse(EXPECTED_BUNDLE)
+
+    with pytest.raises(KeyError) as excinfo:
+        bundle['non existent']
+    assert "neither a valid bundle property nor does it match the id property" in str(excinfo.value)

stix2/v20/bundle.py

@@ -35,3 +35,21 @@ 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):
+        if "objects" in self._inner:
+            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
+        else:
+            raise KeyError("There are no objects in this empty bundle")
+
+    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)

stix2/v21/bundle.py

@@ -33,3 +33,21 @@ 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):
+        if "objects" in self._inner:
+            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
+        else:
+            raise KeyError("There are no objects in this empty bundle")
+
+    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)