xref: /aosp_15_r20/external/tink/python/examples/walkthrough/load_encrypted_keyset_test.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"""Test for load_encrypted_keyset."""
15from absl.testing import absltest
16
17import tink
18
19from tink import aead
20from tink import cleartext_keyset_handle
21
22import load_encrypted_keyset
23from tink.testing import fake_kms
24
25_FAKE_KMS_AEAD_KEYSET = r"""{
26  "key": [
27    {
28      "keyData": {
29        "keyMaterialType": "SYMMETRIC",
30        "typeUrl": "type.googleapis.com/google.crypto.tink.AesGcmKey",
31        "value": "GiBWyUfGgYk3RTRhj/LIUzSudIWlyjCftCOypTr0jCNSLg=="
32      },
33      "keyId": 294406504,
34      "outputPrefixType": "TINK",
35      "status": "ENABLED"
36    }
37  ],
38  "primaryKeyId": 294406504
39}"""
40
41_KEYSET_TO_ENCRYPT = r"""{
42  "key": [
43    {
44      "keyData": {
45        "keyMaterialType": "SYMMETRIC",
46        "typeUrl": "type.googleapis.com/google.crypto.tink.AesGcmKey",
47        "value": "GhD+9l0RANZjzZEZ8PDp7LRW"
48      },
49      "keyId": 1931667682,
50      "outputPrefixType": "TINK",
51      "status": "ENABLED"
52    }
53  ],
54  "primaryKeyId": 1931667682
55}"""
56
57# Encryption of _KEYSET_TO_ENCRYPT with _FAKE_KMS_AEAD_KEYSET with no
58# associated data.
59_ENCRYPTED_KEYSET = r"""{
60  "encryptedKeyset": "ARGMSWi6YHyZ/Oqxl00XSq631a0q2UPmf+rCvCIAggSZrwCmxFF797MpY0dqgaXu1fz2eQ8zFNhlyTXv9kwg1kY6COpyhY/68zNBUkyKX4CharLYfpg1LgRl+6rMzIQa0XDHh7ZDmp1CevzecZIKnG83uDRHxxSv3h8c/Kc="
61}"""
62
63# Fake KMS keys are base64-encoded keysets. This was generated from
64# _FAKE_KMS_AEAD_KEYSET by first serializing it to bytes using a
65# tink.BinaryKeysetWriter, and then encoding it as base64.
66_FAKE_KMS_KEY_URI = (
67    'fake-kms://COiSsYwBEmQKWAowdHlwZS5nb29nbGVhcGlzLmNvbS9nb29nbGUuY3J5cHRvLnR'
68    'pbmsuQWVzR2NtS2V5EiIaIFbJR8aBiTdFNGGP8shTNK50haXKMJ-0I7KlOvSMI1IuGAEQARjok'
69    'rGMASAB')
70
71
72class LoadEncryptedKeysetTest(absltest.TestCase):
73
74  def setUp(self):
75    super().setUp()
76    aead.register()
77    fake_kms.register_client()
78
79  def test_load_encrypted_keyset_fails_if_kms_key_is_invalid(self):
80    with self.assertRaises(tink.TinkError):
81      load_encrypted_keyset.LoadEncryptedKeyset(
82          _ENCRYPTED_KEYSET,
83          kms_key_uri='fake-kms://invalid-kms-key',
84          associated_data=b'')
85
86  def test_load_encrypted_keyset_fails_if_keyset_is_invalid(self):
87    with self.assertRaises(tink.TinkError):
88      load_encrypted_keyset.LoadEncryptedKeyset(
89          'Invalid keyset', _FAKE_KMS_KEY_URI, associated_data=b'')
90
91  def test_load_encrypted_keyset_returns_a_valid_keyset(self):
92    keyset_handle = load_encrypted_keyset.LoadEncryptedKeyset(
93        _ENCRYPTED_KEYSET, _FAKE_KMS_KEY_URI, associated_data=b'')
94
95    # Make sure that we can use this primitive.
96    aead_primitive = keyset_handle.primitive(aead.Aead)
97    plaintext = b'Some plaintext'
98    associated_data = b'Some associated data'
99    ciphertext = aead_primitive.encrypt(plaintext, associated_data)
100    self.assertEqual(
101        aead_primitive.decrypt(ciphertext, associated_data), plaintext)
102
103    # Make sure we can use the loaded keyset to decrypt a ciphertext encrypted
104    # with _KEYSET_TO_ENCRYPT.
105    expected_keyset_handle = cleartext_keyset_handle.read(
106        tink.JsonKeysetReader(_KEYSET_TO_ENCRYPT))
107    expected_aead = expected_keyset_handle.primitive(aead.Aead)
108    self.assertEqual(
109        aead_primitive.decrypt(
110            expected_aead.encrypt(plaintext, associated_data), associated_data),
111        plaintext)
112    # Make sure we can use _KEYSET_TO_ENCRYPT to decrypt the ciphertext produced
113    # by the keyset we loaded.
114    self.assertEqual(
115        expected_aead.decrypt(
116            aead_primitive.encrypt(plaintext, associated_data),
117            associated_data), plaintext)
118
119
120if __name__ == '__main__':
121  absltest.main()
122