#!/usr/bin/python3
import socket
import argparse
import time
import sys
import configparser
import os

# Sends authentication request to MFA server
# Receive either pass or fail response from MFA server
# Returns status to PAM

HEADER_LENGTH = 64
KEY_LENGTH = 64
ACK_LENGTH = 3
DISCONNECT_MESSAGE = "DISCONNECT"
FORMAT = "utf-8"
RESPONSE_LENGTH = 1
PAM_SUCCESS = 0
PAM_AUTH_ERR = 7

def parse_arguments():
    # Parse command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--user",type=str,help="PAM username",required=True)
    parser.add_argument("--service",type=str,help="PAM service",required=True)
    parser.add_argument("--host",type=str,help="PAM hostname")
    parser.add_argument("--config",type=str,help="Path to config file",\
                        default="/etc/mfa/mfa.conf")
    parser.add_argument("--server",type=str,help="MFA server address")
    parser.add_argument("--port",type=str,help="MFA server PAM connection port")
    return parser.parse_args()

def init_connection(mfa_server, pam_port):
    # Attempts to connect to MFA server with provided address and port
    # Repeats connection attempts once per second until timeout is reached
    # Returns the socket if connection was successful or None otherwise
    connection = None
    timeout = 0
    timeout_length = 5
    sleep_length = 1
    while connection == None and timeout < timeout_length:
        try:
            connection = socket.create_connection((mfa_server,pam_port))
            return connection
        except (ConnectionError,ConnectionRefusedError):
            time.sleep(sleep_length)
            timeout += sleep_length
    return None


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 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

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

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

    return server,port


def main():
    authed = "0"
    failed = "1"

    # Get arguments
    args = parse_arguments()
    confparser = read_config(args.config)
    mfa_server,pam_port = get_vars(args,confparser)
    user = args.user
    service = args.service

    # Compile data to send to server
    # Read server and port from config file but allow command line options
    # to override those settings
    if args.server != None:
        mfa_server = args.server
    if args.port != None:
        pam_port = args.port
    # Get hostname if not given on command line
    if args.host == None:
        with open("/etc/hostname") as f:
            hostname = f.read().strip()
    else:
        hostname = args.host
    data = user + "," + hostname + "," + service 
    

    # Initalize connection to MFA server.  Quit if unable to connect.
    connection = init_connection(mfa_server,pam_port)
    if connection == None:
        print(failed)
        sys.exit(1)

    # Send authentication data to MFA server
    data_length = len(data)
    length_msg = str(data_length)
    length_msg += ' ' * (HEADER_LENGTH - len(length_msg))
    connection.send(length_msg.encode(FORMAT))
    connection.send(data.encode(FORMAT))

    # Listen for response from MFA server
    # Response will be 0 for authenticated and 1 for denied
    response = int(connection.recv(RESPONSE_LENGTH).decode(FORMAT))

    # Print success/failure for PAM module
    print(str(response))

if __name__ == '__main__':
    main()