def test_cipher(self):
"""Test AES cipher with all key lengths"""
import test_keys
import key_expander
test_data = test_keys.TestKeys()
for key_size in 128, 192, 256:
test_key_expander = key_expander.KeyExpander(key_size)
test_expanded_key = test_key_expander.expand(
test_data.test_key[key_size])
test_cipher = AESCipher(test_expanded_key)
test_result_ciphertext = test_cipher.cipher_block(
test_data.test_block_plaintext)
self.assertEquals(len([
i for i, j in zip(
test_result_ciphertext,
test_data.test_block_ciphertext_validated[key_size])
if i == j
]),
16,
msg='Test %d bit cipher' % key_size)
test_result_plaintext = test_cipher.decipher_block(
test_data.test_block_ciphertext_validated[key_size])
self.assertEquals(len([
i for i, j in zip(test_result_plaintext,
test_data.test_block_plaintext) if i == j
]),
16,
msg='Test %d bit decipher' % key_size)
def test_mode(self):
# Self test
import key_expander
import aes_cipher
import test_keys
test_data = test_keys.TestKeys()
test_expander = key_expander.KeyExpander(256)
test_expanded_key = test_expander.expand(test_data.test_mode_key)
test_cipher = aes_cipher.AESCipher(test_expanded_key)
test_cfb = CFBMode(test_cipher, 16)
test_cfb.set_iv(test_data.test_mode_iv)
for k in range(4):
self.assertEquals(len([i for i, j in zip(test_data.test_cfb_ciphertext[k],
test_cfb.encrypt_block(test_data.test_mode_plaintext[k])) if
i == j]),
16,
msg='CFB encrypt test block' + str(k))
test_cfb.set_iv(test_data.test_mode_iv)
for k in range(4):
self.assertEquals(len([i for i, j in zip(test_data.test_mode_plaintext[k],
test_cfb.decrypt_block(test_data.test_cfb_ciphertext[k])) if
i == j]),
16,
msg='CFB decrypt test block' + str(k))
def encrypt_file(self: AESdemo, in_file_path: str,
out_file_path: str) -> bool:
password = None #bg# was optional arg
#If a password is provided, generate new salt and create key and iv
if password is not None:
#bg#self.new_salt()
#bg#self.create_key_from_password(password)
pass #bg#
else:
self._salt = None
#If key and iv are not provided are established above, bail out.
if self._key is None or self._iv is None:
return False
#Initialize encryption using key and iv
key_expander_256 = key_expander.KeyExpander(256)
expanded_key = key_expander_256.expand(self._key)
aes_cipher_256 = aes_cipher.AESCipher(expanded_key)
aes_cbc_256 = cbc_mode.CBCMode(aes_cipher_256, 16)
aes_cbc_256.set_iv(self._iv)
#Get filesize of original file for storage in encrypted file
#bg: removed exception handling
filesize = os.stat(in_file_path)[6]
with open(in_file_path, 'rb') as in_file:
with open(out_file_path, 'wb') as out_file:
#Write salt if present
if self._salt is not None:
#bg#out_file.write(self._salt)
pass
#Write filesize of original
out_file.write(struct.pack('L', filesize))
#Encrypt to eof
eof = False
while not eof:
in_data = in_file.read(16)
if len(in_data) == 0:
eof = True
else:
out_data = aes_cbc_256.encrypt_block(
[int(x) for x in bytearray(in_data)])
out_file.write(bytes(out_data))
self._salt = None
return True
def decrypt_file(self: AESdemo, in_file_path: str,
out_file_path: str) -> bool:
password = None #bg# was optional arg.
with open(in_file_path, 'rb') as in_file:
#If a password is provided, generate key and iv using salt from file.
#bg#if password is not None:
#bg# self._salt = in_file.read (32)
#bg# self.create_key_from_password (password)
#Key and iv have not been generated or provided, bail out
if self._key is None or self._iv is None:
return False
#Initialize encryption using key and iv
key_expander_256 = key_expander.KeyExpander(256)
expanded_key = key_expander_256.expand(self._key)
aes_cipher_256 = aes_cipher.AESCipher(expanded_key)
aes_cbc_256 = cbc_mode.CBCMode(aes_cipher_256, 16)
aes_cbc_256.set_iv(self._iv)
#Read original file size
filesize = struct.unpack('L',
in_file.read(struct.calcsize('L')))[0]
#Decrypt to eof
with open(out_file_path, 'wb') as out_file:
eof = False
while not eof:
in_data = in_file.read(16)
if len(in_data) == 0:
eof = True
else:
out_data = aes_cbc_256.decrypt_block(
list(bytearray(in_data)))
#At end of file, if end of original file is within < 16 bytes slice it out.
if filesize - out_file.tell() < 16:
out_file.write(
bytes(out_data[:filesize - out_file.tell()]))
else:
out_file.write(bytes(out_data))
self._salt = None
return True
