Cleaned up AES CTR.

Author: ricmoo

pyscrypt/aesctr.py

@@ -30,6 +30,7 @@ def __init__(self, nbits, initial_value = 1):
         if nbits % 8 != 0: raise ValueError('invalid counter length')
         self._counter = [ 0 ] * (nbits // 8)
 
+        # Initialize the vector with the initial value
         index = len(self._counter) - 1
         while initial_value:
             self._counter[index] = initial_value % 256
@@ -78,21 +79,14 @@ def encrypt(self, plaintext):
         encrypted = [ ]
         for c in plaintext:
             if len(self._remaining_counter) == 0:
-                encrypted_counter = self._aes.encrypt(self._counter(), self._key, len(self._key))
-                self._remaining_counter.extend(encrypted_counter)
+                self._remaining_counter = self._aes.encrypt(self._counter(), self._key, len(self._key))
             encrypted.append(self._remaining_counter.pop(0) ^ ord(c))
 
         return "".join(chr(c) for c in encrypted)
 
     def decrypt(self, crypttext):
-        decrypted = [ ]
-        for c in crypttext:
-            if len(self._remaining_counter) == 0:
-                encrypted_counter = self._aes.encrypt(self._counter(), self._key, len(self._key))
-                self._remaining_counter.extend(encrypted_counter)
-            decrypted.append(self._remaining_counter.pop(0) ^ ord(c))
-
-        return "".join(chr(c) for c in decrypted)
+        # AES-CTR is symetric
+        return self.encrypt(crypttext)
 
 
 if __name__ == '__main__':

pyscrypt/file.py

@@ -58,7 +58,7 @@
 MODE_READ  = 'r'
 MODE_WRITE = 'w'
 
-BLOCK_SIZE = 16
+BLOCK_SIZE = 1024
 
 class InvalidScryptFileFormat(Exception): pass
 
@@ -308,47 +308,36 @@ def read(self, size = None):
             # end-of-file checksum) and decrypt into a decrypted buffer 1 block at a time
             while not self._read_finished:
 
-                # We have enough decrypted bytes (or will after decrypted a few encrypted blocks)
-                if len(self._decrypted_buffer) + len(self._encrypted_buffer) - 32 >= size: break
+                # We have enough decrypted bytes (or will after decrypting the encrypted buffer)
+                available = len(self._decrypted_buffer) + len(self._encrypted_buffer) - 32
+                if available >= size: break
 
-                # Read enough AES-256 blocks with a little extra for the possible final checksum
-                count = int(BLOCK_SIZE * math.ceil(size / BLOCK_SIZE)) + 32
-                data = ''
-                while len(data) < count:
-                  chunk = self._fp.read(count - len(data))
-                  if not chunk: break
-                  data += chunk
+                # Read a little extra for the possible final checksum
+                data = self._fp.read(BLOCK_SIZE)
 
-                self._encrypted_buffer += data
-
-                # We didn't get as much as we wanted... The file must be done
-                if len(data) != count:
+                # No data left; we're done
+                if not data:
                     self._read_finished = True
                     break
 
-        # Decrypt as many of the encrypted blocks as possible (leaving the final check sum)
-        while len(self._encrypted_buffer) >= BLOCK_SIZE + 32:
-            block = self._encrypted_buffer[:BLOCK_SIZE]
-            self._decrypted_buffer += self._crypto.decrypt(block)
-            self._checksumer.update(block)
-            self._encrypted_buffer = self._encrypted_buffer[BLOCK_SIZE:]
+                self._encrypted_buffer += data
+
+        # Decrypt as much of the encrypted data as possible (leaving the final check sum)
+        safe = self._encrypted_buffer[:-32]
+        self._encrypted_buffer = self._encrypted_buffer[-32:]
+        self._decrypted_buffer += self._crypto.decrypt(safe)
+        self._checksumer.update(safe)
 
-        # We read all the bytes, so what is left is possible a little more encrypted bytes and the checksum
+        # We read all the bytes, only the checksum remains
         if self._read_finished:
-            leftover = self._encrypted_buffer[:-32]
-            checksum = self._encrypted_buffer[-32:]
-            if leftover:
-                self._decrypted_buffer += self._crypto.decrypt(leftover)
-                self._checksumer.update(leftover)
-            self._check_final_checksum(checksum)
+            self._check_final_checksum(self._encrypted_buffer)
 
         # Send back the number of bytes requests and remove them from the buffer
         decrypted = self._decrypted_buffer[:size]
         self._decrypted_buffer = self._decrypted_buffer[size:]
 
         return decrypted
 
-
     # Write operations
     def flush(self):
         "Flush the underlying file object's I/O buffer."
@@ -422,13 +411,7 @@ def write(self, str):
         if not self._done_header:
             self._write_header()
 
-        self._decrypted_buffer += str
-
-        output = ''
-        while len(self._decrypted_buffer) >= BLOCK_SIZE:
-            block = self._crypto.encrypt(self._decrypted_buffer[:BLOCK_SIZE])
-            output += block
-            self._decrypted_buffer = self._decrypted_buffer[BLOCK_SIZE:]
-            self._checksumer.update(block)
-        self._fp.write(output)
+        encrypted = self._crypto.encrypt(str)
+        self._checksumer.update(encrypted)
+        self._fp.write(encrypted)
 

tests/run-tests-file.py

@@ -24,7 +24,7 @@
     sf.close()
 
     result = {True: "pass", False: "fail"}[decrypted == plaintext]
-    print "Test Encrypt/Decrypt: text_length=%s result=%s" % (text_length, result)
+    print "Test Encrypt/Decrypt: text_length=%s result=%s valid=%s" % (text_length, result, sf.valid)
 
 
 # Generate some files to make sure the tarsnap scrypt utility can read them