steganography.py

import argparse
import sys
from PIL import Image


def encode_subcommand(args):
    im = Image.open(args.input_image)
    pixels = im.load()

    data = args.data.encode('ascii') + b'\0'

    if len(data) > im.size[0] * im.size[1]:
        print("Image not large enough to store data")
        exit(1)

    for i, b in enumerate(data):
        x, y = i // im.size[1], i % im.size[1]

        pixels[x, y] = (
            (pixels[x, y][0] & 0xf8) + (b >> 5),
            (pixels[x, y][1] & 0xf8) + ((b >> 2) & 0x7),
            (pixels[x, y][2] & 0xfc) + (b & 0x3),
        )

    im.save(args.output_image)


def decode_subcommand(args):
    im = Image.open(args.image)
    pixels = im.load()

    data = []

    for x in range(im.size[0]):
        for y in range(im.size[1]):
            b = ((pixels[x, y][0] & 0x7) << 5) + ((pixels[x, y][1] & 0x7) << 2) + (pixels[x, y][2] & 0x3)
            if b == 0:
                break

            data.append(chr(b))

    print(''.join(data))


def main(argv):
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(required=True)

    encode_parser = subparsers.add_parser("encode")
    encode_parser.add_argument("data")
    encode_parser.add_argument("input_image")
    encode_parser.add_argument("output_image")
    encode_parser.set_defaults(func=encode_subcommand)

    decode_parser = subparsers.add_parser("decode")
    decode_parser.add_argument("image")
    decode_parser.set_defaults(func=decode_subcommand)

    args = parser.parse_args(sys.argv[1:])

    args.func(args)


if __name__ == "__main__":
    main(sys.argv)