Wrap client loop in try catch block

1 files changed, 21 insertions, 14 deletions

diff --git a/client/client.py b/client/client.py
index 0388073..bc4da96 100755
--- a/client/client.py
+++ b/client/client.py
@@ -166,20 +166,27 @@ def main():
     # Main loop
     running = True
     while running:
-        # Receive MFA prompt from server
-        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)
-        # Send answer to MFA server
-        answer_length = len(answer)
-        length_msg = str(answer_length)
-        length_msg += ' ' * (HEADER_LENGTH - len(length_msg))
-        conn.send(length_msg.encode(FORMAT))
-        conn.send(answer.encode(FORMAT))
+        try:
+            # Receive MFA prompt from server
+            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)
+            # Send answer to MFA server
+            length_msg = str(len(answer))
+            # Handle very large input
+            if len(length_msg) > HEADER_LENGTH:
+                die("error: response too big")
+            length_msg += ' ' * (HEADER_LENGTH - len(length_msg))
+            conn.send(length_msg.encode(FORMAT))
+            conn.send(answer.encode(FORMAT))
+        except KeyboardInterrupt:
+            #conn.send(DISCONNECT_MESSAGE.encode(FORMAT))
+            conn.close()
+            die("")