Crypto helper which uses PyCrypto.

crypto.py

import os
import hmac
import hashlib

from Crypto.Cipher import AES


class Crypto(object):
    """Authenticated encryption class
    
    Encryption algorithm: AES-CBC
    Signing algorithm: HMAC-SHA256
    """

    # Taken from http://code.activestate.com/recipes/576980-authenticated-encryption-with-pycrypto/
    # Modified + reviewed by @wolever

    AES_BLOCK_SIZE = 16
    SIG_SIZE = hashlib.sha256().digest_size

    def __init__(self, key_string, key_bytes=32):
        self.keys = self.extract_keys(key_string, key_bytes)
        self.key_bytes = key_bytes

    @classmethod
    def generate_key_string(cls, key_bytes=32):
        key = os.urandom(key_bytes + cls.SIG_SIZE)
        return key.encode("base64").replace("\n", "")

    @classmethod
    def extract_keys(cls, key_string, key_bytes):
        key = key_string.decode("base64")
        assert len(key) == key_bytes + cls.SIG_SIZE
        return (key[:key_bytes], key[key_bytes:])

    @classmethod
    def pad(cls, data, width):
        needed = width - ((len(data) + 1) % width)
        return data + ('\x00' * needed) + chr(needed + 1)

    @classmethod
    def unpad(cls, data):
        padding = ord(data[-1])
        return data[:-padding]

    def encrypt(self, data):
        """encrypt data with AES-CBC and sign it with HMAC-SHA256"""
        aes_key, hmac_key = self.keys
        data = self.pad(data, self.AES_BLOCK_SIZE)
        iv_bytes = os.urandom(self.AES_BLOCK_SIZE)
        cypher = AES.new(aes_key, AES.MODE_CBC, iv_bytes)
        data = iv_bytes + cypher.encrypt(data)
        sig = hmac.new(hmac_key, data, hashlib.sha256).digest()
        return sig + data

    def decrypt(self, data):
        """verify HMAC-SHA256 signature and decrypt data with AES-CBC"""
        aes_key, hmac_key = self.keys
        sig = data[:self.SIG_SIZE]
        data = data[self.SIG_SIZE:]
        if hmac.new(hmac_key, data, hashlib.sha256).digest() != sig:
            raise Exception("message authentication failed")
        iv_bytes = data[:self.AES_BLOCK_SIZE]
        data = data[self.AES_BLOCK_SIZE:]
        cypher = AES.new(aes_key, AES.MODE_CBC, iv_bytes)
        data = cypher.decrypt(data)
        return self.unpad(data)


if __name__ == "__main__":
    # Test padding
    for size in xrange(Crypto.AES_BLOCK_SIZE * 2):
        expected = "a" * size
        padded = Crypto.pad(expected, Crypto.AES_BLOCK_SIZE)
        actual = Crypto.unpad(padded)
        assert expected == actual, "%r != %r" %(padded, actual)
    print "Padding OK."

    key = Crypto.generate_key_string()
    crypto = Crypto(key)
    print "Using key %r" %(key, )
    for size in xrange(Crypto.AES_BLOCK_SIZE * 2):
        expected = "a" * size
        encrypted = crypto.encrypt(expected)
        actual = crypto.decrypt(encrypted)
        assert expected == actual, "%r != %r" %(padded, actual)
    print "Encryption OK."