1# Copyright 2020 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"""PrfSet wrapper.""" 15 16from typing import Type, Mapping 17from tink import core 18from tink.prf import _prf_set 19 20 21class _WrappedPrfSet(_prf_set.PrfSet): 22 """Implements PrfSet for a set of PrfSet primitives.""" 23 24 def __init__(self, primitive_set: core.PrimitiveSet): 25 self._primitive_set = primitive_set 26 27 def primary_id(self) -> int: 28 return self._primitive_set.primary().key_id 29 30 def all(self) -> Mapping[int, _prf_set.Prf]: 31 return { 32 entry.key_id: entry.primitive 33 for entry in self._primitive_set.raw_primitives() 34 } 35 36 def primary(self) -> _prf_set.Prf: 37 return self._primitive_set.primary().primitive 38 39 40class PrfSetWrapper(core.PrimitiveWrapper[_prf_set.Prf, _prf_set.PrfSet]): 41 """A PrimitiveWrapper for the PrfSet primitive. 42 43 The returned primitive works with a keyset (rather than a single key). To sign 44 a message, it uses the primary key in the keyset, and prepends to the 45 signature a certain prefix associated with the primary key. 46 """ 47 48 def wrap(self, primitives_set: core.PrimitiveSet) -> _WrappedPrfSet: 49 return _WrappedPrfSet(primitives_set) 50 51 def primitive_class(self) -> Type[_prf_set.PrfSet]: 52 return _prf_set.PrfSet 53 54 def input_primitive_class(self) -> Type[_prf_set.Prf]: 55 return _prf_set.Prf 56