Using TINKKEY to create EncryptedKeyset for BigQuery

main.py

# 1. Create KEK in KMS
# 2. Use TinkKey to generate DEK and wrap that with KEK (i.,e encryptedKeyset)
# 3. use python tink to encrypt some data with wrappeed DEK keyset
# 4. Use BQ functions to decrypt data using ciphertext in 3 and EncryptedKeyset in 2


# ref

# https://blog.salrashid.dev/articles/2022/bq_kms/
# https://github.com/salrashid123/tink_samples/



# ```bash
# export PROJECT_ID=`gcloud config get-value core/project`
# export GCLOUD_USER=`gcloud config get-value core/account`
# export KEYRING_NAME=bqr
# export LOCATION=us
# export KEY_NAME=k1

# gcloud kms keyrings create $KEYRING_NAME --location=$LOCATION
# gcloud kms keys create $KEY_NAME --keyring=$KEYRING_NAME  --purpose=encryption  --default-algorithm=google-symmetric-encryption  --location=$LOCATION
# ```

# ## for bq delegation
# gcloud kms keys add-iam-policy-binding $KEY_NAME \
#   --keyring=$KEYRING_NAME --location=$LOCATION \
#   --member=user:$GCLOUD_USER \
#   --role=roles/cloudkms.cryptoKeyDecrypterViaDelegation
# '

# ## allow the adc user to use the key for local encryption/decryption

# gcloud kms keys add-iam-policy-binding $KEY_NAME \
#   --keyring=$KEYRING_NAME --location=$LOCATION \
#   --member=user:$GCLOUD_USER \
#   --role=roles/cloudkms.cryptoKeyDecrypter
  
# gcloud kms keys add-iam-policy-binding $KEY_NAME \
#   --keyring=$KEYRING_NAME --location=$LOCATION \
#   --member=user:$GCLOUD_USER \
#   --role=roles/cloudkms.cryptoKeyEncrypter

# # download tinkkey

# wget https://storage.googleapis.com/tinkey/tinkey-1.6.1.tar.gz
# gzip -d tinkey-1.6.1.tar.gz
# tar xvf tinkey-1.6.1.tar
# rm tinkey-1.6.1.tar

# # generate encryptedKeyset 

# ./tinkey create-keyset --key-template AES256_GCM --out-format json \
#   --out encrypted_aead_keyset.json  \
#   --master-key-uri gcp-kms://projects/$PROJECT_ID/locations/us/keyRings/$KEYRING_NAME/cryptoKeys/$KEY_NAME

# cat encrypted_aead_keyset.json  | jq '.'

# {
#   "encryptedKeyset": "CiQA0Rg9S46PMKLyYxjopEUPX/vH7syBeaBu2NGmi3whMZer4UYSlQEAQ8LPmJSlBPfK3yICCAzc9Yrnw35t4vyy08A7ilVvCT0kYRpfX9Eaub4HLlI5rSyBBawsMjtnIrDJ9qsje2QWBj1PjXO3SenajG+y6G6j9G4bb5D4UReeUlZ7vc+1D/YvXws3eE39dgUcBv1PkSrtH9XFsaIg4TiKDu4MLx9wfkSaxINgIO5b+HkbxQhjoPiOKSL4FA==",
#   "keysetInfo": {
#     "primaryKeyId": 2072948758,
#     "keyInfo": [
#       {
#         "typeUrl": "type.googleapis.com/google.crypto.tink.AesGcmKey",
#         "status": "ENABLED",
#         "keyId": 2072948758,
#         "outputPrefixType": "TINK"
#       }
#     ]
#   }
# }


# ### encrypt some data locally with that keyset

# python main.py 
#     AXuOsBZAT2Ssfb0WwGLKjepEznDlQ7U+evJtz7r0aYoNXb1HHzFjvhR5jtHosgNU
#     Juarez Gold DSS


# ### decrypt that data using the encrypted keyset and bq

# bq query --use_legacy_sql=false  "
#   DECLARE
#     KMS_RESOURCE_NAME STRING;
#   DECLARE
#     FIRST_LEVEL_KEYSET BYTES;
#   SET
#     KMS_RESOURCE_NAME = 'gcp-kms://projects/$PROJECT_ID/locations/us/keyRings/$KEYRING_NAME/cryptoKeys/$KEY_NAME';
#   SET
#     FIRST_LEVEL_KEYSET = FROM_BASE64('CiQA0Rg9S46PMKLyYxjopEUPX/vH7syBeaBu2NGmi3whMZer4UYSlQEAQ8LPmJSlBPfK3yICCAzc9Yrnw35t4vyy08A7ilVvCT0kYRpfX9Eaub4HLlI5rSyBBawsMjtnIrDJ9qsje2QWBj1PjXO3SenajG+y6G6j9G4bb5D4UReeUlZ7vc+1D/YvXws3eE39dgUcBv1PkSrtH9XFsaIg4TiKDu4MLx9wfkSaxINgIO5b+HkbxQhjoPiOKSL4FA==');
#   SELECT
#     AEAD.DECRYPT_STRING( 
#         KEYS.KEYSET_CHAIN(KMS_RESOURCE_NAME, FIRST_LEVEL_KEYSET),
#         FROM_BASE64('AXuOsBZAT2Ssfb0WwGLKjepEznDlQ7U+evJtz7r0aYoNXb1HHzFjvhR5jtHosgNU'),
#         'additional_data') as ff
# "


# +-----------------+
# |       ff        |
# +-----------------+
# | Juarez Gold DSS |
# +-----------------+




import base64
import io
import tink
from tink import aead
from tink import tink_config
from tink.proto import tink_pb2
from tink.proto import common_pb2
from tink.integration import gcpkms
from tink import core

from tink import cleartext_keyset_handle

tink_config.register()
aead.register()

key_uri = 'gcp-kms://projects/mineral-minutia-820/locations/us/keyRings/bqr/cryptoKeys/k1'
tink_keyset = 'encrypted_aead_keyset.json'

gcp_client = gcpkms.GcpKmsClient(key_uri=key_uri,credentials_path="")
gcp_aead = gcp_client.get_aead(key_uri)

with open(tink_keyset, 'r') as f:
    json_keyset = f.read()
reader = tink.JsonKeysetReader(json_keyset)

keyset_handle = tink.KeysetHandle.read(reader, gcp_aead)

aead_primitive = keyset_handle.primitive(aead.Aead)

plaintext = b'Juarez Gold DSS'
associated_data = b'additional_data'

ciphertext = aead_primitive.encrypt(plaintext, associated_data)
print(base64.b64encode(ciphertext).decode('utf-8'))

text = aead_primitive.decrypt(ciphertext, associated_data)
print(text)