feat(net): добавление функционала для работы с сетью
This commit is contained in:
@@ -0,0 +1,7 @@
|
|||||||
|
from .controller import NetworkController
|
||||||
|
from .protocol import NetworkMesssage
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"NetworkController",
|
||||||
|
"NetworkMesssage",
|
||||||
|
]
|
||||||
@@ -0,0 +1,150 @@
|
|||||||
|
import socket
|
||||||
|
import struct
|
||||||
|
import threading
|
||||||
|
from typing import Callable, Optional
|
||||||
|
|
||||||
|
import wx
|
||||||
|
|
||||||
|
from .protocol import NetworkMesssage
|
||||||
|
|
||||||
|
OnMessageCallback = Callable[[NetworkMesssage], None]
|
||||||
|
OnNoticeCallback = Callable[[str], None]
|
||||||
|
|
||||||
|
|
||||||
|
def receive_exact(sock: socket.socket, n: int) -> Optional[bytes]:
|
||||||
|
data = bytearray()
|
||||||
|
|
||||||
|
while len(data) < n:
|
||||||
|
packet = sock.recv(n - len(data))
|
||||||
|
if not packet:
|
||||||
|
return None
|
||||||
|
data.extend(packet)
|
||||||
|
|
||||||
|
return bytes(data)
|
||||||
|
|
||||||
|
|
||||||
|
class NetworkController(threading.Thread):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
on_message_callback: OnMessageCallback,
|
||||||
|
on_notice_callback: Optional[OnNoticeCallback] = None,
|
||||||
|
):
|
||||||
|
super().__init__(daemon=True)
|
||||||
|
self.on_message_callback = on_message_callback
|
||||||
|
self.on_notice_callback = on_notice_callback
|
||||||
|
|
||||||
|
self.sock: Optional[socket.socket] = None
|
||||||
|
self.conn: Optional[socket.socket] = None
|
||||||
|
|
||||||
|
self.is_running: bool = False
|
||||||
|
self.is_server: bool = False
|
||||||
|
|
||||||
|
def serve(self, host: str = "0.0.0.0", port: int = 8994) -> None:
|
||||||
|
self.is_server = True
|
||||||
|
|
||||||
|
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||||
|
self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||||
|
self.sock.bind((host, port))
|
||||||
|
self.sock.listen(1)
|
||||||
|
|
||||||
|
self.is_running = True
|
||||||
|
self.start()
|
||||||
|
|
||||||
|
def connect(self, host: str, port: int = 8994) -> bool:
|
||||||
|
self.is_server = False
|
||||||
|
|
||||||
|
try:
|
||||||
|
self.conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||||
|
self.conn.connect((host, port))
|
||||||
|
|
||||||
|
self.is_running = True
|
||||||
|
self._notify("Connected to host")
|
||||||
|
self.start()
|
||||||
|
return True
|
||||||
|
except Exception as e:
|
||||||
|
self._notify(f"Connection failed: {e}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
def send(self, event: str, payload: dict = {}) -> None:
|
||||||
|
if not self.conn or not self.is_running:
|
||||||
|
return
|
||||||
|
|
||||||
|
m = NetworkMesssage(event, payload)
|
||||||
|
|
||||||
|
try:
|
||||||
|
self.conn.sendall(m.encode())
|
||||||
|
except Exception as e:
|
||||||
|
self._notify(f"Send error: {e}")
|
||||||
|
self._close_active_connection()
|
||||||
|
|
||||||
|
def run(self) -> None:
|
||||||
|
HEADER_SIZE = 4
|
||||||
|
|
||||||
|
while self.is_running:
|
||||||
|
if self.is_server and self.sock and not self.conn:
|
||||||
|
self._notify("Waiting for peer to connect...")
|
||||||
|
|
||||||
|
try:
|
||||||
|
self.conn, addr = self.sock.accept()
|
||||||
|
self._notify(f"Peer connected from {addr[0]}:{addr[1]}")
|
||||||
|
except Exception:
|
||||||
|
break
|
||||||
|
|
||||||
|
while self.is_running and self.conn:
|
||||||
|
try:
|
||||||
|
header_bytes = receive_exact(self.conn, HEADER_SIZE)
|
||||||
|
if not header_bytes:
|
||||||
|
self._notify("Peer disconnected")
|
||||||
|
break
|
||||||
|
|
||||||
|
body_size = struct.unpack(">I", header_bytes)[0]
|
||||||
|
|
||||||
|
body_bytes = receive_exact(self.conn, body_size)
|
||||||
|
if not body_bytes:
|
||||||
|
self._notify("Peer disconnected unexpectedly")
|
||||||
|
break
|
||||||
|
|
||||||
|
json_str = body_bytes.decode("utf-8")
|
||||||
|
m = NetworkMesssage.from_json(json_str)
|
||||||
|
|
||||||
|
wx.CallAfter(self.on_message_callback, m)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self._notify(f"Read error: {e}")
|
||||||
|
break
|
||||||
|
|
||||||
|
self._close_active_connection()
|
||||||
|
|
||||||
|
if not self.is_server:
|
||||||
|
break
|
||||||
|
|
||||||
|
self.stop()
|
||||||
|
|
||||||
|
def _close_active_connection(self) -> None:
|
||||||
|
if self.conn:
|
||||||
|
try:
|
||||||
|
self.conn.shutdown(socket.SHUT_RDWR)
|
||||||
|
self.conn.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
self.conn = None
|
||||||
|
|
||||||
|
def _notify(self, notice: str) -> None:
|
||||||
|
if self.on_notice_callback:
|
||||||
|
wx.CallAfter(self.on_notice_callback, notice)
|
||||||
|
|
||||||
|
def stop(self) -> None:
|
||||||
|
self.is_running = False
|
||||||
|
|
||||||
|
self._close_active_connection()
|
||||||
|
|
||||||
|
if self.sock:
|
||||||
|
try:
|
||||||
|
self.sock.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
self.sock = None
|
||||||
|
|
||||||
|
self._notify("Connection closed")
|
||||||
@@ -0,0 +1,46 @@
|
|||||||
|
import itertools
|
||||||
|
import json
|
||||||
|
import struct
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
|
||||||
|
_id_generator = itertools.count(0)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class NetworkMesssage:
|
||||||
|
event: str
|
||||||
|
payload: dict = field(default_factory=dict)
|
||||||
|
id: int = field(default_factory=lambda: next(_id_generator))
|
||||||
|
version: str = "1.0"
|
||||||
|
|
||||||
|
def to_json(self) -> str:
|
||||||
|
return json.dumps(
|
||||||
|
{
|
||||||
|
"version": self.version,
|
||||||
|
"id": self.id,
|
||||||
|
"event": self.event,
|
||||||
|
"payload": self.payload,
|
||||||
|
},
|
||||||
|
ensure_ascii=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_json(cls, json_str: str) -> "NetworkMesssage":
|
||||||
|
data = json.loads(json_str)
|
||||||
|
return cls(
|
||||||
|
event=data.get("event", "unknown"),
|
||||||
|
payload=data.get("payload", {}),
|
||||||
|
id=data.get("id", -1),
|
||||||
|
version=data.get("version", "unknown"),
|
||||||
|
)
|
||||||
|
|
||||||
|
def encode(self) -> bytes:
|
||||||
|
raw_bytes = self.to_json().encode("utf-8")
|
||||||
|
length_prefix = struct.pack(">I", len(raw_bytes))
|
||||||
|
return length_prefix + raw_bytes
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
for i in range(10):
|
||||||
|
message = NetworkMesssage("none")
|
||||||
|
print(message.encode())
|
||||||
Reference in New Issue
Block a user