Merge pull request #264 from khdesai/stix2_issue_257

Stix2 issue 257: make accessing bundles easier
master
Chris Lenk 2019-05-22 09:59:42 -04:00 committed by GitHub
commit 51b2db4fba
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 233 additions and 0 deletions

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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)