Added support for TOTP

Added TOTP as an MFA option. Also made a couple of of minor changes.
Changed all database connections to use with statement. Read some
options from a config file.

2 files changed, 134 insertions, 83 deletions

diff --git a/server/mfac.py b/server/mfac.py
index 779fa44..884c61c 100755
--- a/server/mfac.py
+++ b/server/mfac.py
@@ -6,7 +6,12 @@ import pyotp
 import sys
 
 DB_NAME = "mfa.db"
-KEY_LENGTH = 64
+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)
@@ -14,9 +19,10 @@ def die(msg):
 
 def parse_arguments():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--add-client",action="store_true",help="Add a client")
     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")
@@ -25,59 +31,76 @@ def parse_arguments():
 
     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):
-    conn = sqlite3.connect(DB_NAME)
-    c = conn.cursor()
-    c.execute("SELECT * FROM clients WHERE alias=?",(alias,))
-    client = c.fetchall()
-    conn.close()
+    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_key(alias):
-    CLIENT_ALIAS_INDEX = 0
-    CLIENT_KEY_INDEX = 1
+def get_client_info(alias,info):
     if not alias_exists(alias):
         die("Error: alias does not exist")
     else:
-        conn = sqlite3.connect(DB_NAME)
-        c = conn.cursor()
-        c.execute("SELECT * FROM clients WHERE alias=?",(alias,))
-        client = c.fetchone()
-        conn.close()
-        return str(client[CLIENT_KEY_INDEX])
+        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):
-        key = pyotp.random_base32(length=64)
-        conn = sqlite3.connect(DB_NAME)
-        c = conn.cursor()
-        c.execute("INSERT INTO clients VALUES (?,?)",(alias,key))
-        conn.commit()
-        conn.close()
-        print("key: " + key)
+        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_key(alias)
-        conn = sqlite3.connect(DB_NAME)
-        c = conn.cursor()
-        c.execute("INSERT INTO applications VALUES (?,?,?,?,?)",
-                (username,hostname,service,alias,mfa_methods))
-        conn.commit()
-        conn.close()
+        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():
@@ -90,6 +113,10 @@ def main():
         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):
@@ -98,11 +125,17 @@ def main():
     if args.add_client:
         add_client(args.alias)
     elif args.get_client:
-        key = get_client_key(args.alias)
+        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"))
         
         
 
diff --git a/server/mfad.py b/server/mfad.py
index 7a2fc40..d045e14 100755
--- a/server/mfad.py
+++ b/server/mfad.py
@@ -26,28 +26,54 @@ FORMAT = "utf-8"
 AUTHED = 0
 DENIED = 1
 
+# DB object index constants
+DB_USERNAME_INDEX = 0
+DB_HOSTNAME_INDEX = 1
+DB_SERVICE_INDEX = 2
+DB_ALIAS_INDEX = 3
+DB_MFAMETHODS_INDEX = 4
+
+CLIENT_ALIAS_INDEX = 0
+CLIENT_KEY_INDEX = 1
+CLIENT_SECRET_INDEX = 2
+
 # Stores connected clients as a dictionary with the client key as the dictionary 
 # key and a tuple of (socket,(addr,port)) as the value
 client_connections = dict()
 
 
-def eval_mfa(mfa_methods,client_response):
+def eval_mfa(client_key, mfa_methods, client_response):
+    print("response: " + client_response)
+    print("length: " + str(len(client_response)))
+    print("methods: " + str(mfa_methods))
     # Evaluates MFA and decides if authenticated or denied
     # Returns 0 for authenticated on 1 for denied
     if "push" in mfa_methods and client_response == "allow":
         return AUTHED
     elif "totp" in mfa_methods and len(client_response) == 6:
-        # Only attempt to validate if response is a valid TOTP format
+        # Only attempt to validate if response is a valid TOTP format 
         totp_format = (r'(\d)(\d)(\d)(\d)(\d)(\d)')
         totp_regex = re.compile(totp_format)
         matched = totp_regex.match(client_response)
         if matched:
-            return validate_totp(int(client_response))
+            return validate_totp(client_key, client_response)
     return DENIED
 
 
-def validate_totp(client_response):
-    pass
+def validate_totp(client_key, client_response):
+    secret = ""
+    with sqlite3.connect(DB_NAME) as conn:
+        c = conn.cursor()
+        c.execute("SELECT * FROM clients WHERE key=?",(client_key,))
+        client = c.fetchone()
+        secret = client[CLIENT_SECRET_INDEX]
+    totp = pyotp.TOTP(secret)
+    print("Client Response: " + str(client_response))
+    print("Valid TOTP: " + str(totp.now()))
+    if totp.verify(client_response):
+        return AUTHED
+    else:
+        return DENIED
 
 
 ################################################################################
