imonlinux
1eef4ba917
Includes various compatibility fixes, reversions, and clean-up from the previous development cycle to create a stable base for the MQTT implementation. This ensures the application starts reliably with the pinned aioesphomeapi version.
162 lines
4.9 KiB
Python
162 lines
4.9 KiB
Python
"""Partial ESPHome server implementation."""
|
|
|
|
import asyncio
|
|
import logging
|
|
from abc import abstractmethod
|
|
from collections.abc import Iterable
|
|
from typing import TYPE_CHECKING, List, Optional
|
|
|
|
# pylint: disable=no-name-in-module
|
|
from aioesphomeapi._frame_helper.packets import make_plain_text_packets
|
|
from aioesphomeapi.api_pb2 import ( # type: ignore[attr-defined]
|
|
DisconnectRequest,
|
|
DisconnectResponse,
|
|
HelloRequest,
|
|
HelloResponse,
|
|
PingRequest,
|
|
PingResponse,
|
|
)
|
|
from aioesphomeapi.core import MESSAGE_TYPE_TO_PROTO
|
|
from google.protobuf import message
|
|
|
|
PROTO_TO_MESSAGE_TYPE = {v: k for k, v in MESSAGE_TYPE_TO_PROTO.items()}
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
|
|
class APIServer(asyncio.Protocol):
|
|
|
|
def __init__(self, name: str) -> None:
|
|
self.name = name
|
|
|
|
self._buffer: Optional[bytes] = None
|
|
self._buffer_len: int = 0
|
|
self._pos: int = 0
|
|
self._transport = None
|
|
self._writelines = None
|
|
|
|
@abstractmethod
|
|
def handle_message(self, msg: message.Message) -> Iterable[message.Message]:
|
|
pass
|
|
|
|
def process_packet(self, msg_type: int, packet_data: bytes) -> None:
|
|
msg_class = MESSAGE_TYPE_TO_PROTO[msg_type]
|
|
msg_inst = msg_class.FromString(packet_data)
|
|
|
|
if isinstance(msg_inst, HelloRequest):
|
|
self.send_messages(
|
|
[
|
|
HelloResponse(
|
|
api_version_major=1,
|
|
api_version_minor=10,
|
|
name=self.name,
|
|
)
|
|
]
|
|
)
|
|
return
|
|
|
|
if isinstance(msg_inst, DisconnectRequest):
|
|
self.send_messages([DisconnectResponse()])
|
|
_LOGGER.debug("Disconnect requested")
|
|
if self._transport:
|
|
self._transport.close()
|
|
self._transport = None
|
|
self._writelines = None
|
|
elif isinstance(msg_inst, PingRequest):
|
|
self.send_messages([PingResponse()])
|
|
elif msgs := self.handle_message(msg_inst):
|
|
if isinstance(msgs, message.Message):
|
|
msgs = [msgs]
|
|
|
|
self.send_messages(msgs)
|
|
|
|
def send_messages(self, msgs: List[message.Message]):
|
|
if self._writelines is None:
|
|
return
|
|
|
|
packets = [
|
|
(PROTO_TO_MESSAGE_TYPE[msg.__class__], msg.SerializeToString())
|
|
for msg in msgs
|
|
]
|
|
packet_bytes = make_plain_text_packets(packets)
|
|
self._writelines(packet_bytes)
|
|
|
|
def connection_made(self, transport) -> None:
|
|
self._transport = transport
|
|
self._writelines = transport.writelines
|
|
|
|
def data_received(self, data: bytes):
|
|
if self._buffer is None:
|
|
self._buffer = data
|
|
self._buffer_len = len(data)
|
|
else:
|
|
self._buffer += data
|
|
self._buffer_len += len(data)
|
|
|
|
while self._buffer_len >= 3:
|
|
self._pos = 0
|
|
if (preamble := self._read_varuint()) != 0x00:
|
|
_LOGGER.error("Incorrect preamble: %s", preamble)
|
|
return
|
|
|
|
if (length := self._read_varuint()) == -1:
|
|
_LOGGER.error("Incorrect length: %s", length)
|
|
return
|
|
|
|
if (msg_type := self._read_varuint()) == -1:
|
|
_LOGGER.error("Incorrect message type: %s", msg_type)
|
|
return
|
|
|
|
if length == 0:
|
|
self._remove_from_buffer()
|
|
self.process_packet(msg_type, b"")
|
|
continue
|
|
|
|
if (packet_data := self._read(length)) is None:
|
|
return
|
|
|
|
self._remove_from_buffer()
|
|
self.process_packet(msg_type, packet_data)
|
|
|
|
def _read(self, length: int) -> bytes | None:
|
|
new_pos = self._pos + length
|
|
if self._buffer_len < new_pos:
|
|
return None
|
|
original_pos = self._pos
|
|
self._pos = new_pos
|
|
if TYPE_CHECKING:
|
|
assert self._buffer is not None, "Buffer should be set"
|
|
cstr = self._buffer
|
|
return cstr[original_pos:new_pos]
|
|
|
|
def connection_lost(self, exc):
|
|
self._transport = None
|
|
self._writelines = None
|
|
|
|
def _read_varuint(self) -> int:
|
|
if not self._buffer:
|
|
return -1
|
|
|
|
result = 0
|
|
bitpos = 0
|
|
cstr = self._buffer
|
|
while self._buffer_len > self._pos:
|
|
val = cstr[self._pos]
|
|
self._pos += 1
|
|
result |= (val & 0x7F) << bitpos
|
|
if (val & 0x80) == 0:
|
|
return result
|
|
bitpos += 7
|
|
return -1
|
|
|
|
def _remove_from_buffer(self) -> None:
|
|
end_of_frame_pos = self._pos
|
|
self._buffer_len -= end_of_frame_pos
|
|
if self._buffer_len == 0:
|
|
self._buffer = None
|
|
return
|
|
if TYPE_CHECKING:
|
|
assert self._buffer is not None, "Buffer should be set"
|
|
cstr = self._buffer
|
|
self._buffer = cstr[end_of_frame_pos : self._buffer_len + end_of_frame_pos]
|