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"""A container class for a set of primitives.""" 16 17import collections 18from typing import Generic, List, Type, TypeVar 19 20from tink.proto import tink_pb2 21from tink.core import _crypto_format 22from tink.core import _tink_error 23 24P = TypeVar('P') 25Entry = collections.namedtuple( 26 'Entry', 'primitive, identifier, status, output_prefix_type, key_id') 27 28 29def new_primitive_set(primitive_class): 30 return PrimitiveSet(primitive_class) 31 32 33class PrimitiveSet(Generic[P]): 34 """A container class for a set of primitives. 35 36 PrimitiveSet is an auxiliary class used for supporting key rotation: 37 primitives in a set correspond to keys in a keyset. Users will usually work 38 with primitive instances, which essentially wrap primitive sets. For example 39 an instance of an Aead-primitive for a given keyset holds a set of 40 Aead-primitives corresponding to the keys in the keyset, and uses the set 41 members to do the actual crypto operations: to encrypt data the primary 42 Aead-primitive from the set is used, and upon decryption the ciphertext's 43 prefix determines the id of the primitive from the set. 44 """ 45 46 def __init__(self, primitive_class: Type[P]): 47 self._primitives = {} # Dict[bytes, List[Entry]] 48 self._primary = None 49 self._primitive_class = primitive_class 50 51 def primitive_class(self) -> Type[P]: 52 return self._primitive_class 53 54 def primitive_from_identifier(self, identifier: bytes) -> List[Entry]: 55 """Returns a copy of the list of entries for a given identifier.""" 56 # Copy the list so that if the user modifies the list, it does not affect 57 # the internal data structure. 58 return self._primitives.get(identifier, [])[:] 59 60 def all(self) -> List[List[Entry]]: 61 """Returns a list of copies of all lists of entries in the primitive set.""" 62 return list(entries[:] for entries in self._primitives.values()) 63 64 def primitive(self, key: tink_pb2.Keyset.Key) -> List[Entry]: 65 """Returns a copy of the list of entries for a given key.""" 66 return self.primitive_from_identifier(_crypto_format.output_prefix(key)) 67 68 def raw_primitives(self) -> List[Entry]: 69 """Returns a copy of the list of entries of keys with raw prefix.""" 70 # All raw keys have the same identifier, which is just b''. 71 return self.primitive_from_identifier(_crypto_format.RAW_PREFIX) 72 73 def add_primitive(self, primitive: P, key: tink_pb2.Keyset.Key) -> Entry: 74 """Adds a new primitive and key entry to the set, and returns the entry.""" 75 if not isinstance(primitive, self._primitive_class): 76 raise _tink_error.TinkError( 77 'The primitive is not an instance of {}'.format( 78 self._primitive_class)) 79 identifier = _crypto_format.output_prefix(key) 80 81 entry = Entry(primitive, identifier, key.status, key.output_prefix_type, 82 key.key_id) 83 entries = self._primitives.setdefault(identifier, []) 84 entries.append(entry) 85 return entry 86 87 def set_primary(self, entry: Entry) -> None: 88 self._primary = entry 89 90 def primary(self) -> Entry: 91 return self._primary 92