AwinSimpleP2P/utils.py
2025-06-08 20:17:59 +08:00

105 lines
3.0 KiB
Python

import json
import socket
import struct
import time
def pack_data(action, message=None, data=None):
if data is None:
data = b''
if message is None:
message = {}
# 将字符串编码为UTF-8字节串
action_bytes = action.encode('utf-8')
message_bytes = json.dumps(message).encode('utf-8')
# 使用大端序打包长度信息(4字节无符号整数)
# >I: 大端序无符号整型(4字节)
packed = struct.pack('>I', len(action_bytes))
packed += action_bytes
packed += struct.pack('>I', len(message_bytes))
packed += message_bytes
packed += data # 直接附加二进制数据
return packed
def unpack_data(packed):
# 解析action长度(前4字节)
action_len = struct.unpack('>I', packed[:4])[0]
offset = 4
# 提取action字节并解码
action_bytes = packed[offset:offset + action_len]
action = action_bytes.decode('utf-8')
offset += action_len
# 解析message长度(接下来的4字节)
message_len = struct.unpack('>I', packed[offset:offset + 4])[0]
offset += 4
# 提取message字节并解码
message_bytes = packed[offset:offset + message_len]
message = json.loads(message_bytes.decode('utf-8'))
offset += message_len
# 剩余部分是原始二进制数据
data = packed[offset:]
return action, message, data
def send_pack_tcp(sock, pack):
sock.sendall(pack)
def send_tcp(sock, action, data, message):
packed = pack_data(action, message, data)
sock.sendall(packed)
def send_udp(sock, action, data, message, addr):
if isinstance(sock, int):
sock_ = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock_.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock_.bind(('0.0.0.0', sock))
else:
sock_ = sock
packed = pack_data(action, message, data)
sock_.sendto(packed, addr)
if isinstance(sock, int):
sock_.close()
def send_data_udp(sock, addr, data, message=None):
send_udp(sock, 'data', data, message, addr)
def send_data_tcp(sock, data, message=None):
send_tcp(sock, 'data', data, message)
def send_action_udp(sock, addr, action, message=None):
if isinstance(message, str):
message = {'message': message}
elif isinstance(message, dict):
message['timestamp'] = int(time.time())
else:
try:
str_message = str(message)
message = {'message': str_message}
except Exception as e:
print(f'无法将数据转换为字符串: {e}')
send_udp(sock, action, None, message, addr)
def send_action_tcp(sock, action, message=None):
if isinstance(message, str):
message = {'message': message}
send_tcp(sock, action, None, message)
def recv_tcp(sock):
pack = sock.recv(4096)
action, message, data = unpack_data(pack)
return action, message, data
def recv_udp(sock):
pack, addr = sock.recvfrom(4096)
action, message, data = unpack_data(pack)
return action, message, data, addr