import lib.rsa as rsa
import lib.diffie_hellman as diffie
from hashlib import sha256

BITS = 32
ALPH = {'a': 0, 'b': 1, 'c': 2, 'd': 3, 'e': 4, 'f': 5, 'g': 6, 'h': 7, 'i': 8, 'j': 9, 'k': 10, 'l': 11, 'm': 12, 'n': 13, 'o': 14, 'p': 15, 'q': 16, 'r': 17, 's': 18, 't': 19, 'u': 20, 'v': 21, 'w': 22, 'x': 23, 'y': 24, 'z': 25, ' ': 26}
LEN = BITS // 5

def encode(text):
    sum = 0
    for i in range(len(text)):
        # Codage de la chaîne en base 27
        sum += ALPH[text[i]] * 27 ** i
    return sum

def decode(num):
    text = ""
    while num > 0:
        text += list(ALPH.keys())[num % 27]
        num //= 27
    return text

def encrypt(text, e, n):
    crypt = []
    while len(text) > 0:
        if len(text) >= LEN:
            crypt.append(rsa.encrypt(encode(text[:LEN]), e, n))
            text = text[LEN:]
        else:
            crypt.append(rsa.encrypt(encode(text), e, n))
            text = ""
    return crypt

def decrypt(crypt, d, n):
    text = ""
    for i in range(len(crypt)):
        text += decode(rsa.decrypt(crypt[i], d, n))
    return text

if __name__ == "__main__":
    p, g = diffie.get_p_and_g(BITS)

    a = diffie.get_x(BITS)
    b = diffie.get_x(BITS)

    A = diffie.get_X(a, p, g)
    B = diffie.get_X(b, p, g)
    
    K_a = diffie.get_K(B, a, p)
    K_b = diffie.get_K(A, b, p)

    e, d, n = rsa.get_keys(BITS)

    c_d = d ^ K_a
    h_d = sha256(str(c_d).encode())

    m = "hello world"
    crypt = encrypt(m, e, n)
    print(crypt)
    print(decrypt(crypt, d, n))