From 6b01263feb0ba45a28fb7475b90c523c731d2db2 Mon Sep 17 00:00:00 2001 From: Richard Moore Date: Wed, 7 May 2014 04:05:37 -0400 Subject: [PATCH] Cleaned up AES CTR. --- pyscrypt/aesctr.py | 14 +++------- pyscrypt/file.py | 57 +++++++++++++++-------------------------- tests/run-tests-file.py | 2 +- 3 files changed, 25 insertions(+), 48 deletions(-) diff --git a/pyscrypt/aesctr.py b/pyscrypt/aesctr.py index 790e686..094f716 100644 --- a/pyscrypt/aesctr.py +++ b/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__': diff --git a/pyscrypt/file.py b/pyscrypt/file.py index 45ccc22..105f9a5 100644 --- a/pyscrypt/file.py +++ b/pyscrypt/file.py @@ -58,7 +58,7 @@ MODE_READ = 'r' MODE_WRITE = 'w' -BLOCK_SIZE = 16 +BLOCK_SIZE = 1024 class InvalidScryptFileFormat(Exception): pass @@ -308,39 +308,29 @@ 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] @@ -348,7 +338,6 @@ def read(self, size = None): 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) diff --git a/tests/run-tests-file.py b/tests/run-tests-file.py index 5f05840..5d5f360 100644 --- a/tests/run-tests-file.py +++ b/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