#!/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")
    parser.add_argument("--port",type=int,help="Port to connect to")
    parser.add_argument("--config",type=str,help="Path to config file",\
                        default="/etc/mfa/mfa.conf")
    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 read_config(config_file):
    # Read config file for server and port info
    # Return tuple (server,port)
    server = ""
    port = 0
    with open(config_file) as conf:
        line = None
        while line != "":
            line = conf.readline()
            if line.startswith("server ="):
                server = line.split("=")[1].strip()
            if line.startswith("port ="):
                port = int(line.split("=")[1].strip())
    return (server,port)

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

    # Read server and port from config file but allow command line options
    # to override those settings
    mfa_server, client_port = read_config(args.config)
    if args.server != None:
        mfa_server = args.server
    if args.port != None:
        client_port = args.port

    # 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()