diff --git a/stegano/lsb/lsb.py b/stegano/lsb/lsb.py index 7d12790..ee5ff3a 100755 --- a/stegano/lsb/lsb.py +++ b/stegano/lsb/lsb.py @@ -31,7 +31,7 @@ from stegano import tools def hide( - input_image: Union[str, IO[bytes]], + image: Union[str, IO[bytes]], message: str, encoding: str = "UTF-8", shift: int = 0, @@ -40,91 +40,35 @@ def hide( """Hide a message (string) in an image with the LSB (Least Significant Bit) technique. """ - message_length = len(message) - assert message_length != 0, "message length is zero" + hider = tools.Hider(image, message, encoding, auto_convert_rgb) + width, height = hider.encoded_image.size - img = tools.open_image(input_image) - - if img.mode not in ["RGB", "RGBA"]: - if not auto_convert_rgb: - print("The mode of the image is not RGB. Mode is {}".format(img.mode)) - answer = input("Convert the image to RGB ? [Y / n]\n") or "Y" - if answer.lower() == "n": - raise Exception("Not a RGB image.") - img = img.convert("RGB") - - encoded = img.copy() - width, height = img.size - index = 0 - - message = str(message_length) + ":" + str(message) - message_bits = "".join(tools.a2bits_list(message, encoding)) - message_bits += "0" * ((3 - (len(message_bits) % 3)) % 3) - - npixels = width * height - len_message_bits = len(message_bits) - if len_message_bits > npixels * 3: - raise Exception( - "The message you want to hide is too long: {}".format(message_length) - ) for row in range(height): for col in range(width): if shift != 0: shift -= 1 continue - if index + 3 <= len_message_bits: - # Get the colour component. - pixel = img.getpixel((col, row)) - r = pixel[0] - g = pixel[1] - b = pixel[2] - - # Change the Least Significant Bit of each colour component. - r = tools.setlsb(r, message_bits[index]) - g = tools.setlsb(g, message_bits[index + 1]) - b = tools.setlsb(b, message_bits[index + 2]) - - # Save the new pixel - if img.mode == "RGBA": - encoded.putpixel((col, row), (r, g, b, pixel[3])) - else: - encoded.putpixel((col, row), (r, g, b)) - - index += 3 + if hider.encode_another_pixel(): + hider.encode_pixel((col, row)) else: - img.close() - return encoded + return hider.encoded_image -def reveal(input_image: Union[str, IO[bytes]], encoding: str = "UTF-8", shift: int = 0): +def reveal( + encoded_image: Union[str, IO[bytes]], + encoding: str = "UTF-8", + shift: int = 0, +): """Find a message in an image (with the LSB technique).""" - img = tools.open_image(input_image) - width, height = img.size - buff, count = 0, 0 - bitab = [] - limit = None + revealer = tools.Revealer(encoded_image, encoding) + width, height = revealer.encoded_image.size + for row in range(height): for col in range(width): if shift != 0: shift -= 1 continue - # pixel = [r, g, b] or [r,g,b,a] - pixel = img.getpixel((col, row)) - if img.mode == "RGBA": - pixel = pixel[:3] # ignore the alpha - for color in pixel: - buff += (color & 1) << (tools.ENCODINGS[encoding] - 1 - count) - count += 1 - if count == tools.ENCODINGS[encoding]: - bitab.append(chr(buff)) - buff, count = 0, 0 - if bitab[-1] == ":" and limit is None: - try: - limit = int("".join(bitab[:-1])) - except Exception: - pass - if len(bitab) - len(str(limit)) - 1 == limit: - img.close() - return "".join(bitab)[len(str(limit)) + 1 :] + if revealer.decode_pixel((col, row)): + return revealer.secret_message diff --git a/stegano/lsbset/lsbset.py b/stegano/lsbset/lsbset.py index b2fb75f..fcf77d2 100644 --- a/stegano/lsbset/lsbset.py +++ b/stegano/lsbset/lsbset.py @@ -31,7 +31,7 @@ from stegano import tools def hide( - input_image: Union[str, IO[bytes]], + image: Union[str, IO[bytes]], message: str, generator: Iterator[int], shift: int = 0, @@ -41,75 +41,33 @@ def hide( """Hide a message (string) in an image with the LSB (Least Significant Bit) technique. """ - message_length = len(message) - assert message_length != 0, "message length is zero" + hider = tools.Hider(image, message, encoding, auto_convert_rgb) + width = hider.encoded_image.width - img = tools.open_image(input_image) - - if img.mode not in ["RGB", "RGBA"]: - if not auto_convert_rgb: - print("The mode of the image is not RGB. Mode is {}".format(img.mode)) - answer = input("Convert the image to RGB ? [Y / n]\n") or "Y" - if answer.lower() == "n": - raise Exception("Not a RGB image.") - img = img.convert("RGB") - - encoded = img.copy() - width, height = img.size - index = 0 - - message = str(message_length) + ":" + str(message) - message_bits = "".join(tools.a2bits_list(message, encoding)) - message_bits += "0" * ((3 - (len(message_bits) % 3)) % 3) - - npixels = width * height - len_message_bits = len(message_bits) - if len_message_bits > npixels * 3: - raise Exception( - "The message you want to hide is too long: {}".format(message_length) - ) while shift != 0: next(generator) shift -= 1 - while index + 3 <= len_message_bits: + while hider.encode_another_pixel(): generated_number = next(generator) col = generated_number % width row = int(generated_number / width) - coordinate = (col, row) - r, g, b, *a = encoded.getpixel(coordinate) + hider.encode_pixel((col, row)) - # Change the Least Significant Bit of each colour component. - r = tools.setlsb(r, message_bits[index]) - g = tools.setlsb(g, message_bits[index + 1]) - b = tools.setlsb(b, message_bits[index + 2]) - - # Save the new pixel - if img.mode == "RGBA": - encoded.putpixel(coordinate, (r, g, b, *a)) - else: - encoded.putpixel(coordinate, (r, g, b)) - - index += 3 - - return encoded + return hider.encoded_image def reveal( - input_image: Union[str, IO[bytes]], + encoded_image: Union[str, IO[bytes]], generator: Iterator[int], shift: int = 0, encoding: str = "UTF-8", ): """Find a message in an image (with the LSB technique).""" - img = tools.open_image(input_image) - img_list = list(img.getdata()) - width, height = img.size - buff, count = 0, 0 - bitab = [] - limit = None + revealer = tools.Revealer(encoded_image, encoding) + width = revealer.encoded_image.width while shift != 0: next(generator) @@ -121,22 +79,5 @@ def reveal( col = generated_number % width row = int(generated_number / width) - # pixel = [r, g, b] or [r,g,b,a] - pixel = img.getpixel((col, row)) - if img.mode == "RGBA": - pixel = pixel[:3] # ignore the alpha - - for color in pixel: - buff += (color & 1) << (tools.ENCODINGS[encoding] - 1 - count) - count += 1 - if count == tools.ENCODINGS[encoding]: - bitab.append(chr(buff)) - buff, count = 0, 0 - if bitab[-1] == ":" and limit is None: - if "".join(bitab[:-1]).isdigit(): - limit = int("".join(bitab[:-1])) - else: - raise IndexError("Impossible to detect message.") - - if len(bitab) - len(str(limit)) - 1 == limit: - return "".join(bitab)[len(str(limit)) + 1 :] + if revealer.decode_pixel((col, row)): + return revealer.secret_message diff --git a/stegano/tools.py b/stegano/tools.py index 6ddeac3..78fb9f3 100755 --- a/stegano/tools.py +++ b/stegano/tools.py @@ -115,3 +115,104 @@ def open_image(fname_or_instance: Union[str, IO[bytes]]): return fname_or_instance return Image.open(fname_or_instance) + + +class Hider: + def __init__( + self, + input_image: Union[str, IO[bytes]], + message: str, + encoding: str = "UTF-8", + auto_convert_rgb: bool = False, + ): + self._index = 0 + + message_length = len(message) + assert message_length != 0, "message length is zero" + + image = open_image(input_image) + + if image.mode not in ["RGB", "RGBA"]: + if not auto_convert_rgb: + print("The mode of the image is not RGB. Mode is {}".format(image.mode)) + answer = input("Convert the image to RGB ? [Y / n]\n") or "Y" + if answer.lower() == "n": + raise Exception("Not a RGB image.") + + image = image.convert("RGB") + + self.encoded_image = image.copy() + image.close() + + message = str(message_length) + ":" + str(message) + self._message_bits = "".join(a2bits_list(message, encoding)) + self._message_bits += "0" * ((3 - (len(self._message_bits) % 3)) % 3) + + width, height = self.encoded_image.size + npixels = width * height + self._len_message_bits = len(self._message_bits) + + if self._len_message_bits > npixels * 3: + raise Exception( + "The message you want to hide is too long: {}".format(message_length) + ) + + def encode_another_pixel(self): + return True if self._index + 3 <= self._len_message_bits else False + + def encode_pixel(self, coordinate: tuple): + # Get the colour component. + r, g, b, *a = self.encoded_image.getpixel(coordinate) + + # Change the Least Significant Bit of each colour component. + r = setlsb(r, self._message_bits[self._index]) + g = setlsb(g, self._message_bits[self._index + 1]) + b = setlsb(b, self._message_bits[self._index + 2]) + + # Save the new pixel + if self.encoded_image.mode == "RGBA": + self.encoded_image.putpixel(coordinate, (r, g, b, *a)) + else: + self.encoded_image.putpixel(coordinate, (r, g, b)) + + self._index += 3 + + +class Revealer: + def __init__(self, encoded_image: Union[str, IO[bytes]], encoding: str = "UTF-8"): + self.encoded_image = open_image(encoded_image) + self._encoding_length = ENCODINGS[encoding] + self._buff, self._count = 0, 0 + self._bitab = [] + self._limit = None + self.secret_message = "" + + def decode_pixel(self, coordinate: tuple): + # pixel = [r, g, b] or [r,g,b,a] + pixel = self.encoded_image.getpixel(coordinate) + + if self.encoded_image.mode == "RGBA": + pixel = pixel[:3] # ignore the alpha + + for color in pixel: + self._buff += (color & 1) << (self._encoding_length - 1 - self._count) + self._count += 1 + + if self._count == self._encoding_length: + self._bitab.append(chr(self._buff)) + self._buff, self._count = 0, 0 + + if self._bitab[-1] == ":" and self._limit is None: + if "".join(self._bitab[:-1]).isdigit(): + self._limit = int("".join(self._bitab[:-1])) + else: + raise IndexError("Impossible to detect message.") + + if len(self._bitab) - len(str(self._limit)) - 1 == self._limit: + self.secret_message = "".join(self._bitab)[len(str(self._limit)) + 1 :] + self.encoded_image.close() + + return True + + else: + return False