"""Module for handling the TCP connection with Gateway.""" import asyncio import ssl import sys from typing import Any, Callable, Coroutine, List, Set from .api.frame_creation import frame_from_raw from .api.frames import FrameBase from .config import Config from .exception import PyVLXException from .log import PYVLXLOG from .slip import get_next_slip, is_slip, slip_pack class SlipTokenizer: """Helper class for splitting up binary stream to slip packets.""" def __init__(self) -> None: """Init Tokenizer.""" self.data = bytes() def feed(self, chunk: bytes) -> None: """Feed chunk to tokenizer.""" if not chunk: return self.data += chunk def has_tokens(self) -> bool: """Return True if Tokenizer has tokens.""" return is_slip(self.data) def get_next_token(self) -> bytes | None: """Get next token from Tokenizer.""" slip, self.data = get_next_slip(self.data) return slip class TCPTransport(asyncio.Protocol): """Class for handling asyncio connection transport.""" def __init__( self, frame_received_cb: Callable[[FrameBase], None], connection_lost_cb: Callable[[], None], ): """Init TCPTransport.""" self.frame_received_cb = frame_received_cb self.connection_lost_cb = connection_lost_cb self.tokenizer = SlipTokenizer() def connection_made(self, transport: object) -> None: """Handle sucessful connection.""" PYVLXLOG.debug("Socket connection to KLF 200 opened") def data_received(self, data: bytes) -> None: """Handle data received.""" self.tokenizer.feed(data) while self.tokenizer.has_tokens(): raw = self.tokenizer.get_next_token() assert raw is not None try: frame = frame_from_raw(raw) if frame is not None: self.frame_received_cb(frame) except PyVLXException: PYVLXLOG.error("Error in data_received", exc_info=sys.exc_info()) def connection_lost(self, exc: object) -> None: """Handle lost connection.""" PYVLXLOG.debug("Socket connection to KLF 200 has been lost") self.connection_lost_cb() CallbackType = Callable[[FrameBase], Coroutine[Any, Any, None]] class Connection: """Class for handling TCP connection.""" CONNECT_TIMEOUT = 10.0 def __init__(self, config: Config): """Init TCP connection.""" self.config = config self.transport: asyncio.Transport | None = None self.frame_received_cbs: List[CallbackType] = [] self.connection_closed_cbs: List[Callable[[], Coroutine[Any, Any, None]]] = [] self.connection_opened_cbs: List[Callable[[], Coroutine[Any, Any, None]]] = [] self.connected = False self.connection_counter = 0 self.tasks: Set[asyncio.Task[None]] = set() def __del__(self) -> None: """Destruct connection.""" self.disconnect(notify_callbacks=False) def disconnect(self, notify_callbacks: bool = True) -> None: """Disconnect connection and notify callbacks if specified. Running callbacks only makes sense if loop is still running, so it it can be skipped, mostly for the case of destructor being called during shutdown when loop is already closed. """ if self.transport is not None: self.transport.close() self.transport = None self.connected = False PYVLXLOG.debug("TCP transport closed.") if not notify_callbacks: return try: loop = asyncio.get_running_loop() except RuntimeError: PYVLXLOG.debug("Skipping connection closed callbacks because no event loop is running.") return for connection_closed_cb in self.connection_closed_cbs: task = loop.create_task(connection_closed_cb()) self.tasks.add(task) task.add_done_callback(self.tasks.discard) async def connect(self) -> None: """Connect to gateway via SSL.""" tcp_client = TCPTransport(self.frame_received_cb, connection_lost_cb=self.on_connection_lost) loop = asyncio.get_running_loop() assert self.config.host is not None try: async with asyncio.timeout(self.CONNECT_TIMEOUT): self.transport, _ = await loop.create_connection( lambda: tcp_client, host=self.config.host, port=self.config.port, ssl=self.create_ssl_context(), ) except asyncio.TimeoutError as error: self.transport = None self.connected = False raise PyVLXException( f"Socket connection to KLF 200 timed out after {self.CONNECT_TIMEOUT} seconds" ) from error except (OSError, ssl.SSLError) as error: self.transport = None self.connected = False raise PyVLXException(f"Failed to open socket connection to KLF 200: {error}") from error self.connected = True self.connection_counter += 1 PYVLXLOG.debug( "Number of connections since last HA start: %s", self.connection_counter ) for connection_opened_cb in self.connection_opened_cbs: task = asyncio.create_task(connection_opened_cb()) self.tasks.add(task) task.add_done_callback(self.tasks.remove) def register_frame_received_cb(self, callback: CallbackType) -> None: """Register frame received callback.""" self.frame_received_cbs.append(callback) def unregister_frame_received_cb(self, callback: CallbackType) -> None: """Unregister frame received callback.""" self.frame_received_cbs.remove(callback) def register_connection_closed_cb(self, callback: Callable[[], Coroutine[Any, Any, None]]) -> None: """Register connection closed callback.""" self.connection_closed_cbs.append(callback) def unregister_connection_closed_cb(self, callback: Callable[[], Coroutine[Any, Any, None]]) -> None: """Unregister connection closed callback.""" self.connection_closed_cbs.remove(callback) def register_connection_opened_cb(self, callback: Callable[[], Coroutine[Any, Any, None]]) -> None: """Register connection opened callback.""" self.connection_opened_cbs.append(callback) def unregister_connection_opened_cb(self, callback: Callable[[], Coroutine[Any, Any, None]]) -> None: """Unregister connection opened callback.""" self.connection_opened_cbs.remove(callback) def write(self, frame: FrameBase) -> None: """Write frame to Bus.""" if not isinstance(frame, FrameBase): raise PyVLXException("Frame not of type FrameBase", *type(frame)) PYVLXLOG.debug("SEND: %s", frame) assert self.transport is not None self.transport.write(slip_pack(bytes(frame))) @staticmethod def create_ssl_context() -> ssl.SSLContext: """Create and return SSL Context.""" ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) ssl_context.check_hostname = False ssl_context.verify_mode = ssl.CERT_NONE return ssl_context def frame_received_cb(self, frame: FrameBase) -> None: """Received message.""" PYVLXLOG.debug("REC: %s", frame) for frame_received_cb in self.frame_received_cbs: task = asyncio.create_task(frame_received_cb(frame)) self.tasks.add(task) task.add_done_callback(self.tasks.remove) def on_connection_lost(self) -> None: """Server closed connection.""" self.disconnect()