背景是我23年写过一个udp穿透的方法，那时候手搓了一个小demo，细节见：
udp穿透的方法:https://blog.csdn.net/yeshennet/article/details/135165159
上周公司同事有在讨论这个，我又回头看了下我那时候写的脚本。现在有AI了，提了一下说明，生成了代码如下：
python
# Server (Public IP)
import socket
import threading
clients = {} # {client_id: (ip, port)}
def handle_client(sock):
while True:
try:
data, addr = sock.recvfrom(1024)
msg = data.decode()
if msg.startswith("REGISTER:"):
# Register client
client_id = msg.split(":")[1]
clients[client_id] = addr
print(f"Registered: {client_id} -> {addr}")
sock.sendto(b"OK", addr)
elif msg.startswith("QUERY:"):
# Query peer address
target_id = msg.split(":")[1]
if target_id in clients:
target_addr = clients[target_id]
# Send target address to requester
response = f"ADDR:{target_addr[0]}:{target_addr[1]}"
sock.sendto(response.encode(), addr)
# Notify target that someone wants to connect
notify = f"PEER:{addr[0]}:{addr[1]}"
sock.sendto(notify.encode(), target_addr)
print(f"Connect: {addr} <-> {target_addr}")
else:
sock.sendto(b"NOT_FOUND", addr)
except Exception as e:
print(f"Error: {e}")
def main():
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind(("0.0.0.0", 9999))
print("Server started on port 9999")
handle_client(sock)
if __name__ == "__main__":
main()
python
# Client A/B
import socket
import time
import threading
class UDPClient:
def __init__(self, client_id, server_addr):
self.client_id = client_id
self.server_addr = server_addr
self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
self.sock.bind(("0.0.0.0", 0))
self.sock.settimeout(0.5) # Non-blocking receive with timeout
self.peer_addr = None
self.connected = False
self.running = True
def register(self):
# Register to server
msg = f"REGISTER:{self.client_id}"
self.sock.sendto(msg.encode(), self.server_addr)
print(f"Registered as {self.client_id}")
time.sleep(0.5)
def query_peer(self, peer_id):
# Query peer address from server
msg = f"QUERY:{peer_id}"
self.sock.sendto(msg.encode(), self.server_addr)
print(f"Querying {peer_id}...")
def send_thread(self):
# Thread for sending packets periodically
last_send_time = 0
while self.running:
current_time = time.time()
# Send packet every 1 second if not connected or as keepalive
if self.peer_addr and (current_time - last_send_time >= 1.0):
try:
msg = f"PING:{self.client_id}"
self.sock.sendto(msg.encode(), self.peer_addr)
last_send_time = current_time
if not self.connected:
print(f"Punching hole to {self.peer_addr}...")
except Exception as e:
print(f"Send error: {e}")
time.sleep(0.1) # Small sleep to avoid busy loop
def receive_thread(self):
# Thread for receiving packets
while self.running:
try:
data, addr = self.sock.recvfrom(1024)
msg = data.decode()
if msg.startswith("ADDR:"):
# Got peer address from server
parts = msg.split(":")
peer_ip = parts[1]
peer_port = int(parts[2])
self.peer_addr = (peer_ip, peer_port)
print(f"Got peer address: {self.peer_addr}")
elif msg.startswith("PEER:"):
# Server notified us that someone wants to connect
parts = msg.split(":")
peer_ip = parts[1]
peer_port = int(parts[2])
self.peer_addr = (peer_ip, peer_port)
print(f"Peer wants to connect: {self.peer_addr}")
elif msg.startswith("PING:"):
# Received ping from peer
if not self.connected:
self.connected = True
print(f"*** Hole punched! Connected to {addr} ***")
# Send pong back
pong_msg = f"PONG:{self.client_id}"
self.sock.sendto(pong_msg.encode(), addr)
elif msg.startswith("PONG:"):
# Received pong from peer
if not self.connected:
self.connected = True
print(f"*** Hole punched! Connected to {addr} ***")
elif msg.startswith("MSG:"):
# Normal message from peer
print(f"<< {msg[4:]}")
except socket.timeout:
# Timeout, continue loop
continue
except Exception as e:
if self.running:
print(f"Receive error: {e}")
def run(self, peer_id):
# Start threads
recv_thread = threading.Thread(target=self.receive_thread, daemon=True)
send_thread = threading.Thread(target=self.send_thread, daemon=True)
recv_thread.start()
send_thread.start()
# Register and query peer
self.register()
time.sleep(1)
self.query_peer(peer_id)
# Wait for connection
print("Waiting for connection...")
while not self.connected:
time.sleep(0.5)
# Interactive chat
print("You can now send messages (type 'quit' to exit):")
try:
while True:
user_input = input("> ")
if user_input.lower() == 'quit':
break
if self.peer_addr:
msg = f"MSG:{user_input}"
self.sock.sendto(msg.encode(), self.peer_addr)
except KeyboardInterrupt:
pass
self.running = False
print("Shutting down...")
def main():
import sys
if len(sys.argv) != 4:
print("Usage: python client.py <client_id> <peer_id> <server_ip>")
print("Example: python client.py A B 1.2.3.4")
return
client_id = sys.argv[1]
peer_id = sys.argv[2]
server_ip = sys.argv[3]
server_addr = (server_ip, 9999)
client = UDPClient(client_id, server_addr)
client.run(peer_id)
if __name__ == "__main__":
main()Usage
1. Start server on public IP:
bash
python server.py2. Start Client A (in LAN A):
bash
python client.py A B <server_public_ip>3. Start Client B (in LAN B):
bash
python client.py B A <server_public_ip>Key Features
- Non-blocking: Uses separate threads for sending and receiving
- Auto-retry: Sends UDP packet every 1 second when not connected
- Keepalive: Continues sending packets after connection for NAT mapping maintenance
- Python 3: Compatible with Python 3.x
- No Chinese: All comments and messages in English