AWS KMS + Python Cryptography using Fernet

kmsencrypt.py

#!/usr/bin/env python
"""
kmsencrypt.py

AWS kms + python Cryptography library file encrypt and decrypt

This will perform a file encryption and decryption using AWS KMS for generating a data key 
rather than using the Fernet generate_key function.
Assumes that AWS access key, secret or token have been setup outside using credentials file or envvars

!! WARNING - I am not a security expert so use at your own risk !!
"""
import sys
import base64

import boto3
from cryptography.fernet import Fernet

KEY_ID='alias/my_key'  # <- place you kms keyid or alias here

def main():
    # get a data key from kms
    kms_client = boto3.client('kms')
    data_key_dict = kms_client.generate_data_key(
        KeyId=KEY_ID, KeySpec='AES_256')

    # get the components from kms response
    encrypted_key = base64.b64encode(data_key_dict['CiphertextBlob'])
    master_key_id = data_key_dict['KeyId']
    plain_key = base64.b64encode(data_key_dict['Plaintext'])

    # encrypt file with data key using cryptography.fernet library
    with open("./data.txt", mode='rb') as data_fh:
        cipher = Fernet(plain_key)
        encrypted_data_content = cipher.encrypt(data_fh.read())
    
    # remove sensitive variables
    del plain_key, cipher, data_key_dict
    
    # write content to file
    with open("./data.txt.enc", mode='wb') as encdata_fh:
        encdata_fh.write(encrypted_data_content)

    print("Encryped file...")
    print("enckey={}\nmasterkey={}".format(encrypted_key, master_key_id))

    #--------------------------------------------------------------- 
    # OK, lets decrypt the file. You only have the encrypted key to work with
    #--------------------------------------------------------------- 

    # decrypt the data key using aws kms
    data_key_dict = kms_client.decrypt(CiphertextBlob=base64.b64decode(encrypted_key))
    plain_key = base64.b64encode(data_key_dict['Plaintext'])

    # decrypt the file using plan key and fernet
    cipher = Fernet(plain_key)
    with open("./data.txt.enc", mode='rb') as encdata_fh:
        data = cipher.decrypt(encdata_fh.read())

    # remove the key variables
    del plain_key, cipher, data_key_dict
    
    print("\nDecrypted file...")
    print("The content is as follows:\n{}".format(data.decode()))


if __name__ == '__main__':
    sys.exit(int(main() or 0))