/*
 * Copyright 2021 Google LLC
 *
 * 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.
 */

package cloudsql.tink;

// [START cloud_sql_mysql_cse_key]

import com.google.crypto.tink.Aead;
import com.google.crypto.tink.KeysetHandle;
import com.google.crypto.tink.KmsClients;
import com.google.crypto.tink.aead.AeadConfig;
import com.google.crypto.tink.aead.AeadKeyTemplates;
import com.google.crypto.tink.integration.gcpkms.GcpKmsClient;
import com.google.crypto.tink.proto.KeyTemplate;
import java.security.GeneralSecurityException;

public class CloudKmsEnvelopeAead {

  public static Aead get(String kmsUri) throws GeneralSecurityException {
    AeadConfig.register();
    // Generate a new envelope key template, then generate key material.
    KeyTemplate kmsEnvKeyTemplate = AeadKeyTemplates
        .createKmsEnvelopeAeadKeyTemplate(kmsUri, AeadKeyTemplates.AES128_GCM);
    KeysetHandle keysetHandle = KeysetHandle.generateNew(kmsEnvKeyTemplate);

    // Register the KMS client.
    KmsClients.add(new GcpKmsClient()
        .withDefaultCredentials());

    // Create envelope AEAD primitive from keysetHandle
    return keysetHandle.getPrimitive(Aead.class);
  }
}
// [END cloud_sql_mysql_cse_key]
@@ -684,52 +684,8 @@
 ad;%0A
-import com.google.crypto.tink.KeysetHandle;%0A
 impo
@@ -719,17 +719,16 @@
 msClient
-s
 ;%0Aimport
@@ -821,16 +821,68 @@
 plates;%0A
+import com.google.crypto.tink.aead.KmsEnvelopeAead;%0A
 import c
@@ -940,57 +940,8 @@
 nt;%0A
-import com.google.crypto.tink.proto.KeyTemplate;%0A
 impo
@@ -1126,20 +1126,19 @@
 ();%0A
+%0A
     // 
-Gener
+Cre
 ate 
@@ -1147,373 +1147,184 @@
 new 
-envelope key template, then generate key material.%0A    KeyTemplate kmsEnvKeyTemplate = AeadKeyTemplates%0A        .createKmsEnvelopeAeadKeyTemplate(kmsUri, AeadKeyTemplates.AES128_GCM
+KMS Client%0A    KmsClient client = new GcpKmsClient().withDefaultCredentials(
 );%0A
+%0A
     
-KeysetHandle keysetHandle = KeysetHandle.generateNew(kmsEnvKeyTemplate);%0A%0A    // Register the KMS client.%0A    KmsClients.add(new GcpKmsClient()%0A        .withDefaultCredentials()
+// Create an AEAD primitive using the Cloud KMS key%0A    Aead gcpAead = client.getAead(kmsUri
 );%0A%0A
@@ -1337,16 +1337,19 @@
  Create 
+an 
 envelope
@@ -1367,74 +1367,76 @@
 tive
- from keysetHandle%0A    return keysetHandle.getPrimitive(Aead.class
+%0A    return new KmsEnvelopeAead(AeadKeyTemplates.AES128_GCM, gcpAead
 );%0A 

replace deprecated envelope encryption method in Tink sample (#4838)