Create Hider and Revealer class

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