From 93de151025678a43f07a13c2e193dcba226d56df Mon Sep 17 00:00:00 2001
From: Peter Justin <peter.justin@outlook.com>
Date: Mon, 8 Apr 2019 09:57:50 +0200
Subject: [PATCH] Add option to either pass a path or an Image instance

---
 stegano/exifHeader/exifHeader.py |  8 ++++++--
 stegano/lsb/lsb.py               |  4 ++--
 stegano/lsbset/lsbset.py         |  4 ++--
 stegano/red/red.py               |  5 +++--
 stegano/tools.py                 | 14 ++++++++++++++
 5 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/stegano/exifHeader/exifHeader.py b/stegano/exifHeader/exifHeader.py
index 2abf1d4..24a8d65 100644
--- a/stegano/exifHeader/exifHeader.py
+++ b/stegano/exifHeader/exifHeader.py
@@ -26,7 +26,10 @@ __revision__ = "$Date: 2017/01/18 $"
 __license__ = "GPLv3"
 
 from PIL import Image
+
 import piexif
+from stegano import tools
+
 
 def hide(input_image_file, img_enc, secret_message = None, secret_file = None, img_format = None):
     """Hide a message (string) in an image.
@@ -43,7 +46,7 @@ def hide(input_image_file, img_enc, secret_message = None, secret_file = None, i
     except:
         text = compress(b64encode(secret_message))
 
-    img = Image.open(input_image_file)
+    img = tools.open_image(input_image_file)
 
     if img_format is None:
         img_format = img.format
@@ -66,7 +69,8 @@ def reveal(input_image_file):
     from base64 import b64decode
     from zlib import decompress
 
-    img = Image.open(input_image_file)
+    img = tools.open_image(input_image_file)
+
     try:
         if img.format in ['JPEG', 'TIFF']:
             if 'exif' in img.info:
diff --git a/stegano/lsb/lsb.py b/stegano/lsb/lsb.py
index d92b3ba..8f08391 100755
--- a/stegano/lsb/lsb.py
+++ b/stegano/lsb/lsb.py
@@ -42,7 +42,7 @@ def hide(input_image: Union[str, IO[bytes]],
     message_length = len(message)
     assert message_length != 0, "message length is zero"
 
-    img = Image.open(input_image)
+    img = tools.open_image(input_image)
 
     if img.mode not in ['RGB', 'RGBA']:
         if not auto_convert_rgb:
@@ -96,7 +96,7 @@ def hide(input_image: Union[str, IO[bytes]],
 def reveal(input_image: Union[str, IO[bytes]], encoding='UTF-8'):
     """Find a message in an image (with the LSB technique).
     """
-    img = Image.open(input_image)
+    img = tools.open_image(input_image)
     width, height = img.size
     buff, count = 0, 0
     bitab = []
diff --git a/stegano/lsbset/lsbset.py b/stegano/lsbset/lsbset.py
index 061c21d..74cbd0b 100644
--- a/stegano/lsbset/lsbset.py
+++ b/stegano/lsbset/lsbset.py
@@ -45,7 +45,7 @@ def hide(input_image: Union[str, IO[bytes]],
     message_length = len(message)
     assert message_length != 0, "message length is zero"
 
-    img = Image.open(input_image)
+    img = tools.open_image(input_image)
 
     if img.mode not in ['RGB', 'RGBA']:
         if not auto_convert_rgb:
@@ -105,7 +105,7 @@ def reveal(input_image: Union[str, IO[bytes]],
             encoding: str = 'UTF-8'):
     """Find a message in an image (with the LSB technique).
     """
-    img = Image.open(input_image)
+    img = tools.open_image(input_image)
     img_list = list(img.getdata())
     width, height = img.size
     buff, count = 0, 0
diff --git a/stegano/red/red.py b/stegano/red/red.py
index e429191..0488d90 100755
--- a/stegano/red/red.py
+++ b/stegano/red/red.py
@@ -29,6 +29,7 @@ import sys
 
 from PIL import Image
 from typing import Union, IO
+from stegano import tools
 
 def hide(input_image: Union[str, IO[bytes]], message: str):
     """
@@ -41,7 +42,7 @@ def hide(input_image: Union[str, IO[bytes]], message: str):
     message_length = len(message)
     assert message_length != 0, "message message_length is zero"
     assert message_length < 255, "message is too long"
-    img = Image.open(input_image)
+    img = tools.open_image(input_image)
     # Use a copy of image to hide the text in
     encoded = img.copy()
     width, height = img.size
@@ -70,7 +71,7 @@ def reveal(input_image: Union[str, IO[bytes]]):
     hidden message characters (ASCII values).
     The red value of the first pixel is used for message_length of string.
     """
-    img = Image.open(input_image)
+    img = tools.open_image(input_image)
     width, height = img.size
     message = ""
     index = 0
diff --git a/stegano/tools.py b/stegano/tools.py
index a9b9cc5..48f4d39 100755
--- a/stegano/tools.py
+++ b/stegano/tools.py
@@ -29,6 +29,9 @@ import base64
 import itertools
 from typing import List, Iterator, Tuple, Union
 from functools import reduce
+from typing import IO, Iterator, List, Tuple, Union
+
+from PIL import Image
 
 ENCODINGS = {
     'UTF-8': 8,
@@ -99,3 +102,14 @@ def base642binary(b64_fname: str) -> bytes:
     """
     b64_fname += '==='
     return base64.b64decode(b64_fname)
+
+def open_image(fname_or_instance: Union[str, IO[bytes]]):
+    """Opens a Image and returns it.
+
+    :param fname_or_instance: Can either be the location of the image as a
+                              string or the Image.Image instance itself.
+    """
+    if isinstance(fname_or_instance, Image.Image):
+        return fname_or_instance
+
+    return Image.open(fname_or_instance)