from common import *
from hashlib import sha256
from AES import AESCipher
import socket
import os
import base64
import time

HOST = "127.0.0.1"
PORT = AS_PORT

def process_message(message):
    message = message.split("||")
    if message[0] == "register":
        body = message[1].split("<>")
        username = body[0]
        password = body[1]
        password = sha256(password.encode()).hexdigest()
        key = os.urandom(32)
        print(f"Registering user [{username}]")
        return register_user(username, password, key)
    elif message[0] == "login":
        body = message[1].split("<>")
        username = body[0]
        password = body[1]
        password = sha256(password.encode()).hexdigest()
        return login(username, password)
    elif message[0] == "request":
        """
        M1 = [ID_C + {ID_S + T_R + N1}Kc]

        T_c_tgs = {ID_C + T_R + K_c_tgs}K_tgs
        M2 = [{K_c_tgs + N_1}Kc + T_c_tgs]
        """
        body = message[1].split("<>")

        ID_C = body[0]
        Kc = get_user_key(ID_C)
        AES = AESCipher(Kc)

        inner_message = AES.decrypt(body[1])
        inner_message = inner_message.split("<>")
        T_R = inner_message[1]
        N1 = inner_message[2]

        # generate m2
        K_tgs = get_ktgs()
        M2_AES = AESCipher(K_tgs)
        K_c_tgs = os.urandom(16)
        T_c_tgs = f"{ID_C}<>{T_R}<>{K_c_tgs}"
        T_c_tgs = M2_AES.encrypt(T_c_tgs)

        M2 = f"{K_c_tgs}<>{N1}"
        M2 = AES.encrypt(M2)
        M2 = f"{M2}<>{T_c_tgs}"
        return M2
    else:
        print("Invalid message")

def register_user(username, password, key):
    with open("users.data", "a") as f:
        f.write(f"{username}<>{password}<>{base64.b64encode(key)}\n")
    return f"success<>{key}"

def login(username, password):
    with open("users.data", "r", newline='\n') as f:
        for line in f:
            line = line.strip()
            line = line.replace("\n", "")
            user, pw, key = line.split("<>")
            if user == username and pw == password:
                key = eval(key)
                key = base64.b64decode(key)
                print(f"Login succeeded [{username}]")
                return f"success<>{key}"
        print(f"Login failed! [{username}]")
        return "failure"

def get_user_key(username):
    with open("users.data", "r", newline='\n') as f:
        for line in f:
            line = line.strip()
            line = line.replace("\n", "")
            user, pw, key = line.split("<>")
            if user == username:
                key = eval(key)
                key = base64.b64decode(key)
                return key
    return None

def get_ktgs():
    message = f"getkey"
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, TGS_PORT))
        s.sendall(message.encode())

        data = s.recv(4096)
        data = data.decode()
        key = eval(data)
        key = base64.b64decode(key)
        return key

def main():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind((HOST, PORT))
        s.listen()
        while True:
            conn, addr = s.accept()
            with conn:
                data = conn.recv(4096)
                data = data.decode()

                response = process_message(data)

                conn.sendall(response.encode())

if __name__ == "__main__":
    main()