Advertisement
Najeebsk

IMAGE-STEGANOGRAPHY.py

Feb 12th, 2024
735
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 6.96 KB | None | 0 0
  1. import os
  2. import numpy as np
  3. from imageio import imread, imwrite
  4.  
  5. import argparse
  6.  
  7. max_value = 255 # max uint value per pixel per channel
  8. header_len = 4*8 # uint32 bit length
  9.  
  10. def read_image(img_path):
  11.     """
  12.        Reads an image from file and flattens it.
  13.        Args:
  14.            img_path    path to the image
  15.        Returns:
  16.            ndarray     numpy array containing the image in a flat shape
  17.            ndarray     shape of the read image before flattening
  18.    """
  19.     img = np.array(imread(img_path), dtype=np.uint8)
  20.     orig_shape = img.shape
  21.     return img.flatten(), orig_shape
  22.  
  23. def write_image(img_path, img_data, shape):
  24.     """
  25.        Writes an image to a path from a flat numpy array, usig the shape provided.
  26.        Args:
  27.            img_path    path were to save the image
  28.            img_data    numpy array containing the image (flat)
  29.            shape       shape of the image to be saved
  30.    """
  31.     img_data = np.reshape(img_data, shape)
  32.     imwrite(img_path, img_data)
  33.  
  34. def bytes2array(byte_data):
  35.     """
  36.        Converts byte data to a bit array (numpy array, dtype=np.uint8).
  37.        Args:
  38.            byte_data   the byte data
  39.        Returns:
  40.            ndarray     a numpy array of the single bits that composed the byte data
  41.    """
  42.     byte_array = np.frombuffer(byte_data, dtype=np.uint8)
  43.     return np.unpackbits(byte_array)
  44.  
  45. def array2bytes(bit_array):
  46.     """
  47.        Converts a bit array (numpy array, dtype=np.uint8) to byte data.
  48.        Args:
  49.            bit_array   the bit array
  50.        Returns:
  51.            bytes       the byte data
  52.    """
  53.     byte_array = np.packbits(bit_array)
  54.     return byte_array.tobytes()
  55.  
  56. def read_file(file_path):
  57.     """
  58.        Reads a file as a bit array (numpy array, dtype=np.uint8)
  59.        Args:
  60.            file_path   path to the file
  61.        Returns:
  62.            ndarray     the bit array
  63.    """
  64.     file_bytes = open(file_path, "rb").read()
  65.     return bytes2array(file_bytes)
  66.  
  67. def write_file(file_path, file_bit_array):
  68.     """
  69.        Writes a file to a path from a bit array (numpy array, dtype=np.uint8).
  70.        Args:
  71.            file_path       path to the file
  72.            file_bit_array  the bit array of the file
  73.    """
  74.     bytes_data = array2bytes(file_bit_array)
  75.     f = open(file_path, 'wb')
  76.     f.write(bytes_data)
  77.     f.close()
  78.  
  79. def encode_data(image, file_data):
  80.     """
  81.        Encodes the file data onto the image
  82.        Args:
  83.            image       the original image numpy array (flat)
  84.            file_data   the file data (bit array)
  85.        Returns:
  86.            ndarray     the encoded image as a numpy array
  87.    """
  88.     or_mask = file_data
  89.     and_mask = np.zeros_like(or_mask)
  90.     and_mask = (and_mask + max_value - 1) + or_mask
  91.     res = np.bitwise_or(image, or_mask)
  92.     res = np.bitwise_and(res, and_mask)
  93.     return res
  94.  
  95. def decode_data(encoded_data):
  96.     """
  97.        Decodes the data from an image
  98.        Args:
  99.            encoded_data    the encoded image as numpy array
  100.        Returns:
  101.            ndarray         the bit array containig the file bits
  102.    """
  103.     out_mask = np.ones_like(encoded_data)
  104.     output = np.bitwise_and(encoded_data, out_mask)
  105.     return output
  106.  
  107. def _main(args):
  108.     """Main fuction of the script"""
  109.     if args.image is not None and args.file is not None:
  110.         if args.encode:
  111.             img_path = args.image
  112.             file_path = args.file
  113.             if not os.path.isfile(img_path):
  114.                 print("Image file does not exist")
  115.                 return
  116.             if not os.path.isfile(file_path):
  117.                 print("File does not exist")
  118.                 return
  119.  
  120.             output_path = args.output
  121.             extension = os.path.splitext(output_path)[1][1:]
  122.             if extension == '':  # if no extension, append png
  123.                 output_path = output_path + '.png'
  124.             elif extension != 'png':  # replace the wrong extension with png
  125.                 li = output_path.rsplit(extension, 1)
  126.                 output_path = 'png'.join(li)
  127.  
  128.             image, shape_orig = read_image(img_path)
  129.             file = read_file(file_path)
  130.             file_len = file.shape[0]
  131.             len_array = np.array([file_len], dtype=np.uint32).view(np.uint8)
  132.             len_array = np.unpackbits(len_array)
  133.             img_len = image.shape[0]
  134.  
  135.             if file_len >= img_len - header_len:  # 4 bytes are used to store file length
  136.                 print("File too big, error")
  137.                 return
  138.             else:  #  Insert padding. Using random padding, otherwise values would all be even if padding with zeros (could be noticed in histogram).
  139.                 tmp = file
  140.                 file = np.random.randint(2, size=img_len, dtype=np.uint8)
  141.                 file[header_len:header_len+file_len] = tmp
  142.                 # file = np.pad(file, (header_len,img_len - file_len - header_len), 'constant', constant_values=(0, 0))
  143.  
  144.             file[:header_len] = len_array
  145.             encoded_data = encode_data(image, file)
  146.  
  147.             write_image(output_path, encoded_data, shape_orig)
  148.             print("Image encoded")
  149.             return
  150.  
  151.         if args.decode:
  152.             img_path = args.image
  153.             if not os.path.isfile(img_path):
  154.                 print("Image file does not exist")
  155.                 return
  156.             file_path = args.file
  157.             encoded_data, shape_orig = read_image(img_path)
  158.             data = decode_data(encoded_data)
  159.             el_array = np.packbits(data[:header_len])
  160.             extracted_len = el_array.view(np.uint32)[0]
  161.             data = data[header_len:extracted_len+header_len]
  162.             write_file(file_path, data)
  163.             print("Image decoded")
  164.             return
  165.  
  166.         print("Error, no action specified!")
  167.         return
  168.  
  169.     print("Error, image or file not specified")
  170.  
  171. if __name__ == '__main__':
  172.     parser = argparse.ArgumentParser(
  173.         description='Conceal small files inside a PNG image and extract them back')
  174.     group = parser.add_mutually_exclusive_group()
  175.     group.add_argument(
  176.         '-e',
  177.         '--encode',
  178.         help='If present the script will conceal the file in the image and produce a new encoded image',
  179.         action="store_true")
  180.     group.add_argument(
  181.         '-d',
  182.         '--decode',
  183.         help='If present the script will decode the concealed data in the image and produce a new file with this data',
  184.         action="store_true")
  185.     parser.add_argument(
  186.         '-i',
  187.         '--image',
  188.         help='Path to an image to use for concealing or file extraction')
  189.     parser.add_argument(
  190.         '-f',
  191.         '--file',
  192.         help='Path to the file to conceal or to extract')
  193.     parser.add_argument(
  194.         '-o',
  195.         '--output',
  196.         help='Path where to save the encoded image. Specify only the file name, or use .png extension; png extension will be added automatically',
  197.         default='encoded.png')
  198.  
  199.     _main(parser.parse_args())
  200.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement