Implement GCP KMS integration in Python, without wrapping C++ version.

PiperOrigin-RevId: 521439181
Change-Id: If410287106953071e43390c82c3e050b61d5a4f8

kokoro/gcp_ubuntu/bazel/run_tests.sh

@@ -36,6 +36,9 @@ readonly GITHUB_ORG="https://github.com/tink-crypto"
 
 ./kokoro/testutils/copy_credentials.sh "testdata" "all"
 
+# TODO(b/276277854) It is not clear why this is needed.
+pip3 install google-cloud-kms==2.15.0 --user
+
 TINK_PY_MANUAL_TARGETS=()
 # These tests require valid credentials to access KMS services.
 if [[ -n "${KOKORO_ROOT:-}" ]]; then

kokoro/macos_external/bazel/run_tests.sh

@@ -35,8 +35,10 @@ readonly GITHUB_ORG="https://github.com/tink-crypto"
   "${GITHUB_ORG}/tink-cc-gcpkms"
 
 ./kokoro/testutils/copy_credentials.sh "testdata" "all"
-# Install protobuf pip packages.
+
+# TODO(b/276277854) It is not clear why this is needed.
 pip3 install protobuf==3.20.3 --user
+pip3 install google-cloud-kms==2.15.0 --user
 
 TINK_PY_MANUAL_TARGETS=()
 # These tests require valid credentials to access KMS services.

kokoro/macos_external/examples/bazel_kms/run_tests.sh

@@ -47,10 +47,9 @@ cp "examples/WORKSPACE" "examples/WORKSPACE.bak"
 ./kokoro/testutils/replace_http_archive_with_local_repository.py \
   -f "examples/WORKSPACE" -t "${TINK_BASE_DIR}"
 
-# Install protobuf pip packages.
-# TODO(b/253216420): Investigate why these need to be installed instead on
-# MacOS, but not on GCP Ubuntu.
+# TODO(b/276277854) It is not clear why this is needed.
 pip3 install protobuf==3.20.3 --user
+pip3 install google-cloud-kms==2.15.0 --user
 pip3 install google-cloud-storage==2.5.0 --user
 
 # All manual test targets except *test_package ones.

requirements.in

@@ -1,3 +1,6 @@
 absl-py==1.3.0
 protobuf==4.21.9
 boto3==1.26.89
+google-auth==2.16.2
+google-api-core==2.11.0
+google-cloud-kms==2.16.1

tink/integration/gcpkms/BUILD.bazel

@@ -20,8 +20,10 @@ py_library(
     deps = [
         "//tink/aead",
         "//tink/aead:_kms_aead_key_manager",
-        "//tink/cc/pybind:tink_bindings",
         "//tink/core",
+        requirement("google-auth"),
+        requirement("google-api-core"),
+        requirement("google-cloud-kms"),
     ],
 )
 
@@ -69,6 +71,7 @@ py_test(
         "//tink:cleartext_keyset_handle",
         "//tink:tink_python",
         "//tink/aead",
+        "//tink/aead:_kms_aead_key_manager",
         "//tink/testing:helper",
         requirement("absl-py"),
     ],

tink/integration/gcpkms/_gcp_kms_client.py

@@ -15,14 +15,53 @@
 
 from typing import Optional
 
+from google.api_core import exceptions as core_exceptions
+from google.cloud import kms_v1
+from google.oauth2 import service_account
+
 from tink import aead
 from tink import core
 from tink.aead import _kms_aead_key_manager
-from tink.cc.pybind import tink_bindings
 
 GCP_KEYURI_PREFIX = 'gcp-kms://'
 
 
+class _GcpKmsAead(aead.Aead):
+  """Implements the Aead interface for GCP KMS."""
+
+  def __init__(
+      self, client: kms_v1.KeyManagementServiceClient, name: str
+  ) -> None:
+    self.client = client
+    self.name = name
+
+  def encrypt(self, plaintext: bytes, associated_data: bytes) -> bytes:
+    try:
+      response = self.client.encrypt(
+          request=kms_v1.types.service.EncryptRequest(
+              name=self.name,
+              plaintext=plaintext,
+              additional_authenticated_data=associated_data,
+          )
+      )
+      return response.ciphertext
+    except core_exceptions.GoogleAPIError as e:
+      raise core.TinkError(e)
+
+  def decrypt(self, ciphertext: bytes, associated_data: bytes) -> bytes:
+    try:
+      response = self.client.decrypt(
+         request=kms_v1.types.service.DecryptRequest(
+             name=self.name,
+             ciphertext=ciphertext,
+             additional_authenticated_data=associated_data
+         )
+      )
+      return response.plaintext
+    except core_exceptions.GoogleAPIError as e:
+      raise core.TinkError(e)
+
+
 class GcpKmsClient(_kms_aead_key_manager.KmsClient):
   """Basic GCP client for AEAD."""
 
