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, 27 insertions, 4 deletions

diff --git a/client/client.py b/client/client.py
index abbc9de..b2429b6 100755
--- a/client/client.py
+++ b/client/client.py
@@ -14,8 +14,10 @@ FORMAT = "utf-8"
 
 def parse_arguments():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--server",type=str,help="IP of MFA Server",required=True)
-    parser.add_argument("--port",type=int,help="Port to connect to",required=True)
+    parser.add_argument("--server",type=str,help="IP of MFA Server")
+    parser.add_argument("--port",type=int,help="Port to connect to")
+    parser.add_argument("--config",type=str,help="Path to config file",\
+                        default="/etc/mfa/mfa.conf")
     parser.add_argument("--key",type=str,help="Client connection key",required=True)
     return parser.parse_args()
 
@@ -50,13 +52,34 @@ def init_connection(mfa_server, client_port, client_key):
     return connection
 
 
+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 main():
     # Get arguments, exit if unable to connect
     args = parse_arguments()
-    mfa_server = args.server
-    client_port = args.port
     client_key = args.key
 
+    # Read server and port from config file but allow command line options
+    # to override those settings
+    mfa_server, client_port = read_config(args.config)
+    if args.server != None:
+        mfa_server = args.server
+    if args.port != None:
+        client_port = args.port
+
     # Exit if invalid key is provided
     if len(client_key) != KEY_LENGTH:
         print("invalid key")