Use PKCS7 padding conventions.

This makes SimpleAES compatible with many of the other AES
implementations out there, most notably CryptoJS.  It uses PKCS7-style
padding for the byte stream, which is a revertible, unambiguous padding
that avoids data loss while not requiring length announcements in the
stream header.

This is the default pad-encoding for CryptoJS and possibly many more
algorithms out there.

SimpleAES/__init__.py

@@ -4,7 +4,8 @@
 import base64
 from StringIO import StringIO
 from Crypto.Cipher import AES
-from .version import VERSION
+from version import VERSION
+from padding import PKCS7Encoder
 
 __title__ = 'SimpleAES'
 __version__ = VERSION
@@ -29,20 +30,21 @@ def encrypt(self, string):
         iv = _random_noise(16)
         aes = AES.new(self._key, AES.MODE_CBC, iv)
 
-        fin = StringIO(string)
+        # Make a PKCS7-padded string, which is guaranteed to be decodeable
+        # without data loss
+        encoder = PKCS7Encoder(aes.block_size)
+        fin = StringIO(encoder.encode(string))
         fout = StringIO()
         try:
             # Fixed-length data encoded in the encrypted string, first
-            fout.write(struct.pack('<Q', len(string)))
+            # fout.write(struct.pack('<Q', len(string)))
             fout.write(iv)
 
             while True:
                 chunk = fin.read(self.chunksize)
                 chunk_len = len(chunk)
                 if chunk_len == 0:
                     break  # done
-                elif chunk_len % 16 != 0:
-                    chunk += _random_noise(16 - chunk_len % 16)
                 fout.write(aes.encrypt(chunk))
             cipherbytes = fout.getvalue()
         finally:
@@ -55,9 +57,7 @@ def decrypt(self, cipherbytes):
         """Decrypts a string using AES-256."""
         fin = StringIO(cipherbytes)
         fout = StringIO()
-
         try:
-            input_size = struct.unpack('<Q', fin.read(struct.calcsize('Q')))[0]
             iv = fin.read(16)
             aes = AES.new(self._key, AES.MODE_CBC, iv)
 
@@ -67,15 +67,14 @@ def decrypt(self, cipherbytes):
                     break  # done
                 fout.write(aes.decrypt(chunk))
 
-            # truncate any padded random noise
-            fout.truncate(input_size)
-
             text = fout.getvalue()
         finally:
             fin.close()
             fout.close()
 
-        return text
+        # Unpad the padded string
+        encoder = PKCS7Encoder(aes.block_size)
+        return encoder.decode(text)
 
     def base64_encrypt(self, string):
         """Encrypts a string using AES-256, but returns the result
@@ -90,6 +89,44 @@ def base64_decrypt(self, ciphertext):
         plaintext = self.decrypt(cipherbytes)
         return plaintext
 
+    def convert(self, cipherbytes):
+        """Converts an old-style ciphertext into a new, PKCS7-padded, ciphertext."""
+        decrypted = self.decrypt_compat(cipherbytes)
+        encrypted = self.encrypt(decrypted)
+        assert self.decrypt(encrypted) == decrypted
+        return encrypted
+
+    def decrypt_compat(self, cipherbytes):
+        """Decrypts a string that's encoded with a SimpleAES version < 1.0.
+        To convert a ciphertext to the new-style algo, use:
+
+            aes = SimpleAES('my secret')
+            aes.convert(legacy_ciphertext)
+        """
+        fin = StringIO(cipherbytes)
+        fout = StringIO()
+
+        try:
+            input_size = struct.unpack('<Q', fin.read(struct.calcsize('Q')))[0]
+            iv = fin.read(16)
+            aes = AES.new(self._key, AES.MODE_CBC, iv)
+
+            while True:
+                chunk = fin.read(self.chunksize)
+                if len(chunk) == 0:
+                    break  # done
+                fout.write(aes.decrypt(chunk))
+
+            # truncate any padded random noise
+            fout.truncate(input_size)
+
+            text = fout.getvalue()
+        finally:
+            fin.close()
+            fout.close()
+
+        return text
+
 
 __all__ = ['SimpleAES']
 

SimpleAES/padding.py

@@ -0,0 +1,29 @@
+from binascii import hexlify, unhexlify
+
+
+class InvalidBlockSizeError(Exception):
+    """Raised for invalid block sizes"""
+    pass
+
+
+class PKCS7Encoder():
+    """Technique for padding a string as defined in RFC 2315, section 10.3,
+    note #2.
+    """
+    def __init__(self, block_size=16):
+        if block_size < 1 or block_size > 99:
+            raise InvalidBlockSizeError('The block size must be between 1 ' \
+                    'and 99')
+        self.block_size = block_size
+
+    def encode(self, text):
+        text_length = len(text)
+        amount_to_pad = self.block_size - (text_length % self.block_size)
+        if amount_to_pad == 0:
+            amount_to_pad = self.block_size
+        pad = unhexlify('%02d' % amount_to_pad)
+        return text + pad * amount_to_pad
+
+    def decode(self, text):
+        pad = int(hexlify(text[-1]))
+        return text[:-pad]