Fix up error handling

filecrypt.py

@@ -8,6 +8,7 @@ import os
 
 FILE_ID_LENGTH = 22
 TOKEN_LENGTH = 22
+HEADER_LENGTH = 56
 
 def encrypt_file(filename_in, chunk_size=1000000//16):
     file_id = secrets.token_urlsafe(16)
@@ -24,6 +25,7 @@ def encrypt_file(filename_in, chunk_size=1000000//16):
         fout.write(token_tag) # 16 bytes
         fout.write(auth_secret) # 16 bytes
         fout.write(data_nonce) # 8 bytes
+        assert fout.tell() == HEADER_LENGTH
 
         cipher = AES.new(key, AES.MODE_CTR,
                 initial_value=struct.pack('<Q', 0), nonce=data_nonce)
@@ -37,7 +39,7 @@ def encrypt_file(filename_in, chunk_size=1000000//16):
     return file_id, token
 
 def payload_size(path):
-    return os.stat(path).st_size - 56
+    return os.stat(path).st_size - HEADER_LENGTH
 
 def decrypt_generator(filename, token, seek=0, end=None, chunk_size=1000000//16):
     with open(filename, 'rb') as fin:
@@ -45,27 +47,30 @@ def decrypt_generator(filename, token, seek=0, end=None, chunk_size=1000000//16)
         token_tag = fin.read(16)
         auth_secret = fin.read(16)
         data_nonce = fin.read(8)
+        assert fin.tell() == HEADER_LENGTH
 
-        ciphertext = base64.b64decode(token + '='*((3-len(token)%3)%3))
-        token_cipher = AES.new(auth_secret, AES.MODE_GCM, nonce=token_nonce)
-        key = token_cipher.decrypt_and_verify(ciphertext, token_tag)
+    ciphertext = base64.b64decode(token + '='*((3-len(token)%3)%3))
+    token_cipher = AES.new(auth_secret, AES.MODE_GCM, nonce=token_nonce)
+    key = token_cipher.decrypt_and_verify(ciphertext, token_tag)
 
-        # Use token cipher's block size assuming the two are the same
-        cipher = AES.new(key, AES.MODE_CTR,
-                initial_value=struct.pack('>Q', seek // token_cipher.block_size), nonce=data_nonce)
-        print('Seeking to position', seek, 'iv', seek // token_cipher.block_size, 'pos', seek - seek %
-                cipher.block_size, 'offset', seek % cipher.block_size)
-        fin.seek(seek - seek % cipher.block_size, 1)
-        offset = seek % cipher.block_size
+    def generator():
+        with open(filename, 'rb') as fin:
+            fin.seek(HEADER_LENGTH)
+            # Use token cipher's block size assuming the two are the same
+            cipher = AES.new(key, AES.MODE_CTR,
+                    initial_value=struct.pack('>Q', seek // token_cipher.block_size), nonce=data_nonce)
+            fin.seek(seek - seek % cipher.block_size, 1)
+            offset = seek % cipher.block_size
 
-        to_send = end - seek + 1 if end else None
-        block = fin.read(cipher.block_size*chunk_size)
-        while block and (to_send is None or to_send > 0):
-            data = cipher.decrypt(block)
-            out = data[offset:to_send]
-            yield out
-            if to_send is not None:
-                to_send -= len(out)
-            offset = 0
+            to_send = end - seek + 1 if end else None
             block = fin.read(cipher.block_size*chunk_size)
+            while block and (to_send is None or to_send > 0):
+                data = cipher.decrypt(block)
+                out = data[offset:to_send]
+                yield out
+                if to_send is not None:
+                    to_send -= len(out)
+                offset = 0
+                block = fin.read(cipher.block_size*chunk_size)
+    return generator()