From ceddd0cb74f4404b8d743c3fce011c2c9fda4f32 Mon Sep 17 00:00:00 2001 From: Sam Chudnick Date: Sun, 11 Jun 2023 08:14:22 -0400 Subject: Wrap client loop in try catch block --- client/client.py | 35 +++++++++++++++++++++-------------- 1 file changed, 21 insertions(+), 14 deletions(-) (limited to 'client') 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("") -- cgit v1.2.3