xref: /aosp_15_r20/external/tink/python/examples/walkthrough/write_keyset.py (revision e7b1675dde1b92d52ec075b0a92829627f2c52a5)
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"""Example to showcase how to create a keyset."""
15# [START tink_walkthrough_write_keyset]
16from typing import TextIO
17
18import tink
19from tink import aead
20
21
22def GetKmsAead(kms_kek_uri: str) -> aead.Aead:
23  """Returns an AEAD primitive from a KMS Key Encryption Key URI."""
24  # To obtain a primitive that uses the KMS to encrypt/decrypt we simply create
25  # keyset from the appropriate template and get an AEAD primitive from it.
26  template = aead.aead_key_templates.create_kms_aead_key_template(kms_kek_uri)
27  kms_aead_keyset_handle = tink.new_keyset_handle(template)
28  return kms_aead_keyset_handle.primitive(aead.Aead)
29
30
31def WriteEncryptedKeyset(keyset_handle: tink.KeysetHandle,
32                         text_io_stream: TextIO,
33                         kms_kek_uri: str,
34                         associated_data: bytes = b'') -> None:
35  """Encrypts keyset_hanlde with a KMS and writes it to text_io_stream as JSON.
36
37  The keyset is encrypted with a KMS using the KMS key kms_kek_uri.
38
39  Prerequisites:
40    - Register AEAD implementations of Tink.
41    - Register a KMS client that can use kms_kek_uri.
42    - Create a keyset and obtain a handle to it.
43
44  Args:
45    keyset_handle: Keyset to write.
46    text_io_stream: I/O stream where writng the Keyset to.
47    kms_kek_uri: URI of the KMS key to use to encrypt the keyset.
48    associated_data: Associated data to which tie the ciphertext.
49
50  Raises:
51    tink.TinkError in case of errors.
52  """
53  keyset_handle.write_with_associated_data(
54      tink.JsonKeysetWriter(text_io_stream), GetKmsAead(kms_kek_uri),
55      associated_data)
56
57
58# [END tink_walkthrough_write_keyset]
59