cli: make burn terminal bidirectional

Author: eorry-bit

src/defib/cli/app.py

@@ -400,24 +400,16 @@ def on_sigint(*_: object) -> None:
             if output == "human":
                 console.print("[dim]--- Terminal mode (Ctrl-C to exit) ---[/dim]")
 
-            stop = False
-
-            def on_sigint(*_: object) -> None:
-                nonlocal stop
-                stop = True
-
-            signal.signal(signal.SIGINT, on_sigint)
-
+            from defib.cli.terminal import run_raw_terminal
             try:
-                while not stop:
-                    try:
-                        data = await transport.read(256, timeout=0.1)
-                        _sys.stdout.buffer.write(data)
-                        _sys.stdout.buffer.flush()
-                    except Exception:
-                        pass
+                await run_raw_terminal(
+                    transport,
+                    _sys.stdin.buffer,
+                    _sys.stdout.buffer,
+                )
+            except KeyboardInterrupt:
+                pass
             finally:
-                signal.signal(signal.SIGINT, signal.SIG_DFL)
                 if output == "human":
                     console.print("\n[dim]--- Terminal closed ---[/dim]")
 

src/defib/cli/terminal.py

@@ -0,0 +1,151 @@
+"""Bidirectional raw terminal bridge for the burn command."""
+
+from __future__ import annotations
+
+import asyncio
+import importlib
+import os
+import signal
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import BinaryIO
+
+from defib.transport.base import Transport, TransportTimeout
+
+
+@contextmanager
+def _raw_terminal(fd: int) -> Iterator[None]:
+    """Disable canonical input, translations, and local echo."""
+    if os.name != "posix" or not os.isatty(fd):
+        yield
+        return
+
+    import termios
+    import tty
+
+    previous = termios.tcgetattr(fd)
+    tty.setraw(fd)
+    try:
+        yield
+    finally:
+        termios.tcsetattr(fd, termios.TCSADRAIN, previous)
+
+
+async def _pump_posix_stdin(
+    transport: Transport,
+    fd: int,
+    stop: asyncio.Event,
+) -> None:
+    loop = asyncio.get_running_loop()
+    queue: asyncio.Queue[bytes] = asyncio.Queue()
+
+    def on_readable() -> None:
+        try:
+            data = os.read(fd, 1024)
+        except OSError:
+            data = b""
+        if not data:
+            loop.remove_reader(fd)
+        queue.put_nowait(data)
+
+    loop.add_reader(fd, on_readable)
+    try:
+        while not stop.is_set():
+            try:
+                data = await asyncio.wait_for(queue.get(), timeout=0.1)
+            except TimeoutError:
+                continue
+            if not data:
+                stop.set()
+                return
+            if b"\x03" in data:
+                before_sigint, _, _ = data.partition(b"\x03")
+                if before_sigint:
+                    await transport.write(before_sigint)
+                stop.set()
+                return
+            await transport.write(data)
+    finally:
+        loop.remove_reader(fd)
+
+
+async def _pump_windows_stdin(
+    transport: Transport,
+    stop: asyncio.Event,
+) -> None:
+    msvcrt = importlib.import_module("msvcrt")
+
+    while not stop.is_set():
+        if not msvcrt.kbhit():
+            await asyncio.sleep(0.01)
+            continue
+        char = msvcrt.getwch()
+        if char == "\x03":
+            stop.set()
+            return
+        if char in ("\x00", "\xe0"):
+            msvcrt.getwch()
+            continue
+        await transport.write(char.encode())
+
+
+async def _pump_transport(
+    transport: Transport,
+    stdout: BinaryIO,
+    stop: asyncio.Event,
+) -> None:
+    while not stop.is_set():
+        try:
+            data = await transport.read(256, timeout=0.1)
+        except TransportTimeout:
+            continue
+        if not data:
+            stop.set()
+            return
+        stdout.write(data)
+        stdout.flush()
+
+
+async def _bridge_terminal(
+    transport: Transport,
+    stdin: BinaryIO,
+    stdout: BinaryIO,
+    stop: asyncio.Event,
+) -> None:
+    if os.name == "nt":
+        stdin_task = asyncio.create_task(_pump_windows_stdin(transport, stop))
+    else:
+        stdin_task = asyncio.create_task(_pump_posix_stdin(transport, stdin.fileno(), stop))
+    output_task = asyncio.create_task(_pump_transport(transport, stdout, stop))
+    stop_task = asyncio.create_task(stop.wait())
+    tasks = {stdin_task, output_task, stop_task}
+
+    done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+    stop.set()
+    for task in pending:
+        task.cancel()
+    await asyncio.gather(*pending, return_exceptions=True)
+
+    for task in done:
+        if task is not stop_task:
+            task.result()
+
+
+async def run_raw_terminal(
+    transport: Transport,
+    stdin: BinaryIO,
+    stdout: BinaryIO,
+) -> None:
+    """Bridge stdin and transport until EOF, disconnect, or Ctrl-C."""
+    stop = asyncio.Event()
+    previous_handler = signal.getsignal(signal.SIGINT)
+
+    def on_sigint(*_: object) -> None:
+        stop.set()
+
+    signal.signal(signal.SIGINT, on_sigint)
+    try:
+        with _raw_terminal(stdin.fileno()):
+            await _bridge_terminal(transport, stdin, stdout, stop)
+    finally:
+        signal.signal(signal.SIGINT, previous_handler)

