diff --git a/stegano/lsb/lsb.py b/stegano/lsb/lsb.py index 917e5ce..95215d7 100755 --- a/stegano/lsb/lsb.py +++ b/stegano/lsb/lsb.py @@ -44,7 +44,8 @@ def hide(input_image_file, message, auto_convert_rgb=False): assert message_length != 0, "message length is zero" img = Image.open(input_image_file) - if img.mode != 'RGB': + + 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)) @@ -71,7 +72,10 @@ def hide(input_image_file, message, auto_convert_rgb=False): if index + 3 <= len_message_bits : # Get the colour component. - (r, g, b) = img.getpixel((col, row)) + 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]) @@ -79,7 +83,12 @@ def hide(input_image_file, message, auto_convert_rgb=False): b = tools.setlsb(b, message_bits[index+2]) # Save the new pixel - encoded.putpixel((col, row), (r, g , b)) + if img.mode == 'RGBA': + updated_pixel = (r, g, b, pixel[3]) + else: + updated_pixel = (r, g, b) + + encoded.putpixel((col, row), updated_pixel) index += 3 else: @@ -102,8 +111,11 @@ def reveal(input_image_file): for row in range(height): for col in range(width): - # color = [r, g, b] - for color in img.getpixel((col, row)): + # 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)<<(7-count) count += 1 if count == 8: diff --git a/tests/sample-files/transparent.png b/tests/sample-files/transparent.png new file mode 100644 index 0000000..d2fb11d Binary files /dev/null and b/tests/sample-files/transparent.png differ diff --git a/tests/test_lsb.py b/tests/test_lsb.py index 97a965d..c464724 100644 --- a/tests/test_lsb.py +++ b/tests/test_lsb.py @@ -49,6 +49,17 @@ class TestLSB(unittest.TestCase): self.assertEqual(message, clear_message) + def test_with_transparent_png(self): + messages_to_hide = ["a", "foo", "Hello World!", ":Python:"] + + for message in messages_to_hide: + secret = lsb.hide("./tests/sample-files/transparent.png", message) + secret.save("./image.png") + + clear_message = lsb.reveal("./image.png") + + self.assertEqual(message, clear_message) + def test_with_text_file(self): text_file_to_hide = "./tests/sample-files/lorem_ipsum.txt" with open(text_file_to_hide) as f: