xref: /aosp_15_r20/external/tink/python/tink/testing/keyset_builder.py (revision e7b1675dde1b92d52ec075b0a92829627f2c52a5) 
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"""Python implementation of a KeysetBuilder."""
15
16import io
17import random
18from tink.proto import tink_pb2
19import tink
20from tink import cleartext_keyset_handle
21
22_MAX_INT32 = 4294967295  # 2^32-1
23
24
25def _new_key_data(key_template: tink_pb2.KeyTemplate) -> tink_pb2.KeyData:
26  return tink.core.Registry.new_key_data(key_template)
27
28
29def _generate_unused_key_id(keyset: tink_pb2.Keyset) -> int:
30  while True:
31    key_id = random.randint(1, _MAX_INT32)
32    if key_id not in {key.key_id for key in keyset.key}:
33      return key_id
34
35
36def raw_template(template: tink_pb2.KeyTemplate) -> tink_pb2.KeyTemplate:
37  template_copy = tink_pb2.KeyTemplate()
38  template_copy.CopyFrom(template)
39  template_copy.output_prefix_type = tink_pb2.RAW
40  return template_copy
41
42
43def legacy_template(template: tink_pb2.KeyTemplate) -> tink_pb2.KeyTemplate:
44  template_copy = tink_pb2.KeyTemplate()
45  template_copy.CopyFrom(template)
46  template_copy.output_prefix_type = tink_pb2.LEGACY
47  return template_copy
48
49
50class KeysetBuilder:
51  """A KeysetBuilder provides convenience methods for managing Keysets.
52
53  It provides methods for adding, disabling, enabling, or deleting keys.
54  The validity of the keyset is checked when creating a keyset_handle.
55  """
56
57  def __init__(self, keyset_proto: tink_pb2.Keyset):
58    self._keyset = keyset_proto
59
60  def keyset_handle(self) -> tink.KeysetHandle:
61    keyset_copy = tink_pb2.Keyset()
62    keyset_copy.CopyFrom(self._keyset)
63    return cleartext_keyset_handle.from_keyset(keyset_copy)
64
65  def keyset(self) -> bytes:
66    return self._keyset.SerializeToString()
67
68  def public_keyset(self) -> bytes:
69    public_handle = self.keyset_handle().public_keyset_handle()
70    public_keyset = io.BytesIO()
71    writer = tink.BinaryKeysetWriter(public_keyset)
72    cleartext_keyset_handle.write(writer, public_handle)
73    return public_keyset.getvalue()
74
75  def add_new_key(self, key_template: tink_pb2.KeyTemplate) -> int:
76    """Generates a new key, adds it to the keyset, and returns its ID."""
77    new_key = self._keyset.key.add()
78    new_key.key_data.CopyFrom(_new_key_data(key_template))
79    new_key.status = tink_pb2.ENABLED
80    new_key_id = _generate_unused_key_id(self._keyset)
81    new_key.key_id = new_key_id
82    new_key.output_prefix_type = key_template.output_prefix_type
83    return new_key_id
84
85  def set_primary_key(self, key_id: int) -> None:
86    """Sets a key as primary."""
87    for key in self._keyset.key:
88      if key.key_id == key_id:
89        self._keyset.primary_key_id = key_id
90        return
91    raise tink.TinkError('key not found: %d' % key_id)
92
93  def enable_key(self, key_id: int) -> None:
94    """Enables a key."""
95    for key in self._keyset.key:
96      if key.key_id == key_id:
97        key.status = tink_pb2.ENABLED
98        return
99    raise tink.TinkError('key not found: %d' % key_id)
100
101  def disable_key(self, key_id: int) -> None:
102    """Disables a key."""
103    for key in self._keyset.key:
104      if key.key_id == key_id:
105        key.status = tink_pb2.DISABLED
106        return
107    raise tink.TinkError('key not found: %d' % key_id)
108
109  def delete_key(self, key_id: int) -> None:
110    """Deletes a key."""
111    for key in self._keyset.key:
112      if key.key_id == key_id:
113        self._keyset.key.remove(key)
114        return
115    raise tink.TinkError('key not found: %d' % key_id)
116
117
118def from_keyset(keyset: bytes) -> KeysetBuilder:
119  """Return a KeysetBuilder for a Keyset copied from a KeysetHandle."""
120  keyset_proto = tink_pb2.Keyset.FromString(keyset)
121  return KeysetBuilder(keyset_proto)
122
123
124def from_keyset_handle(keyset_handle: tink.KeysetHandle) -> KeysetBuilder:
125  """Return a KeysetBuilder for a Keyset copied from a KeysetHandle."""
126  keyset_buffer = io.BytesIO()
127  cleartext_keyset_handle.write(
128      tink.BinaryKeysetWriter(keyset_buffer), keyset_handle)
129  return from_keyset(keyset_buffer.getvalue())
130
131
132def new_keyset_builder() -> KeysetBuilder:
133  return KeysetBuilder(tink_pb2.Keyset())
134