From dc73c10e433ddd790792ee06f03ce8cdd18b59b0 Mon Sep 17 00:00:00 2001
From: recyvan <e877586809@163.com>
Date: Tue, 26 Nov 2024 15:48:19 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=BA=86=E4=B8=80=E4=B8=8B?=
 =?UTF-8?q?=E4=B8=8D=E7=9F=A5=E5=90=8Dbug?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 recv/recv_net.py     | 34 ++++++++++++++++++++--------------
 sender/sender_net.py | 30 ++++++++++++++++--------------
 2 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/recv/recv_net.py b/recv/recv_net.py
index 4d9ca7a..0498554 100644
--- a/recv/recv_net.py
+++ b/recv/recv_net.py
@@ -28,11 +28,11 @@ from sender import sender
 #     pass
 
 class recver_net():
-    def __init__(self,port,client_host,client_port):
+    def __init__(self, port, client_host, client_port):
         self.client_host = client_host
         self.client_port = client_port
         self.letter = Letter()
-        self.port=port
+        self.port = port
         self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.server.bind(('0.0.0.0', self.port))
         self.server.listen(10)
@@ -51,18 +51,21 @@ class recver_net():
 
     def send_thread(self, conn):
         while True:
-            flag: int = int(input(""))
-            if flag == 0:
-                data = input("[-] Enter data to send: ")
-                msg = {"flag": 0, "data": data}
-                msg = json.dumps(msg).encode("utf-8")
-                self.send_data(conn, msg)
-            elif flag == 1:
-                print("[+] Calling file transfer module...")
-                # file_name = input("Enter file name to send: ")
-                self.send_file(conn)
-            else:
-                print("[x] Error,plesae enter 0 to send data, or 1 to send file.")
+            try:
+                flag: int = int(input())
+                if flag == 0:
+                    data = input("[-] Enter data to send: ")
+                    msg = {"flag": 0, "data": data}
+                    msg = json.dumps(msg).encode("utf-8")
+                    self.send_data(conn, msg)
+                elif flag == 1:
+                    print("[+] Calling file transfer module...")
+                    # file_name = input("Enter file name to send: ")
+                    self.send_file(conn)
+                else:
+                    print("[x] Error,plesae enter 0 to send data, or 1 to send file.")
+            except ValueError:
+                print("[x] Error, please enter a valid number.")
 
     def send_data(self, conn, data):
         conn.sendall(data)
@@ -105,6 +108,7 @@ class recver_net():
     def run(self):
         threading.Thread(target=self.main).start()
 
+
 def input_verify():
     while True:
         try:
@@ -128,6 +132,8 @@ def input_verify():
         else:
             print("Invalid address format. Please enter in the format 'host:port'.")
     return port, client_host, client_port
+
+
 # test
 if __name__ == '__main__':
     port, client_host, client_port = input_verify()
diff --git a/sender/sender_net.py b/sender/sender_net.py
index b99df16..60a2af7 100644
--- a/sender/sender_net.py
+++ b/sender/sender_net.py
@@ -50,18 +50,22 @@ class sender_net():
 
     def send_thread(self, conn):
         while True:
-            flag: int = int(input())
-            if flag == 0:
-                data = input("[-] Enter data to send: ")
-                msg = {"flag": 0, "data": data}
-                msg = json.dumps(msg).encode("utf-8")
-                self.send_data(conn, msg)
-            elif flag == 1:
-                print("[+] Calling file transfer module...")
-                # file_name = input("Enter file name to send: ")
-                self.send_file(conn)
-            else:
-                print("[x] Error,plesae enter 0 to send data, or 1 to send file.")
+            try:
+                flag: int = int(input())
+                if flag == 0:
+                    data = input("[-] Enter data to send: ")
+                    msg = {"flag": 0, "data": data}
+                    msg = json.dumps(msg).encode("utf-8")
+                    self.send_data(conn, msg)
+                elif flag == 1:
+                    print("[+] Calling file transfer module...")
+                    # file_name = input("Enter file name to send: ")
+                    self.send_file(conn)
+                else:
+                    print("[x] Error,plesae enter 0 to send data, or 1 to send file.")
+            except ValueError:
+                print("[x] Error, please enter a valid number.")
+
 
     def send_data(self, conn, msg):
         conn.sendall(msg)
@@ -134,8 +138,6 @@ def input_verify():
 
 #test
 if __name__ == '__main__':
-
-
     port, client_host, client_port = input_verify()
     a = sender_net(port, client_host, client_port)
     a.run()