def update_key_policy()

in functions/source/update_custommanaged_cmk_policy.py [0:0]


def update_key_policy(key_id, arn_credentials, reuse_key_for_cluster_volumes):
    global databricks_s3bucketandEBS_statement
    global databricks_s3bucketonly_statement
    current_key_policy = json.loads(get_key_policy(key_id))
    statements = current_key_policy['Statement']
    print('no. of statements in current policy are : {}'.format(len(statements)))

    # update the kms policy statement
    if reuse_key_for_cluster_volumes == 'True':
        databricks_s3bucketandEBS_statement3['Principal']['AWS'] = str(arn_credentials)
        statements.append(databricks_s3bucketandEBS_statement1)
        statements.append(databricks_s3bucketandEBS_statement2)
        statements.append(databricks_s3bucketandEBS_statement3)
    else:
        statements.append(databricks_s3bucketonly_statement1)
        statements.append(databricks_s3bucketonly_statement2) 

    print('no. of statements in new policy are : {}'.format(len(statements)))
    current_key_policy['Statement'] = statements
    print('new policy is : {}'.format(json.dumps(current_key_policy)))  
    
    # Update new statements in current key policy
    response = client.put_key_policy(
        KeyId=key_id,
        PolicyName='default',
        Policy=json.dumps(current_key_policy)
    )