9cc1f1d58f
- Align protocol documentation (PROTOCOL.md, README.md, VERSIONING.md) - Go: add nonce test, update communicator - Kotlin: update Communicator, TcpClient, TcpServer, add TLS test - Python: update all modules, add certificate test resources - TypeScript: update communicator, tcp/ws clients and servers, add tests - Dart: update communicator, heartbeat mixin, and tests
95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import ssl
|
|
from contextlib import suppress
|
|
from typing import Any
|
|
|
|
from google.protobuf.message import DecodeError
|
|
from websockets.exceptions import ConnectionClosed
|
|
|
|
try:
|
|
from websockets.asyncio.client import connect as ws_connect
|
|
except ImportError: # websockets 12.x legacy API
|
|
from websockets import connect as ws_connect
|
|
|
|
from toki_socket.base_client import BaseClient
|
|
from toki_socket.communicator import ParserMap, type_name_of
|
|
from toki_socket.packets.message_common_pb2 import HeartBeat, PacketBase
|
|
|
|
|
|
class WsClient(BaseClient):
|
|
def __init__(
|
|
self,
|
|
ws: Any,
|
|
interval_sec: int,
|
|
wait_sec: int,
|
|
parser_map: ParserMap,
|
|
) -> None:
|
|
super().__init__(interval_sec, wait_sec, self._do_close_impl)
|
|
self._ws = ws
|
|
self._init_base()
|
|
self.communicator.initialize(self, parser_map)
|
|
self.communicator.set_write_error_handler(lambda _: asyncio.create_task(self._on_disconnected()))
|
|
self.communicator.add_listener(
|
|
type_name_of(HeartBeat),
|
|
lambda _: asyncio.create_task(self._on_heartbeat()),
|
|
)
|
|
self._read_task = asyncio.create_task(self._read_loop())
|
|
asyncio.create_task(self._send_heartbeat())
|
|
|
|
async def write_packet(self, base: PacketBase) -> None:
|
|
await self._ws.send(base.SerializeToString())
|
|
|
|
async def wait_closed(self) -> None:
|
|
await self.communicator.wait_closed()
|
|
|
|
async def _do_close_impl(self) -> None:
|
|
with suppress(Exception):
|
|
await self._ws.close()
|
|
|
|
async def _read_loop(self) -> None:
|
|
try:
|
|
async for message in self._ws:
|
|
if not isinstance(message, bytes):
|
|
continue
|
|
base = PacketBase()
|
|
base.ParseFromString(message)
|
|
self.communicator.on_received_data(
|
|
base.typeName,
|
|
base.data,
|
|
base.nonce,
|
|
base.responseNonce,
|
|
)
|
|
asyncio.create_task(self._send_heartbeat())
|
|
except (ConnectionClosed, ConnectionError, OSError, DecodeError):
|
|
pass
|
|
finally:
|
|
await self._on_disconnected()
|
|
|
|
|
|
async def connect_ws(
|
|
host: str,
|
|
port: int,
|
|
path: str,
|
|
interval_sec: int,
|
|
wait_sec: int,
|
|
parser_map: ParserMap,
|
|
) -> WsClient:
|
|
uri = f"ws://{host}:{port}{path}"
|
|
ws = await ws_connect(uri)
|
|
return WsClient(ws, interval_sec, wait_sec, parser_map)
|
|
|
|
|
|
async def connect_wss(
|
|
host: str,
|
|
port: int,
|
|
path: str,
|
|
ssl_context: ssl.SSLContext,
|
|
interval_sec: int,
|
|
wait_sec: int,
|
|
parser_map: ParserMap,
|
|
) -> WsClient:
|
|
uri = f"wss://{host}:{port}{path}"
|
|
ws = await ws_connect(uri, ssl=ssl_context)
|
|
return WsClient(ws, interval_sec, wait_sec, parser_map)
|