Create Hider and Revealer class

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)
+    hider = tools.Hider(image, message, encoding, auto_convert_rgb)
-    assert message_length != 0, "message length is zero"
+    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.
+            if hider.encode_another_pixel():
-                pixel = img.getpixel((col, row))
+                hider.encode_pixel((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
             else:
-                img.close()
+                return hider.encoded_image
-                return encoded
 
 
-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)
+    revealer = tools.Revealer(encoded_image, encoding)
-    width, height = img.size
+    width, height = revealer.encoded_image.size
-    buff, count = 0, 0
+
-    bitab = []
-    limit = None
     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:
+            if revealer.decode_pixel((col, row)):
-                img.close()
+                return revealer.secret_message
-                return "".join(bitab)[len(str(limit)) + 1 :]

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)
+    hider = tools.Hider(image, message, encoding, auto_convert_rgb)
-    assert message_length != 0, "message length is zero"
+    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.
+    return hider.encoded_image
-        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
 
 
 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)
+    revealer = tools.Revealer(encoded_image, encoding)
-    img_list = list(img.getdata())
+    width = revealer.encoded_image.width
-    width, height = img.size
-    buff, count = 0, 0
-    bitab = []
-    limit = None
 
     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]
+        if revealer.decode_pixel((col, row)):
-        pixel = img.getpixel((col, row))
+            return revealer.secret_message
-        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 :]

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