PROBLEM

As a data engineer, we deal with many PII and PHI information. Once data is loaded to DataWareHouse we might leave the data files in same path or move to archive path and left untouched. Question here is how to leave the sensitive information in a open buckets on cloud. In order to maintain the security standards, it is advised to encrypt the files once used. We have to enable security protocols on those files.

DESIGN GOALS

In order to demo I have used AWS KMS and AWS Lambda to solve this problem.

SOLUTION - STEPS

Create an AWS Lambda - Function Example - poc-encrypt-decrypt-s3
Add policy to the Lambda role that has a read, list and write access to KMS.
Add layers to have smart-open and cryptography packages to run lambda function.

Create layer using AWS Cloudshell

ROOT_DIR="lambda_layers/python/lib/python3.7/site-packages/"

ZIP_FILE="aws_s3_security_layer.zip"

LAYER_NAME="aws_s3_security_layer"

cd ~

rm -rf ${ROOT_DIR}*

mkdir -p ${ROOT_DIR}

python3 -m venv venv

source venv/bin/activate

# 1) install the dependencies in the desired folder

python3 -m pip install cffi  -t  ${ROOT_DIR}

pip3 install smart-open cryptography  -t  ${ROOT_DIR}.

cd  ${ROOT_DIR}

rm -r *dist-info __pycache__

cd -

deactivate

# 2) Zip the lambda_layers folder

cd lambda_layers

zip -r ${ZIP_FILE} *

# 3) publish layer

aws lambda publish-layer-version \
    --layer-name ${LAYER_NAME} \
    --license-info "MIT" \
    --compatible-runtimes python3.6 python3.7 python3.8 python3.9  \
    --zip-file fileb://${ZIP_FILE} \
    --compatible-architectures "arm64" "x86_64"


# aws lambda publish-layer-version --layer-name my-layer --description "My layer"  \
# --license-info "MIT" --content S3Bucket=lambda-layers-us-east-2-123456789012,S3Key=layer.zip \
# --compatible-runtimes python3.6 python3.7 python3.8 \
# --compatible-architectures "arm64" "x86_64"

# https://docs.aws.amazon.com/lambda/latest/dg/configuration-layers.html

echo "-*- Done -*-"

AWS KMS Library files

Save as lib_kms.py in AWS Lambda code

import boto3
import logging
import base64
from cryptography.fernet import Fernet
from smart_open import open

NUM_BYTES_FOR_LEN = 4

def retrieve_cmk(desc):
    """Retrieve an existing KMS CMK based on its description

    :param desc: Description of CMK specified when the CMK was created
    :return Tuple(KeyId, KeyArn) where:
        KeyId: CMK ID
        KeyArn: Amazon Resource Name of CMK
    :return Tuple(None, None) if a CMK with the specified description was
    not found
    """

    # Retrieve a list of existing CMKs
    # If more than 100 keys exist, retrieve and process them in batches
    kms_client = boto3.client('kms')
    try:
        response = kms_client.list_keys()
    except Exception as e:
        logging.error(e)
        return None, None

    done = False
    while not done:
        for cmk in response['Keys']:
            # Get info about the key, including its description
            try:
                key_info = kms_client.describe_key(KeyId=cmk['KeyArn'])
            except Exception as e:
                logging.error(e)
                return None, None

            # Is this the key we're looking for?
            if key_info['KeyMetadata']['Description'] == desc:
                return cmk['KeyId'], cmk['KeyArn']

        # Are there more keys to retrieve?
        if not response['Truncated']:
            # No, the CMK was not found
            logging.debug('A CMK with the specified description was not found')
            done = True
        else:
            # Yes, retrieve another batch
            try:
                response = kms_client.list_keys(Marker=response['NextMarker'])
            except Exception as e:
                logging.error(e)
                return None, None

    # All existing CMKs were checked and the desired key was not found
    return None, None


