Compress folder, encrypt, and upload to s3

glacify.py

import os
import sys
import subprocess
from datetime import date
from contextlib import contextmanager
import logging
from cStringIO import StringIO
import boto


logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


@contextmanager
def compress(folder):
    folder = folder.rstrip('/')
    tar = subprocess.Popen(
        ['tar', 'czf', '-', os.path.basename(folder)],
        cwd=os.path.dirname(folder),
        stdout=subprocess.PIPE,
    )
    openssl = subprocess.Popen(
        ['openssl', 'aes-256-cbc', '-salt', '-pass', 'env:GLACIFY_AES_KEY'],
        stdin=tar.stdout,
        stdout=subprocess.PIPE,
    )
    try:
        yield openssl.stdout
    except:
        openssl.kill()
        tar.kill()
        raise
    finally:
        openssl.wait()
        tar.wait()


def iter_chunks(file_stream, block_size=10485760):
    while True:
        block = file_stream.read(block_size)
        if not block:
            break
        yield block


def upload(bucket_name, file_name, file_stream, auth):
    s3 = boto.connect_s3(*auth)
    bucket = s3.get_bucket(bucket_name)
    logger.info("begin multipart upload: %r", file_name)
    multipart = bucket.initiate_multipart_upload(file_name)
    try:
        for number, chunk in enumerate(iter_chunks(file_stream), 1):
            logger.info("uploading chunk %d (%d)", number, len(chunk))
            multipart.upload_part_from_file(StringIO(chunk), number)
    except:
        logger.info("aborting")
        multipart.cancel_upload()
        raise
    else:
        logger.info("success!")
        multipart.complete_upload()


def main():
    bucket_name = os.environ['GLACIFY_BUCKET']
    auth = (os.environ['GLACIFY_AWS_ID'], os.environ['GLACIFY_AWS_KEY'])
    (folder, basename) = sys.argv[1:]
    file_name = "%s-%s.tgz.aes" % (basename, date.today().isoformat())
    with compress(folder) as tar_file:
        upload(bucket_name, file_name, tar_file, auth)


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    main()