1# Copyright 2019 Google LLC 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""Tink Registry.""" 16 17from typing import Any, Tuple, Type, TypeVar 18 19from tink.proto import tink_pb2 20from tink.core import _key_manager 21from tink.core import _primitive_set 22from tink.core import _primitive_wrapper 23from tink.core import _tink_error 24 25P = TypeVar('P') 26 27 28class Registry: 29 """A global container of key managers. 30 31 Registry maps supported key types to a corresponding KeyManager object, 32 which 'understands' the key type (i.e., the KeyManager can instantiate the 33 primitive corresponding to given key, or can generate new keys of the 34 supported key type). Keeping KeyManagers for all primitives in a single 35 Registry (rather than having a separate KeyManager per primitive) enables 36 modular construction of compound primitives from 'simple' ones, 37 e.g., AES-CTR-HMAC AEAD encryption uses IND-CPA encryption and a MAC. 38 39 Registry is initialized at startup, and is later used to instantiate 40 primitives for given keys or keysets. 41 """ 42 43 _key_managers = {} # type: dict[str, Tuple[_key_manager.KeyManager, bool]] 44 _wrappers = {} # type: dict[Type, _primitive_wrapper.PrimitiveWrapper] 45 46 @classmethod 47 def reset(cls) -> None: 48 """Resets the registry.""" 49 cls._key_managers = {} 50 cls._wrappers = {} 51 52 @classmethod 53 def _key_manager_internal( 54 cls, type_url: str) -> Tuple[_key_manager.KeyManager, bool]: 55 """Returns a key manager, new_key_allowed pair for the given type_url.""" 56 if type_url not in cls._key_managers: 57 raise _tink_error.TinkError( 58 'No manager for type {} has been registered.'.format(type_url)) 59 return cls._key_managers[type_url] 60 61 @classmethod 62 def key_manager(cls, type_url: str) -> _key_manager.KeyManager: 63 """Returns a key manager for the given type_url and primitive_class. 64 65 Args: 66 type_url: Key type string 67 68 Returns: 69 A KeyManager object 70 """ 71 key_mgr, _ = cls._key_manager_internal(type_url) 72 return key_mgr 73 74 @classmethod 75 def register_key_manager(cls, 76 key_manager: _key_manager.KeyManager, 77 new_key_allowed: bool = True) -> None: 78 """Tries to register a key_manager for the given key_manager.key_type(). 79 80 Args: 81 key_manager: A KeyManager object 82 new_key_allowed: If new_key_allowed is true, users can generate new keys 83 with this manager using Registry.new_key() 84 """ 85 key_managers = cls._key_managers 86 type_url = key_manager.key_type() 87 primitive_class = key_manager.primitive_class() 88 89 if not key_manager.does_support(type_url): 90 raise _tink_error.TinkError( 91 'The manager does not support its own type {}.'.format(type_url)) 92 93 if type_url in key_managers: 94 existing, existing_new_key = key_managers[type_url] 95 if (type(existing) != type(key_manager) or # pylint: disable=unidiomatic-typecheck 96 existing.primitive_class() != primitive_class): 97 raise _tink_error.TinkError( 98 'A manager for type {} has been already registered.'.format( 99 type_url)) 100 else: 101 if not existing_new_key and new_key_allowed: 102 raise _tink_error.TinkError( 103 ('A manager for type {} has been already registered ' 104 'with forbidden new key operation.').format(type_url)) 105 key_managers[type_url] = (existing, new_key_allowed) 106 else: 107 key_managers[type_url] = (key_manager, new_key_allowed) 108 109 @classmethod 110 def primitive(cls, key_data: tink_pb2.KeyData, primitive_class: Type[P]) -> P: 111 """Creates a new primitive for the key given in key_data. 112 113 It looks up a KeyManager identified by key_data.type_url, 114 and calls manager's primitive(key_data) method. 115 116 Args: 117 key_data: KeyData object 118 primitive_class: The expected primitive class 119 120 Returns: 121 A primitive for the given key_data 122 Raises: 123 Error if primitive_class does not match the registered primitive class. 124 """ 125 key_mgr = cls.key_manager(key_data.type_url) 126 if key_mgr.primitive_class() != primitive_class: 127 raise _tink_error.TinkError( 128 'Wrong primitive class: type {} uses primitive {}, and not {}.' 129 .format(key_data.type_url, key_mgr.primitive_class().__name__, 130 primitive_class.__name__)) 131 return key_mgr.primitive(key_data) 132 133 @classmethod 134 def new_key_data(cls, key_template: tink_pb2.KeyTemplate) -> tink_pb2.KeyData: 135 """Generates a new key for the specified key_template.""" 136 key_mgr, new_key_allowed = cls._key_manager_internal( 137 key_template.type_url) 138 139 if not new_key_allowed: 140 raise _tink_error.TinkError( 141 'KeyManager for type {} does not allow for creation of new keys.' 142 .format(key_template.type_url)) 143 144 return key_mgr.new_key_data(key_template) 145 146 @classmethod 147 def public_key_data(cls, 148 private_key_data: tink_pb2.KeyData) -> tink_pb2.KeyData: 149 """Generates a new key for the specified key_template.""" 150 if (private_key_data.key_material_type != 151 tink_pb2.KeyData.ASYMMETRIC_PRIVATE): 152 raise _tink_error.TinkError('The keyset contains a non-private key') 153 key_mgr = cls.key_manager(private_key_data.type_url) 154 if not isinstance(key_mgr, _key_manager.PrivateKeyManager): 155 raise _tink_error.TinkError( 156 'manager for key type {} is not a PrivateKeyManager' 157 .format(private_key_data.type_url)) 158 return key_mgr.public_key_data(private_key_data) 159 160 @classmethod 161 def register_primitive_wrapper( 162 cls, wrapper: _primitive_wrapper.PrimitiveWrapper) -> None: 163 """Tries to register a PrimitiveWrapper. 164 165 Args: 166 wrapper: A PrimitiveWrapper object. 167 Raises: 168 TinkError if a different wrapper has already been registered for the same 169 Primitive. 170 """ 171 if (wrapper.primitive_class() in cls._wrappers and 172 type(cls._wrappers[wrapper.primitive_class()]) != type(wrapper)): # pylint: disable=unidiomatic-typecheck 173 raise _tink_error.TinkError( 174 'A wrapper for primitive {} has already been added.'.format( 175 wrapper.primitive_class().__name__)) 176 wrapped = wrapper.wrap( 177 _primitive_set.PrimitiveSet(wrapper.input_primitive_class())) 178 if not isinstance(wrapped, wrapper.primitive_class()): 179 raise _tink_error.TinkError( 180 'Wrapper for primitive {} generates incompatible primitive of type {}' 181 .format(wrapper.primitive_class().__name__, 182 type(wrapped).__name__)) 183 cls._wrappers[wrapper.primitive_class()] = wrapper 184 185 @classmethod 186 def input_primitive_class(cls, primitive_class: Any) -> Any: 187 """Returns the primitive class that gets wrapped into primitive_class. 188 189 Args: 190 primitive_class: Class of the output primitive of a wrapper. 191 Returns: 192 the primitive class that gets wrapped. This needs to be the type used 193 in the primitive set. 194 Raises: 195 TinkError if no wrapper for this primitive class is registered. 196 """ 197 if primitive_class not in cls._wrappers: 198 raise _tink_error.TinkError( 199 'No PrimitiveWrapper registered for primitive {}.' 200 .format(primitive_class.__name__)) 201 wrapper = cls._wrappers[primitive_class] 202 return wrapper.input_primitive_class() 203 204 @classmethod 205 def wrap(cls, 206 primitive_set: _primitive_set.PrimitiveSet, 207 primitive_class: Type[P]) -> P: 208 """Wraps a set of primitives into a single primitive. 209 210 Args: 211 primitive_set: A PrimitiveSet object. 212 primitive_class: Class of the output primitive. 213 Returns: 214 A primitive of type primitive_class that wraps the primitives in 215 primitive_set. 216 Raises: 217 TinkError if no wrapper for this primitive class is registered or the type 218 of the primitives in primitive_set are don't match the 219 input_primitive_class of the wrapper. 220 """ 221 if primitive_class not in cls._wrappers: 222 raise _tink_error.TinkError( 223 'No PrimitiveWrapper registered for primitive {}.' 224 .format(primitive_class.__name__)) 225 wrapper = cls._wrappers[primitive_class] 226 if primitive_set.primitive_class() != wrapper.input_primitive_class(): 227 raise _tink_error.TinkError( 228 'Wrapper for primitive {} wraps type {}, but the primitive_set' 229 'has type {}' 230 .format(wrapper.primitive_class().__name__, 231 wrapper.input_primitive_class().__name__, 232 primitive_set.primitive_class().__name__)) 233 return wrapper.wrap(primitive_set) 234