@@ -64,28 +90,22 @@ def get_client_key(username,hostname,service):
     # This is done by checking the PAM request against a preconfigured 
     # database mapping request info (username,hostname,etc...) to clients
     # Returns a tuple consisting of the key and approved MFA methods
-    DB_USERNAME_INDEX = 0
-    DB_HOSTNAME_INDEX = 1
-    DB_SERVICE_INDEX = 2
-    DB_ALIAS_INDEX = 3
-    DB_MFAMETHODS_INDEX = 4
-
-    CLIENT_ALIAS_INDEX = 0
-    CLIENT_KEY_INDEX = 1
-
-    conn = sqlite3.connect(DB_NAME)
-    c = conn.cursor()
-    c.execute("""SELECT * FROM applications WHERE username=? AND hostname=?
-                 AND service=?""",(username,hostname,service))
-    application = c.fetchone()
-    # Return None if no results found
-    if application == None:
-        return application
-
-    alias = application[DB_ALIAS_INDEX]
-    c.execute("SELECT * FROM clients WHERE alias=?",(alias,))
-    client = c.fetchone()
-    conn.close()
+
+    application = None
+    client = None
+    with sqlite3.connect(DB_NAME) as conn:
+        c = conn.cursor()
+        c.execute("""SELECT * FROM applications WHERE username=? AND hostname=?
+                     AND service=?""",(username,hostname,service))
+        application = c.fetchone()
+        # Return None if no results found
+        if application == None:
+            return application
+
+        alias = application[DB_ALIAS_INDEX]
+        c.execute("SELECT * FROM clients WHERE alias=?",(alias,))
+        client = c.fetchone()
+
     client_key = client[CLIENT_KEY_INDEX]
     methods = application[DB_MFAMETHODS_INDEX]
     return (client_key,methods)
@@ -129,11 +149,12 @@ def prompt_client(client_key, user, host, service, methods, timeout=10):
 
 def validate_client(client_key):
     # Validates a client
-    conn = sqlite3.connect(DB_NAME)
-    c = conn.cursor()
-    c.execute("SELECT * FROM clients WHERE key=?",(client_key,))
-    client = c.fetchall()
-    conn.close()
+    client = None
+    with sqlite3.connect(DB_NAME) as conn:
+        c = conn.cursor()
+        c.execute("SELECT * FROM clients WHERE key=?",(client_key,))
+        client = c.fetchall()
+
     if len(client) == 0:
         # No client matches provided key, invalid
         return False
@@ -172,7 +193,7 @@ def handle_pam(conn, addr):
 
     # Correlate request to client
     client_key,mfa_methods = get_client_key(user,host,service)
-    mfa_methods = mfa_methods.split(',')
+    mfa_methods = mfa_methods.split(' ')
     if client_key == None:
         print("No applications found for user="+user+" host="+host+" service="+service)
         conn.send(str(DENIED).encode(FORMAT))
@@ -182,8 +203,7 @@ def handle_pam(conn, addr):
     response = prompt_client(client_key,user,host,service,mfa_methods)
 
     # Evaluate Response
-    auth_type = "push"
-    decision = eval_mfa(auth_type, response)
+    decision = eval_mfa(client_key, mfa_methods, response)
 
     # Return response to PAM module
     # Respone will either be 0 for authenticated and 1 for denied
@@ -209,22 +229,20 @@ def listen_pam(addr, port):
 ################################################################################
 
 def create_db():
-    conn = sqlite3.connect(DB_NAME)
-    c = conn.cursor()
-    c.execute("""CREATE TABLE applications (
-                username text,
-                hostname text,
-                service text,
-                client_key text,
-                mfa_methods text
-            )""")
-    conn.commit()
-    c.execute("""CREATE TABLE clients (
-                alias,
-                key
-             )""")
-    conn.commit()
-    conn.close()
+    with sqlite3.connect(DB_NAME) as conn:
+        c = conn.cursor()
+        c.execute("""CREATE TABLE applications (
+                    username text,
+                    hostname text,
+                    service text,
+                    client_key text,
+                    mfa_methods text
+                )""")
+        c.execute("""CREATE TABLE clients (
+                    alias text,
+                    key text,
+                    totp_secret text
+                 )""")
 
 
 def main():