Python Cryptography ExampleLast weekend I wrote some code to help people use PyCrypto in a canonical fashion to symmetrically encrypt and decrypt strings and files. By canonical I mean I've followed Colin Percival's advice in Cryptographic Right Answers.The code isn't very verbose, and I'd appreciate any comments or questions about the code, the tests, and the documentation. After I get this review I'm going to submit this to Security Stack Exchange for a design review.Source code and a subset of tests follow:crypto.py:# ---------------------------------------------------------------------------# Copyright (c) 2012 Asim Ihsan (asim dot ihsan at gmail dot com)# Distributed under the MIT/X11 software license, see the accompanying# file license.txt or http://www.opensource.org/licenses/mit-license.php.# ---------------------------------------------------------------------------import osimport sysimport structimport cStringIO as StringIOimport bz2from Crypto.Cipher import AESfrom Crypto.Hash import SHA256, HMACfrom Crypto.Protocol.KDF import PBKDF2# ----------------------------------------------------------------------------#   Constants.# ----------------------------------------------------------------------------# Length of salts in bytes.salt_length_in_bytes = 16# Hash function to use in general.hash_function = SHA256# PBKDF pseudo-random function. Used to mix a password and a salt.# See Crypto\Protocol\KDF.pypbkdf2_prf = lambda p, s: HMAC.new(p, s, hash_function).digest()# PBKDF count, number of iterations.pbkdf2_count = 1000# PBKDF derived key length.pbkdf2_dk_len = 32# ----------------------------------------------------------------------------class HMACIsNotValidException(Exception):    passclass InvalidFormatException(Exception):    def __init__(self, reason):        self.reason = reason    def __str__(self):        return repr(self.reason)class CTRCounter:     Callable class that returns an iterating counter for PyCrypto    AES in CTR mode.    def __init__(self, nonce):         Initialize the counter object.        @nonce      An 8-byte binary string.                assert(len(nonce)==8)        self.nonce = nonce        self.cnt = 0    def __call__(self):         Return the next 16 byte counter, as a binary string.         right_half = struct.pack('>Q', self.cnt)        self.cnt += 1        return self.nonce + right_halfdef encrypt_string(plaintext, key, compress=False):    plaintext_obj = StringIO.StringIO(plaintext)    ciphertext_obj = StringIO.StringIO()    encrypt_file(plaintext_obj, key, ciphertext_obj, compress=compress)    return ciphertext_obj.getvalue()def decrypt_string(ciphertext, key):    plaintext_obj = StringIO.StringIO()    ciphertext_obj = StringIO.StringIO(ciphertext)    decrypt_file(ciphertext_obj, key, plaintext_obj)    return plaintext_obj.getvalue()def decrypt_file(ciphertext_file_obj,                 key,                 plaintext_file_obj,                 chunk_size=4096):    # ------------------------------------------------------------------------    #   Unpack the header values from the ciphertext.    # ------------------------------------------------------------------------    header_format = >HHHHHQ?H    header_size = struct.calcsize(header_format)    header_string = ciphertext_file_obj.read(header_size)    try:        header = struct.unpack(header_format, header_string)    except struct.error:        raise InvalidFormatException(Header is invalid.)    pbkdf2_count = header[0]    pbkdf2_dk_len = header[1]    password_salt_size = header[2]    nonce_size = header[3]    hmac_salt_size = header[4]    ciphertext_size = header[5]    compress = header[6]    hmac_size = header[7]    # ------------------------------------------------------------------------    #   Unpack everything except the ciphertext and HMAC, which are the    #   last two strings in the ciphertext file.    # ------------------------------------------------------------------------    encrypted_string_format = ''.join([>,                                       %ss % (password_salt_size, ) ,                                       %ss % (nonce_size, ),                                       %ss % (hmac_salt_size, )])    encrypted_string_size = struct.calcsize(encrypted_string_format)    body_string = ciphertext_file_obj.read(encrypted_string_size)    try:        body = struct.unpack(encrypted_string_format, body_string)    except struct.error:        raise InvalidFormatException(Start of body is invalid.)    password_salt = body[0]    nonce = body[1]    hmac_salt = body[2]    # ------------------------------------------------------------------------    # ------------------------------------------------------------------------    #   Prepare the HMAC with everything except the ciphertext.    #    #   Notice we do not HMAC the ciphertext_size, just like the encrypt    #   stage.    # ------------------------------------------------------------------------    hmac_password_derived = PBKDF2(password = key,                                   salt = hmac_salt,                                   dkLen = pbkdf2_dk_len,                                   count = pbkdf2_count,                                   prf = pbkdf2_prf)    elems_to_hmac = [str(pbkdf2_count),                     str(pbkdf2_dk_len),                     str(len(password_salt)),                     password_salt,                     str(len(nonce)),                     nonce,                     str(len(hmac_salt)),                     hmac_salt]    hmac_object = HMAC.new(key = hmac_password_derived,                           msg = ''.join(elems_to_hmac),                           digestmod = hash_function)    # ------------------------------------------------------------------------    # ------------------------------------------------------------------------    #   First pass: stream in the ciphertext object into the HMAC object    #   and verify that the HMAC is correct.    #    #   Notice we don't need to decompress anything here even if compression    #   is in use. We're using Encrypt-Then-MAC.    # ------------------------------------------------------------------------    ciphertext_file_pos = ciphertext_file_obj.tell()    ciphertext_bytes_read = 0    while True:        bytes_remaining = ciphertext_size - ciphertext_bytes_read        current_chunk_size = min(bytes_remaining, chunk_size)        ciphertext_chunk = ciphertext_file_obj.read(current_chunk_size)        if ciphertext_chunk == '':            break        ciphertext_bytes_read += len(ciphertext_chunk)        hmac_object.update(ciphertext_chunk)    if ciphertext_bytes_read != ciphertext_size:        raise InvalidFormatException(first pass ciphertext_bytes_read %s != ciphertext_size %s % (ciphertext_bytes_read, ciphertext_size))    # the rest of the file is the HMAC.    hmac = ciphertext_file_obj.read()    if len(hmac) != hmac_size:        raise InvalidFormatException(len(hmac) %s != hmac_size %s % (len(hmac), hmac_size))    hmac_calculated = hmac_object.digest()    if hmac != hmac_calculated:        raise HMACIsNotValidException    # ------------------------------------------------------------------------    # ------------------------------------------------------------------------    #   Second pass: stream in the ciphertext object and decrypt it into the    #   plaintext object.    # ------------------------------------------------------------------------    cipher_password_derived = PBKDF2(password = key,                                     salt = password_salt,                                     dkLen = pbkdf2_dk_len,                                     count = pbkdf2_count,                                     prf = pbkdf2_prf)    cipher_ctr = AES.new(key = cipher_password_derived,                         mode = AES.MODE_CTR,                         counter = CTRCounter(nonce))    ciphertext_file_obj.seek(ciphertext_file_pos, os.SEEK_SET)    ciphertext_bytes_read = 0    if compress:        decompressor = bz2.BZ2Decompressor()    while True:        bytes_remaining = ciphertext_size - ciphertext_bytes_read        current_chunk_size = min(bytes_remaining, chunk_size)        ciphertext_chunk = ciphertext_file_obj.read(current_chunk_size)        end_of_file = ciphertext_chunk == ''        ciphertext_bytes_read += len(ciphertext_chunk)        plaintext_chunk = cipher_ctr.decrypt(ciphertext_chunk)        if compress:            try:                decompressed = decompressor.decompress(plaintext_chunk)            except EOFError:                decompressed =             plaintext_chunk = decompressed        plaintext_file_obj.write(plaintext_chunk)        if end_of_file:            break    if ciphertext_bytes_read != ciphertext_size:        raise InvalidFormatException(second pass ciphertext_bytes_read %s != ciphertext_size %s % (ciphertext_bytes_read, ciphertext_size))    # ------------------------------------------------------------------------def encrypt_file(plaintext_file_obj,                 key,                 ciphertext_file_obj,                 chunk_size = 4096,                 compress = False):    # ------------------------------------------------------------------------    #   Prepare input key.    # ------------------------------------------------------------------------    password_salt = os.urandom(salt_length_in_bytes)    cipher_password_derived = PBKDF2(password = key,                                     salt = password_salt,                                     dkLen = pbkdf2_dk_len,                                     count = pbkdf2_count,                                     prf = pbkdf2_prf)    # ------------------------------------------------------------------------    # ------------------------------------------------------------------------    #   Prepare cipher object.    # ------------------------------------------------------------------------    nonce_size = 8    nonce = os.urandom(nonce_size)    cipher_ctr = AES.new(key = cipher_password_derived,                         mode = AES.MODE_CTR,                         counter = CTRCounter(nonce))    # ------------------------------------------------------------------------    # ------------------------------------------------------------------------    #   Prepare HMAC object, and hash what we have so far.    #    #   Notice that we do not HMAC the size of the ciphertext. We don't    #   know how big it'll be until we compress it, if we do, and we can't    #   compress it without reading it into memory. So let the HMAC of    #   the ciphertext itself do.    # ------------------------------------------------------------------------    hmac_salt = os.urandom(salt_length_in_bytes)    hmac_password_derived = PBKDF2(password = key,                                   salt = hmac_salt,                                   dkLen = pbkdf2_dk_len,                                   count = pbkdf2_count,                                   prf = pbkdf2_prf)    elems_to_hmac = [str(pbkdf2_count),                     str(pbkdf2_dk_len),                     str(len(password_salt)),                     password_salt,                     str(len(nonce)),                     nonce,                     str(len(hmac_salt)),                     hmac_salt]    hmac_object = HMAC.new(key = hmac_password_derived,                           msg = ''.join(elems_to_hmac),                           digestmod = hash_function)    # ------------------------------------------------------------------------    # ------------------------------------------------------------------------    #   Write in what we have so far into the output, ciphertext file.    #    #   Given that the plaintext may be compressed we don't know what    #   it's final length will be without compressing it, and we can't    #   do this without reading it all into memory. Hence let's    #   put 0 as the ciphertext length for now and fill it in after.    # ------------------------------------------------------------------------    header_format = ''.join([>,                             H,                               # PBKDF2 count                             H,                               # PBKDF2 derived key length                             H,                               # Length of password salt                             H,                               # Length of CTR nonce                             H,                               # Length of HMAC salt                             Q,                               # Length of ciphertext                             ?,                               # Is compression used?                             H,                               # Length of HMAC                             %ss % (len(password_salt), ) ,   # Password salt                             %ss % (nonce_size, ),            # CTR nonce                             %ss % (len(hmac_salt), )])       # HMAC salt    header = struct.pack(header_format,                         pbkdf2_count,                         pbkdf2_dk_len,                         len(password_salt),                         len(nonce),                         len(hmac_salt),                         0,                                     # This is the ciphertext size, wrong for now.                         compress,                         hmac_object.digest_size,                         password_salt,                         nonce,                         hmac_salt)    ciphertext_file_obj.write(header)    # ------------------------------------------------------------------------    # ------------------------------------------------------------------------    #   Stream in the input file and stream out ciphertext into the    #   ciphertext file.    # ------------------------------------------------------------------------    ciphertext_size = 0    if compress:        compressor = bz2.BZ2Compressor()    while True:        plaintext_chunk = plaintext_file_obj.read(chunk_size)        end_of_file = plaintext_chunk == ''        if compress:            if end_of_file:                compressed = compressor.flush()            else:                compressed = compressor.compress(plaintext_chunk)            plaintext_chunk = compressed        ciphertext_chunk = cipher_ctr.encrypt(plaintext_chunk)        ciphertext_size += len(ciphertext_chunk)        ciphertext_file_obj.write(ciphertext_chunk)        hmac_object.update(ciphertext_chunk)        if end_of_file:            break    # ------------------------------------------------------------------------    # ------------------------------------------------------------------------    #   Write the HMAC to the ciphertext file.    # ------------------------------------------------------------------------    hmac = hmac_object.digest()    ciphertext_file_obj.write(hmac)    # ------------------------------------------------------------------------    # ------------------------------------------------------------------------    #   Go back to the header and update the ciphertext size.    #    #   Notice that we capture the header such that the last unpacked    #   element of the struct is the unsigned long long of the ciphertext    #   length.    # ------------------------------------------------------------------------    # Read in.    ciphertext_file_obj.seek(0, os.SEEK_SET)    header_format = >HHHHHQ    header_size = struct.calcsize(header_format)    header = ciphertext_file_obj.read(header_size)    # Modify.    header_elems = list(struct.unpack(header_format, header))    header_elems[-1] = ciphertext_size    # Write out.    header = struct.pack(header_format, *header_elems)    ciphertext_file_obj.seek(0, os.SEEK_SET)    ciphertext_file_obj.write(header)    # ------------------------------------------------------------------------test_crypto.py:#!/usr/bin/env python2.7# ---------------------------------------------------------------------------# Copyright (c) 2012 Asim Ihsan (asim dot ihsan at gmail dot com)# Distributed under the MIT/X11 software license, see the accompanying# file license.txt or http://www.opensource.org/licenses/mit-license.php.# ---------------------------------------------------------------------------import osimport sysimport tempfileimport cStringIO as StringIOsrc_path = os.path.abspath(os.path.join(__file__, os.pardir, os.pardir, src))assert(os.path.isdir(src_path))sys.path.append(src_path)from utilities import cryptofrom nose.tools import raises, assert_false, assert_true, assert_not_equal, assert_equal, assert_lessfrom nose.plugins.skip import SkipTest, Skip_multiprocess_can_split_ = Truedef test_encrypt_then_decrypt_string():    plaintext = this is some text    key = this is my key    ciphertext = crypto.encrypt_string(plaintext, key)    plaintext_after = crypto.decrypt_string(ciphertext, key)    assert_equal(plaintext_after, plaintext)@raises(crypto.HMACIsNotValidException)def test_encrypt_then_alter_raises_exception():    plaintext = this is some text    key = this is my key    ciphertext = crypto.encrypt_string(plaintext, key)    ciphertext = ciphertext[:-len(plaintext)] + '\0' * len(plaintext)    plaintext_after = crypto.decrypt_string(ciphertext, key)@raises(crypto.InvalidFormatException)def test_decrypt_then_damage_raises_exception():    plaintext = this is some text    key = this is my key    ciphertext = crypto.encrypt_string(plaintext, key)    ciphertext = ciphertext[:len(ciphertext)-5]    plaintext_after = crypto.decrypt_string(ciphertext, key)@raises(crypto.HMACIsNotValidException)def test_decrypt_with_wrong_password_raises_exception():    plaintext = this is some text    key = this is my key    ciphertext = crypto.encrypt_string(plaintext, key)    another_key = this is my key 2    plaintext_after = crypto.decrypt_string(ciphertext, another_key)def test_encrypt_then_decrypt_empty_string():    plaintext =     key = this is my key    ciphertext = crypto.encrypt_string(plaintext, key)    plaintext_after = crypto.decrypt_string(ciphertext, key)    assert_equal(plaintext_after, plaintext)def test_compressed_encrypt_then_decrypt_string():    plaintext = X * 4096    key = this is my key    ciphertext = crypto.encrypt_string(plaintext, key, compress=True)    assert_less(len(ciphertext), len(plaintext) / 10)    plaintext_after = crypto.decrypt_string(ciphertext, key)    assert_equal(plaintext, plaintext_after)def test_compressed_encrypt_then_decrypt_random_string():    plaintext = os.urandom(1024 * 1024)    key = this is my key    ciphertext = crypto.encrypt_string(plaintext, key, compress=True)    plaintext_after = crypto.decrypt_string(ciphertext, key)    assert_equal(plaintext, plaintext_after)def _test_encrypt_then_decrypt_file(plaintext_size,                                    chunk_size,                                    wrong_password=False,                                    alter_file=False,                                    truncate_header=False,                                    truncate_body=False,                                    compress=False):    tempfile.tempdir = os.path.join(__file__, os.pardir)    plaintext_file = tempfile.NamedTemporaryFile(delete=False)    ciphertext_file = tempfile.NamedTemporaryFile(delete=False)    plaintext_after_file = tempfile.NamedTemporaryFile(delete=False)    plaintext_filepath = plaintext_file.name    ciphertext_filepath = ciphertext_file.name    plaintext_after_filepath = plaintext_after_file.name    plaintext_file.close()    ciphertext_file.close()    plaintext_after_file.close()    try:        key = this is my key        # --------------------------------------------------------------------        #   Write plaintext to file.        # --------------------------------------------------------------------        with open(plaintext_filepath, wb) as f:            cnt = 0            while cnt < plaintext_size:                current_chunk_size = min(4096, plaintext_size - cnt)                f.write(X * current_chunk_size)                cnt += current_chunk_size        # --------------------------------------------------------------------        # --------------------------------------------------------------------        #   Encrypt plaintext file to ciphertext file.        #        #   Notice that the output, ciphertext file requires read and        #   write access.        # --------------------------------------------------------------------        with open(plaintext_filepath, rb) as f_in:            with open(ciphertext_filepath, rb+) as f_out:                crypto.encrypt_file(f_in,                                    key,                                    f_out,                                    chunk_size=chunk_size,                                    compress=compress)        # --------------------------------------------------------------------        # --------------------------------------------------------------------        #   If wrong password then let's adjust the password.        # --------------------------------------------------------------------        if wrong_password:            key = this is my key 2        # --------------------------------------------------------------------        # --------------------------------------------------------------------        #   If alter file then let's alter the file.        #   If truncate_body then let's skip the last ten bytes from the        #   file.        #   If truncate_header then let's skip the first ten bytes from the        #   file.        # --------------------------------------------------------------------        if alter_file:            with open(ciphertext_filepath, rb+) as f:                f.seek(-10, os.SEEK_END)                f.write('\0' * 10)        if truncate_header or truncate_body:            with open(ciphertext_filepath, rb) as f:                contents = f.read()            if truncate_header:                with open(ciphertext_filepath, wb) as f:                    f.write(contents[10:])            else:                with open(ciphertext_filepath, wb) as f:                    f.write(contents[:-10])        # --------------------------------------------------------------------        # --------------------------------------------------------------------        #   Decrypt ciphertext file to another plaintext file.        # --------------------------------------------------------------------        with open(ciphertext_filepath, rb) as f_in:            with open(plaintext_after_filepath, wb) as f_out:                crypto.decrypt_file(f_in, key, f_out, chunk_size=chunk_size)        # --------------------------------------------------------------------        with open(plaintext_filepath, rb) as f_original:            with open(plaintext_after_filepath, rb) as f_after:                while True:                    f_original_chunk = f_original.read(chunk_size)                    f_after_chunk = f_after.read(chunk_size)                    assert_equal(f_original_chunk, f_after_chunk)                    if f_original_chunk == '':                        break    finally:        for (file_obj, filepath) in [(plaintext_file, plaintext_filepath),                                     (ciphertext_file, ciphertext_filepath),                                     (plaintext_after_file, plaintext_after_filepath)]:            os.remove(filepath)def test_encrypt_then_decrypt_file_normal():    _test_encrypt_then_decrypt_file(plaintext_size = 17,                                    chunk_size = 4096)
Canonical Python symmetric cryptography example
python;programming challenge;security;cryptography