Implemented CBC and ECB file encryption and corrected tests to work for file encryption as well.

Signed-off-by: Gabriel Lindeblad <Gabriel.lindeblad@icloud.com>

Author: Glindeb

src/AES_Python/main.py

@@ -4,13 +4,16 @@
 be used in any other use case than educational and no security is guaranteed for data
 encrypted or decrypted using this tool.
 """
+import os
 
 # Imported libraries
 import numpy as np  # Used for arrays and mathematical operations.
 import galois  # Used for GF(2^8) multiplication in mix columns operation.
 from numpy.typing import NDArray  # Used for type hinting numpy arrays.
 from typing import Any  # Used for type hinting __getattr__ function.
 from secrets import token_bytes  # Used for generating random key if needed.
+from os.path import getsize  # Used to acquire size of files.
+from os import remove  # Used to remove files.
 
 
 class AES:
@@ -138,7 +141,7 @@ def get(self, item: str) -> Any:
         :param item: Attribute to be retrieved. Valid attributes (running_mode, key, iv).
         :return: Attribute value.
         """
-        if item:
+        if item in ["running_mode", "key", "iv"]:
             return self.__dict__[f"_{item}"]
         else:
             raise AttributeError(f"No attribute <{item}> exists!")
@@ -176,8 +179,17 @@ def enc(self, *, data_string: str = "", file_path: str = "",
                 return self.__cbc_enc(data_string=data_string, keys=self.key_expand(key), iv=iv)
             else:
                 raise NotImplementedError(f"{running_mode} is not supported!")
+        elif file_path:
+            if running_mode == "ECB":
+                self.__ecb_enc(file_path=file_path, keys=self.key_expand(key))
+                return ""
+            elif running_mode == "CBC":
+                self.__cbc_enc(file_path=file_path, keys=self.key_expand(key), iv=iv)
+                return ""
+            else:
+                raise NotImplementedError(f"{running_mode} is not supported!")
         else:
-            raise NotImplementedError("File encryption is not implemented yet...")
+            raise RuntimeWarning("No file or string was give...")
 
     def dec(self, *, data_string: str = "", file_path: str = "",
             running_mode: str = "", key: str = "", iv: str = "") -> str:
@@ -212,8 +224,17 @@ def dec(self, *, data_string: str = "", file_path: str = "",
                 return self.__cbc_dec(data_string=data_string, keys=self.key_expand(key), iv=iv)
             else:
                 raise NotImplementedError(f"{running_mode} is not supported!")
+        elif file_path:
+            if running_mode == "ECB":
+                self.__ecb_dec(file_path=file_path, keys=self.key_expand(key))
+                return ""
+            elif running_mode == "CBC":
+                self.__cbc_dec(file_path=file_path, keys=self.key_expand(key), iv=iv)
+                return ""
+            else:
+                raise NotImplementedError(f"{running_mode} is not supported!")
         else:
-            raise NotImplementedError("File encryption is not implemented yet...")
+            raise RuntimeWarning("No file or string was give...")
 
     @classmethod
     def __ecb_enc(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[np.uint8]) -> str:
@@ -234,7 +255,7 @@ def __ecb_enc(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[
                 enc: str = "".join(cls.vec_chr(cls.__enc_schedule(raw, keys).flatten().astype(np.uint8)))
                 output_string += enc
 
-            extra = len(data_string) % 16   # Calculates length of final data block
+            extra = len(data_string) % 16  # Calculates length of final data block
             result: str = ""
 
             if extra != 0:  # If last data block not integer multiple of 16 adds extra padding
@@ -244,8 +265,33 @@ def __ecb_enc(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[
                 result = "".join(cls.vec_chr(cls.__enc_schedule(raw.reshape(4, 4), keys).flatten().astype(np.uint8)))
 
             return output_string + result
+        elif file_path:
+            file_size = getsize(file_path)
+
+            with open(f"{file_path}.enc", 'wb') as output, open(file_path, 'rb') as data:
+                for i in range(int(file_size / 16)):
+                    raw = np.array([i for i in data.read(16)], dtype=np.uint8).reshape(4, 4)
+                    out = bytes((cls.__enc_schedule(raw, keys).flatten()).tolist())
+                    output.write(out)
+
+                if file_size % 16 != 0:
+                    raw_l = [i for i in data.read()]
+                    raw_l, length = cls.__add_padding(raw_l)
+
+                    out = bytes((cls.__enc_schedule(np.array(raw_l).reshape(4, 4), keys).flatten()).tolist())
+                    identifier = bytes((cls.__enc_schedule(np.array([0 for i in range(15)] + [length]).reshape(4, 4),
+                                                           keys).flatten()).tolist())
+
+                    output.write(out + identifier)
+                else:
+                    identifier = bytes((cls.__enc_schedule(np.array([0 for i in range(16)]).reshape(4, 4),
+                                                           keys).flatten()).tolist())
+                    output.write(identifier)
+
+            remove(file_path)
+            return ""
         else:
-            raise NotImplementedError
+            raise RuntimeError("No string or file path received...?")
 
     @classmethod
     def __ecb_dec(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[np.uint8]) -> str:
@@ -268,8 +314,30 @@ def __ecb_dec(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[
                 output_string += dec
 
             return output_string
+        elif file_path:
+            file_size = getsize(file_path)
+            file_name = file_path[:-4]
+
+            with open(f"{file_name}", 'wb') as output, open(file_path, 'rb') as data:
+                for i in range(int(file_size / 16) - 2):
+                    raw = np.array([i for i in data.read(16)]).reshape(4, 4)
+                    result = bytes((cls.__dec_schedule(raw, keys).flatten()).tolist())
+                    output.write(result)
+
+                data_final = np.array([i for i in data.read(16)]).reshape(4, 4)
+                identifier = np.array([i for i in data.read()]).reshape(4, 4)
+
+                data_dec = (cls.__dec_schedule(data_final, keys).flatten()).tolist()
+                id_l = (cls.__dec_schedule(identifier, keys).flatten()).tolist()
+
+                result = bytes(cls.__remove_padding(data_dec, id_l))
+
+                output.write(result)
+
+            remove(file_path)
+            return ""
         else:
-            raise NotImplementedError
+            raise RuntimeError("No string or file path received...?")
 
     @classmethod
     def __cbc_enc(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[np.uint8], iv: str) -> str:
@@ -289,7 +357,7 @@ def __cbc_enc(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[
             for i in range(len(data_string) // 16):  # Encryption cycle, skips last if not integer multiple of 16 bytes
                 raw: NDArray[np.uint8] = np.array([ord(i) for i in data_string[(i * 16): ((i + 1) * 16)]],
                                                   dtype=np.uint8).reshape(4, 4)
-                enc = np.bitwise_xor(raw, enc_array)    # Xor operation with previous encrypted block or iv
+                enc = np.bitwise_xor(raw, enc_array)  # Xor operation with previous encrypted block or iv
                 enc_array = cls.__enc_schedule(enc, keys)
                 output_string += "".join(cls.vec_chr(enc_array.flatten().astype(np.uint8)))
 
@@ -306,8 +374,35 @@ def __cbc_enc(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[
                 result = "".join(cls.vec_chr(cls.__enc_schedule(temp_array, keys).flatten().astype(np.uint8)))
 
             return output_string + result
+        elif file_path:
+            file_size = getsize(file_path)
+
+            with open(f"{file_path}.enc", 'wb') as output, open(file_path, 'rb') as data:
+                for i in range(int(file_size / 16)):
+                    raw = np.array([i for i in data.read(16)]).reshape(4, 4)
+                    raw = np.bitwise_xor(raw, enc_array)
+                    enc_array = cls.__enc_schedule(raw, keys)
+                    output.write(bytes((enc_array.flatten()).tolist()))
+
+                if file_size % 16 != 0:
+                    final = [i for i in data.read()]
+                    final, length = cls.__add_padding(final)
+
+                    raw = np.bitwise_xor(np.array(final).reshape(4, 4), enc_array)
+                    enc_array = cls.__enc_schedule(raw, keys)
+
+                    identifier = np.bitwise_xor(np.array([0 for i in range(15)] + [length]).reshape(4, 4), enc_array)
+                    identifier = cls.__enc_schedule(identifier, keys)
+
+                    output.write(bytes((enc_array.flatten()).tolist() + (identifier.flatten()).tolist()))
+                else:
+                    identifier = np.bitwise_xor(np.array([0 for i in range(16)]).reshape(4, 4), enc_array)
+                    id_bytes = bytes(((cls.__enc_schedule(identifier, keys)).flatten()).tolist())
+                    output.write(id_bytes)
+            remove(file_path)
+            return ""
         else:
-            raise NotImplementedError
+            raise RuntimeError("No string or file path received...?")
 
     @classmethod
     def __cbc_dec(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[np.uint8], iv: str) -> str:
@@ -325,7 +420,7 @@ def __cbc_dec(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[
             output_string: str = ""
 
             for i in range(len(data_string) // 16):  # Decryption cycle
-                raw: NDArray[np.uint8] = np.array(   # Reads in input string 16 bytes at a time
+                raw: NDArray[np.uint8] = np.array(  # Reads in input string 16 bytes at a time
                     [ord(i) for i in data_string[(i * 16): ((i + 1) * 16)]], dtype=np.uint8).reshape(4, 4)
 
                 temp_array = cls.__dec_schedule(raw, keys)
@@ -336,8 +431,44 @@ def __cbc_dec(cls, *, data_string: str = "", file_path: str = "", keys: NDArray[
                 output_string += dec
 
             return output_string
+        elif file_path:
+            file_size = getsize(file_path)
+            file_name = file_path[:-4]
+
+            with open(f"{file_name}", 'wb') as output, open(file_path, 'rb') as data:
+                if int(file_size / 16) - 3 >= 0:
+                    vector = np.array([i for i in data.read(16)]).reshape(4, 4)
+                    raw = cls.__dec_schedule(vector, keys)
+                    result = np.bitwise_xor(raw, dec_array)
+                    output.write(bytes((result.flatten()).tolist()))
+
+                    for i in range(int(file_size / 16) - 3):
+                        raw = np.array([i for i in data.read(16)]).reshape(4, 4)
+                        result = cls.__dec_schedule(raw, keys)
+                        result = np.bitwise_xor(result, vector)
+                        vector = raw
+                        output.write(bytes((result.flatten()).tolist()))
+                else:
+                    vector = dec_array
+
+                data_pice = np.array([i for i in data.read(16)]).reshape(4, 4)
+                vector_1, identifier = data_pice, np.array([i for i in data.read()]).reshape(4, 4)
+
+                result = cls.__dec_schedule(data_pice, keys)
+                identifier = cls.__dec_schedule(identifier, keys)
+
+                identifier = np.bitwise_xor(identifier, vector_1)
+                data_pice = np.bitwise_xor(result, vector)
+
+                result_bytes = bytes(cls.__remove_padding((data_pice.flatten()).tolist(),
+                                                          (identifier.flatten()).tolist()))
+
+                output.write(result_bytes)
+
+            remove(file_path)
+            return ""
         else:
-            raise NotImplementedError
+            raise RuntimeError("No string or file path received...?")
 
     @staticmethod
     def key_gen(length: int = 16) -> str:
@@ -533,3 +664,21 @@ def __mix_columns(matrix: NDArray[np.uint8], shift: int) -> NDArray[np.uint8]:
                 result[j, i] = temp
 
         return result.transpose()
+
+    @staticmethod
+    def __add_padding(data: list[int]) -> tuple[list[int], int]:
+        """Adds a padding to ensure a bloke size of 16 bytes."""
+        length = 16 - len(data)
+        for i in range(length):
+            data.append(0)
+        return data, length
+
+    @staticmethod
+    def __remove_padding(data: list[int], identifier: list[int]) -> list[int]:
+        """Removes the padding from the data using information from identifier."""
+        if identifier[-1] == 0:
+            return data
+        elif 0 < identifier[-1] < 16:
+            return data[:-identifier[-1]]
+        else:
+            raise ValueError('Invalid padding')

temp/test.py

@@ -1,25 +1,25 @@
 import numpy as np
+import os
 from AES_Python import AES
 
 aes_test = AES(running_mode="CBC",
                key="2b7e151628aed2a6abf7158809cf4f3c",
-               iv="000102030405060708090A0B0C0D0E0F"
+               iv="000102030405060708090a0b0c0d0e0f"
                )
 
-print(aes_test, "\n")
+data = b'\x1b\x16\x86:\xb9*w\xc5)"\xe4\xe9D\\\xf1\xee\x8b\x03\xcc\xe7\x0c~\xba7\xcf\x0f\x9c\x16dM$\xe9\x91\xef\xc3\xa6\xd2\xf0\xcd\xc2\xee\x86\xf0\x90\x8a]\x87\xf5R\xe2.c\xd4\xc6T\xdc\xe0#\xa7X\x8b_\x81\x04'
+file_name = "tmp.txt"
+expected = b'1234567890123456789012345678901234567890'
 
-data_raw = '6bc1bee22e409f96e93d7e117393172a'
+with open(f"{file_name}.enc", "wb") as file:
+    file.write(data)
 
-data = "".join([chr(i) for i in bytes.fromhex(data_raw)])
+aes_test.dec(file_path=f"{file_name}.enc")
 
-print("Original data:", data)
+with open(f"{file_name}", "rb") as file:
+    result = file.read()
 
-enc_data = aes_test.enc(data_string=data)
+os.remove(f"{file_name}")
 
-print("Encrypted data:", bytearray([ord(i) for i in enc_data]).hex())
-
-print("Expected:", "7649abac8119b246cee98e9b12e9197d")
-
-dec_data = aes_test.dec(data_string=enc_data)
-
-print("Decrypted data:", dec_data)
+print("Result:", result)
+print("Expected:", expected)

tests/test_file_dec.py

@@ -3,7 +3,6 @@
 from AES_Python import AES
 
 
-@pytest.mark.skip(reason="Not correctly implemented yet")
 @pytest.mark.parametrize("data,key,file_name,expected", [
     # 128 bit
     (b'|\x94\x18\xcf\x1c\xf0\xef\xa0\xff\xa4\xbb\xe9\xd8\x8am\xa40f\xe4\x1eg\x9d\x88\xb8\xef\xeb{=J\xf3\xf6\xc1',
@@ -28,7 +27,7 @@ def test_file_dec_ecb(data, key, file_name, expected):
         file.write(data)
 
     ecb = AES(key=key, running_mode="ECB")
-    ecb.dec(file_path=file_name)
+    ecb.dec(file_path=f"{file_name}.enc")
 
     with open(file_name, "rb") as file:
         result = file.read()
@@ -38,7 +37,6 @@ def test_file_dec_ecb(data, key, file_name, expected):
     assert result == expected
 
 
-@pytest.mark.skip(reason="Not correctly implemented yet")
 @pytest.mark.parametrize("data,key,file_name,iv,expected", [
     # 128 bit
     (b'\xe4\xa7\x0e\xbd\x84\xfa\xf5\xd8`\xb8\xa1\x10\x0b~\xadh\x89Feso\xc5~_|\xe9\x1bG\xd9*\\\x81',
@@ -64,8 +62,8 @@ def test_file_dec_cbc(data, key, file_name, iv, expected):
     with open(f"{file_name}.enc", "wb") as file:
         file.write(data)
 
-    cbc = AES(key=key, running_mode="CBC")
-    cbc.dec(file_path=file_name)
+    cbc = AES(key=key, iv=iv, running_mode="CBC")
+    cbc.dec(file_path=f"{file_name}.enc")
 
     with open(file_name, "rb") as file:
         result = file.read()

tests/test_file_enc.py

@@ -3,7 +3,6 @@
 from AES_Python import AES
 
 
-@pytest.mark.skip(reason="Not correctly implemented yet")
 @pytest.mark.parametrize("data,key,file_name,expected", [
     # 128 bit
     (b'1234567890', "2b7e151628aed2a6abf7158809cf4f3c", "tmp.txt",
@@ -36,7 +35,6 @@ def test_file_enc_ecb(data, key, file_name, expected):
     assert result == expected
 
 
-@pytest.mark.skip(reason="Not correctly implemented yet")
 @pytest.mark.parametrize("data,key,file_name,iv,expected", [
     # 128 bit
     (b'1234567890', "2b7e151628aed2a6abf7158809cf4f3c", "tmp.txt",