Better error handling if server disconnects

1 files changed, 8 insertions, 1 deletions

diff --git a/client/client.py b/client/client.py
index 1c7e155..70d85a0 100755
--- a/client/client.py
+++ b/client/client.py
@@ -14,6 +14,10 @@ DISCONNECT_MESSAGE = "BYE"
 FORMAT = "utf-8"
 import configparser
 
+def die(msg):
+    print(msg)
+    sys.exit(1)
+
 def parse_arguments():
     parser = argparse.ArgumentParser()
     parser.add_argument("--server",type=str,help="IP of MFA Server")
@@ -113,7 +117,10 @@ def main():
     running = True
     while running:
         # Receive MFA prompt from server
-        prompt_len = int(conn.recv(HEADER_LENGTH).decode(FORMAT))
+        header = conn.recv(HEADER_LENGTH).decode(FORMAT) 
+        if header == "":
+            die("error: lost connection to server")
+        prompt_len = int(header)
         prompt = conn.recv(prompt_len).decode(FORMAT)
         # Ask user for response
         answer = prompt_user(prompt)