1 files changed, 41 insertions, 23 deletions

diff --git a/client/client.py b/client/client.py
index b2429b6..1c7e155 100755
--- a/client/client.py
+++ b/client/client.py
@@ -4,6 +4,7 @@ import socket
 import time
 import argparse
 import sys
+import os
 
 HEADER_LENGTH = 64
 KEY_LENGTH = 64
@@ -11,6 +12,7 @@ DISCONNECT_LENGTH = ACK_LENGTH = 3
 ACK_MESSAGE = "ACK"
 DISCONNECT_MESSAGE = "BYE"
 FORMAT = "utf-8"
+import configparser
 
 def parse_arguments():
     parser = argparse.ArgumentParser()
@@ -18,7 +20,7 @@ def parse_arguments():
     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)
+    parser.add_argument("--key",type=str,help="Client connection key")
     return parser.parse_args()
 
 def prompt_user(prompt):
@@ -53,32 +55,48 @@ def init_connection(mfa_server, client_port, client_key):
 
 
 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)
+    parser = configparser.ConfigParser(inline_comment_prefixes="#")
+    parser.read(config_file)
+    return parser
+
+
+def get_vars(args,confparser):
+    if not os.path.exists(args.config):
+        print("Unable to open config file")
+        sys.exit(1)
+
+    server = None
+    port = None
+    key = None
+
+    # Set values from config file first
+    if confparser.has_section("client"):
+        server = confparser.get("client","server",fallback=None) 
+        port = confparser.get("client","port",fallback=None) 
+        key = confparser.get("client","key",fallback=None)
+        
+    # Let command line args overwrite any values
+    if args.server:
+        server = args.server
+    if args.port:
+        port = args.port
+    if args.key:
+        key = args.key
+
+    # Exit if any value is null
+    if None in [server,port,key]:
+        print("error: one or more items unspecified")
+        sys.exit(1)
+
+    return server,port,key
+
 
 def main():
     # Get arguments, exit if unable to connect
     args = parse_arguments()
-    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
+    confparser = read_config(args.config)
+
+    mfa_server,client_port,client_key = get_vars(args,confparser)
 
     # Exit if invalid key is provided
     if len(client_key) != KEY_LENGTH: