Minor updates for readability improvements.

encrypt.py

@@ -6,7 +6,7 @@
 from pkcs import PKCS7
 from feistel import FeistelNetwork
 from modes import ECB, CBC, CTR
-from iterators import file_block_iterator, eof_signal_iterator
+from iterators import file_block_iterator
 
 """
 The Mike Encryption Standard
@@ -66,16 +66,12 @@ def main():
         raise ValueError("Mode of operation {} is not recognised".format(args.mode))
 
     # Encrypt or decrypt
-    if args.encrypt:
-        file_blocks = file_block_iterator(args.input_file, cipher.block_size)
-        with open(args.output_file, 'wb') as f:
-            for ciphertext in mode.encrypt(file_blocks):
-                f.write(ciphertext)
-    else:
-        file_blocks = file_block_iterator(args.input_file, cipher.block_size)
-        with open(args.output_file, 'wb') as f:
-            for plaintext in mode.decrypt(file_blocks):
-                f.write(plaintext)
+    operation = mode.encrypt if args.encrypt else mode.decrypt
+
+    file_blocks = file_block_iterator(args.input_file, cipher.block_size)
+    with open(args.output_file, 'wb') as f:
+        for block in operation(file_blocks):
+            f.write(block)
 
 
 if __name__ == "__main__":

feistel.py

@@ -1,3 +1,4 @@
+import sys
 import hashlib
 
 """ Class that implements a Feistel Cipher in ECB mode
@@ -22,27 +23,31 @@ def generate_round_keys(self, key, round_count, block_size):
 
         # [Prekey: 64] [R x Subkeys: 32] [Postkey: 64]
         keys = {
-            "prekey" : key_data[0:block_size],
-            "roundkeys" : [key_data[block_size + (half_block_size * a): block_size + half_block_size + (half_block_size * a)] for a in range(round_count)],
+            "prekey" : key_data[:block_size],
+            "roundkeys" : lambda r : key_data[block_size + (half_block_size * r): block_size + half_block_size + (half_block_size * r)],
             "postkey" : key_data[-block_size:],
         }
 
         return keys
 
     def _xor(self, a, b):
-        a_i = int.from_bytes(a, "little")
-        b_i = int.from_bytes(b, "little")
-        return (a_i ^ b_i).to_bytes(len(a), "little")
+        byteorder = sys.byteorder
+        a_i = int.from_bytes(a, byteorder)
+        b_i = int.from_bytes(b, byteorder)
+        return (a_i ^ b_i).to_bytes(len(a), byteorder)
 
-
-    def _reverse(self, block):
+    def _split(self, block):
         L = block[:self.block_size//2]
         R = block[self.block_size//2:]
+
+        return L, R
+
+    def _reverse(self, block):
+        L, R = self._split(block)
         return R+L
 
     def round(self, block, round_key):
-        L = block[:self.block_size//2]
-        R = block[self.block_size//2:]
+        L, R = self._split(block)
 
         Rprime = self._xor(R, round_key)
         sha256 = hashlib.sha256()
@@ -59,7 +64,7 @@ def encrypt_block(self, block):
 
         # Rounds
         for i in range(self.round_count):
-            block = self.round(block, self.keys["roundkeys"][i])
+            block = self.round(block, self.keys["roundkeys"](i))
 
         # Reverse half-block
         block = self._reverse(block)
@@ -75,8 +80,8 @@ def decrypt_block(self, block):
         block = self._xor(block, self.keys["postkey"])
 
         # Rounds
-        for i in range(self.round_count - 1, -1, -1):
-            block = self.round(block, self.keys["roundkeys"][i])
+        for i in reversed(range(self.round_count)):
+            block = self.round(block, self.keys["roundkeys"](i))
 
         # Reverse half-block
         block = self._reverse(block)

iterators.py

@@ -13,7 +13,7 @@ def file_block_iterator(path, block_size):
         data = f.read(block_size)
         while data != b"":
             yield data
-            data = f.read(64)
+            data = f.read(block_size)
 
 def list_block_iterator(message, block_size):
     for idx in range (0, len(message), block_size):

modes.py

@@ -1,20 +1,15 @@
 from iterators import eof_signal_iterator
-
+from utils import Utils
 """ Classes implementing modes of encryption:
 
 * We can extend this, currently only ECB, CBC and CTR is supported
 
 """
 class ModeOfOperation():
-    def __init__(self, cipher, iv = None, nonce = None):
+    def __init__(self, cipher):
         self.cipher = cipher
-        self.iv = iv
-        self.nonce = nonce
         self.block_size = cipher.block_size
 
-    def _xor(self, block1, block2):
-        return bytes([a ^ b for a, b in zip(block1, block2)])
-
 class ECB(ModeOfOperation):
     def __init__(self, cipher, padding_scheme):
         super(ECB, self).__init__(cipher)
@@ -57,8 +52,9 @@ def decrypt(self, block_iterator):
 class CBC(ModeOfOperation):
 
     def __init__(self, cipher, iv, padding_scheme):
-        super(CBC, self).__init__(cipher=cipher, iv=iv)
+        super(CBC, self).__init__(cipher=cipher)
         #initialize cipher_block with value None
+        self.iv = iv
         self.cipher_block = None
         self.padding_scheme = padding_scheme
 
@@ -74,22 +70,22 @@ def encrypt(self, block_iterator):
 
             if not eof:
                 self.cipher_block = self.cipher.encrypt_block(
-                    self._xor(data, self.cipher_block)
+                    Utils.xor(data, self.cipher_block)
                 )
             else:
                 #executed only once, for the last block of the file
                 block = data if not eof else self.padding_scheme.apply(data)
                 if len(block) == self.block_size:
                     self.cipher_block = self.cipher.encrypt_block(
-                        self._xor(block, self.cipher_block)
+                        Utils.xor(block, self.cipher_block)
                     )
                 elif len(block) == 2 * self.block_size:
                     last_block = self.cipher.encrypt_block(
-                        self._xor(block[:self.block_size], self.cipher_block)
+                        Utils.xor(block[:self.block_size], self.cipher_block)
                     )
                     #This will append an entire block of padding (??)
                     self.cipher_block = self.cipher.encrypt_block(
-                        self._xor(block[self.block_size:], last_block)
+                        Utils.xor(block[self.block_size:], last_block)
                     )
                     #set the cipher_block variable to be last block of real data
                     #prepended to the extra block of padding
@@ -108,7 +104,7 @@ def decrypt(self, block_iterator):
         self.cipher_block, eof = next(eof_iterator)
 
         for data, eof in eof_iterator:
-            plaintext = self._xor(
+            plaintext = Utils.xor(
                 self.cipher.decrypt_block(data), self.cipher_block
             )
             self.cipher_block = data
@@ -149,7 +145,7 @@ def encrypt(self, block_iterator):
             if eof and len(data) < self.block_size:
                 xor_block = xor_block[:len(data)]
 
-            ciphertext = self._xor(data, xor_block)
+            ciphertext = Utils.xor(data, xor_block)
             counter += 1
 
             yield ciphertext

utils.py

@@ -0,0 +1,4 @@
+class Utils:
+    @staticmethod
+    def xor(block1, block2):
+        return bytes([a ^ b for a, b in zip(block1, block2)])