tmux-amq/src/tamq/ptytap.py

from __future__ import annotations

import fcntl
import os
import pty
import re
import select
import signal
import sys
import termios
import tty
from collections.abc import Callable, Sequence

from .broker import InputBroker


ENHANCED_ENTER = re.compile(
    rb"\x1b(?:\[13(?:;[0-9:]+)*u|\[27;[0-9:]+;13~|OM)"
)
TERMINAL_INPUT_SEQUENCE = re.compile(
    rb"\x1b(?:"
    rb"\[[0-?]*[ -/]*[@-~]"
    rb"|O."
    rb"|P.*?(?:\x1b\\)"
    rb"|\].*?(?:\x07|\x1b\\)"
    rb")",
    re.DOTALL,
)


def observed_line(raw: bytes) -> str:
    """Return typed text without terminal replies or local editing controls."""
    decoded = TERMINAL_INPUT_SEQUENCE.sub(b"", raw).decode(
        "utf-8", errors="replace"
    )
    text: list[str] = []
    for character in decoded:
        if character in ("\b", "\x7f"):
            if text:
                text.pop()
        elif character == "\t" or character >= " ":
            text.append(character)
    return "".join(text)


class PtyTap:
    """Full-duplex PTY proxy. Input is observed, never rewritten."""

    def __init__(self, command: Sequence[str], broker: InputBroker, *, on_line: Callable[[str], None] | None = None):
        self.command = list(command)
        self.broker = broker
        self.on_line = on_line

    def _observe_input(self, buffer: bytearray, data: bytes) -> None:
        buffer.extend(data)
        while True:
            boundaries = [
                (index, index + 1)
                for marker in (b"\r", b"\n")
                if (index := buffer.find(marker)) >= 0
            ]
            if match := ENHANCED_ENTER.search(buffer):
                boundaries.append(match.span())
            if not boundaries:
                return
            start, end = min(boundaries)
            raw = bytes(buffer[:start])
            del buffer[:end]
            line = observed_line(raw)
            routed = self.broker.inspect_line(line)
            if routed and self.on_line:
                self.on_line(line)

    def run(self) -> int:
        stdin_fd = sys.stdin.fileno()
        stdout_fd = sys.stdout.fileno()
        master, slave = pty.openpty()
        copy_winsize(stdin_fd, slave)
        pid = os.fork()
        if pid == 0:
            try:
                os.close(master)
                os.setsid()
                fcntl.ioctl(slave, termios.TIOCSCTTY, 0)
                for fd in (0, 1, 2):
                    os.dup2(slave, fd)
                if slave > 2:
                    os.close(slave)
                os.execvp(self.command[0], self.command)
            except BaseException as exc:
                os.write(2, f"tamq tap: cannot start {self.command[0]}: {exc}\n".encode())
                os._exit(127)
        os.close(slave)
        input_buffer = bytearray()
        saved_terminal = termios.tcgetattr(stdin_fd) if os.isatty(stdin_fd) else None
        saved_handlers: dict[int, signal.Handlers] = {}
        input_open = True

        def resize(_signum=None, _frame=None) -> None:
            copy_winsize(stdin_fd, master)

        def forward(signum, _frame) -> None:
            try:
                os.killpg(pid, signum)
            except ProcessLookupError:
                pass

        status = 0
        try:
            if saved_terminal is not None:
                tty.setraw(stdin_fd)
            saved_handlers[signal.SIGWINCH] = signal.signal(signal.SIGWINCH, resize)
            for signum in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP):
                saved_handlers[signum] = signal.signal(signum, forward)
            resize()
            while True:
                sources = [master]
                if input_open:
                    sources.append(stdin_fd)
                try:
                    readable, _, _ = select.select(sources, [], [])
                except InterruptedError:
                    continue
                if master in readable:
                    try:
                        data = os.read(master, 65536)
                    except OSError:
                        break
                    if not data:
                        break
                    write_all(stdout_fd, data)
                if input_open and stdin_fd in readable:
                    data = os.read(stdin_fd, 65536)
                    if not data:
                        input_open = False
                        try:
                            os.killpg(pid, signal.SIGHUP)
                        except ProcessLookupError:
                            pass
                        continue
                    write_all(master, data)
                    self._observe_input(input_buffer, data)
        except BaseException:
            try:
                os.killpg(pid, signal.SIGHUP)
            except ProcessLookupError:
                pass
            raise
        finally:
            for signum, handler in saved_handlers.items():
                signal.signal(signum, handler)
            if saved_terminal is not None:
                termios.tcsetattr(stdin_fd, termios.TCSAFLUSH, saved_terminal)
            os.close(master)
            try:
                _, status = os.waitpid(pid, 0)
            except ChildProcessError:
                status = 0
        return os.waitstatus_to_exitcode(status)


def copy_winsize(source_fd: int, target_fd: int) -> bool:
    """Copy terminal rows/columns and pixel dimensions when a source tty exists."""
    try:
        size = fcntl.ioctl(source_fd, termios.TIOCGWINSZ, b"\0" * 8)
        fcntl.ioctl(target_fd, termios.TIOCSWINSZ, size)
        return True
    except OSError:
        return False


def write_all(fd: int, data: bytes) -> None:
    view = memoryview(data)
    while view:
        written = os.write(fd, view)
        if written == 0:
            raise OSError("terminal write returned zero bytes")
        view = view[written:]