From 41132171e065081736c82ee283c64309e30baa9b Mon Sep 17 00:00:00 2001
From: Sam Chudnick <sam@chudnick.com>
Date: Sun, 3 Jul 2022 05:45:04 -0400
Subject: Better error handling if server disconnects

---
 client/client.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

(limited to 'client')

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)
-- 
cgit v1.2.3