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.

1 files changed, 19 insertions, 5 deletions

diff --git a/pam/pam.py b/pam/pam.py
index 28450ee..5a2fee8 100755
--- a/pam/pam.py
+++ b/pam/pam.py
@@ -22,6 +22,11 @@ def parse_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):
@@ -43,7 +48,7 @@ def init_connection(mfa_server, pam_port):
     return None
 
 
-def read_config(config_file="/etc/mfa/mfa.conf"):
+def read_config(config_file):
     # Read config file for server and port info
     # Return tuple (server,port)
     server = ""
@@ -68,10 +73,19 @@ def main():
     service = args.service
 
     # Compile data to send to server
-    mfa_server, pam_port = read_config()
-    hostname = None
-    with open("/etc/hostname") as f:
-        hostname = f.read().strip()
+    # Read server and port from config file but allow command line options
+    # to override those settings
+    mfa_server, pam_port = read_config(args.config)
+    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