#!/usr/bin/env python3

import argparse
import sqlite3
import pyotp
import sys

DB_NAME = "mfa.db"
CLIENT_KEY_LENGTH = 64
TOTP_KEY_LENGTH = 24

CLIENT_ALIAS_INDEX = 0
CLIENT_KEY_INDEX = 1
CLIENT_TOTP_INDEX = 2

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

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--alias",type=str,help="Alias for new client")

    parser.add_argument("--add-client",action="store_true",help="Add a client")

    parser.add_argument("--add-app",action="store_true",help="Add an application")
    parser.add_argument("--user",type=str,help="Application username")
    parser.add_argument("--host",type=str,help="Application hostname")
    parser.add_argument("--service",type=str,help="Application service name")
    parser.add_argument("--methods",type=str,nargs="+",help="Allowed MFA methods")

    parser.add_argument("--get-client",action="store_true",help="Get a client key")

    parser.add_argument("--update-totp",action="store_true",help="Update a TOTP")

    parser.add_argument("--get-totp",action="store_true",help="Get client's TOTP")

    return parser.parse_args()


def alias_exists(alias):
    client = None
    with sqlite3.connect(DB_NAME) as conn:
        c = conn.cursor()
        c.execute("SELECT * FROM clients WHERE alias=?",(alias,))
        client = c.fetchall()
    if len(client) == 0:
        return False
    elif len(client) == 1:
        return True


def get_client_info(alias,info):
    if not alias_exists(alias):
        die("Error: alias does not exist")
    else:
        client = None
        with sqlite3.connect(DB_NAME) as conn:
            c = conn.cursor()
            c.execute("SELECT * FROM clients WHERE alias=?",(alias,))
            client = c.fetchone()
        if info == "key":
            return str(client[CLIENT_KEY_INDEX])
        elif info == "totp":
            return str(client[CLIENT_TOTP_INDEX])


def update_totp(alias):
    if alias_exists(alias):
        totp_secret = pyotp.random_base32(TOTP_KEY_LENGTH)
        with sqlite3.connect(DB_NAME) as conn:
            c = conn.cursor()
            c.execute("UPDATE clients SET totp_secret=? WHERE alias=?",\
                    (totp_secret,alias))
        print("totp secret: " + totp_secret)
    else:
        die("error: alias does not exist")


def add_client(alias):
    if not alias_exists(alias):
        client_key = pyotp.random_base32(CLIENT_KEY_LENGTH)
        totp_secret = pyotp.random_base32(TOTP_KEY_LENGTH)
        with sqlite3.connect(DB_NAME) as conn:
            c = conn.cursor()
            c.execute("INSERT INTO clients VALUES (?,?,?)",\
                    (alias,client_key,totp_secret))

        print("client key: " + client_key)
        print("totp secret: " + totp_secret)
        print("uri: " + pyotp.TOTP(totp_secret).provisioning_uri(alias+"@mfad"))
    else:
        die("Error: alias already used")

def add_app(username, hostname, service, alias, mfa_methods):
    if not alias_exists(alias):
        die("Error: alias does not exist")
    else:
        client_key = get_client_info(alias, "key")
        with sqlite3.connect(DB_NAME) as conn:
            c = conn.cursor()
            c.execute("INSERT INTO applications VALUES (?,?,?,?,?)",
                    (username,hostname,service,alias,mfa_methods))


def main():
    args = parse_arguments()
    # Sanity checks
    if (args.add_client and args.add_app) or (args.add_client and args.get_client) \
            or (args.get_client and args.add_app):
        die("Error: cannot specify multiple actions")
    if args.add_client and args.alias == None:
        die("Error: must specify alias to provision a client")
    if args.get_client and args.alias == None:
        die("Error: no alias specified")
    if args.update_totp and args.alias == None:
        die("Error: no alias specified")
    if args.get_totp and args.alias == None:
        die("Error: no alias specified")
    if args.add_app and (args.user == None or args.host == None or \
                         args.service == None or args.alias == None \
                         or args.methods == None):
        die("Error: --add-app requires all of --user,--host,--service,--alias,--methods")
        
    if args.add_client:
        add_client(args.alias)
    elif args.get_client:
        key = get_client_info(args.alias,"key")
        print(key)
    elif args.add_app:
        methods = " ".join(args.methods)
        add_app(args.user,args.host,args.service,args.alias,methods)
    elif args.update_totp:
        update_totp(args.alias)
    elif args.get_totp:
        secret = get_client_info(args.alias,"totp")
        print(secret)
        print(pyotp.TOTP(secret).provisioning_uri(args.alias+"@mfad"))
        
        


if __name__ == '__main__':
    main()