• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2022 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"""A container to store precomputed keys."""
15
16import textwrap
17from typing import Dict
18
19from tink.proto import tink_pb2
20from util import key_util
21
22
23class TestKeysContainer():
24  """Container for test keys."""
25
26  _map: Dict[str, tink_pb2.Keyset.Key]
27
28  def __init__(self):
29    self._map = {}
30
31  def add_key(self, template: str, key: str) -> None:
32    """Adds a new key to the list of precomputed keys.
33
34    The arguments need to be in the format produced by key_util.text_format,
35    but can be additionally indented and have whitespace (it needs to be in the
36    format after calling textwrap.dedent() and strip()).
37
38    The key will be stored in a map keyed by the template, unless a key-value
39    pair with this template as key was previously inserted in a call to
40    'add_key'.
41
42    Args:
43      template: A key template in the format created by key_util.text_format,
44        possibly indented, and with additional spaces at the beginning and the
45        end.
46      key: A key corresponding to the template in the format created by
47        key_util.text_format, possibly indented, and with additional spaces.
48    """
49
50    dedented_template = textwrap.dedent(template).strip()
51    dedented_key = textwrap.dedent(key).strip()
52    parsed_template = tink_pb2.KeyTemplate()
53    # We parse to check the correctness of the formatting
54    key_util.parse_text_format(dedented_template, parsed_template)
55
56    parsed_key = tink_pb2.Keyset.Key()
57    key_util.parse_text_format(dedented_key, parsed_key)
58    if dedented_template in self._map:
59      raise ValueError('Template already present')
60    self._map[dedented_template] = parsed_key
61
62  def get_key(self, template: tink_pb2.KeyTemplate) -> tink_pb2.Keyset.Key:
63    """Returns a previously stored key for this template."""
64
65    template_text_format = key_util.text_format(template)
66    return self._map[template_text_format]
67