Stegosaurus — let's steganograpy together.

stego.py

#!/usr/bin/python

import argparse
import hashlib
import itertools
import struct
import sys
import zlib

from PIL import Image

from Crypto.Cipher import AES   # pip install pycrypto
from Crypto import Random


def conceal(im, plaintext, password):
    """
    Compresses and encrypts data, and hides it in an image.
    """

    # Compress data and pad it out
    payload = zlib.compress(plaintext, 9)
    paylen = len(payload)
    if paylen % AES.block_size:
        payload += '\x00' * (AES.block_size - paylen % AES.block_size)

    # See if it will fit.
    blocks = im.size[0] // 4 * im.size[1] // 4 - 3 # Two meta blocks, one IV
    datablocks = len(payload) // AES.block_size

    if blocks < datablocks:
        raise Exception("Too much data for this image! " +
                        "Can store %d blocks, got %d." % (blocks, datablocks))

    # Encrypt our data
    key = hashlib.sha256(password).digest()
    iv = Random.new().read(AES.block_size)
    ciphertext = AES.new(key, AES.MODE_CFB, iv).encrypt(payload)

    # Meta blocks
    md5hash = hashlib.md5(password).digest()
    lengths = struct.pack("!QQ", datablocks, paylen)

    return insert(im, md5hash + lengths + iv + ciphertext)

def reveal(im, password):
    """
    Extracts concealed message from an image.
    """
    data = extract(im)

    # Meta blocks
    md5hash = data.next()
    datablocks, datalen = struct.unpack("!QQ", data.next())

    # Validate password
    if md5hash != hashlib.md5(password).digest():
        raise ValueError("Not the right password!")
    key = hashlib.sha256(password).digest()

    # Decrypt and unpad
    iv = data.next()
    ciphertext = ''.join(itertools.islice(data, datablocks))
    plaintext = AES.new(key, AES.MODE_CFB, iv).decrypt(ciphertext)[:datalen]

    return zlib.decompress(plaintext)


def insert(im, data):
    """Actually places data in an image."""

    pix = im.load()

    data = list(data)
    data.reverse()

    for i in range(im.size[1] // 4):
        for j in range(im.size[0] // 4):
            if not data:
                break
            for x in range(4):
                for y in range(4):
                    byte = ord(data.pop())
                    x_c = j * 4 + x
                    y_c = i * 4 + y
                    r, g, b, a = pix[x_c, y_c]
                    pix[x_c, y_c] = (r & 0xfc | ((byte & 0b11000000) >> 6),
                                     g & 0xfc | ((byte & 0b00110000) >> 4),
                                     b & 0xfc | ((byte & 0b00001100) >> 2),
                                     a & 0xfc |  (byte & 0b00000011))
    return im

def extract(im):
    """Extracts raw data from image, one block at a time."""

    pix = im.load()

    for i in range(im.size[1] // 4):
        for j in range(im.size[0] // 4):
            block = []
            for x in range(4):
                for y in range(4):
                    r, g, b, a = pix[j * 4 + x, i * 4 + y]
                    byte = 0
                    byte |= (r & 0x3) << 6
                    byte |= (g & 0x3) << 4
                    byte |= (b & 0x3) << 2
                    byte |= (a & 0x3)
                    block.append(chr(byte))
            yield ''.join(block)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Stegosaurus -- let's steganograpy together."
    )
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--conceal', '-c', action='store_true', dest='conceal',
                       help='conceal data in an image')
    group.add_argument('--reveal', '-r', action='store_false', dest='conceal',
                       help='reveal data stored in an image')
    parser.add_argument('--password', '-p', metavar='PASS', default='',
                        help='password to encrypt your message')
    parser.add_argument('--output', '-o', metavar='FILE',
                        help='output filename (omit for in-place (conceal ' +
                             'mode) or stdout (reveal mode))')
    parser.add_argument('--message', '-m', metavar='FILE',
                        type=open, default=sys.stdin,
                        help='[conceal mode only] file containing your ' +
                             'message (omit for stdin)')
    parser.add_argument('image', metavar='IMG', help='the vessel image')
    args = parser.parse_args()

    try:
        im = Image.open(args.image).convert('RGBA')
    except IOError:
        print "Can't open %s. Please review." % args.image
        sys.exit(1)

    if args.conceal:
        im = conceal(im, args.message.read(), args.password)
        im.save(args.image if args.output is None else args.output, 'PNG')
        print 'Done.'
    else:
        message = reveal(im, args.password)
        if args.output is None:
            print message
        else:
            with open(args.output, 'w') as o:
                o.write(message)
            print 'Done.'