ohmqtt 0.0.3__py3-none-any.whl

ohmqtt/client.py

@@ -0,0 +1,266 @@
+from __future__ import annotations
+
+import ssl
+import threading
+from typing import Final, Iterable, Sequence
+import weakref
+
+from .connection import Address, ConnectParams, Connection, MessageHandlers
+from .logger import get_logger
+from .mqtt_spec import MQTTReasonCode
+from .packet import MQTTAuthPacket
+from .property import MQTTAuthProps, MQTTConnectProps, MQTTPublishProps, MQTTWillProps
+from .persistence.base import PublishHandle
+from .session import Session
+from .subscriptions import Subscriptions, SubscribeCallback, SubscribeHandle, UnsubscribeHandle, RetainPolicy
+from .topic_alias import AliasPolicy
+
+logger: Final = get_logger("client")
+
+
+class Client:
+    """High level interface for the MQTT client."""
+    __slots__ = (
+        "__weakref__",
+        "_thread",
+        "connection",
+        "session",
+        "subscriptions",
+    )
+
+    def __init__(self, db_path: str = "", *, db_fast: bool = False) -> None:
+        self._thread: threading.Thread | None = None
+        message_handlers = MessageHandlers()
+        with message_handlers as handlers:
+            self.connection = Connection(handlers)
+            self.subscriptions = Subscriptions(handlers, self.connection, weakref.ref(self))
+            self.session = Session(handlers, self.connection, db_path=db_path, db_fast=db_fast)
+            handlers.register(MQTTAuthPacket, self.handle_auth)
+
+    def __enter__(self) -> Client:
+        self.start_loop()
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.shutdown()
+
+    def connect(
+        self,
+        address: str,
+        *,
+        client_id: str = "",
+        clean_start: bool = False,
+        connect_timeout: float | None = None,
+        reconnect_delay: int = 0,
+        keepalive_interval: int = 0,
+        tcp_nodelay: bool = True,
+        tls_context: ssl.SSLContext | None = None,
+        tls_hostname: str = "",
+        will_topic: str = "",
+        will_payload: bytes = b"",
+        will_qos: int = 0,
+        will_retain: bool = False,
+        will_properties: MQTTWillProps | None = None,
+        connect_properties: MQTTConnectProps | None = None,
+    ) -> None:
+        """Connect to the broker."""
+        _address = Address(address)
+        params = ConnectParams(
+            address=_address,
+            client_id=client_id,
+            clean_start=clean_start,
+            connect_timeout=connect_timeout,
+            reconnect_delay=reconnect_delay,
+            keepalive_interval=keepalive_interval,
+            tcp_nodelay=tcp_nodelay,
+            tls_context=tls_context,
+            tls_hostname=tls_hostname,
+            will_topic=will_topic,
+            will_payload=will_payload,
+            will_qos=will_qos,
+            will_retain=will_retain,
+            will_properties=will_properties if will_properties is not None else MQTTWillProps(),
+            connect_properties=connect_properties if connect_properties is not None else MQTTConnectProps(),
+        )
+        self.session.set_params(params)
+        self.connection.connect(params)
+
+    def disconnect(self) -> None:
+        """Disconnect from the broker."""
+        self.connection.disconnect()
+
+    def shutdown(self) -> None:
+        """Shutdown the client and close the connection."""
+        self.connection.shutdown()
+
+    def publish(
+        self,
+        topic: str,
+        payload: bytes,
+        *,
+        qos: int = 0,
+        retain: bool = False,
+        properties: MQTTPublishProps | None = None,
+        alias_policy: AliasPolicy = AliasPolicy.NEVER,
+    ) -> PublishHandle:
+        """Publish a message to a topic."""
+        properties = properties if properties is not None else None
+        return self.session.publish(
+            topic,
+            payload,
+            qos=qos,
+            retain=retain,
+            properties=properties,
+            alias_policy=alias_policy,
+        )
+
+    def subscribe(
+        self,
+        topic_filter: str,
+        callback: SubscribeCallback,
+        max_qos: int = 2,
+        *,
+        share_name: str | None = None,
+        no_local: bool = False,
+        retain_as_published: bool = False,
+        retain_policy: RetainPolicy = RetainPolicy.ALWAYS,
+        sub_id: int | None = None,
+        user_properties: Sequence[tuple[str, str]] | None = None,
+    ) -> SubscribeHandle | None:
+        """Subscribe to a topic filter with a callback.
+
+        If the client is connected, returns a handle which can be used to unsubscribe from the topic filter
+        or wait for the subscription to be acknowledged.
+
+        If the client is not connected, returns None."""
+        return self.subscriptions.subscribe(
+            topic_filter,
+            callback,
+            max_qos=max_qos,
+            share_name=share_name,
+            no_local=no_local,
+            retain_as_published=retain_as_published,
+            retain_policy=retain_policy,
+            sub_id=sub_id,
+            user_properties=user_properties,
+        )
+
+    def unsubscribe(
+        # This method must have the same signature as the subscribe method.
+        # This lets us match the unsubscribe to the subscribe with the same args.
+        self,
+        topic_filter: str,
+        callback: SubscribeCallback,
+        max_qos: int = 2,
+        *,
+        share_name: str | None = None,
+        no_local: bool = False,
+        retain_as_published: bool = False,
+        retain_policy: RetainPolicy = RetainPolicy.ALWAYS,
+        sub_id: int | None = None,
+        user_properties: Iterable[tuple[str, str]] | None = None,
+        unsub_user_properties: Iterable[tuple[str, str]] | None = None,
+    ) -> UnsubscribeHandle | None:
+        """Unsubscribe from a topic filter.
+
+        If the client is connected, returns a handle which can be used to wait for the unsubscription to be acknowledged.
+
+        If the client is not connected, returns None."""
+        return self.subscriptions.unsubscribe(
+            topic_filter,
+            callback,
+            max_qos=max_qos,
+            share_name=share_name,
+            no_local=no_local,
+            retain_as_published=retain_as_published,
+            retain_policy=retain_policy,
+            sub_id=sub_id,
+            user_properties=user_properties,
+            unsub_user_properties=unsub_user_properties,
+        )
+
+    def auth(
+        self,
+        *,
+        authentication_method: str | None = None,
+        authentication_data: bytes | None = None,
+        reason_string: str | None = None,
+        user_properties: Sequence[tuple[str, str]] | None = None,
+        reason_code: MQTTReasonCode = MQTTReasonCode.Success,
+    ) -> None:
+        """Send an AUTH packet to the broker."""
+        properties = MQTTAuthProps()
+        if authentication_method is not None:
+            properties.AuthenticationMethod = authentication_method
+        if authentication_data is not None:
+            properties.AuthenticationData = authentication_data
+        if reason_string is not None:
+            properties.ReasonString = reason_string
+        if user_properties is not None:
+            properties.UserProperty = user_properties
+        packet = MQTTAuthPacket(
+            reason_code=reason_code,
+            properties=properties,
+        )
+        self.connection.send(packet)
+
+    def wait_for_connect(self, timeout: float | None = None) -> None:
+        """Wait for the client to connect to the broker.
+
+        Raises TimeoutError if the timeout is exceeded."""
+        if not self.connection.wait_for_connect(timeout):
+            raise TimeoutError("Waiting for connection timed out")
+
+    def wait_for_disconnect(self, timeout: float | None = None) -> None:
+        """Wait for the client to disconnect from the broker.
+
+        Raises TimeoutError if the timeout is exceeded."""
+        if not self.connection.wait_for_disconnect(timeout):
+            raise TimeoutError("Waiting for disconnection timed out")
+
+    def wait_for_shutdown(self, timeout: float | None = None) -> None:
+        """Wait for the client to disconnect and finalize.
+
+        Raises TimeoutError if the timeout is exceeded."""
+        if not self.connection.wait_for_shutdown(timeout):
+            raise TimeoutError("Waiting for disconnection timed out")
+
+    def start_loop(self) -> None:
+        """Start the client state machine in a separate thread."""
+        if self._thread is not None:
+            raise RuntimeError("Connection loop already started")
+        self._thread = threading.Thread(target=self.loop_forever, daemon=True)
+        self._thread.start()
+
+    def loop_once(self, max_wait: float | None = 0.0) -> None:
+        """Run a single iteration of the MQTT client loop.
+
+        If max_wait is 0.0 (the default), this call will not block.
+
+        If max_wait is None, this call will block until the next event.
+
+        Any other numeric max_wait value may block for maximum that amount of time in seconds."""
+        self.connection.loop_once(max_wait)
+
+    def loop_forever(self) -> None:
+        """Run the MQTT client loop.
+
+        This will run until the client is stopped or shutdown.
+        """
+        self.connection.loop_forever()
+
+    def loop_until_connected(self, timeout: float | None = None) -> bool:
+        """Run the MQTT client loop until the client is connected to the broker.
+
+        If a timeout is provided, the loop will give up after that amount of time.
+
+        Returns True if the client is connected, False if the timeout was reached."""
+        return self.connection.loop_until_connected(timeout)
+
+    def is_connected(self) -> bool:
+        """Check if the client is connected to the broker."""
+        return self.connection.is_connected()
+
+    def handle_auth(self, packet: MQTTAuthPacket) -> None:
+        """Callback for an AUTH packet from the broker."""
+        logger.debug("Got an AUTH packet")

