stegano.py

#!/usr/bin/env python
"""Quickly written steganography tool."""
# Arguments(d): input.png input.bin output.png
# Arguments(e): input.png output.bin(stdout if nothing)
import sys
from PIL import Image

def load(image):
    if image.format != 'PNG':
        raise ValueError('File format should be PNG')
    if image.mode != 'RGBA':
        raise ValueError('Image mode should be RGBA')
    return image.load()

def putbits(byte, bits: '0bXX'):
    bits &= 0b00000011
    byte &= 0b11111100
    return byte | bits

def getbits(byte):
    return byte & 0b11

def encode(image, data):
    pix = load(image)
    image_size = image.width * image.height
    if len(data) + 8 > image_size:
        raise IndexError('Image is %sx%spx, data is %s bytes' % (image.width, image.height, len(data)))
    size = len(data).to_bytes(length=4, byteorder='big')
    data = size + data
    x, y, i = 0, 0, 0
    while i < len(data):
        pixel = pix[x, y]
        byte = data[i]
        newpix = tuple((putbits(pixel[k], byte >> ((3-k)*2)) for k in range(4)))
        pix[x, y] = newpix
        x+=1
        if x == image.width:
            x = 0
            y += 1
        i += 1
    return image


def decode(image):
    pix = load(image)
    header = True
    size = 4
    data = []
    x, y = 0, 0
    while len(data) < size:
        pixel = pix[x, y]
        bits = (getbits(pixel[k]) << ((3-k)*2) for k in range(4))
        data.append(sum(bits))

        x+=1
        if x == image.width:
            x = 0
            y += 1
        if header and len(data) == 4:
            header = False
            size = int.from_bytes(data, byteorder='big')
            data = []
    return bytes(data)


def main():
    if not 1 < len(sys.argv) <= 4:
        print('Encoding args: input.png input.bin output.png')
        print('Decoding args: input.png [output.bin]')
    elif len(sys.argv) <= 3:
        image = Image.open(sys.argv[1])
        data = decode(image)
        out = sys.stdout.buffer
        if len(sys.argv) == 3:
            out = open(sys.argv[2], 'wb')
        out.write(data)
        out.flush()
    else:
        image = Image.open(sys.argv[1])
        data = open(sys.argv[2], 'rb').read()
        output = encode(image, data)
        output.save(sys.argv[3])


if __name__ == '__main__':
    main()