def create_cmk(desc='Customer Master Key'):
    """Create a KMS Customer Master Key

    The created CMK is a Customer-managed key stored in AWS KMS.

    :param desc: key description
    :return Tuple(KeyId, KeyArn) where:
        KeyId: AWS globally-unique string ID
        KeyArn: Amazon Resource Name of the CMK
    :return Tuple(None, None) if error
    """

    # Create CMK
    kms_client = boto3.client('kms')
    try:
        response = kms_client.create_key(Description=desc)
    except Exception as e:
        logging.error(e)
        return None, None

    # Return the key ID and ARN
    return response['KeyMetadata']['KeyId'], response['KeyMetadata']['Arn']


def create_data_key(cmk_id, key_spec='AES_256'):
    """Generate a data key to use when encrypting and decrypting data

    :param cmk_id: KMS CMK ID or ARN under which to generate and encrypt the
    data key.
    :param key_spec: Length of the data encryption key. Supported values:
        'AES_128': Generate a 128-bit symmetric key
        'AES_256': Generate a 256-bit symmetric key
    :return Tuple(EncryptedDataKey, PlaintextDataKey) where:
        EncryptedDataKey: Encrypted CiphertextBlob data key as binary string
        PlaintextDataKey: Plaintext base64-encoded data key as binary string
    :return Tuple(None, None) if error
    """

    # Create data key
    kms_client = boto3.client('kms')
    try:
        response = kms_client.generate_data_key(KeyId=cmk_id, KeySpec=key_spec)
    except Exception as e:
        logging.error(e)
        return None, None

    # Return the encrypted and plaintext data key
    return response['CiphertextBlob'], base64.b64encode(response['Plaintext'])


def encrypt_file(filename, cmk_id):
    """Encrypt a file using an AWS KMS CMK

    A data key is generated and associated with the CMK.
    The encrypted data key is saved with the encrypted file. This enables the
    file to be decrypted at any time in the future and by any program that
    has the credentials to decrypt the data key.
    The encrypted file is saved to <filename>.encrypted
    Limitation: The contents of filename must fit in memory.

    :param filename: File to encrypt
    :param cmk_id: AWS KMS CMK ID or ARN
    :return: True if file was encrypted. Otherwise, False.
    """

    # Read the entire file into memory
    try:
        with open(filename, 'rb') as file:
            file_contents = file.read()
    except IOError as e:
        logging.error(e)
        return False

    # Generate a data key associated with the CMK
    # The data key is used to encrypt the file. Each file can use its own
    # data key or data keys can be shared among files.
    # Specify either the CMK ID or ARN
    data_key_encrypted, data_key_plaintext = create_data_key(cmk_id)
    if data_key_encrypted is None:
        return False
    logging.info('Created new AWS KMS data key')

    # Encrypt the file
    f = Fernet(data_key_plaintext)
    file_contents_encrypted = f.encrypt(file_contents)

    # Write the encrypted data key and encrypted file contents together
    try:
        with open(filename + '.encrypted', 'wb') as file_encrypted:
            file_encrypted.write(len(data_key_encrypted).to_bytes(NUM_BYTES_FOR_LEN,
                                                                  byteorder='big'))
            file_encrypted.write(data_key_encrypted)
            file_encrypted.write(file_contents_encrypted)
    except IOError as e:
        logging.error(e)
        return False

    # For the highest security, the data_key_plaintext value should be wiped
    # from memory. Unfortunately, this is not possible in Python. However,
    # storing the value in a local variable makes it available for garbage
    # collection.
    return True


def decrypt_data_key(data_key_encrypted):
    """Decrypt an encrypted data key

    :param data_key_encrypted: Encrypted ciphertext data key.
    :return Plaintext base64-encoded binary data key as binary string
    :return None if error
    """

    # Decrypt the data key
    kms_client = boto3.client('kms')
    try:
        response = kms_client.decrypt(CiphertextBlob=data_key_encrypted)
    except Exception as e:
        logging.error(e)
        return None

    # Return plaintext base64-encoded binary data key
    return base64.b64encode((response['Plaintext']))


