from PIL import Image
import sys
import types

def tobits(s):
    result = []
    for c in s:
        bits = bin(ord(c))[2:]
        bits = '00000000'[len(bits):] + bits
        result.extend([int(b) for b in bits])
    return result

def frombits(bits):
    chars = []
    for b in range(len(bits) / 8):
        byte = bits[b*8:(b+1)*8]
        chars.append(chr(int(''.join([str(bit) for bit in byte]), 2)))
    return ''.join(chars)

def clear_last_two_bits(num):
    #black and white
    if type(num) is types.IntType:
        return num - (num % 4)
    #RGB
    if type(num) is types.TupleType:
        R, G, B = num
        return (R - (R % 4),G - (G % 4),B - (B % 4))

def encode_message(pixel, message):
    msg = 2 * message[0] + message[1]

    #black and white
    if type(pixel) is types.IntType:
        return pixel + msg
    #RGB
    if type(pixel) is types.TupleType:
        R, G, B = pixel
        if R >= G and R >= B:
            return (R + msg, G, B)
        if G >= R and G >= B:
            return (R, G + msg, B)
        if B >= G and B >= R:
            return (R, G, B + msg)


def decode_message(pixel):
    if type(pixel) is types.IntType:
        return "{0:02b}".format(pixel%4)
    #RGB
    if type(pixel) is types.TupleType:
        R, G, B = pixel
        if R >= G and R >= B:
            return  "{0:02b}".format(R%4)
        if G >= R and G >= B:
            return  "{0:02b}".format(G%4)
        if B >= G and B >= R:
            return  "{0:02b}".format(B%4)

IMAGE_PATH = "./stegano.bmp"
MESSAGE = "KYN_2016_Steganography"
MESG_SIZE = sys.getsizeof(MESSAGE)
encoding = 'utf-8'

im = Image.open(IMAGE_PATH)

pixels = im.load()

width, height = im.size

number_of_pixels = width * height
number_of_bytes = (number_of_pixels * 4) / 8
if MESG_SIZE > number_of_bytes:
    print ("Image too small for the message")

message = tobits(MESSAGE)

length = len(message)

extracted_msg = ""

for x in range(0,width):
    for y in range(0,height):
        bits_encoded = 2 *(x*height + y)
        if bits_encoded < len(message):
            sub_message = message[bits_encoded:bits_encoded + 2]
        else:
            sub_message = 2*[0]

        pixels[x,y] = clear_last_two_bits(pixels[x,y])
        pixels[x,y] = encode_message(pixels[x,y], sub_message)


for x in range(0,width):
    for y in range(0,height):
        extracted_msg = extracted_msg + decode_message(pixels[x,y])

im.save('./stegano2.bmp')

message_str = frombits(message)
message2 = frombits(extracted_msg)

print (message_str)
print (message2)