ohmqtt/connection/__init__.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from typing import Final
+
+from .address import Address as Address
+from .fsm import FSM
+from .fsm import InvalidStateError as InvalidStateError
+from .handlers import MessageHandlers as MessageHandlers
+from .states import (
+    ClosingState,
+    ClosedState,
+    ConnectingState,
+    ConnectedState,
+    ReconnectWaitState,
+    ShutdownState,
+)
+from .types import ConnectParams as ConnectParams
+from .types import StateEnvironment, ReceivablePacketT, SendablePacketT
+from ..error import MQTTError
+from ..logger import get_logger
+
+logger: Final = get_logger("connection")
+
+
+class Connection:
+    """Interface for the MQTT connection."""
+    __slots__ = ("_handlers", "fsm")
+
+    def __init__(self, handlers: MessageHandlers) -> None:
+        state_env = StateEnvironment(packet_callback=self.handle_packet)
+        self.fsm = FSM(env=state_env, init_state=ClosedState, error_state=ShutdownState)
+        self._handlers = handlers
+
+    def handle_packet(self, packet: ReceivablePacketT) -> None:
+        """Handle incoming packets by routing them to registered handlers."""
+        logger.debug("<--- %s", packet)
+        exceptions = self._handlers.handle(packet)
+        if exceptions:
+            if any(True for exc in exceptions if isinstance(exc, MQTTError)):
+                # If there is an MQTTError, raise it.
+                raise next(exc for exc in exceptions if isinstance(exc, MQTTError))
+            # Otherwise. raise the first exception.
+            raise exceptions[0]
+
+    def can_send(self) -> bool:
+        """Check if the connection is in a state where data can be sent."""
+        with self.fsm.lock:
+            state = self.fsm.get_state()
+            return state is ConnectedState
+
+    def send(self, packet: SendablePacketT) -> None:
+        """Send data to the connection."""
+        with self.fsm.lock:
+            if not self.can_send():
+                state = self.fsm.get_state()
+                raise InvalidStateError(f"Cannot send data in state {state.__name__}")
+            logger.debug("---> %s", packet)
+            data = packet.encode()
+            self.fsm.env.write_buffer.extend(data)
+            self.fsm.selector.interrupt()
+
+    def connect(self, params: ConnectParams) -> None:
+        """Connect to the MQTT broker."""
+        with self.fsm.lock:
+            self.fsm.set_params(params)
+            self.fsm.request_state(ConnectingState)
+            self.fsm.selector.interrupt()
+
+    def disconnect(self) -> None:
+        """Disconnect from the MQTT broker."""
+        with self.fsm.lock:
+            self.fsm.request_state(ClosingState)
+            self.fsm.selector.interrupt()
+
+    def shutdown(self) -> None:
+        """Shutdown the connection."""
+        with self.fsm.lock:
+            self.fsm.request_state(ShutdownState)
+            self.fsm.selector.interrupt()
+
+    def is_connected(self) -> bool:
+        """Check if the connection is established."""
+        return self.fsm.get_state() == ConnectedState
+
+    def wait_for_connect(self, timeout: float | None = None) -> bool:
+        """Wait for the connection to be established.
+
+        Returns True if the connection is established, False if the timeout is reached."""
+        return self.fsm.wait_for_state((ConnectedState,), timeout)
+
+    def wait_for_disconnect(self, timeout: float | None = None) -> bool:
+        """Wait for the connection to be closed.
+
+        Returns True if the connection is closed, False if the timeout is reached."""
+        return self.fsm.wait_for_state((ClosedState, ShutdownState, ReconnectWaitState), timeout)
+
+    def wait_for_shutdown(self, timeout: float | None = None) -> bool:
+        """Wait for the connection to be closed and finalized.
+
+        Returns True if the connection is closed, False if the timeout is reached."""
+        return self.fsm.wait_for_state((ShutdownState,), timeout)
+
+    def loop_once(self, max_wait: float | None = 0.0) -> None:
+        """Run a single iteration of the state machine.
+
+        If max_wait is None, wait indefinitely. Otherwise, wait for the specified time."""
+        self.fsm.loop_once(max_wait)
+
+    def loop_forever(self) -> None:
+        """Run the state machine until the connection is closed."""
+        self.fsm.loop_until_state((ShutdownState,))
+
+    def loop_until_connected(self, timeout: float | None = None) -> bool:
+        """Run the state machine until the connection is established.
+
+        Returns True if the connection is established, False otherwise."""
+        return self.fsm.loop_until_state((ConnectedState,), timeout)

