# coding: utf-8

import base64
import Image
import sys

""" This is a simple steganography script
Base64-encoded message will be hidden in a image file
"""

def  bin6bits(n):
  """ Convert given decimal number to 6-bit binary
  >>> bin6bits(10)
  '001010'
  """
  return bin(n)[2:].zfill(6)

def c2b(c):
  """ Convert given character to 6-bit binary
  >>> c2b('0')
  '000000'
  """
  code = ord(c)
  ascii_range = lambda head, tail: ord(head) <= code and code <= ord(tail)
  if ascii_range('0', '9'):
    return bin6bits(code - ord('0'))
  elif ascii_range('A', 'Z'):
    return bin6bits(code - ord('A') + 10)
  elif ascii_range('a', 'z'):
    return bin6bits(code - ord('a') + 10 + 26)
  elif c == '+':
    return bin6bits(62)
  elif c == '/':
    return bin6bits(63)

def hide(c, r, g, b):
  """ Hide given character to last-2bits of r/g/b each
  """
  cb = c2b(c)
  rb, gb, bb = bin6bits(r), bin6bits(g), bin6bits(b)
  rb = rb[:-2] + cb[0:2]
  gb = gb[:-2] + cb[2:4]
  bb = bb[:-2] + cb[4:6]
  return int(rb, 2), int(gb, 2), int(bb, 2)

def encode(src_txt, src_img, dst_img):
  """ Create new image which has hidden message(src_txt)
  """
  txt = open(src_txt).read()
  b64_txt = base64.b64encode(txt)

  img = Image.open(src_img)
  w, h = img.size

  n = len(b64_txt.replace('=',''))
  n_padding = len(b64_txt) - n
  d = len(str(n))

  if (2 + d + n) > w * h:
    print 'Picture size is too small, or text length is too long...'
    quit()

  pxl = img.load()

  # The first pixel shows the number of padding (=)
  r, g, b = pxl[0, 0]
  pxl[0, 0] = hide(str(n_padding), r, g, b)

  # The second pixel indicates digit number of text's length
  r, g, b = pxl[1, 0]
  pxl[1, 0] = hide(str(d), r, g, b)

  # From the third pixel, write the text length
  for i in range(2, d+2):
    c = str(n)[i-2]
    y = i / w
    x = i % w
    r, g, b = pxl[x, y]
    pxl[x, y] = hide(c, r, g, b)

  # After the length information, memorize the text
  for i in range(n):
    c = b64_txt[i]
    if c == '=': break
    y = (2 + d + i) / w
    x = (2 + d + i) % w
    r, g, b = pxl[x, y]
    pxl[x, y] = hide(c, r, g, b)

  img.save(dst_img, 'PNG')
  print '[success] saved as %s\n%d message length + %d padding' % (dst_img, n, n_padding)

def main():
  encode('source.txt', 'sushi.png', 'yogurt.png')

if __name__ == '__main__':
  import doctest
  doctest.testmod()
  main()