From a0f6f24d7c2f3b3e246658aea2847014510cc002 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Bonhomme?= Date: Sun, 22 May 2016 00:13:29 +0200 Subject: [PATCH] Check if the generator is implemented in Stegano. --- stegano/lsb/lsb.py | 2 +- stegano/lsbset/lsbset.py | 8 ++++++-- tests/test_lsbset.py | 8 ++++++-- 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/stegano/lsb/lsb.py b/stegano/lsb/lsb.py index 5b6c3d8..3cb3d77 100755 --- a/stegano/lsb/lsb.py +++ b/stegano/lsb/lsb.py @@ -64,7 +64,7 @@ def hide(input_image_file, message, auto_convert_rgb=False): 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: {}).". \ + raise Exception("The message you want to hide is too long: {}". \ format(message_length)) for row in range(height): for col in range(width): diff --git a/stegano/lsbset/lsbset.py b/stegano/lsbset/lsbset.py index 1969a29..ffceadd 100644 --- a/stegano/lsbset/lsbset.py +++ b/stegano/lsbset/lsbset.py @@ -66,9 +66,13 @@ def hide(input_image_file, message, generator_function, auto_convert_rgb=False): 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 (%s > %s).""" % (len_message_bits, npixels * 3)) + raise Exception("The message you want to hide is too long: {}".\ + format(message_length)) - generator = getattr(generators, generator_function)() + try: + generator = getattr(generators, generator_function)() + except AttributeError as e: + raise e while index + 3 <= len_message_bits : generated_number = next(generator) diff --git a/tests/test_lsbset.py b/tests/test_lsbset.py index 82da78c..03282a0 100644 --- a/tests/test_lsbset.py +++ b/tests/test_lsbset.py @@ -41,7 +41,6 @@ class TestLSBSet(unittest.TestCase): def test_hide_and_reveal(self): messages_to_hide = ["a", "foo", "Hello World!", ":Python:"] - for message in messages_to_hide: secret = lsbset.hide("./tests/sample-files/Lenna.png", message, "eratosthenes") @@ -53,7 +52,6 @@ class TestLSBSet(unittest.TestCase): def test_hide_and_reveal_with_bad_generator(self): message_to_hide = "Hello World!" - secret = lsbset.hide("./tests/sample-files/Lenna.png", message_to_hide, "eratosthenes") secret.save("./image.png") @@ -61,6 +59,12 @@ class TestLSBSet(unittest.TestCase): with self.assertRaises(IndexError): clear_message = lsbset.reveal("./image.png", "identity") + def test_with_unknown_generator(self): + message_to_hide = "Hello World!" + with self.assertRaises(AttributeError): + secret = lsbset.hide("./tests/sample-files/Lenna.png", + message_to_hide, "eratosthene") + def tearDown(self): try: os.unlink("./image.png")