certificates.py

#!/usr/bin/env python
"""
Run cumanage operations in an AWS hosted stack.
"""

import json
import sys
import shellish
import boto3
from . import base
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Util import Counter
from pwgen import pwgen

AWS_REGIONS = [
    'us-east-2',
    'us-west-2',
]
SALT = 'LTzkm1w/p3ReDm9kmfmnwQ=='
ENCODING = 'utf-8'
AES_KEY_BYTES = 32


def __derive_key__(password, salt=None):
    '''Return key using PBKDF2'''
    if not salt:
        salt = SALT
    return PBKDF2(
        password,
        salt.encode(ENCODING),
        AES_KEY_BYTES
    )


def __encrypt_data__(password, salt, data):
    '''Encrypt given data using password and PBKDF2 key stretching'''
    key = __derive_key__(password, salt)
    aes_cipher = AES.new(key, AES.MODE_CTR, counter=Counter.new(128))
    return aes_cipher.encrypt(data.encode(ENCODING))


class CertificateCommand(base.HostingCommand):
    '''Manage Client Certificates'''
    name = 'certificate'

    def setup_args(self, parser):
        self.add_subcommand(UploadCertificateCommand)


class UploadCertificateCommand(base.HostingCommand):
    '''Upload Client Certificates to AWS'''

    name = 'upload'

    def setup_args(self, parser):
        self.add_argument('--secret-id', action='store', required=True,
                          help='Path for secret storage')
        self.add_argument('--certificate-file', action='store',
                          help='Certificate file to encrypt and upload')
        self.add_argument('--output', '-o', action='store',
                          help='Name of output file if specified')

    def run(self, args):

        def read_file(filename):
            with open(filename, 'r') as fh:
                return fh.read()

        secret_id = args.secret_id
        if args.certificate_file:
            certificate = read_file(args.certificate_file)
        else:
            certificate = sys.stdin.read()
        salt = pwgen(16)
        private_key = pwgen(50)

        encrypted_certificate = __encrypt_data__(private_key,
                                                 salt,
                                                 certificate)

        if args.output:
            with open(args.output, 'wb') as output:
                output.write(encrypted_certificate)
        else:
            print(encrypted_certificate)

        shellish.vtmlprint(
            '<green>Encrypted and Uploaded Certificate</green>',
            file=sys.stderr
        )

        secret_arns = update_secrets_value(secret_id, salt, private_key)

        shellish.vtmlprint(
            '<green>Created/Updated Encryption Keys: %s</green>' % (
                ', '.join(secret_arns)),
            file=sys.stderr
        )


def update_secrets_value(secret_id, salt, private_key):
    '''Update new secrets value for `secret_id`'''
    arns = []
    for region in AWS_REGIONS:
        client = boto3.client('secretsmanager', region_name=region)
        secret_dictionary = get_secret_value(client, secret_id)

        secret_dictionary['salt'] = salt
        secret_dictionary['password'] = private_key

        arn = put_secret_value(client, secret_id, secret_dictionary)
        arns.append(arn)

    return arns


def get_secret_value(client, secret_id):
    '''Return the JSON data for the provided `secret_id`

    If there is no secret for the provided `secret_id` return an empty
    dictionary.
    '''
    try:
        response = client.get_secret_value(SecretId=secret_id)
        return json.loads(response['SecretString'])
    except client.exceptions.ResourceNotFoundException:
        return {}


def put_secret_value(client, secret_id, secrets):
    '''Create or Update secret value of the provided `secret_id`'''
    json_secrets = json.dumps(secrets)
    if get_secret_value(client, secret_id) == {}:
        response = client.create_secret(
            Name=secret_id,
            SecretString=json_secrets
        )
    else:
        response = client.put_secret_value(
            SecretId=secret_id,
            SecretString=json.dumps(secrets)
        )
    return response['ARN']