diff --git a/stix2/registration.py b/stix2/registration.py index bc07995..e3aa5e5 100644 --- a/stix2/registration.py +++ b/stix2/registration.py @@ -68,13 +68,11 @@ def _register_marking(new_marking, version=DEFAULT_VERSION): if not re.match(PREFIX_21_REGEX, prop_name): raise ValueError("Property name '%s' must begin with an alpha character." % prop_name) - if version: - v = 'v' + version.replace('.', '') - else: - # Use default version (latest) if no version was provided. - v = 'v' + DEFAULT_VERSION.replace('.', '') + class_maps = registry.get_stix2_class_maps( + version or DEFAULT_VERSION + ) - OBJ_MAP_MARKING = registry.STIX2_OBJ_MAPS[v]['markings'] + OBJ_MAP_MARKING = class_maps['markings'] if mark_type in OBJ_MAP_MARKING.keys(): raise DuplicateRegistrationError("STIX Marking", mark_type) OBJ_MAP_MARKING[mark_type] = new_marking @@ -130,13 +128,11 @@ def _register_observable(new_observable, version=DEFAULT_VERSION): "is not a ListProperty containing ReferenceProperty." % prop_name, ) - if version: - v = 'v' + version.replace('.', '') - else: - # Use default version (latest) if no version was provided. - v = 'v' + DEFAULT_VERSION.replace('.', '') + class_maps = registry.get_stix2_class_maps( + version or DEFAULT_VERSION + ) - OBJ_MAP_OBSERVABLE = registry.STIX2_OBJ_MAPS[v]['observables'] + OBJ_MAP_OBSERVABLE = class_maps['observables'] if new_observable._type in OBJ_MAP_OBSERVABLE.keys(): raise DuplicateRegistrationError("Cyber Observable", new_observable._type) OBJ_MAP_OBSERVABLE[new_observable._type] = new_observable @@ -182,8 +178,6 @@ def _register_observable_extension( if not re.match(PREFIX_21_REGEX, prop_name): raise ValueError("Property name '%s' must begin with an alpha character." % prop_name) - v = 'v' + version.replace('.', '') - try: observable_type = observable._type except AttributeError: @@ -192,8 +186,9 @@ def _register_observable_extension( "created with the @CustomObservable decorator.", ) - OBJ_MAP_OBSERVABLE = registry.STIX2_OBJ_MAPS[v]['observables'] - EXT_MAP = registry.STIX2_OBJ_MAPS[v]['observable-extensions'] + class_maps = registry.get_stix2_class_maps(version) + OBJ_MAP_OBSERVABLE = class_maps['observables'] + EXT_MAP = class_maps['observable-extensions'] try: if ext_type in EXT_MAP[observable_type].keys(): diff --git a/stix2/registry.py b/stix2/registry.py index 6cb6cd8..692fe3e 100644 --- a/stix2/registry.py +++ b/stix2/registry.py @@ -26,3 +26,18 @@ def _collect_stix2_mappings(): elif re.match(r'^stix2\.v2[0-9]\.common$', name) and is_pkg is False: mod = importlib.import_module(name, str(top_level_module.__name__)) STIX2_OBJ_MAPS[ver]['markings'] = mod.OBJ_MAP_MARKING + + +def get_stix2_class_maps(stix_version): + """ + Get the stix2 class mappings for the given STIX version. + + :param stix_version: A STIX version as a string + :return: The class mappings. This will be a dict mapping from some general + category name, e.g. "object" to another mapping from STIX type + to a stix2 class. + """ + stix_vid = "v" + stix_version.replace(".", "") + cls_maps = STIX2_OBJ_MAPS[stix_vid] + + return cls_maps diff --git a/stix2/utils.py b/stix2/utils.py index 94921a7..76fbea0 100644 --- a/stix2/utils.py +++ b/stix2/utils.py @@ -340,21 +340,6 @@ def _stix_type_of(value): return type_ -def _get_stix2_class_maps(stix_version): - """ - Get the stix2 class mappings for the given STIX version. - - :param stix_version: A STIX version as a string - :return: The class mappings. This will be a dict mapping from some general - category name, e.g. "object" to another mapping from STIX type - to a stix2 class. - """ - stix_vid = "v" + stix_version.replace(".", "") - cls_maps = mappings.STIX2_OBJ_MAPS[stix_vid] - - return cls_maps - - def is_sdo(value, stix_version=stix2.DEFAULT_VERSION): """ Determine whether the given object, type, or ID is/is for an SDO. @@ -369,7 +354,7 @@ def is_sdo(value, stix_version=stix2.DEFAULT_VERSION): # Eventually this needs to be moved into the stix2 library (and maybe # improved?); see cti-python-stix2 github issue #450. - cls_maps = _get_stix2_class_maps(stix_version) + cls_maps = mappings.get_stix2_class_maps(stix_version) type_ = _stix_type_of(value) result = type_ in cls_maps["objects"] and type_ not in { "relationship", "sighting", "marking-definition", "bundle", @@ -389,7 +374,7 @@ def is_sco(value, stix_version=stix2.DEFAULT_VERSION): :return: True if the type of the given value is an SCO type; False if not """ - cls_maps = _get_stix2_class_maps(stix_version) + cls_maps = mappings.get_stix2_class_maps(stix_version) type_ = _stix_type_of(value) result = type_ in cls_maps["observables"] @@ -425,7 +410,7 @@ def is_object(value, stix_version=stix2.DEFAULT_VERSION): :return: True if the type of the given value is a valid STIX type with respect to the given STIX version; False if not """ - cls_maps = _get_stix2_class_maps(stix_version) + cls_maps = mappings.get_stix2_class_maps(stix_version) type_ = _stix_type_of(value) result = type_ in cls_maps["observables"] or type_ in cls_maps["objects"]