Merge pull request #264 from khdesai/stix2_issue_257
Stix2 issue 257: make accessing bundles easiermaster
commit
51b2db4fba
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue