import sys, socket, select, getpass
import hashlib, random, string
from GetRSA import getRSA
from friedMess import *

PRIME_MINIMUM = 100
HASHBYTES = 4

G_username = ""
G_password = ""
G_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
G_rsakey = []

KEY_PUBLIC = 0
KEY_PRIVATE = 1
KEY_N = 2

STATUS_NOTFOUND = "404"
STATUS_UNAUTHORIZED  = "403"
STATUS_FOUND = "302"
STATUS_OK = "200"

ServerKey_Public = 0
ServerKey_N = 0

G_listCommands = ""
for line in open("commands.txt", 'r'):
    G_listCommands += line

def printFriedMessLogo():
    print("\n" * 50)
    print("    ______    ____     ____    ______    ____     __  ___    ______   _____   _____")
    print("   / ____/   / __ \   /  _/   / ____/   / __ \   /  |/  /   / ____/  / ___/  / ___/")
    print("  / /_      / /_/ /   / /    / __/     / / / /  / /|_/ /   / __/     \__ \   \__ \ ")
    print(" / __/     / _, _/  _/ /    / /___    / /_/ /  / /  / /   / /___    ___/ /  ___/ / ")
    print("/_/       /_/ |_|  /___/   /_____/   /_____/  /_/  /_/   /_____/   /____/  /____/  ")
    print("\n")

def printPreChar():
    print("> ", end="")
    sys.stdout.flush()

def sendMessage(message):
    message = str(message)
    print("Sending: " + message)
    G_socket.send(message.encode("ASCII"))

def formatMessage(command, recipient, message):
    nonce = generateNonce(16)
    payload = G_username + "," + command + "," + recipient + "," + message
    h_message = hashSHA(nonce + message + G_password)
    h_message = encodeRSA(h_message, ServerKey_Public, ServerKey_N)
    message = nonce + "," + str(h_message) + "," + payload
    return message

def generateNonce(length):
    return ''.join([random.choice(string.ascii_letters + string.digits) for n in range(length)])

def authenticateUser():
    nonce = generateNonce(16)
    message = G_username + ",AUTH," + str(G_rsakey[KEY_PUBLIC])+" "+str(G_rsakey[KEY_N]) + ","
    h_message = hashlib.sha256()
    h_message.update((nonce + message + G_password).encode("UTF-8"))
    h_message = str(int(h_message.hexdigest(), 16))[:HASHBYTES]
    message = nonce + "," + str(h_message) + "," + message
    sendMessage(message)

def validateMessage(incHash, salt, message, RSAK, RSAN):
    valid = False
    incHash = encodeRSA(incHash, RSAK, RSAN)
    checkHash = hashSHA(message+salt)
    if(str(checkHash) == str(incHash)):
        valid = True
    else:
        valid = False
    return valid

def validCommand(command):
    if(command in G_listCommands):
        return True
    return False

def chat_client():
    global G_username
    global G_password
    global G_socket
    global G_rsakey
    global ServerKey_Public
    global ServerKey_N

    if(len(sys.argv) < 3) :
        print('Usage : python3 client.py <host> <port>')
        sys.exit()

    HOST = sys.argv[1]
    PORT = int(sys.argv[2])

    G_socket.settimeout(2)

    # Try to connect to the chat server
    try :
        G_socket.connect((HOST, PORT))
    except Exception as ex :
        print('Unable to connect' + str(ex))
        sys.exit()

    printFriedMessLogo()

    G_rsakey = getRSA(PRIME_MINIMUM, PRIME_MINIMUM)

    print("Please enter your zID (zXXXXXXX) and password (password will not echo):\n")

    loginSuccess = False

    while(loginSuccess == False):
        G_username = input("zID: ")
        G_password = getpass.getpass("Password: ")

        authenticateUser()
        socket_list = [sys.stdin, G_socket]
        ready_to_read,ready_to_write,in_error = select.select(socket_list , [], [])
        for sock in ready_to_read:
            if sock == G_socket:
                data = sock.recv(4096).decode("ASCII")
                if(data == STATUS_OK):
                    loginSuccess = True
                    print("Login Successful!")
                    printFriedMessLogo()
                    print("Welcome back " + G_username + "!")
                    print("Your connection is protected against replay attacks and data alteration\n")
                elif(data == STATUS_UNAUTHORIZED):
                    print("Incorrect zID or password\n")
    while(1):
        socket_list = [sys.stdin, G_socket]

        # Get the list sockets which are readable
        ready_to_read,ready_to_write,in_error = select.select(socket_list , [], [])

        for sock in ready_to_read:
            if sock == G_socket:
                # incoming message from remote server, s
                data = sock.recv(4096).decode("ASCII")
                if not data :
                    print('\nDisconnected from chat server')
                    sys.exit()
                else :
                    data = data.split(",")
                    h_incMessage = data[0]
                    sender = data[1]
                    command = data[2]
                    recipient = data[3]
                    incMessage = data[4]
                    if(validateMessage(h_incMessage, G_password, sender+","+command+","+recipient+","+incMessage, G_rsakey[KEY_PRIVATE], G_rsakey[KEY_N])):
                        if(command == "ALLOCKEYS"):
                            ServerKeys = incMessage.split(" ")
                            ServerKey_Public = int(ServerKeys[0])
                            ServerKey_N = int(ServerKeys[1])
                    printPreChar()

            else :
                # user entered a message
                msg = sys.stdin.readline()
                command = msg.strip().split(" ", 1)[0]
                if(validCommand(command)):
                    if(command == "whoelse"):
                        msg = formatMessage("whoelse", "", "")
                    sendMessage(msg)
                else:
                    print("Invalid command. Open the assignment to see list of commands")
                printPreChar()

if __name__ == "__main__":

    sys.exit(chat_client())