from Crypto.Cipher import DES
import  matplotlib.pyplot as plt
import time

def DESEncrypt(key,plaintext):
    cipher = DES.new(key,DES.MODE_ECB)
    ciphertext = cipher.encrypt(plaintext)
    return ciphertext

def DESDecrypt(key,ciphertext):
    cipher = DES.new(key,DES.MODE_ECB)
    plaintext = cipher.decrypt(ciphertext)
    return plaintext

#source
#https://stackoverflow.com/questions/21017698/converting-int-to-bytes-in-python-3
def int_to_bytes(x: int) -> bytes:
    return x.to_bytes((x.bit_length() + 7) // 8, 'big', signed = False)
#source
#https://stackoverflow.com/questions/21017698/converting-int-to-bytes-in-python-3
def int_from_bytes(xbytes: bytes) -> int:
    return int.from_bytes(xbytes, 'big',signed = False)

def intArrayToByteArray(intArray):
    tempByteArray = bytearray()
    for j in range(len(intArray)):
        if intArray[j] == 0:
            tempByteArray += b'0'
        tempByteArray += int_to_bytes(intArray[j])

    return tempByteArray



keyTable = []

start_time = time.time()

#Creates a table with all possible keys.
for firstByte in range(0,128):
    for secondByte in range(0,128):
        for thirdByte in range(0,64):
            key = [0] * 8
            key[0] = firstByte << 1
            key[1] = secondByte << 1
            key[2] = thirdByte << 1
            tempArray = intArrayToByteArray(key)
            keyTable.append(tempArray)



print("Key Table Gen. Finished: Exec. Time: %s" % (time.time() - start_time))


#byte1 = [1,1,1,1,1,1,1,0] # 0 (parity bit)
byte1_int = 254
byte2_int = 254
byte3_int = 126
key1  = bytearray(b'')
key1.append(byte1_int)
key1.append(byte2_int)
key1.append(byte3_int)
for i in range(5):
    key1 += b'0'
print("key1 bytearray: {} \nkey1 hex: {}".format(key1,key1.hex()))

key2  = bytearray(b'')
byte1_int_key2 = 170
byte2_int_key2 = 170
byte3_int_key2 = 42
key2.append(byte1_int_key2)
key2.append(byte2_int_key2)
key2.append(byte3_int_key2)
for i in range(5):
    key2 += b'0'

print("key2 bytearray: {} \nkey2 hex: {}".format(key2,key2.hex()))

plaintext_orig = bytearray(b'01234567')
print("Plaintext: ",plaintext_orig.hex())
ciphertext = DESEncrypt(key2,DESEncrypt(key1,plaintext_orig))

EncTable = dict()
for key in keyTable:
    tempCipher = DESEncrypt(key, plaintext_orig)
    EncTable[tempCipher] = key

print("All ciphertext calculated!")

restoredKeyPair = []

for key in keyTable:
    tempDecryptedValue = DESDecrypt(key, ciphertext)
    if tempDecryptedValue in EncTable.keys():
        key1 = EncTable[tempDecryptedValue]
        key2 = key
        restoredKeyPair.append((key1, key2))
        key1 = 0
        key2 = 0

print("EXECUTION DONE --- TIME: --- %s seconds ---" % (time.time() - start_time))
print("RETRIEVED KEY PAIR: ",restoredKeyPair)
print("RETRIVEVED KEY PAIR: ",restoredKeyPair[0][0].hex(),restoredKeyPair[0][1].hex())