From c9dda3da004d14b9b47754540d0260909f922c15 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Bonhomme?= <cedric@cedricbonhomme.org>
Date: Thu, 19 May 2016 07:30:50 +0200
Subject: [PATCH] Improved 'red' module and added unit tests for the 'red'
 module.

---
 stegano/red.py    | 24 ++++++++++++-----
 tests/test_red.py | 66 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_red.py

diff --git a/stegano/red.py b/stegano/red.py
index 92080f3..006b8c3 100755
--- a/stegano/red.py
+++ b/stegano/red.py
@@ -28,18 +28,18 @@ import sys
 
 from PIL import Image
 
-def hide(img, message):
+def hide(input_image_file, message):
     """
     Hide a message (string) in an image.
 
     Use the red portion of a pixel (r, g, b) tuple to
     hide the message string characters as ASCII values.
-    The red value of the first pixel is used for length of string.
+    The red value of the first pixel is used for length of the string.
     """
     length = len(message)
-    # Limit length of message to 255
-    if length > 255:
-        return False
+    assert length != 0, "lenght of message is null"
+    assert length < 255, "message is too long"
+    img = Image.open(input_image_file)
     # Use a copy of image to hide the text in
     encoded = img.copy()
     width, height = img.size
@@ -59,7 +59,7 @@ def hide(img, message):
             index += 1
     return encoded
 
-def reveal(img):
+def reveal(input_image_file):
     """
     Find a message in an image.
 
@@ -67,6 +67,7 @@ def reveal(img):
     hidden message characters (ASCII values).
     The red value of the first pixel is used for length of string.
     """
+    img = Image.open(input_image_file)
     width, height = img.size
     message = ""
     index = 0
@@ -81,6 +82,17 @@ def reveal(img):
             index += 1
     return message
 
+def write(image, output_image_file):
+    """
+    """
+    try:
+        image.save(output_image_file)
+    except Exception as e:
+        # If hide() returns an error (Too long message).
+        print(e)
+    finally:
+        image.close()
+
 if __name__ == '__main__':
     # Point of entry in execution mode.
     from optparse import OptionParser
diff --git a/tests/test_red.py b/tests/test_red.py
new file mode 100644
index 0000000..a8e034d
--- /dev/null
+++ b/tests/test_red.py
@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+#-*- coding: utf-8 -*-
+
+# Stéganô - Stéganô is a basic Python Steganography module.
+# Copyright (C) 2010-2016  Cédric Bonhomme - https://www.cedricbonhomme.org
+#
+# For more information : https://github.com/cedricbonhomme/Stegano
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>
+
+__author__ = "Cedric Bonhomme"
+__version__ = "$Revision: 0.1 $"
+__date__ = "$Date: 2016/05/19 $"
+__license__ = "GPLv3"
+
+import os
+import unittest
+
+from stegano import red
+
+class TestRed(unittest.TestCase):
+
+    def test_hide_empty_message(self):
+        """
+        Test hiding the empty string.
+        """
+        with self.assertRaises(AssertionError):
+            secret = red.hide("./examples/pictures/Lenna.png", "")
+
+    def test_hide_and_reveal(self):
+        messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
+
+        for message in messages_to_hide:
+            secret = red.hide("./examples/pictures/Lenna.png", message)
+            secret.save("./image.png")
+
+            clear_message = red.reveal("./image.png")
+
+            self.assertEqual(message, clear_message)
+
+    def test_with_too_long_message(self):
+        with open("./examples/lorem_ipsum.txt") as f:
+            message = f.read()
+        with self.assertRaises(AssertionError):
+            red.hide("./examples/pictures/Lenna.png", message)
+
+    def tearDown(self):
+        try:
+            os.unlink("./image.png")
+        except:
+            pass
+
+
+if __name__ == '__main__':
+    unittest.main()