1111

recv/recv_net.py

@@ -1,33 +1,12 @@
-import hashlib
 import json
-import selectors
 import socket
 import threading
 from time import sleep
-
-import select
-
 from entity.Letter import Letter, json_to_obj
 import recv
 from sender import sender
 
-
+class RecverNet:
-# from entity.Letter import Letter
-
-
-# def main():
-#     # 用户输入各种数据填充letter字段
-#     # 获取用户的公私钥对进行签名
-#     # 使用对方的公钥进行加密
-#     # 发送信件
-#     pass
-#
-#
-# def sendLetter(letter: Letter, target="192.168.195.162:8426"):
-#     # 向目标ip和端口发送指定的信件
-#     pass
-
-class recver_net():
     def __init__(self, port, client_host, client_port):
         self.client_host = client_host
         self.client_port = client_port
@@ -36,106 +15,114 @@ class recver_net():
         self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.server.bind(('0.0.0.0', self.port))
         self.server.listen(10)
+        self.is_running = True
 
     def recver_thread(self):
-        while True:
+        while self.is_running:
-            conn, addr = self.server.accept()
+            try:
-            data = conn.recv(10240)
+                conn, addr = self.server.accept()
-            msg = json.loads(data.decode())
+                print(f"\n[*] Connection from {addr}")
-            if int(msg['flag']) == 0:
+                with conn:
-                print("[+] haved received data:" + msg["data"])
+                    while self.is_running:
-            elif int(msg['flag']) == 1:
+                        data = conn.recv(10240)
-                self.recv_file(msg['data'])
+                        if not data:
-            else:
+                            print("\n[x] Disconnected from client.")
-                print("[x] Error")
+                            break
+                        msg = json.loads(data.decode())
+                        self.handle_message(msg)
+            except Exception as e:
+                print(f"\n[x] Error in recver_thread: {e}")
+
+    def handle_message(self, msg):
+        if msg['flag'] == 0:
+            print("\n[+] Received data:", msg["data"])
+        elif msg['flag'] == 1:
+            self.recv_file(msg['data'])
+        else:
+            print("\n[x] Error: Unknown flag")
 
     def send_thread(self, conn):
-        while True:
+        while self.is_running:
             try:
-                flag: int = int(input())
+                flag = input("[*] Enter 0 to send data, 1 to send file: ")
-                if flag == 0:
+                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:
+                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.")
+                    print("[x] Invalid input. Please enter 0 or 1.")
-            except ValueError:
+            except Exception as e:
-                print("[x] Error, please enter a valid number.")
+                print(f"[x] Error in send_thread: {e}")
+                self.is_running = False
 
-    def send_data(self, conn, data):
+    def send_data(self, conn, msg):
-        conn.sendall(data)
+        conn.sendall(json.dumps(msg).encode("utf-8"))
 
     def send_file(self, conn):
-        self.letter: Letter = sender.main()
+        self.letter = sender.main()
-        letter = self.letter.to_dict()
+        letter_data = self.letter.to_dict()
-        msg = {"flag": 1, "data": letter}
+        msg = {"flag": 1, "data": letter_data}
-        msg = json.dumps(msg).encode("utf-8")
+        self.send_data(conn, msg)
-        conn.sendall(msg)
-
-    def recv_data(self, conn):
-        con, addr = conn.accept()
-        data = con.recv(1024)
-        if not data:
-            return None
-        # print(data.decode())
-        return data.decode()
 
     def recv_file(self, data):
-        letter: Letter = json_to_obj(data)
+        letter = json_to_obj(data)
         recv.handleLetter(letter)
 
-    def main(self):
+    def connect(self):
-
         while True:
             try:
-                client = socket.socket()  # 定义协议类型,相当于生命socket类型,同时生成socket连接对象
+                client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                 client.connect((self.client_host, self.client_port))
-                print("[*] Connected...")
+                return client
-                print("[*] if you want to send data, enter 0, if you want to send file, enter 1.")
-                sleep(1)
-                break
             except socket.error:
-                print("[*] Waiting for ...")
+                print("[*] Connection failed. Retrying...")
-        # 加入线程
+                sleep(2)
-        threading.Thread(target=self.recver_thread, ).start()
+
-        threading.Thread(target=self.send_thread, args=(client,)).start()
+    def main(self):
+        client = self.connect()
+        threading.Thread(target=self.recver_thread, daemon=True).start()
+        self.send_thread(client)
+        while self.is_running==False:
+            self.is_running=True
+            client=self.connect()
+            self.send_thread(client)
+            sleep(1)
 
     def run(self):
-        threading.Thread(target=self.main).start()
+        self.main()
 
 
 def input_verify():
     while True:
         try:
-            port = int(input("[*] Enter to the listen port: "))
+            port = int(input("[*] Enter the listen port: "))
-            break
+            if 0 <= port <= 65535:
+                break
+            else:
+                print("[x] Port must be between 0 and 65535.")
         except ValueError:
-            print("[x] Error, please enter a valid port number.")
+            print("[x] Invalid port number. Please try again.")
 
     while True:
-        addr = input("[*] Enter the address to connect to(127.0.0.1:8424): ")
+        addr = input("[*] Enter the address to connect to (e.g., 127.0.0.1:8424): ")
         if ':' in addr:
             client_host, client_port = addr.split(":")
             if client_host and client_port.isdigit():
                 client_port = int(client_port)
-                if 0 <= client_port <= 65535:  # 检查端口范围
+                if 0 <= client_port <= 65535:
-                    break  # 输入有效，跳出循环
+                    return port, client_host, client_port
                 else:
-                    print("Port must be between 0 and 65535.")
+                    print("[x] Port must be between 0 and 65535.")
             else:
-                print("Invalid address format. Please enter in the format 'host:port'.")
+                print("[x] Invalid address format.")
         else:
-            print("Invalid address format. Please enter in the format 'host:port'.")
+            print("[x] Invalid address format.")
-    return port, client_host, client_port
-
 
-# test
+# Test
 if __name__ == '__main__':
     port, client_host, client_port = input_verify()
-    a = recver_net(port, client_host, client_port)
+    recver_net = RecverNet(port, client_host, client_port)
-    a.run()
+    recver_net.run()