Started working on the invalid cases

1 files changed, 27 insertions, 5 deletions

diff --git a/cryptography/fernet.py b/cryptography/fernet.py
index 064aceec..880d96f0 100644
--- a/cryptography/fernet.py
+++ b/cryptography/fernet.py
@@ -10,6 +10,10 @@ from cryptography.hazmat.primitives.hmac import HMAC
 from cryptography.hazmat.primitives.block import BlockCipher, ciphers, modes
 
 
+class InvalidToken(Exception):
+    pass
+
+
 class Fernet(object):
     def __init__(self, key):
         super(Fernet, self).__init__()
@@ -23,6 +27,9 @@ class Fernet(object):
         return self._encrypt_from_parts(data, current_time, iv)
 
     def _encrypt_from_parts(self, data, current_time, iv):
+        if isinstance(data, six.text_type):
+            raise TypeError("Unicode-objects must be encoded before encryption")
+
         padder = padding.PKCS7(ciphers.AES.block_size).padder()
         padded_data = padder.update(data) + padder.finalize()
         encryptor = BlockCipher(
@@ -41,28 +48,43 @@ class Fernet(object):
         )
 
     def decrypt(self, data, ttl=None, current_time=None):
-        # TODO: whole function is a giant hack job with no error checking
+        if isinstance(data, six.text_type):
+            raise TypeError("Unicode-objects must be encoded before decryption")
+
         if current_time is None:
             current_time = int(time.time())
-        data = base64.urlsafe_b64decode(data)
+
+        try:
+            data = base64.urlsafe_b64decode(data)
+        except TypeError:
+            raise InvalidToken
+
         assert six.indexbytes(data, 0) == 0x80
         timestamp = data[1:9]
         iv = data[9:25]
         ciphertext = data[25:-32]
         if ttl is not None:
             if struct.unpack(">Q", timestamp)[0] + ttl < current_time:
-                raise ValueError
+                raise InvalidToken
         h = HMAC(self.signing_key, digestmod=hashes.SHA256)
         h.update(data[:-32])
         hmac = h.digest()
+
         if not constant_time_compare(hmac, data[-32:]):
-            raise ValueError
+            raise InvalidToken
+
         decryptor = BlockCipher(
             ciphers.AES(self.encryption_key), modes.CBC(iv)
         ).decryptor()
         plaintext_padded = decryptor.update(ciphertext) + decryptor.finalize()
         unpadder = padding.PKCS7(ciphers.AES.block_size).unpadder()
-        return unpadder.update(plaintext_padded) + unpadder.finalize()
+
+        unpadded = unpadder.update(plaintext_padded)
+        try:
+            unpadded += unpadder.finalize()
+        except ValueError:
+            raise InvalidToken
+        return unpadded
 
 
 def constant_time_compare(a, b):