Update client.py

MP1/client.py

@@ -22,17 +22,21 @@ def send_message(sock, message):
     sock.sendall(header)
     sock.sendall(e_data)
 
-def receive_message(sock):
-    header = sock.recv(10)
-    if not header:
-        return None
-    message_length = int(header.decode('utf-8'))
+def receive_exact(sock, num_bytes):
     received_data = b""
-    while len(received_data) < message_length:
-        chunk = sock.recv(65536)
+    while len(received_data) < num_bytes:
+        chunk = sock.recv(num_bytes - len(received_data))
         if not chunk:
             break
         received_data += chunk
+    return received_data
+
+def receive_message(sock):
+    header = receive_exact(sock, 10)
+    if not header:
+        return None
+    message_length = int(header.decode('utf-8'))
+    received_data = receive_exact(sock, message_length)
     return received_data.decode('utf-8')
 
 def grep_all_ips(ips, cmd):