diff --git a/SimpleAES/__init__.py b/SimpleAES/__init__.py index de2bb23..0f40702 100644 --- a/SimpleAES/__init__.py +++ b/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,11 +30,14 @@ 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(' 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]