tests/test_terminal.py

@@ -0,0 +1,80 @@
+"""Tests for the interactive raw terminal bridge."""
+
+import asyncio
+import io
+import os
+import signal
+
+import pytest
+
+from defib.cli.terminal import run_raw_terminal
+from defib.transport.mock import MockTransport
+
+
+pytestmark = pytest.mark.skipif(os.name != "posix", reason="PTY tests require POSIX")
+
+
+async def _wait_until(predicate, timeout: float = 1.0) -> None:
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + timeout
+    while loop.time() < deadline:
+        if predicate():
+            return
+        await asyncio.sleep(0.01)
+    raise AssertionError("condition not met before timeout")
+
+
+@pytest.mark.asyncio
+async def test_pty_input_is_forwarded_and_transport_output_is_printed() -> None:
+    import pty
+    import termios
+
+    master_fd, slave_fd = pty.openpty()
+    stdin = os.fdopen(slave_fd, "rb", buffering=0)
+    stdout = io.BytesIO()
+    transport = MockTransport()
+    transport.enqueue_rx(b"hisilicon # ")
+
+    task = asyncio.create_task(run_raw_terminal(transport, stdin, stdout))
+    try:
+        await _wait_until(lambda: not termios.tcgetattr(stdin.fileno())[3] & termios.ECHO)
+        os.write(master_fd, b"help\r")
+        await _wait_until(lambda: b"help\r" in transport.all_tx_data)
+        await _wait_until(lambda: b"hisilicon # " in stdout.getvalue())
+        os.write(master_fd, b"\x03")
+        await asyncio.wait_for(task, timeout=1.0)
+    finally:
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+        os.close(master_fd)
+        stdin.close()
+
+
+@pytest.mark.asyncio
+async def test_sigint_stops_bridge_and_restores_pty() -> None:
+    import pty
+    import termios
+
+    master_fd, slave_fd = pty.openpty()
+    stdin = os.fdopen(slave_fd, "rb", buffering=0)
+    stdout = io.BytesIO()
+    transport = MockTransport()
+    original_attrs = termios.tcgetattr(stdin.fileno())
+    original_handler = signal.getsignal(signal.SIGINT)
+
+    task = asyncio.create_task(run_raw_terminal(transport, stdin, stdout))
+    try:
+        await _wait_until(lambda: not termios.tcgetattr(stdin.fileno())[3] & termios.ECHO)
+
+        signal.raise_signal(signal.SIGINT)
+        await asyncio.wait_for(task, timeout=1.0)
+
+        assert termios.tcgetattr(stdin.fileno()) == original_attrs
+        assert signal.getsignal(signal.SIGINT) == original_handler
+    finally:
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+        os.close(master_fd)
+        stdin.close()