from common import *
import os
import socket
from AES import AESCipher
from random import randint

HOST = "127.0.0.1"
CLIENT_ID = None
KEY = None

def register_user(username, password):
    global CLIENT_ID, KEY
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, AS_PORT))
        key = os.urandom(32)
        message = f"register||{username},{password},{key}"
        s.sendall(message.encode())
        data = s.recv(4096)

        data = data.decode()
        if data == "success":
            print("Login succeeded!")
            CLIENT_ID = username
            KEY = key

def login(username, password):
    global CLIENT_ID, KEY
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, AS_PORT))
        message = f"login||{username},{password}"
        s.sendall(message.encode())
        data = s.recv(4096)

        data = data.decode()
        data = data.split(",")
        if data[0] == "success":
            print(f"Login succeeded!")
            CLIENT_ID = username
            KEY = data[1]

def print_something():
    global CLIENT_ID, KEY

    if CLIENT_ID is None:
        print("You must register or login first!")
        return

    AES = AESCipher(C_User_Key)

    ID_C = CLIENT_ID
    ID_S = "PRINTER"
    T_R = int(input("Enter the ticket lifetime (minutes): "))
    T_R = T_R * 60
    N1 = randint(0, 1000000)

    M2 = None
    M4 = None
    M6 = None

    M1_inner = f"{ID_S},{T_R},{N1}"
    M1_inner = AES.encrypt(M1_inner)
    M1 = f"{ID_C},{M1_inner}"
    message = f"request||{M1}"
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, AS_PORT))
        print('sending m1', message)
        s.sendall(message.encode())

        data = s.recv(4096)
        data = data.decode()
        print('received m2', data)
        M2 = data

    '''
    M2 = [{K_c_tgs + N_1}Kc + T_c_tgs]
    M3 = [{ID_C + ID_S + T_R + N2}K_c_tgs + T_c_tgs]
    '''
    M2 = M2.split(",")
    M2_inner = AES.decrypt(M2[0])
    M2_inner = M2_inner.split(",")

    K_c_tgs = eval(M2_inner[0])
    N1 = M2_inner[1]
    T_c_tgs = M2[1]

    M2_AES = AESCipher(K_c_tgs)
    N2 = randint(0, 1000000)

    M3_inner = f"{ID_C},{ID_S},{T_R},{N2}"
    M3_inner = M2_AES.encrypt(M3_inner)
    M3 = f"{M3_inner},{T_c_tgs}"
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, TGS_PORT))
        message = f"request||{M3}"
        print('sending m3', message)
        s.sendall(message.encode())

        data = s.recv(4096)
        data = data.decode()
        print('received m4', data)
        M4 = data

    '''
        M4 = [{K_c_s + T_A + N2}K_c_tgs + T_c_s]
            T_c_s = {ID_C + T_A + K_c_s}K_s
        M5 = [{ID_C + T_A + S_R + N3}K_c_s + T_c_s]
    '''
    M4 = M4.split(",")
    M4_AES = AESCipher(K_c_tgs)
    M4_inner = M4_AES.decrypt(M4[0])
    M4_inner = M4_inner.split(",")
    print('M4_inner', M4_inner)

    K_c_s = eval(M4_inner[0])
    T_A = M4_inner[1]
    N2 = M4_inner[2]

    T_c_s = M4[1]

    M5_AES = AESCipher(K_c_s)
    N3 = randint(0, 1000000)

    M5_inner = f"{ID_C},{T_A},{ID_S},{N3}"
    M5_inner = M5_AES.encrypt(M5_inner)
    M5 = f"{M5_inner},{T_c_s}"

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, SERVICE_PORT))
        message = f"request||{M5}"
        print('sending m5', message)
        s.sendall(message.encode())

        data = s.recv(4096)
        data = data.decode()
        print('received m6', data)
        M6 = data

    '''
    M6 = [{Resposta, N3}K_c_s]
    '''
    M6 = M6.split(",")
    M6_AES = AESCipher(K_c_s)
    M6_inner = M6_AES.decrypt(M6[0])
    M6_inner = M6_inner.split(",")
    resposta = M6_inner[0]
    N3 = M6_inner[1]

    print(f"Resposta: [{resposta}]")

    print("Finished!!!")


def main():
    global CLIENT_ID
    while True:
        # os.system("clear")

        print(f"Hello world! Logged as [{CLIENT_ID}]")
        print("1. Register")
        print("2. Login")
        if CLIENT_ID is not None:
            print("3. Print something")
        print("0. Exit")

        choice = input("Enter your choice: ")

        if choice == "1":
            username = input("Enter your username: ")
            password = input("Enter your password: ")
            register_user(username, password)
        elif choice == "2":
            username = input("Enter your username: ")
            password = input("Enter your password: ")
            login(username, password)
        elif choice == "3":
            print_something()
        elif choice == "0":
            break
        else:
            print("Invalid choice")

if __name__ == "__main__":
    main()