encrypt and decrypt with kms, boto3 and pyCrypto. Python2 compatible

kms.py

#!/usr/bin/env python
import base64
from Crypto import Random
import boto3
from Crypto.Cipher import AES


PAD = lambda s: s + (32 - len(s) % 32) * ' '


def get_arn(aws_data):
    return 'arn:aws:kms:{region}:{account_number}:key/{key_id}'.format(**aws_data)


# encrypt the plaintext message using a kms key
def encrypt_data(aws_data, plaintext_message):

    # get a kms client
    kms_client = boto3.client(
        'kms',
        region_name=aws_data['region'])

    # generate a temporary encryption key
    data_key = kms_client.generate_data_key(
        KeyId=aws_data['key_id'],
        KeySpec='AES_128')

    # get both the plaintext and encrypted parts of the key
    cipher_text_blob = data_key.get('CiphertextBlob')
    plaintext_key = data_key.get('Plaintext')

    # create a random IV (init vector), used in the encryption.
    # the iv can be made public, but must be created with CSPRNG
    iv = Random.new().read(AES.block_size) 
    cypher = AES.new(plaintext_key, AES.MODE_CBC, iv)

    # encrypt with the plain key, and then prefix with iv and base64 the whole thing
    encrypted_data = base64.b64encode(iv + cypher.encrypt(PAD(plaintext_message)))

    # return the encrypted message along with the encrpyted key (in base64)
    # need to preserve both of these data elements in order to decrypt
    return encrypted_data, base64.b64encode(cipher_text_blob)


# decrypt the encrypted data using the encrpyted key and kms
def decrypt_data(aws_data, encrypted_data, cipher_text_blob):
    kms_client = boto3.client(
        'kms',
        region_name=aws_data['region'])

    # convert the encrypted key to plaintext key
    decrypted_key = kms_client.decrypt(CiphertextBlob=base64.b64decode(cipher_text_blob)).get('Plaintext')

    # un-base64 the encrypted data, seperate the iv from the message
    enc = base64.b64decode(encrypted_data)
    iv = enc[:AES.block_size]

    # decrypt the message using the plaintext key and the iv
    cypher = AES.new(decrypted_key, AES.MODE_CBC, iv)
    return cypher.decrypt(enc[AES.block_size:]).rstrip()


def main():
    # Add your account number / region / KMS Key ID here.
    aws_data = {
        'region': 'us-east-1',
        'account_number': 'xxxxx',
        'key_id': 'yyyyy',
    }

    # And your super secret message to envelope encrypt...
    plaintext = 'Hello, World!'

    # Store encrypted_data & cipher_text_blob in your persistent storage. You will need them both later.
    encrypted_data, cipher_text_blob = encrypt_data(aws_data, plaintext)
    print 'encrypted_data: %s' % encrypted_data
    print 'cipher_text_blob: %s.' % cipher_text_blob

    # # Later on when you need to decrypt, get these from your persistent storage.
    decrypted_data = decrypt_data(aws_data, encrypted_data, cipher_text_blob)
    print(decrypted_data)

if __name__ == '__main__':
    main()