78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
import socket
|
|
import threading
|
|
import sys
|
|
|
|
class TyNetworkManager:
|
|
def __init__(self, on_packet_received):
|
|
self.sock = None
|
|
self.on_packet_received = on_packet_received
|
|
self.is_running = False
|
|
|
|
def connect(self, host: str, port: int) -> bool:
|
|
try:
|
|
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
self.sock.connect((host, port))
|
|
self.is_running = True
|
|
threading.Thread(target=self._read_loop, daemon=True).start()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def send_packet(self, packet: bytes) -> bool:
|
|
if not self.sock or not self.is_running:
|
|
return False
|
|
try:
|
|
self.sock.sendall(packet)
|
|
return True
|
|
except Exception:
|
|
self.disconnect()
|
|
return False
|
|
|
|
def register_id(self, my_id: int):
|
|
header = bytes([0x01, 0x03]) + my_id.to_bytes(3, 'big') + (0).to_bytes(3, 'big')
|
|
meta = b'\x00' * 12 + (0).to_bytes(4, 'big')
|
|
self.send_packet(header + meta)
|
|
|
|
def disconnect(self):
|
|
self.is_running = False
|
|
if self.sock:
|
|
try:
|
|
self.sock.close()
|
|
except Exception:
|
|
pass
|
|
self.sock = None
|
|
|
|
def _read_loop(self):
|
|
while self.is_running:
|
|
try:
|
|
header = self._recv_exact(8)
|
|
if not header:
|
|
break
|
|
|
|
meta = self._recv_exact(16)
|
|
if not meta:
|
|
break
|
|
|
|
iv = meta[0:12]
|
|
payload_len = int.from_bytes(meta[12:16], 'big')
|
|
|
|
payload = b""
|
|
if payload_len > 0:
|
|
payload = self._recv_exact(payload_len)
|
|
if not payload:
|
|
break
|
|
|
|
self.on_packet_received(header, iv, payload)
|
|
|
|
except Exception:
|
|
break
|
|
self.disconnect()
|
|
|
|
def _recv_exact(self, length: int) -> bytes:
|
|
data = b""
|
|
while len(data) < length:
|
|
chunk = self.sock.recv(length - len(data))
|
|
if not chunk:
|
|
return b""
|
|
data += chunk
|
|
return data |