Change approach to allow for custom objects

master
Emmanuelle Vargas-Gonzalez 2018-11-14 14:35:22 -05:00
parent de73935d2a
commit c80f39ceed
1 changed files with 5 additions and 5 deletions

View File

@ -44,14 +44,14 @@ def _add(store, stix_data=None, allow_custom=True, version=None):
# Map ID directly to the object, if it is a marking. Otherwise, # Map ID directly to the object, if it is a marking. Otherwise,
# map to a family, so we can track multiple versions. # map to a family, so we can track multiple versions.
if _is_marking(stix_obj): if _is_marking(stix_obj):
store._data[stix_obj.id] = stix_obj store._data[stix_obj["id"]] = stix_obj
else: else:
if stix_obj.id in store._data: if stix_obj.id in store._data:
obj_family = store._data[stix_obj.id] obj_family = store._data[stix_obj["id"]]
else: else:
obj_family = _ObjectFamily() obj_family = _ObjectFamily()
store._data[stix_obj.id] = obj_family store._data[stix_obj["id"]] = obj_family
obj_family.add(stix_obj) obj_family.add(stix_obj)
@ -64,8 +64,8 @@ def _is_marking(obj_or_id):
:return: True if a marking definition, False otherwise. :return: True if a marking definition, False otherwise.
""" """
if isinstance(obj_or_id, _STIXBase): if isinstance(obj_or_id, (_STIXBase, dict)):
id_ = obj_or_id.id id_ = obj_or_id["id"]
else: else:
id_ = obj_or_id id_ = obj_or_id