Steganography tool that uses LSB technique with pseudo-random pixel selection

script.py

import numpy as np
from PIL import Image
import sys, argparse, random

EOF_MARKER = '$eof!'
confs = {'RGB':[0,3], 'RGBA':[1,4]}

def calculate_min_size(number_of_bits):
    """Calculates the minimum image size needed to contain number_of_bits"""
    min_pxls = number_of_bits//3 # 1 pixel = 3 bits stored
    return min_pxls

def get_num_rand(used_pixels, num_of_pixels):
    n = random.randint(0, num_of_pixels)
    while n in used_pixels:
        n = random.randint(0, num_of_pixels)
    used_pixels.add(n)
    return n

def hide_data(seed, input_data, output_file, carrier):
    """
    Takes a seed, some data, a filename for the output image and a carrier. 
    This function hides data inside the carrier using LSB technique
    """
    byte_written = 0
    img = Image.open(carrier, 'r')
    width, height = img.size
    matrix = np.array(list(img.getdata()))
    generator = random.seed(seed)
    used_pixels = set()
    # Get the configuration: is RGB or RGBA?
    conf = confs[img.mode]
    num_of_pixels = matrix.size//conf[1]
    print("Hiding: ",input_data)
    # Append a marker. When this marker is encountered, no more hidden data can be found next.
    input_data += EOF_MARKER
    binary_enc = "".join([format(ord(ch), "08b") for ch in input_data])
    min_size = calculate_min_size(len(binary_enc))
    print("Minimum size needed is {} pixels and the carrier has: {}".format(min_size, num_of_pixels))
    if min_size >= num_of_pixels:
        print("ERROR: The image is not big enough to carry the data with the specified passphrase")
        sys.exit(1)
    start_pixel = get_num_rand(used_pixels, num_of_pixels)
    while byte_written != len(input_data):
        bit_i = 0
        while bit_i != 8:
            px = matrix[start_pixel]
            # colors: Red, Green and Blue
            for c in range(conf[0], conf[1]):
                if bit_i == 8:
                    break
                # Because of Least Significant Bit, we want to modify the last bit of every color
                color = matrix[start_pixel][c]
                lsb = color&1
                # Here, we just use bit manipulation to modify the last bit of the color number
                if lsb != int(binary_enc[(byte_written*8)+bit_i]):
                    color = color>>1 # erase last bit
                    color = color<<1 # zero last bit
                    if lsb == 0: # it means that byte[bit_i]=1, so I need to encode 1
                        color = color|1
                    matrix[start_pixel][c] = color
                bit_i += 1
            start_pixel = get_num_rand(used_pixels, num_of_pixels)
        byte_written += 1
        bit_i = 0
        start_pixel = get_num_rand(used_pixels, num_of_pixels)

    out_img = Image.fromarray(np.uint8(matrix.reshape(height, width, conf[1])), img.mode)
    out_img.save(output_file)
    print("All done!")

def retrieve_data(seed, input_file):
    """
    Takes a passhprase and an image. Outputs the hidden data inside the image, if any.
    """
    img = Image.open(input_file, 'r')
    width, height = img.size
    matrix = np.array(list(img.getdata()))
    # Get the configuration: is RGB or RGBA?
    conf = confs[img.mode]
    num_of_pixels = matrix.size//conf[1]
    generator = random.seed(seed)
    used_pixels = set()
    start_pixel = get_num_rand(used_pixels, num_of_pixels)
    bit_i = 7
    byte = 0
    message = ""
    end = False
    while (end == False):
        while (bit_i >= 0):
            px = matrix[start_pixel]
            # colors: Red, Green and Blue
            for c in range(conf[0], conf[1]):
                if bit_i >= 0:
                    # We are getting the LSB of the pixel color, and then we shift it to the left accordingly
                    byte += (px[c]&1)<<bit_i
                    bit_i -= 1
                else:
                    break
            start_pixel = get_num_rand(used_pixels, num_of_pixels)
        if start_pixel>=num_of_pixels:
            break
        # decoded 1 byte
        message += chr(byte)
        # have I encountered the eof_marker? If yes, the decoding process is done
        if message[-len(EOF_MARKER):] == EOF_MARKER:
            end = True
        byte = 0
        bit_i = 7
        start_pixel = get_num_rand(used_pixels, num_of_pixels)
    
    if end == False:
        print("Nothing found in this image")
    else:
        print("All done!")
        print("The hidden message is: ")
        print(message[:len(message)-len(EOF_MARKER)])
    

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", help="The mode in which you want to run the program: enc or dec")
    parser.add_argument("seed", help="The seed used to encode/decode data.")
    parser.add_argument("--input_data", help="The pathname of the data you want to hide. Must be ASCII.")
    parser.add_argument("--input_image", help="The pathname of the image that contains hidden data")
    parser.add_argument("--output", help="Output filename")
    parser.add_argument("--carrier", help="The pathname of the image that will be used as carrier")
    args = parser.parse_args()
    if args.mode == 'enc':
        if args.input_data == None or args.output == None or args.carrier == None:
            print("ERROR: Argument missing!")
            print(parser.print_help())
            sys.exit(1)
        hide_data(args.seed, args.input_data, args.output, args.carrier)
    elif args.mode == 'dec':
        if args.input_image == None:
            print("ERROR: Argument missing!")
            print(parser.print_help())
            sys.exit(1)
        retrieve_data(args.seed, args.input_image)

if __name__ == '__main__':
    main()

Loving This! Great Work!