ui_cli_manager.ui_cli_manager

  1"""ui_cli_manager.py - Simple Connection Manager for UI/TCP connections"""
  2from __future__ import annotations
  3from queue import Empty
  4from collections import deque
  5from dataclasses import dataclass
  6from typing import Callable, Any
  7import shlex
  8import socket
  9
 10from micronetcode.utils import logger, JournalMsgType, SOCKET_RECV_SIZE, SOCKET_RECV_MAXSIZE
 11from micronetcode.message import SourceType, Message
 12from micronetcode.connection_manager import ConnectionManager
 13from micronetcode.codec import Codec
 14
 15class ShlexASCIICodec(Codec):
 16    """Basic ASCII/shlex (stdlib) based codec. Matches the initial implementation with ascii/newline data."""
 17    def __init__(self, socket_recv_size: int | None = None, socket_recv_maxsize: int | None = None):
 18        super().__init__()
 19        self.socket_recv_maxsize = socket_recv_maxsize or SOCKET_RECV_MAXSIZE
 20        self.socket_recv_size = socket_recv_size or SOCKET_RECV_SIZE
 21
 22    def recv(self, sock: socket.socket) -> bytes | None:
 23        res = bytearray()
 24        while True:
 25            buf = sock.recv(self.socket_recv_size)
 26            if len(buf) == 0: # This means it closed connection
 27                if len(res) == 0:
 28                    return None
 29                return bytes(res)
 30            res.extend(buf)
 31            if res[-1] == ord("\n"):
 32                return bytes(res)
 33            if len(res) > self.socket_recv_maxsize:
 34                raise ValueError(f"Too much data set: {len(res)=} vs {self.socket_recv_maxsize=}")
 35
 36    def decode(self, data: bytes) -> list[Message]:
 37        lines = data.decode("ascii").strip().split("\n")
 38        return [Message(source_type=SourceType.SOCKET, data=line) for line in lines]
 39
 40    def encode(self, resp: Any) -> bytes:
 41        if not isinstance(resp, str):
 42            raise TypeError(f"Expected str, got {type(resp)}")
 43        return resp.encode("ascii")
 44
 45@dataclass
 46class CLICommand:
 47    """dataclass encapsulating the CLI command, its arguments and the source Message (socket, stdin, etc.)"""
 48    message: Message
 49    command: str
 50    args: tuple[Any, ...]
 51    respond: Callable[[str], None] # a callback to respond to this CLI message created by UICLIManager
 52
 53class UICLIManager:
 54    """
 55    UI-CLI manager. Wrapper on top of ConnectionManager + Channel a list commands
 56    Parameters:
 57    - host The host where the socket listens to (e.g. "0.0.0.0")
 58    - port The port where the socket listens to
 59    - cli_commands A dictionary of all valid cli commands and the expected argument count of each (no type validation)
 60    - journal_length An optional parameter that if >0 makes the underlying channel store up to these many messages
 61    - script_lines A list of scripting commands to be executed first before any TCP message
 62    """
 63
 64    def __init__(self, connection_manager: ConnectionManager, cli_commands: dict[str, int],
 65                 script_lines: list[str] | None = None):
 66        if "help" in cli_commands:
 67            raise ValueError("'help' cannot be part of cli_commands, it is auto-inserted")
 68        self.connection_manager = connection_manager
 69        self.cli_commands = cli_commands
 70        self.script_lines = deque(script_lines or [])
 71        # (command, reply) journal of every SCRIPT-source response. The reliable way to observe
 72        # scripted replies (they have no socket client to reply to) — used by the integration tests.
 73        self.script_responses: list[tuple[str, str]] = []
 74
 75    @classmethod
 76    def with_ascii(cls, host: str, port: int, cli_commands: dict[str, int], script_lines: list[str] | None = None,
 77                   max_connections: int | None = None, journal_length: int | None = None,
 78                   channel_queue_maxsize: int | None = None, socket_timeout_s: float | None = None) -> UICLIManager:
 79        """Convenience wrapper to instantiate UICLIManager with a TCP-based connection manager and ASCII codec"""
 80        connection_manager = ConnectionManager(host, port, codec=ShlexASCIICodec(), max_connections=max_connections,
 81                                               journal_length=journal_length, socket_timeout_s=socket_timeout_s,
 82                                               channel_queue_maxsize=channel_queue_maxsize)
 83        return UICLIManager(connection_manager, cli_commands, script_lines)
 84
 85    def start(self):
 86        """Starts the underlying connection manager thread"""
 87        self.connection_manager.start()
 88
 89    def get_cli_command(self) -> CLICommand | None:
 90        """
 91        Returns a command + its valid args if the CLI provided a valid command, None otherwise. No exceptions.
 92        Ran in main thread context.
 93        """
 94
 95        try:
 96            if len(self.script_lines) > 0: # get data from script lines first
 97                message = Message(source_type=SourceType.SCRIPT, data=self.script_lines.popleft())
 98            else:
 99                # Note: this must be nowait so the UI doesn't lag as it's called from the main loop
100                message: Message = self.connection_manager.get_one_message()
101
102            cli_command, err = self._handle_one_message(message)
103            if err is not None:
104                logger.error(err)
105                self._send_response(message=message, data=err)
106                return None
107            return cli_command
108
109        except Empty:
110            return None
111
112    # Private helper functions
113
114    def _handle_one_message(self, message: Message) -> tuple[CLICommand | None, str | None]:
115        """Given one message (source_type, payload, [source]), turn it into a CLICommand(msg, cmd, args) or err"""
116        if not isinstance(message.data, str):
117            raise TypeError(f"Message {message} doesn't have data as string: {type(message.data)}. This is an error. "
118                            f"The connection's Codec ({self.connection_manager.codec}) must make it string.")
119
120        try:
121            cmd_args = shlex.split(message.data)
122        except ValueError as e:
123            return None, f"Bad quoting: {e}"
124
125        if not cmd_args or len(cmd_args) == 0 or cmd_args[0] == "#": # empty or commented out command
126            return None, f"Comment or empty line: {message.data}"
127
128        command, args = cmd_args[0], cmd_args[1: ]
129
130        if command == "help":
131            return None, f"Supported commands: {list(self.cli_commands)}"
132
133        if command not in self.cli_commands:
134            return None, f"Unknown command '{command}'. Supported: {list(self.cli_commands)}"
135
136        if len(args) != self.cli_commands[command]:
137            return None, f"{cmd_args=}, Expected args: {self.cli_commands[command]}, got {len(args)}"
138
139        respond_cb = lambda data: self._send_response(message=message, data=data)
140        return CLICommand(message=message, command=command, args=args, respond=respond_cb), None
141
142    def _send_response(self, message: Message, data: str):
143        """Validate + newline + route by source. Shared by cli_cmd.respond() and the error path."""
144        if not isinstance(data, str):
145            raise TypeError(f"Expected 'str', got {type(data)}")
146        if len(data) == 0:
147            raise ValueError("Cannot respond with empty string")
148        if data[-1] != "\n":
149            data += "\n"
150        if message.client is not None:
151            message.client.channel.main2tcp.put(data, timeout=self.connection_manager.socket_timeout_s)
152        else: # for script/stdin etc: journal the response for tests (no socket client to reply to)
153            logger.debug(f"{JournalMsgType.RESP} - {message=} - {data=}")
154            self.script_responses.append((message.data.rstrip("\n"), data.rstrip("\n")))