#!/usr/bin/env python3

import socket
import time
import argparse
import sys

HEADER_LENGTH = 64
KEY_LENGTH = 64
DISCONNECT_LENGTH = ACK_LENGTH = 3
ACK_MESSAGE = "ACK"
DISCONNECT_MESSAGE = "BYE"
FORMAT = "utf-8"

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--server",type=str,help="IP of MFA Server",required=True)
    parser.add_argument("--port",type=int,help="Port to connect to",required=True)
    parser.add_argument("--key",type=str,help="Client connection key",required=True)
    return parser.parse_args()

def prompt_user(prompt):
    # Prompt user for input and return input
    print(prompt)
    result = input("> ")
    return result


def init_connection(mfa_server, client_port, client_key):
    # Attempts to connect to MFA server with provided address,port, and key.
    # Repeats attempt once a seconds until timeout is reached.
    # Returns socket or None if unable to connect
    connection = None
    timeout = 0
    timeout_length = 5
    sleep_length = 1
    while connection == None and timeout < timeout_length:
        try:
            connection = socket.create_connection((mfa_server,client_port))
            connection.send(client_key.encode(FORMAT))
            response = connection.recv(ACK_LENGTH).decode(FORMAT)
            if response == ACK_MESSAGE:
                print("connected to mfa server")
            elif response == DISCONNECT_MESSAGE:
                print("server terminated connection")
                sys.exit(1)
        except ConnectionError:
            time.sleep(sleep_length)
            timeout += sleep_length
    return connection


def main():
    # Get arguments, exit if unable to connect
    args = parse_arguments()
    mfa_server = args.server
    client_port = args.port
    client_key = args.key

    # Exit if invalid key is provided
    if len(client_key) != KEY_LENGTH:
        print("invalid key")
        sys.exit(1)

    # Open connection to server
    conn = init_connection(mfa_server,client_port,client_key)
    if conn == None:
        print("timed out attempting to connect to server")
        sys.exit(1)

    # Main loop
    running = True
    while running:
        # Receive MFA prompt from server
        prompt_len = int(conn.recv(HEADER_LENGTH).decode(FORMAT))
        prompt = conn.recv(prompt_len).decode(FORMAT)
        # Ask user for response
        answer = prompt_user(prompt)
        # Send answer to MFA server
        answer_length = len(answer)
        length_msg = str(answer_length)
        length_msg += ' ' * (HEADER_LENGTH - len(length_msg))
        conn.send(length_msg.encode(FORMAT))
        conn.send(answer.encode(FORMAT))




if __name__ == '__main__':
    main()