#!/usr/bin/env python3

import socket
import time
import argparse
import sys
import os

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

def die(msg):
    print(msg)
    sys.exit(1)

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")
    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):
    parser = configparser.ConfigParser(inline_comment_prefixes="#")
    parser.read(config_file)
    return parser


def get_vars(args,confparser):
    if not os.path.exists(args.config):
        print("Unable to open config file")
        sys.exit(1)

    server = None
    port = None
    key = None

    # Set values from config file first
    if confparser.has_section("client"):
        server = confparser.get("client","server",fallback=None) 
        port = confparser.get("client","port",fallback=None) 
        key = confparser.get("client","key",fallback=None)
        
    # Let command line args overwrite any values
    if args.server:
        server = args.server
    if args.port:
        port = args.port
    if args.key:
        key = args.key

    # Exit if any value is null
    if None in [server,port,key]:
        print("error: one or more items unspecified")
        sys.exit(1)

    return server,port,key


def main():
    # Get arguments, exit if unable to connect
    args = parse_arguments()
    confparser = read_config(args.config)

    mfa_server,client_port,client_key = get_vars(args,confparser)

    # 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
        header = conn.recv(HEADER_LENGTH).decode(FORMAT) 
        if header == "":
            die("error: lost connection to server")
        prompt_len = int(header)
        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()