Add Python Streaming AEAD key templates.
PiperOrigin-RevId: 267139847
diff --git a/python/streaming_aead/streaming_aead_key_templates.py b/python/streaming_aead/streaming_aead_key_templates.py
new file mode 100644
index 0000000..1576ab9
--- /dev/null
+++ b/python/streaming_aead/streaming_aead_key_templates.py
@@ -0,0 +1,105 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Pre-generated KeyTemplates for StreamingAead.
+
+
+Currently, these templates cannot be used to generate keysets, but they can be
+used to generate individual keys.
+In the future, it will be possible to use these templates to generate a new
+tink_pb2.Keyset with tink_pb2.KeysetHandle. To generate a new keyset that
+contains a single aes_ctr_hmac_streaming_pb2.AesCtrHmacStreamingKey, one can do:
+handle = keyset_handle.KeysetHandle(
+ streaming_aead_key_templates.AES256_CTR_HMAC_SHA256_4KB).
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import google_type_annotations
+from __future__ import print_function
+
+from tink.proto import aes_ctr_hmac_streaming_pb2
+from tink.proto import aes_gcm_hkdf_streaming_pb2
+from tink.proto import common_pb2
+from tink.proto import tink_pb2
+
+_AES_GCM_HKDF_STREAMING_KEY_TYPE_URL = (
+ 'type.googleapis.com/google.crypto.tink.AesGcmHkdfStreamingKey')
+_AES_CTR_HMAC_STREAMING_KEY_TYPE_URL = (
+ 'type.googleapis.com/google.crypto.tink.AesCtrHmacStreamingKey')
+
+
+def create_aes_gcm_hkdf_streaming_key_template(
+ aes_key_size: int, hash_type: common_pb2.HashType, derived_key_size: int,
+ ciphertext_segment_size: int) -> tink_pb2.KeyTemplate:
+ """Creates an AES GCM HKDF Streaming KeyTemplate, and fills in its values."""
+ key_format = aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingKeyFormat()
+ key_format.key_size = aes_key_size
+ key_format.params.hkdf_hash_type = hash_type
+ key_format.params.derived_key_size = derived_key_size
+ key_format.params.ciphertext_segment_size = ciphertext_segment_size
+
+ key_template = tink_pb2.KeyTemplate()
+ key_template.value = key_format.SerializeToString()
+ key_template.type_url = _AES_GCM_HKDF_STREAMING_KEY_TYPE_URL
+ key_template.output_prefix_type = tink_pb2.RAW
+ return key_template
+
+
+def create_aes_ctr_hmac_streaming_key_template(
+ aes_key_size: int, hkdf_hash_type: common_pb2.HashType,
+ derived_key_size: int, mac_hash_type: common_pb2.HashType, tag_size: int,
+ ciphertext_segment_size: int) -> tink_pb2.KeyTemplate:
+ """Creates an AES CTR HMAC Streaming KeyTemplate, and fills in its values."""
+ key_format = aes_ctr_hmac_streaming_pb2.AesCtrHmacStreamingKeyFormat()
+ key_format.key_size = aes_key_size
+
+ key_format.params.ciphertext_segment_size = ciphertext_segment_size
+ key_format.params.derived_key_size = derived_key_size
+ key_format.params.hkdf_hash_type = hkdf_hash_type
+
+ key_format.params.hmac_params.hash = mac_hash_type
+ key_format.params.hmac_params.tag_size = tag_size
+
+ key_template = tink_pb2.KeyTemplate()
+ key_template.value = key_format.SerializeToString()
+ key_template.type_url = _AES_CTR_HMAC_STREAMING_KEY_TYPE_URL
+ key_template.output_prefix_type = tink_pb2.RAW
+ return key_template
+
+
+AES128_GCM_HKDF_4KB = create_aes_gcm_hkdf_streaming_key_template(
+ aes_key_size=16,
+ hash_type=common_pb2.HashType.SHA256,
+ derived_key_size=16,
+ ciphertext_segment_size=4096)
+
+AES256_GCM_HKDF_4KB = create_aes_gcm_hkdf_streaming_key_template(
+ aes_key_size=32,
+ hash_type=common_pb2.HashType.SHA256,
+ derived_key_size=32,
+ ciphertext_segment_size=4096)
+
+AES128_CTR_HMAC_SHA256_4KB = create_aes_ctr_hmac_streaming_key_template(
+ aes_key_size=16,
+ hkdf_hash_type=common_pb2.HashType.SHA256,
+ derived_key_size=16,
+ mac_hash_type=common_pb2.HashType.SHA256,
+ tag_size=32,
+ ciphertext_segment_size=4096)
+
+AES256_CTR_HMAC_SHA256_4KB = create_aes_ctr_hmac_streaming_key_template(
+ aes_key_size=32,
+ hkdf_hash_type=common_pb2.HashType.SHA256,
+ derived_key_size=32,
+ mac_hash_type=common_pb2.HashType.SHA256,
+ tag_size=32,
+ ciphertext_segment_size=4096)
diff --git a/python/streaming_aead/streaming_aead_key_templates_test.py b/python/streaming_aead/streaming_aead_key_templates_test.py
new file mode 100644
index 0000000..14e031f
--- /dev/null
+++ b/python/streaming_aead/streaming_aead_key_templates_test.py
@@ -0,0 +1,141 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for tink.python.streaming_aead_key_templates."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tink.proto import aes_ctr_hmac_streaming_pb2
+from tink.proto import aes_gcm_hkdf_streaming_pb2
+from tink.proto import common_pb2
+from tink.proto import tink_pb2
+from tink.python.streaming_aead import streaming_aead_key_templates
+
+
+class StreamingAeadKeyTemplatesTest(absltest.TestCase):
+
+ def test_aes128_gcm_hkdf_4kb(self):
+ template = streaming_aead_key_templates.AES128_GCM_HKDF_4KB
+ self.assertEqual(
+ 'type.googleapis.com/google.crypto.tink.AesGcmHkdfStreamingKey',
+ template.type_url)
+ self.assertEqual(tink_pb2.RAW, template.output_prefix_type)
+ key_format = aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingKeyFormat()
+ key_format.ParseFromString(template.value)
+
+ self.assertEqual(16, key_format.key_size)
+ self.assertEqual(common_pb2.HashType.SHA256,
+ key_format.params.hkdf_hash_type)
+ self.assertEqual(16, key_format.params.derived_key_size)
+ self.assertEqual(4096, key_format.params.ciphertext_segment_size)
+
+ def test_aes256_gcm_hkdf_4kb(self):
+ template = streaming_aead_key_templates.AES256_GCM_HKDF_4KB
+ self.assertEqual(
+ 'type.googleapis.com/google.crypto.tink.AesGcmHkdfStreamingKey',
+ template.type_url)
+ self.assertEqual(tink_pb2.RAW, template.output_prefix_type)
+ key_format = aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingKeyFormat()
+ key_format.ParseFromString(template.value)
+
+ self.assertEqual(32, key_format.key_size)
+ self.assertEqual(common_pb2.HashType.SHA256,
+ key_format.params.hkdf_hash_type)
+ self.assertEqual(32, key_format.params.derived_key_size)
+ self.assertEqual(4096, key_format.params.ciphertext_segment_size)
+
+ def test_aes128_ctr_hmac_sha256_4kb(self):
+ template = streaming_aead_key_templates.AES128_CTR_HMAC_SHA256_4KB
+ self.assertEqual(
+ 'type.googleapis.com/google.crypto.tink.AesCtrHmacStreamingKey',
+ template.type_url)
+ self.assertEqual(tink_pb2.RAW, template.output_prefix_type)
+ key_format = aes_ctr_hmac_streaming_pb2.AesCtrHmacStreamingKeyFormat()
+ key_format.ParseFromString(template.value)
+
+ self.assertEqual(16, key_format.key_size)
+ self.assertEqual(common_pb2.HashType.SHA256,
+ key_format.params.hkdf_hash_type)
+ self.assertEqual(16, key_format.params.derived_key_size)
+ self.assertEqual(common_pb2.HashType.SHA256,
+ key_format.params.hmac_params.hash)
+ self.assertEqual(32, key_format.params.hmac_params.tag_size)
+ self.assertEqual(4096, key_format.params.ciphertext_segment_size)
+
+ def test_aes256_ctr_hmac_sha256_4kb(self):
+ template = streaming_aead_key_templates.AES256_CTR_HMAC_SHA256_4KB
+ self.assertEqual(
+ 'type.googleapis.com/google.crypto.tink.AesCtrHmacStreamingKey',
+ template.type_url)
+ self.assertEqual(tink_pb2.RAW, template.output_prefix_type)
+ key_format = aes_ctr_hmac_streaming_pb2.AesCtrHmacStreamingKeyFormat()
+ key_format.ParseFromString(template.value)
+
+ self.assertEqual(32, key_format.key_size)
+ self.assertEqual(common_pb2.HashType.SHA256,
+ key_format.params.hkdf_hash_type)
+ self.assertEqual(32, key_format.params.derived_key_size)
+ self.assertEqual(common_pb2.HashType.SHA256,
+ key_format.params.hmac_params.hash)
+ self.assertEqual(32, key_format.params.hmac_params.tag_size)
+ self.assertEqual(4096, key_format.params.ciphertext_segment_size)
+
+ def test_create_aes_gcm_hkdf_streaming_key_template(self):
+ # Intentionally using 'weird' or invalid values for parameters,
+ # to test that the function correctly puts them in the resulting template.
+ template = streaming_aead_key_templates.create_aes_gcm_hkdf_streaming_key_template(
+ aes_key_size=42,
+ hash_type=common_pb2.HashType.SHA1,
+ derived_key_size=76,
+ ciphertext_segment_size=64,
+ )
+ self.assertEqual(
+ 'type.googleapis.com/google.crypto.tink.AesGcmHkdfStreamingKey',
+ template.type_url)
+ self.assertEqual(tink_pb2.RAW, template.output_prefix_type)
+ key_format = aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingKeyFormat()
+ key_format.ParseFromString(template.value)
+ self.assertEqual(42, key_format.key_size)
+ self.assertEqual(common_pb2.HashType.SHA1, key_format.params.hkdf_hash_type)
+ self.assertEqual(76, key_format.params.derived_key_size)
+ self.assertEqual(64, key_format.params.ciphertext_segment_size)
+
+ def test_create_aes_ctr_hmac_streaming_key_template(self):
+ # Intentionally using 'weird' or invalid values for parameters,
+ # to test that the function correctly puts them in the resulting template.
+ template = streaming_aead_key_templates.create_aes_ctr_hmac_streaming_key_template(
+ aes_key_size=42,
+ hkdf_hash_type=common_pb2.HashType.SHA1,
+ derived_key_size=76,
+ mac_hash_type=common_pb2.HashType.UNKNOWN_HASH,
+ tag_size=39,
+ ciphertext_segment_size=64,
+ )
+ self.assertEqual(
+ 'type.googleapis.com/google.crypto.tink.AesCtrHmacStreamingKey',
+ template.type_url)
+ self.assertEqual(tink_pb2.RAW, template.output_prefix_type)
+ key_format = aes_ctr_hmac_streaming_pb2.AesCtrHmacStreamingKeyFormat()
+ key_format.ParseFromString(template.value)
+ self.assertEqual(42, key_format.key_size)
+ self.assertEqual(common_pb2.HashType.SHA1, key_format.params.hkdf_hash_type)
+ self.assertEqual(76, key_format.params.derived_key_size)
+ self.assertEqual(common_pb2.HashType.UNKNOWN_HASH,
+ key_format.params.hmac_params.hash)
+ self.assertEqual(39, key_format.params.hmac_params.tag_size)
+ self.assertEqual(64, key_format.params.ciphertext_segment_size)
+
+
+if __name__ == '__main__':
+ absltest.main()