def decrypt_file(filename):
    """Decrypt a file encrypted by encrypt_file()

    The encrypted file is read from <filename>.encrypted
    The decrypted file is written to <filename>.decrypted

    :param filename: File to decrypt
    :return: True if file was decrypted. Otherwise, False.
    """

    # Read the encrypted file into memory
    try:
        with open(filename + '.encrypted', 'rb') as file:
            file_contents = file.read()
    except IOError as e:
        logging.error(e)
        return False

    # The first NUM_BYTES_FOR_LEN bytes contain the integer length of the
    # encrypted data key.
    # Add NUM_BYTES_FOR_LEN to get index of end of encrypted data key/start
    # of encrypted data.
    data_key_encrypted_len = int.from_bytes(file_contents[:NUM_BYTES_FOR_LEN],
                                            byteorder='big') \
                             + NUM_BYTES_FOR_LEN
    data_key_encrypted = file_contents[NUM_BYTES_FOR_LEN:data_key_encrypted_len]

    # Decrypt the data key before using it
    data_key_plaintext = decrypt_data_key(data_key_encrypted)
    if data_key_plaintext is None:
        return False

    # Decrypt the rest of the file
    f = Fernet(data_key_plaintext)
    file_contents_decrypted = f.decrypt(file_contents[data_key_encrypted_len:])

    # Write the decrypted file contents
    try:
        with open(filename + '.decrypted', 'wb') as file_decrypted:
            file_decrypted.write(file_contents_decrypted)
    except IOError as e:
        logging.error(e)
        return False

    # The same security issue described at the end of encrypt_file() exists
    # here, too, i.e., the wish to wipe the data_key_plaintext value from
    # memory.
    return True

Lambda Function

Save as lambda_function.py in AWS Lambda code

import os
import sys
import json
import logging


root = logging.getLogger()

if root.handlers:
    for handler in root.handlers:
        root.removeHandler(handler)

stream_handler = logging.StreamHandler(sys.stdout)

log_args = {
    "level": logging.INFO,
    "format": "%(asctime)s [%(levelname)s] %(name)s - %(message)s",
    "datefmt": '%d-%b-%y %H:%M' ,
    "handlers": [stream_handler]
}

logging.basicConfig(**log_args)

from lib_kms import *

file_name = os.path.splitext(os.path.basename(__file__))[0]

def lambda_handler(event, context):

    logger = logging.getLogger(f"{file_name}.__main__")

    desc_kms = 'poc-kms-encrypt-decrypt'

    filename = 's3://sample-bucket/testData.csv'

    KeyId, KeyArn = retrieve_cmk(desc_kms)

    logger.info(f'[*] Retrieve Customer Master Key (CMK) = KeyId {KeyId}, KeyArn {KeyArn}')

    if not KeyId and not KeyArn:
        logger.info('[+] Creating new master key')
        KeyId, KeyArn = create_cmk(desc_kms)
        logger.info(f'[*] Create Master Key = KeyId {KeyId}, KeyArn {KeyArn}')

    if not KeyId and not KeyArn:
        logger.error('[-] Customer Master Key (CMK) cannot be none')
        raise

    EncryptedDataKey, PlaintextDataKey = create_data_key(cmk_id = KeyId, key_spec='AES_256')
    logger.info(f'[*] Create Data Key = EncryptedDataKey {EncryptedDataKey}, PlaintextDataKey {PlaintextDataKey}')

    if encrypt_file(filename, cmk_id = KeyId):
        logger.info('[+] file encrypted')
    else:
        logger.error('[-] file  is not encrypted')
        raise

    if decrypt_file(filename ):
        logger.info('[+] file decrypted')
    else:
        logger.error('[-] file  is not decrypted')
        raise


    # TODO implement
    return {
        'statusCode': 200,
        'body': json.dumps('Demo Success !')
    }

Once you have successfully give the s3 file URI and try running it will create the encrypted file and decrypted file in the S3 bucket.

Advantage of using the smart-open library we can encrypt or decrypt without downloading to a local machine. Encryption and decryption using AWS KMS is done on bytes streaming from S3 using smart-open.

GitHub Repo - AWS-S3-encrypt-decrypt-using-AWS-KMS