herosdevices.core.bus.telnet.telnet - Heros Devices 0.11.1 documentation

Skip to content
Back to top

# Source code for herosdevices.core.bus.telnet.telnet

    """This module provides a class for managing telnet connections and a context manager for handling telnet sessions."""

    import time
    from collections.abc import Iterator
    from contextlib import contextmanager
    from typing import Literal, overload

    from heros.helper import log

    from .telnetlib import IAC, NOP, Telnet

    TELNET_DELAYS_DEFAULT: dict = {"write": 0.001, "read_echo": 0.001}

    class TelnetConnection:
        """A class to manage telnet communication connections.

        This class provides functionality to handle telnet connections including opening/closing connections, reading
        data, and writing data.

        Args:
            address: The address of the telnet socket, something like /dev/ttyUSB0.
            port: Port the telnet server is listening on
            read_line_termination: character that terminates a line when reading from the device.
            write_line_termination: character that terminates a line when wrtiting to the device.
            keep_alive: Flag indicating whether to keep the connection open between operations.
            delays: Dictionary containing delay times for in between telnet operations.
                Default telnet delays for telnet devices. Available keys are:

                    * "write": Time to wait after writing a command to the device.
                    * "read_echo": Time to wait before reading a response from the device.

                :py:data:`herosdevices.core.bus.telnet.TELNET_DELAYS_DEFAULT` sets the default delays.
        """

        def __init__(
            self,
            address: str,
            port: int = 23,
            timeout: float = 1.0,
            read_line_termination: bytes = b"\n",
            write_line_termination: bytes = b"",
            keep_alive: bool = True,
            delays: dict | None = None,
        ) -> None:
            self.address = address
            self.port = port
            self.timeout = timeout
            self.read_line_termination = read_line_termination
            self.write_line_termination = write_line_termination
            self.connection = Telnet()
            self.keep_alive = keep_alive
            self.delays = TELNET_DELAYS_DEFAULT | delays if delays else TELNET_DELAYS_DEFAULT

        def check_alive(self) -> bool:
            """Check if the telnet connection is alive."""
            try:
                if self.connection.sock:  # this way I've taken care of problem if the .close() was called
                    self.connection.sock.send(IAC + NOP)  # notice the use of send instead of sendall
                    return True
            except OSError:
                pass
            return False

        @contextmanager
        def operation(self) -> Iterator[None]:
            """Context manager for handling telnet connection operations.

            Ensures the telnet connection is open before performing operations and closes it afterward
            if :code:`self.keep_alive` is False.

            Yields:
                Yields control back to the caller for performing operations within the context.
            """
            if not self.check_alive():
                self.connection.open(self.address, self.port, self.timeout)
            try:
                yield
            finally:
                if not self.keep_alive:
                    self.connection.close()

        def wait(self, operation: str) -> None:
            """Introduce a (synchronous) delay based on the specified operation type.

            Args:
                operation: The operation type. For possible types see :code:`TELNET_DELAYS_DEFAULT`.
            """
            time.sleep(self.delays[operation])

        def read(self) -> str | None:
            """
            Read all available data from the telnet connection and decodes it into a string.

            Returns:
                The decoded data as string, or None if an error occurs.
            """
            with self.operation():
                try:
                    read = self.connection.read_very_eager()
                    return read.decode("ascii")
                except Exception as e:  # noqa: BLE001
                    log.error(
                        "Error reading from telnet connection at %s: %s",
                        self.address,
                        e,
                    )
                    return None

        def read_line(self) -> str | None:
            """Read a single line from the telnet connection.

            Returns:
                The decoded line as string, or None if an error occurs.
            """
            with self.operation():
                try:
                    read = self.connection.read_until(self.read_line_termination)
                    return read.decode("ascii")
                except Exception as e:  # noqa: BLE001
                    log.error(
                        "Error reading from telnet connection at %s: %s",
                        self.address,
                        e,
                    )
                    return None

        @overload
        def write(self, message: str, read_echo: Literal[True] = True, read_line: bool = True) -> str: ...
        @overload
        def write(self, message: str, read_echo: Literal[False], read_line: bool = True) -> None: ...
        @overload
        def write(self, message: str, read_echo: bool = False, read_line: bool = True) -> str | None: ...

        def write(self, message: str, read_echo: bool = False, read_line: bool = True) -> str | None:
            """Write a message to the telnet connection.

            The `self.write_line_termination` is automatically appended to the message if it is not already present.

            Args:
                message: The message to be written to the telnet connection.
                read_echo: If True, reads back the echo after writing. Defaults to False.
                read_line: If True, data is read until `self.read_line_termination` occurs in the data.
                    Otherwise all available data is read.

            Returns:
                If read_echo is True, returns the echo read from the connection as string; otherwise returns None.
            """
            with self.operation():
                self.connection.read_very_lazy()
                enc_msg = message.encode("ascii")
                if not enc_msg.endswith(self.write_line_termination):
                    enc_msg += self.write_line_termination
                self.connection.write(enc_msg)
                self.wait("write")
                if read_echo:
                    if read_line:
                        read = self.read_line()
                    else:
                        self.wait("read_echo")
                        read = self.read()
                    return read
                return None