Google Cloud KMS Integration

backend:
  # Service account key with KMS encrypt/decrypt permissions
  serviceAccountSecretName: "kms-svs-account.json"
  gcpServiceAccountFileEnabled: yes
  gcpEnableKms: yes
  gcloudKmsCryptoKey: "key-name"
  gcloudKmsKeyRing: "test-ring"
  # Project where KMS keys lives
  gcloudProject: "onprem-test-214916"

get_ipython().run_line_magic('pip', 'install -U google-cloud-storage comet_ml google-cloud-kms')

import os
os.environ["COMET_URL_OVERRIDE"] = "https://your-comet-host/clientlib/"

print("ENV", os.environ)

def get_username_and_email():
    # TODO: Customize me
    return ("lothiraldan2", "[email protected]")

def get_project_name():
    # TODO: Customize me
    return "KMS-Integration"

def get_s3_file_name(username):
    return "%(username)s-token.encrypted" % {"username": username}

def get_s3_bucket_and_region():
    # TODO: Customize me
    return ("comet-ml-thales-dev", "global") # Return a specific Region or None for the default one

def get_gcp_kms_settings():
    # TODO: Customize me
    return ("cruise-214915", "global", "gcloud-kms-comet", "kms-cometml-key")
from google.cloud import storage
import io
from google.cloud import kms

def get_api_key(username, email, s3_bucket, region_name, project_id, location_id, key_ring_id, key_id):
    storage_client = storage.Client()
    bucket = storage_client.bucket(s3_bucket)

    # First get the S3 filename
    s3_filename = get_s3_file_name(username)

    # Download the content from S3
    blob = bucket.blob(s3_filename)
    encrypted_data = blob.download_as_string()

    # And decrypt it
    kms_client = kms.KeyManagementServiceClient()

    key_name = kms_client.crypto_key_path(project_id, location_id, key_ring_id, key_id)

    decrypt_response = kms_client.decrypt(
        request={'name': key_name, 'ciphertext': encrypted_data})

    return decrypt_response.plaintext.decode("utf-8")
from comet_ml.exceptions import CometRestApiException
from comet_ml.connection import get_comet_api_client


def ensure_user_account_exists(email, username, send_email=True):
    """ Ensure that an user account given an email exists.
    Args:
        email: string, the user account email
        username: string, the user account username
        send_email: bool, optional. In case an user account is created, should they receive the
            welcome emails.
    Returns:
        In case an user account has been created, ensure_user_account_exists will returns:
            {'cometUserName': 'USERNAME',
             'token': 'q4ddnUELTwMtJG8CCUfuFnKu3',
             'apiKey': 'LzqXBBaBcaeLZFnYqk54xg2cA'}
        In case an user account already exists for this email, ensure_user_account_exists will returns `True`
        In case of errors, ensure_user_account_exists will raises comet_ml.CometRestApiException
    """
    auth_api_client = get_comet_api_client()

    try:
        new_account = auth_api_client.create_user(
            email, username, "cmle-integration", send_email=send_email
        )

        return new_account
    except CometRestApiException as exception:
        if exception.response.status_code == 400:
            if exception.safe_json_response and exception.safe_json_response.get("msg") == "Please use a different email.":
                return True

        raise
import comet_ml

def get_comet_experiment():
    # Get config
    username, email = get_username_and_email()
    project_name = get_project_name()
    s3_bucket, region = get_s3_bucket_and_region()
    project_id, location_id, key_ring_id, key_id = get_gcp_kms_settings()

    # Ensure the user account exits
    account_existing = ensure_user_account_exists(email, username, send_email=False)

    comet_api_key = get_api_key(username, email, s3_bucket, region, project_id, location_id, key_ring_id, key_id)

    return comet_ml.Experiment(comet_api_key, project_name=project_name)
exp = get_comet_experiment()
exp.add_tag("Working")
exp.end()