From a61344a8aafeaca7831a36f4ec78095d9af5f634 Mon Sep 17 00:00:00 2001 From: "Desai, Kartikey H" Date: Tue, 14 May 2019 13:48:54 -0400 Subject: [PATCH 1/4] 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) From 45d3020518bb97b67235d5e41a9ea5d032450b42 Mon Sep 17 00:00:00 2001 From: "Desai, Kartikey H" Date: Fri, 17 May 2019 14:21:35 -0500 Subject: [PATCH 2/4] 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) From 86790a736f1169eb28b82a734ffc3cf5645602ed Mon Sep 17 00:00:00 2001 From: "Desai, Kartikey H" Date: Mon, 20 May 2019 15:29:01 -0500 Subject: [PATCH 3/4] Fixes #257 --- stix2/test/v20/test_bundle.py | 39 ++++++++++++ stix2/test/v21/test_bundle.py | 117 +++++++++++++++++++++++++++++----- stix2/v20/bundle.py | 18 ++++++ stix2/v21/bundle.py | 11 ++-- 4 files changed, 164 insertions(+), 21 deletions(-) diff --git a/stix2/test/v20/test_bundle.py b/stix2/test/v20/test_bundle.py index 72523bb..992cde9 100644 --- a/stix2/test/v20/test_bundle.py +++ b/stix2/test/v20/test_bundle.py @@ -236,3 +236,42 @@ 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 + + +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) diff --git a/stix2/test/v21/test_bundle.py b/stix2/test/v21/test_bundle.py index 3a8e996..4f33fb5 100644 --- a/stix2/test/v21/test_bundle.py +++ b/stix2/test/v21/test_bundle.py @@ -183,23 +183,6 @@ 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['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(): @@ -224,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) diff --git a/stix2/v20/bundle.py b/stix2/v20/bundle.py index 76386ef..dd2bb2c 100644 --- a/stix2/v20/bundle.py +++ b/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) diff --git a/stix2/v21/bundle.py b/stix2/v21/bundle.py index 77b78c1..352cf3b 100644 --- a/stix2/v21/bundle.py +++ b/stix2/v21/bundle.py @@ -35,10 +35,13 @@ class Bundle(_STIXBase): super(Bundle, self).__init__(**kwargs) def get_obj(self, obj_uuid): - 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 + 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: From ce86db2a126fbb02c50b8146b8dfcce1ca561caf Mon Sep 17 00:00:00 2001 From: "Desai, Kartikey H" Date: Mon, 20 May 2019 15:36:35 -0500 Subject: [PATCH 4/4] Fixes #257 --- stix2/test/v20/test_bundle.py | 58 +++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/stix2/test/v20/test_bundle.py b/stix2/test/v20/test_bundle.py index 992cde9..806421a 100644 --- a/stix2/test/v20/test_bundle.py +++ b/stix2/test/v20/test_bundle.py @@ -246,6 +246,64 @@ def test_bundle_obj_id_found(): 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)