1 files changed, 69 insertions, 51 deletions

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