Skip to content

Commit

Permalink
Cleaned up AES CTR.
Browse files Browse the repository at this point in the history
  • Loading branch information
ricmoo committed May 7, 2014
1 parent 306be14 commit 6b01263
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 48 deletions.
14 changes: 4 additions & 10 deletions pyscrypt/aesctr.py
Expand Up @@ -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
Expand Down Expand Up @@ -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__':
Expand Down
57 changes: 20 additions & 37 deletions pyscrypt/file.py
Expand Up @@ -58,7 +58,7 @@
MODE_READ = 'r'
MODE_WRITE = 'w'

BLOCK_SIZE = 16
BLOCK_SIZE = 1024

class InvalidScryptFileFormat(Exception): pass

Expand Down Expand Up @@ -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."
Expand Down Expand Up @@ -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)

2 changes: 1 addition & 1 deletion tests/run-tests-file.py
Expand Up @@ -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
Expand Down

0 comments on commit 6b01263

Please sign in to comment.