Zeroji

stegano.py

Sep 10th, 2018
151
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #!/usr/bin/env python
  2. """Quickly written steganography tool."""
  3. # Arguments(d): input.png input.bin output.png
  4. # Arguments(e): input.png output.bin(stdout if nothing)
  5. import sys
  6. from PIL import Image
  7.  
  8. def load(image):
  9.     if image.format != 'PNG':
  10.         raise ValueError('File format should be PNG')
  11.     if image.mode != 'RGBA':
  12.         raise ValueError('Image mode should be RGBA')
  13.     return image.load()
  14.  
  15. def putbits(byte, bits: '0bXX'):
  16.     bits &= 0b00000011
  17.     byte &= 0b11111100
  18.     return byte | bits
  19.  
  20. def getbits(byte):
  21.     return byte & 0b11
  22.  
  23. def encode(image, data):
  24.     pix = load(image)
  25.     image_size = image.width * image.height
  26.     if len(data) + 8 > image_size:
  27.         raise IndexError('Image is %sx%spx, data is %s bytes' % (image.width, image.height, len(data)))
  28.     size = len(data).to_bytes(length=4, byteorder='big')
  29.     data = size + data
  30.     x, y, i = 0, 0, 0
  31.     while i < len(data):
  32.         pixel = pix[x, y]
  33.         byte = data[i]
  34.         newpix = tuple((putbits(pixel[k], byte >> ((3-k)*2)) for k in range(4)))
  35.         pix[x, y] = newpix
  36.         x+=1
  37.         if x == image.width:
  38.             x = 0
  39.             y += 1
  40.         i += 1
  41.     return image
  42.  
  43.  
  44. def decode(image):
  45.     pix = load(image)
  46.     header = True
  47.     size = 4
  48.     data = []
  49.     x, y = 0, 0
  50.     while len(data) < size:
  51.         pixel = pix[x, y]
  52.         bits = (getbits(pixel[k]) << ((3-k)*2) for k in range(4))
  53.         data.append(sum(bits))
  54.  
  55.         x+=1
  56.         if x == image.width:
  57.             x = 0
  58.             y += 1
  59.         if header and len(data) == 4:
  60.             header = False
  61.             size = int.from_bytes(data, byteorder='big')
  62.             data = []
  63.     return bytes(data)
  64.  
  65.  
  66. def main():
  67.     if not 1 < len(sys.argv) <= 4:
  68.         print('Encoding args: input.png input.bin output.png')
  69.         print('Decoding args: input.png [output.bin]')
  70.     elif len(sys.argv) <= 3:
  71.         image = Image.open(sys.argv[1])
  72.         data = decode(image)
  73.         out = sys.stdout.buffer
  74.         if len(sys.argv) == 3:
  75.             out = open(sys.argv[2], 'wb')
  76.         out.write(data)
  77.         out.flush()
  78.     else:
  79.         image = Image.open(sys.argv[1])
  80.         data = open(sys.argv[2], 'rb').read()
  81.         output = encode(image, data)
  82.         output.save(sys.argv[3])
  83.  
  84.  
  85. if __name__ == '__main__':
  86.     main()
RAW Paste Data