from Crypto.Cipher import AES
from Crypto import Random

class NodeKM:

    def __init__(self, k1, k2, k3, iv):
        self.key_ecb = k1
        self.key_cbc = k2
        self.key_g = k3
        self.iv = iv

    def get_key(self, mode):
        if mode == "ECB":
            cipher = AES.new(self.key_g, AES.MODE_ECB)
            return cipher.encrypt(self.key_ecb)

        if mode == "CBC":
            cipher = AES.new(self.key_g, AES.MODE_CBC, self.iv)
            return cipher.encrypt(self.key_cbc)


class NodeA:
    def __init__(self, k, iv):
        self.key_g = k
        self.iv = iv

    def send_to_B_operating_mode(self, mode):
        global node_b

        self.mode = mode
        node_b.set_mode(mode)

    def get_key_from_KM(self, mode):
        global node_km

        if mode == "ECB":
            cipher = AES.new(self.key_g, AES.MODE_ECB)
        if mode == "CBC":
            cipher = AES.new(self.key_g, AES.MODE_CBC, self.iv)

        self.key_m = cipher.decrypt(node_km.get_key(mode))

    def encrypt(self):
        global node_b

        if mode == "ECB":
            cipher = AES.new(self.key_m, AES.MODE_ECB)
        if mode == "CBC":
            cipher = AES.new(self.key_m, AES.MODE_CBC, self.iv)

        enc_file = cipher.encrypt(_pad(open('data.txt', 'rb').read()))
        enc_key = cipher.encrypt(_pad(self.key_m))
        enc_mode = cipher.encrypt(_pad(self.mode))

        print('enc_file: ', enc_file)
        print('enc_key: ', enc_key)
        print('enc_mode: ', enc_mode)

        print('\ndec_file: ', node_b.decrypt(enc_file))
        print('dec_key: ', node_b.decrypt(enc_key))
        print('dec_mode: ', node_b.decrypt(enc_mode))

class NodeB:
    def __init__(self, k, iv):
        self.key_g = k
        self.iv = iv

    def set_mode(self, mode):
        global node_km
        global node_a

        self.mode = mode

        if mode == "ECB":
            cipher = AES.new(self.key_g, AES.MODE_ECB)
            self.key_m = cipher.decrypt(node_km.get_key(mode))

        if mode == "CBC":
            cipher = AES.new(self.key_g, AES.MODE_CBC, iv)
            self.key_m = cipher.decrypt(node_km.get_key(mode))

    def decrypt(self, content):
        if self.mode == "ECB":
            cipher = AES.new(self.key_m, AES.MODE_ECB)
        if self.mode == "CBC":
            cipher = AES.new(self.key_m, AES.MODE_CBC, self.iv)

        return _unpad(cipher.decrypt(content))



def fill_with_space(key):
    while len(key) < 16:
        key += "."
    return key


def _pad(s):
    if type(s) is bytes:
        s = s.decode("utf-8")
    return s + (AES.block_size - len(s) % AES.block_size) * chr(AES.block_size - len(s) % AES.block_size)

def _unpad(s):
    return s[:-ord(s[len(s)-1:])]

mode = "ECB"

key_ecb = fill_with_space("ecb_key")
key_cbc = fill_with_space("cbc_key")
key_general = fill_with_space("general_key")
iv = Random.new().read(AES.block_size)

node_km = NodeKM(key_ecb, key_cbc, key_general, iv)
node_a = NodeA(key_general, iv)
node_b = NodeB(key_general, iv)

node_a.send_to_B_operating_mode(mode)
node_a.get_key_from_KM(mode)
node_a.encrypt()

print("\nNode A")
print(f"mode: {node_a.mode}")
print(f"key: {node_a.key_g}")
print(f"key mode: {node_a.key_m}")

print("\nNode B")
print(f"mode: {node_b.mode}")
print(f"key: {node_b.key_g}")
print(f"key mode: {node_b.key_m}")