| # Copyright 2019 Google LLC. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| """A container class for a set of primitives.""" |
| |
| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import google_type_annotations |
| from __future__ import print_function |
| |
| import collections |
| from typing import Generic, List, Type, TypeVar |
| |
| from tink.proto import tink_pb2 |
| from tink.python.core import crypto_format |
| from tink.python.core import tink_error |
| |
| P = TypeVar('P') |
| Entry = collections.namedtuple( |
| 'Entry', 'primitive, identifier, status, output_prefix_type') |
| |
| |
| def new_primitive_set(primitive_class): |
| return PrimitiveSet(primitive_class) |
| |
| |
| class PrimitiveSet(Generic[P]): |
| """A container class for a set of primitives. |
| |
| PrimitiveSet is an auxiliary class used for supporting key rotation: |
| primitives in a set correspond to keys in a keyset. Users will usually work |
| with primitive instances, which essentially wrap primitive sets. For example |
| an instance of an Aead-primitive for a given keyset holds a set of |
| Aead-primitives corresponding to the keys in the keyset, and uses the set |
| members to do the actual crypto operations: to encrypt data the primary |
| Aead-primitive from the set is used, and upon decryption the ciphertext's |
| prefix determines the id of the primitive from the set. |
| """ |
| |
| def __init__(self, primitive_class: Type[P]): |
| # map from identifier to a list of Entry |
| self._primitives = {} |
| self._primary = None |
| self._primitive_class = primitive_class |
| |
| def primitive_class(self) -> Type[P]: |
| return self._primitive_class |
| |
| def primitive_from_identifier(self, identifier: bytes) -> List[P]: |
| """Returns a copy of the list of entries for a given identifier.""" |
| # Copy the list so that if the user modifies the list, it does not affect |
| # the internal data structure. |
| return self._primitives.get(identifier, [])[:] |
| |
| def primitive(self, key: tink_pb2.Keyset.Key) -> List[P]: |
| """Returns a copy of the list of primitives for a given identifier.""" |
| return self.primitive_from_identifier(crypto_format.output_prefix(key)) |
| |
| def raw_primitives(self) -> List[P]: |
| """Returns a copy of the list of primitives for a given identifier.""" |
| # All raw keys have the same identifier, which is just b''. |
| return self.primitive_from_identifier(crypto_format.RAW_PREFIX) |
| |
| def add_primitive(self, primitive: P, key: tink_pb2.Keyset.Key) -> Entry: |
| """Adds a new primitive and key entry to the set, and returns the entry.""" |
| if not isinstance(primitive, self._primitive_class): |
| raise tink_error.TinkError( |
| 'The primitive is not an instance of {}'.format( |
| self._primitive_class)) |
| identifier = crypto_format.output_prefix(key) |
| |
| entry = Entry(primitive, identifier, key.status, key.output_prefix_type) |
| entries = self._primitives.setdefault(identifier, []) |
| entries.append(entry) |
| return entry |
| |
| def set_primary(self, entry: Entry) -> None: |
| self._primary = entry |
| |
| def primary(self) -> Entry: |
| return self._primary |