AES using ECB Demo using Python

Here's a little python program that I've knocked together inspired by this article, to demonstrate the serious gotcha involved with using ECB (Electronic Code Book) as a cipher mode.

This code uses the PIL (Python Imaging Library) (I used v1.1.6) and PyCrypto (I used the precompiled binary from here).

These modules are designed for Python 2.7, so that's what I used.

The code:

import Image
import sys
import os
from Crypto.Cipher import AES

    IV_SIZE = 16
    BLOCK_SIZE = 16

    def check_args():
        try:
            if (len(sys.argv) != 4):
                raise Exception()
            elif (not os.path.isfile(sys.argv[1])):
                raise Exception("Input file must exist")
            elif (not sys.argv[3] in ['CBC', 'ECB']):
                raise Exception("Block cipher mode should be ECB or CBC")
            return (sys.argv[1], sys.argv[2], sys.argv[3])
        except Exception, ex:
            print "Usage:", sys.argv[0], \
              "full_path_to_input_image full_path_to_output_image ECB|CBC"
            if len(ex.args) > 0:
                print "--" + str(ex)
            sys.exit(1)


    def encrypt(input_filename, output_filename, cipher_mode):
        """Encrypt an image file and write out the results as a JPEG."""

        input_image = Image.open(input_filename)

        # Key must be one of 16/24/32 bytes in length.
        key = "0123456789ABCDEF"
        mode = AES.MODE_CBC if cipher_mode == 'CBC' else AES.MODE_ECB
        iv = os.urandom(IV_SIZE)

        aes = AES.new(key, mode, iv)

        image_string = input_image.tostring()
        # The input string must be padded to the input block size.
        image_padding_length = BLOCK_SIZE - len(image_string) % BLOCK_SIZE
        image_string += image_padding_length * "~"

        # generate the encrypted image string
        encrypted = aes.encrypt(image_string)

        # create an image from the encrypted string
        encrypted_img = Image.frombuffer("RGB", input_image.size, encrypted, 'raw',
                                         "RGB", 0, 1)

        # create and save the output image
        encrypted_img.save(output_filename, 'JPEG')

        print("Encrypted using AES in " + cipher_mode + " mode and saved to \"" +
               output_filename + "\"!")


    if __name__ == "__main__":
        args = check_args()
        encrypt(*args)

And the test image, and resulting encrypted image, using first ECB then CBC (Cipher Block Chaining).

Input:

test-input Encrypted with ECB: test-ecb Encrypted with CBC: test-cbc

As you can see, ECB outputs the same encrypted data for equivalent input blocks, so the black and white areas are clearly distinguishable. Since CBC has feedback from the previous block's encryption into the next block, the equivalent input blocks are no longer encrypted to the same value, so we see noise.