ohmqtt/connection/address.py

@@ -0,0 +1,91 @@
+from dataclasses import dataclass
+import socket
+from typing import Final, Mapping
+from urllib.parse import urlparse, ParseResult
+
+from ..platform import AF_UNIX, HAS_AF_UNIX
+
+
+DEFAULT_PORTS: Final[Mapping[str, int]] = {
+    "mqtt": 1883,
+    "mqtts": 8883,
+    "unix": 0,
+}
+
+
+def is_ipv6(hostname: str) -> bool:
+    """Check if the hostname is an IPv6 address."""
+    try:
+        socket.inet_pton(socket.AF_INET6, hostname)
+        return True
+    except socket.error:
+        return False
+
+
+def _get_family(parsed: ParseResult) -> socket.AddressFamily:
+    """Get the address family based on the parsed URL scheme."""
+    if HAS_AF_UNIX and parsed.scheme == "unix":
+        return AF_UNIX
+    if parsed.scheme in ("mqtt", "mqtts"):
+        if not parsed.hostname:
+            raise ValueError("Hostname is required for mqtt and mqtts schemes")
+        if is_ipv6(parsed.hostname):
+            return socket.AF_INET6
+        return socket.AF_INET
+    raise ValueError(f"Unsupported scheme: {parsed.scheme}")
+
+
+@dataclass(slots=True, init=False, frozen=True, repr=False)
+class Address:
+    scheme: str
+    family: socket.AddressFamily
+    host: str
+    port: int
+    username: str | None
+    password: str | None
+
+    def __init__(self, address: str = "") -> None:
+        """Parse the address string into family, host, port, username, and password."""
+        # Special case: empty address is allowed, but slots will be empty.
+        if not address:
+            return
+        if address.startswith("unix:"):
+            if not HAS_AF_UNIX:
+                raise ValueError("Unix socket support is not available on this platform")
+        elif "//" not in address:
+            # urlparse may choke on some network address we wish to support, unless we guarantee a //.
+            address = "//" + address
+        parsed = urlparse(address, scheme="mqtt")
+        object.__setattr__(self, "scheme", parsed.scheme)
+        object.__setattr__(self, "family", _get_family(parsed))
+        object.__setattr__(self, "host", parsed.hostname or parsed.path)
+        if not self.host:
+            raise ValueError("No path in address")
+        if HAS_AF_UNIX and self.family == AF_UNIX and self.host == "/":
+            raise ValueError("'/' is not a valid Unix socket path")
+        object.__setattr__(self, "port", parsed.port if parsed.port is not None else DEFAULT_PORTS[parsed.scheme])
+        object.__setattr__(self, "username", parsed.username)
+        object.__setattr__(self, "password", parsed.password)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the address."""
+        if not hasattr(self, "scheme"):
+            # Handle the empty case.
+            return "Address()"
+        userpw = ""
+        if self.username:
+            userpw = self.username
+        if self.password:
+            userpw += ":<hidden>"
+        if userpw:
+            host = f"{userpw}@{self.host}"
+        else:
+            host = self.host
+        if self.scheme != "unix":
+            host = f"{host}:{self.port}"
+        return f"Address({self.scheme}://{host})"
+
+    @property
+    def use_tls(self) -> bool:
+        """Check if the address uses TLS."""
+        return getattr(self, "scheme", None) == "mqtts"

ohmqtt/connection/decoder.py

@@ -0,0 +1,132 @@
+from dataclasses import dataclass, field
+import socket
+import ssl
+from typing import NamedTuple, Final
+
+from ..error import MQTTError
+from ..logger import get_logger
+from ..mqtt_spec import MQTTReasonCode
+from ..packet import decode_packet_from_parts, MQTTPacket
+
+logger: Final = get_logger("connection.decoder")
+
+
+class VarintDecodeResult(NamedTuple):
+    """Result of decoding a variable length integer, in part or whole.
+
+    This state can be used to resume decoding if the socket doesn't have enough data."""
+    value: int
+    multiplier: int
+    complete: bool
+
+
+InitVarintDecodeState: Final = VarintDecodeResult(0, 1, False)
+
+
+class ClosedSocketError(Exception):
+    """Exception raised when the socket is closed."""
+
+
+class WantReadError(Exception):
+    """Indicates that the socket is not ready for reading."""
+
+
+@dataclass(slots=True)
+class IncrementalDecoder:
+    """Incremental decoder for MQTT messages coming in from a socket.
+
+    Attributes:
+        head: The first byte of the packet.
+        length: Variable integer decoding state of the packet length.
+        data: The remaining data of the packet.
+    """
+    head: int = field(default=-1, init=False)
+    length: VarintDecodeResult = field(default=InitVarintDecodeState, init=False)
+    data: bytearray = field(init=False, default_factory=bytearray)
+
+    def reset(self) -> None:
+        """Reset the decoder state."""
+        self.head = -1
+        self.length = InitVarintDecodeState
+        self.data.clear()
+
+    def _recv_one_byte(self, sock: socket.socket | ssl.SSLSocket) -> int:
+        """Receive one byte from the socket.
+
+        Raises WantReadError if the socket is not ready for reading.
+
+        Raises ClosedSocketError if the socket is closed."""
+        try:
+            data = sock.recv(1)
+        except (BlockingIOError, ssl.SSLWantReadError, ssl.SSLWantWriteError):
+            raise WantReadError("Socket not ready for reading")
+        if not data:
+            raise ClosedSocketError("Socket closed")
+        return data[0]
+
+    def _extract_head(self, sock: socket.socket | ssl.SSLSocket) -> None:
+        """Extract the head byte of the next packet from the socket, if needed."""
+        if self.head != -1:
+            return
+        self.head = self._recv_one_byte(sock)
+
+    def _extract_length(self, sock: socket.socket | ssl.SSLSocket) -> None:
+        """Incrementally decode a variable length integer from a socket, if needed.
+
+        Raises WantReadError if the socket is not ready for reading."""
+        # See ohmqtt.serialization.decode_varint for a cleaner implementation.
+        assert self.head != -1  # We shouldn't be here unless we have a head byte.
+        if self.length.complete:
+            return
+        result = self.length.value
+        mult = self.length.multiplier
+        try:
+            while mult < 0x200000:  # This magic is the mult value after pulling 4 bytes.
+                byte = self._recv_one_byte(sock)
+                result += byte % 0x80 * mult
+                if byte < 0x80:
+                    # We have the complete varint.
+                    self.length = VarintDecodeResult(result, mult, True)
+                    return
+                mult *= 0x80
+            raise MQTTError("Varint overflow", MQTTReasonCode.MalformedPacket)
+        except WantReadError:
+            # Not done yet, the socket is neither closed nor ready for reading.
+            # Save the partial state and return.
+            self.length = VarintDecodeResult(result, mult, False)
+            raise
+
+    def _extract_data(self, sock: socket.socket | ssl.SSLSocket) -> None:
+        """Extract all data after the packet length from the socket, if needed."""
+        assert self.length.complete  # We shouldn't be here unless we have a complete length.
+        while len(self.data) < self.length.value:
+            data = sock.recv(self.length.value - len(self.data))
+            if not data:
+                raise ClosedSocketError("Socket closed")
+            self.data.extend(data)
+
+    def decode(self, sock: socket.socket | ssl.SSLSocket) -> MQTTPacket | None:
+        """Decode a single packet straight from the socket.
+
+        Returns None if the socket doesn't have enough data for us to decode a packet.
+
+        Raises ClosedSocketError if the socket is closed."""
+        try:
+            self._extract_head(sock)
+            self._extract_length(sock)
+            self._extract_data(sock)
+        except (BlockingIOError, ssl.SSLWantReadError, ssl.SSLWantWriteError, WantReadError):
+            # If the socket is open but doesn't have enough data for us, we need to wait for more.
+            return None
+        except OSError as exc:
+            raise ClosedSocketError("Socket closed") from exc
+
+        # We have a complete packet, decode it and clear the read buffer.
+        packet_head = self.head
+        packet_data = memoryview(self.data)
+        packet_data.toreadonly()
+        try:
+            return decode_packet_from_parts(packet_head, packet_data)
+        finally:
+            packet_data.release()
+            self.reset()