Little script to demonstrate different cipher modes applied on images

piccrypt.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# author: Timo Schmid
# copyright: 2016
"""
Script to demonstrate different cipher modes
with based on an image
"""
from __future__ import print_function
import sys
import imghdr
import struct
from functools import partial

try:
    from Crypto.Cipher import AES
    from Crypto.Util import RFC1751, Counter, randpool
except ImportError:
    print("PyCrypto required")
    sys.exit(1)

try:
    from PIL import Image
except ImportError:
    print("PILLOW required for full image support. Fallback to BMP")
    Image = None


def pad(data, size=16):
    """
    Pad data based on PKCS#7
    """
    pad_byte = size - (len(data) % size)
    return data + bytes([pad_byte] * pad_byte)


def process(fp):
    """
    Encrypt BMP body with different modes
    """
    if imghdr.what(fp) != 'bmp':
        print("Currently only BITMAP supported")
        sys.exit(1)

    fp.seek(0xA)
    offset = struct.unpack('<I', fp.read(4))[0]
    fp.seek(0x1C)
    bitcount = struct.unpack('<H', fp.read(2))[0]
    if bitcount < 16:
        print("BITMAP uses indexed colors. Results might be undesired")

    fp.seek(0x1E)
    compression = struct.unpack('<I', fp.read(4))[0]
    if compression != 0:
        print("BITMAP uses an unsupported compression")
        sys.exit(1)

    modes = ('ECB', 'CFB', 'ECB', 'OFB', 'CTR')
    pool = randpool.RandomPool()

    for mode in modes:
        fname, ext = fp.name.rsplit('.', 1)
        fname = '{0}.{1}.{2}'.format(
            fname,
            mode.lower(),
            ext
        )
        with open(fname, 'wb') as outf:
            fp.seek(0)
            outf.write(fp.read(offset))

            key = pool.get_bytes(AES.key_size[0])
            print("Mode:", mode)
            print("Key:", RFC1751.key_to_english(key))

            kwargs = {
                'key': key,
                'mode': getattr(AES, "MODE_" + mode),
            }
            if mode in ('CBC', 'CFB', 'OFB'):
                kwargs['IV'] = iv = pool.get_bytes(AES.block_size)
                print("IV:", RFC1751.key_to_english(iv))
            elif mode == 'CTR':
                kwargs['counter'] = Counter.new(AES.block_size * 8)

            cipher = AES.new(**kwargs)

            for chunk in iter(partial(fp.read, AES.block_size), b''):
                if len(chunk) < AES.block_size:
                    chunk = pad(chunk)
                outf.write(cipher.encrypt(chunk))


def process2(source_image):
    """
    Encrypt image data based on different modes
    """
    srcimg = Image.open(source_image)

    modes = ('ECB', 'CFB', 'ECB', 'OFB', 'CTR')
    pool = randpool.RandomPool()

    for mode in modes:
        fname, ext = source_image.rsplit('.', 1)
        fname = '{0}.{1}.{2}'.format(
            fname,
            mode.lower(),
            ext
        )
        key = pool.get_bytes(AES.key_size[0])
        print("Mode:", mode)
        print("Key:", RFC1751.key_to_english(key))

        kwargs = {
            'key': key,
            'mode': getattr(AES, "MODE_" + mode),
        }
        if mode in ('CBC', 'CFB', 'OFB'):
            kwargs['IV'] = iv = pool.get_bytes(AES.block_size)
            print("IV:", RFC1751.key_to_english(iv))
        elif mode == 'CTR':
            kwargs['counter'] = Counter.new(AES.block_size * 8)

        cipher = AES.new(**kwargs)
        data = cipher.encrypt(pad(srcimg.tobytes()))
        dstimg = Image.frombytes(srcimg.mode, srcimg.size, data)
        dstimg.save(fname)


def main():
    """
    Entrypoint
    """
    if len(sys.argv) != 2:
        print("usage: {0} <image>".format(sys.argv[0]))
        sys.exit(1)

    if Image is None:
        with open(sys.argv[1], 'rb') as fp:
            process(fp)
    else:
        process2(sys.argv[1])


if __name__ == "__main__":
    main()