@@ -45,15 +84,20 @@ def __init__(
     """
 
     if not key_uri:
-      self._key_uri = ''
+      self._key_uri = None
     elif key_uri.startswith(GCP_KEYURI_PREFIX):
       self._key_uri = key_uri
     else:
       raise core.TinkError('Invalid key_uri.')
     if not credentials_path:
       credentials_path = ''
-    # Use the C++ GCP KMS client
-    self.cc_client = tink_bindings.GcpKmsClient(self._key_uri, credentials_path)
+    if not credentials_path:
+      self._client = kms_v1.KeyManagementServiceClient()
+      return
+    credentials = service_account.Credentials.from_service_account_file(
+        credentials_path
+    )
+    self._client = kms_v1.KeyManagementServiceClient(credentials=credentials)
 
   def does_support(self, key_uri: str) -> bool:
     """Returns true iff this client supports KMS key specified in 'key_uri'.
@@ -64,9 +108,10 @@ def does_support(self, key_uri: str) -> bool:
     Returns:
       A boolean value which is true if the key is supported and false otherwise.
     """
-    return self.cc_client.does_support(key_uri)
+    if not self._key_uri:
+      return key_uri.startswith(GCP_KEYURI_PREFIX)
+    return key_uri == self._key_uri
 
-  @core.use_tink_errors
   def get_aead(self, key_uri: str) -> aead.Aead:
     """Returns an Aead-primitive backed by KMS key specified by 'key_uri'.
 
@@ -76,8 +121,10 @@ def get_aead(self, key_uri: str) -> aead.Aead:
     Returns:
       An Aead object.
     """
-
-    return aead.AeadCcToPyWrapper(self.cc_client.get_aead(key_uri))
+    if not key_uri.startswith(GCP_KEYURI_PREFIX):
+      raise core.TinkError('Invalid key_uri.')
+    key_id = key_uri[len(GCP_KEYURI_PREFIX) :]
+    return _GcpKmsAead(self._client, key_id)
 
   @classmethod
   def register_client(

tink/integration/gcpkms/_gcp_kms_client_integration_test.py

@@ -56,6 +56,35 @@ def test_encrypt_decrypt(self):
     ciphertext = gcp_aead.encrypt(plaintext, associated_data)
     self.assertEqual(plaintext, gcp_aead.decrypt(ciphertext, associated_data))
 
+  def test_decrypt_with_wrong_ad_fails(self):
+    gcp_client = gcpkms.GcpKmsClient(KEY_URI, CREDENTIAL_PATH)
+    gcp_aead = gcp_client.get_aead(KEY_URI)
+
+    ciphertext = gcp_aead.encrypt(b'plaintext', b'associated_data')
+    with self.assertRaises(tink.TinkError):
+      gcp_aead.decrypt(ciphertext, b'wrong_associated_data')
+
+  def test_decrypt_with_wrong_key_fails(self):
+    gcp_client = gcpkms.GcpKmsClient(None, CREDENTIAL_PATH)
+    gcp_aead1 = gcp_client.get_aead(KEY_URI)
+    gcp_aead2 = gcp_client.get_aead(KEY2_URI)
+
+    ciphertext1 = gcp_aead1.encrypt(b'plaintext', b'associated_data')
+    ciphertext2 = gcp_aead2.encrypt(b'plaintext', b'associated_data')
+
+    # First, verify that both key URIs work.
+    self.assertEqual(
+        b'plaintext', gcp_aead1.decrypt(ciphertext1, b'associated_data')
+    )
+    self.assertEqual(
+        b'plaintext', gcp_aead2.decrypt(ciphertext2, b'associated_data')
+    )
+
+    with self.assertRaises(tink.TinkError):
+      gcp_aead2.decrypt(ciphertext1, b'associated_data')
+    with self.assertRaises(tink.TinkError):
+      gcp_aead1.decrypt(ciphertext2, b'associated_data')
+
   def test_encrypt_decrypt_localized_uri(self):
     gcp_client = gcpkms.GcpKmsClient(LOCAL_KEY_URI, CREDENTIAL_PATH)
     gcp_aead = gcp_client.get_aead(LOCAL_KEY_URI)

tink/integration/gcpkms/_gcp_kms_client_test.py

@@ -56,7 +56,7 @@ def test_client_empty_key_uri(self):
     self.assertEqual(gcp_client.does_support(gcp_key), True)
 
   def test_client_invalid_path(self):
-    with self.assertRaises(ValueError):
+    with self.assertRaises(FileNotFoundError):
       gcpkms.GcpKmsClient(None, CREDENTIAL_PATH + 'corrupted')