Fixed compatibiltiy with Python 3

SimpleAES/__init__.py

@@ -3,17 +3,25 @@
 import base64
 import hashlib
 import warnings
-from StringIO import StringIO
+
+try:
+    from StringIO import StringIO
+except ImportError:
+    from io import StringIO
+
 from Crypto.Cipher import AES
 from .version import VERSION
 from .exceptions import EncryptionError, DecryptionError
+from sys import version_info as python_version
 
 __title__ = 'SimpleAES'
 __version__ = VERSION
 __author__ = 'Vincent Driessen'
 __license__ = 'BSD'
 __copyright__ = 'Copyright 2012 Vincent Driessen'
 
+python_3 = python_version >= (3, 0)
+
 
 def check_output(cmd, input_=None, *popenargs, **kwargs):
     """Custom variant of stdlib's check_output(), but with stdin feed support."""
@@ -32,7 +40,8 @@ def check_output(cmd, input_=None, *popenargs, **kwargs):
 
 
 def _random_noise(len):
-    return ''.join(chr(random.randint(0, 0xFF)) for i in range(len))
+    s = ''.join(chr(random.randint(0, 0xFF)) for i in range(len))
+    return s.encode() if python_3 else s
 
 
 class SimpleAES(object):
@@ -41,6 +50,8 @@ def __init__(self, password):
         self._password = password
 
     def encrypt(self, string):
+        if python_3:
+            string = string.encode()
         """Encrypts a string using AES-256."""
         try:
             envvar = hashlib.sha256(_random_noise(16)).hexdigest()
@@ -55,16 +66,17 @@ def encrypt(self, string):
 
     def decrypt(self, b64_ciphertext, legacy='auto'):
         """Decrypts a string using AES-256."""
-        if legacy is True or (legacy == 'auto' and
-                              not b64_ciphertext.startswith('U2Fsd')):
+        c = b'U2Fsd' if python_3 else 'U2Fsd'
+        if legacy is True or (legacy == 'auto' and not b64_ciphertext.startswith(c)):
             return self._legacy_decrypt(b64_ciphertext)
 
         try:
             envvar = hashlib.sha256(_random_noise(16)).hexdigest()
+            endline = b'\n' if python_3 else '\n'
             plaintext = check_output([
                 'openssl', 'enc', '-d', '-aes-256-cbc', '-a', '-pass',
                 'env:{0}'.format(envvar)],
-                input_=b64_ciphertext + '\n',
+                input_=b64_ciphertext + endline,
                 env={envvar: self._password})
             return plaintext
         except:
@@ -128,4 +140,5 @@ def base64_decrypt(self, ciphertext):
             ciphertext = aes.encrypt(input)
             text = aes.decrypt(ciphertext)
             assert text == input
-    print 'All OK'
+    print('All OK')
+