Skip to main content

max / alloy

16.3 KB · 450 lines History Blame Raw
1 #!/usr/bin/env python3
2 """Read a full-screen TUI over a pty, and type at it.
3
4 `ssh_pty.py` captures the bytes a TUI draws. This reads the SCREEN, and the
5 difference is the whole reason a scripted wizard walk is possible at all.
6
7 ratatui redraws only the cells that changed. "step 1 of 6" becoming
8 "step 2 of 6" puts a single `2` on the wire behind a cursor move, so a script
9 watching the byte stream for the new title sees it exactly once — on the first
10 frame, where everything is written — and then silently never again. Every
11 subsequent screen is a diff against a screen the script never reconstructed.
12 Grepping the stream is not a weaker version of reading the screen; it is a
13 different thing that happens to work on frame one.
14
15 So this keeps a grid and applies the escapes to it. `Screen.text()` is what is
16 on the screen right now, which is a question with an answer, and
17 `Session.wait_for` is "wait until the screen says this".
18
19 The escape subset is crossterm's output and nothing more: absolute and relative
20 cursor moves, the two erases, and printable text. SGR, the alternate-screen
21 toggles, OSC and DCS are parsed only far enough to be skipped, because they
22 carry nothing a script reads. A sequence this does not know is skipped rather
23 than printed, which is the failure that would otherwise fill the grid with
24 `[38;5;` and look like the application misbehaving.
25
26 Not a terminal emulator. No scrollback, no autowrap, no character sets, no
27 line-drawing translation. A full-screen TUI positions every run it draws, so
28 none of that is reachable from here; if something starts depending on it, the
29 honest fix is to say so rather than to grow this file toward xterm.
30
31 Use:
32
33 from tui import Session
34 s = Session(["ssh", ...])
35 s.wait_for(r"step 1 of 6", 60)
36 s.send("alloytest\r")
37 print(s.screen.text())
38 """
39
40 import codecs
41 import fcntl
42 import os
43 import pty
44 import re
45 import select
46 import signal
47 import socket
48 import struct
49 import sys
50 import termios
51 import time
52
53 ROWS, COLS = 40, 120
54
55
56 class Timeout(RuntimeError):
57 """What the screen said when the wait ran out, so a failure is readable."""
58
59 def __init__(self, pattern, seconds, screen):
60 super().__init__(
61 "waited %gs for %r; the screen said:\n\n%s\n" % (seconds, pattern, screen)
62 )
63 self.screen = screen
64
65
66 class Screen:
67 """A character grid, and the escapes needed to keep it current."""
68
69 def __init__(self, rows=ROWS, cols=COLS):
70 self.rows, self.cols = rows, cols
71 self.grid = [[" "] * cols for _ in range(rows)]
72 self.row = self.col = 0
73 # Incremental, because a read can split a multi-byte character and a
74 # per-chunk decode would put a replacement character in the grid.
75 self._decode = codecs.getincrementaldecoder("utf-8")("replace")
76 # Holds a partial escape across reads for the same reason.
77 self._buf = ""
78
79 def text(self):
80 """The screen as lines, right-trimmed. Rows are padded, and trailing
81 blanks on every line would make every pattern need to allow for them."""
82 return "\n".join("".join(row).rstrip() for row in self.grid)
83
84 def feed(self, data):
85 self._buf += self._decode.decode(data)
86 s, i, n = self._buf, 0, len(self._buf)
87 while i < n:
88 ch = s[i]
89 if ch == "\x1b":
90 nxt = self._escape(s, i)
91 if nxt is None:
92 break # incomplete; the rest of it is in the next read
93 i = nxt
94 continue
95 i += 1
96 if ch == "\r":
97 self.col = 0
98 elif ch == "\n":
99 self._newline()
100 elif ch == "\b":
101 self.col = max(0, self.col - 1)
102 elif ch == "\t":
103 self.col = min(self.cols - 1, (self.col // 8 + 1) * 8)
104 elif ch >= " ":
105 self._put(ch)
106 # Everything else is a control character a TUI does not use.
107 self._buf = s[i:]
108
109 # ---- the grid ----
110
111 def _put(self, ch):
112 self.grid[self.row][self.col] = ch
113 # Clamped rather than wrapped. A TUI positions every run it draws, so
114 # the right margin is never reached by accident, and wrapping would
115 # scroll the screen out from under a script on a full-width line.
116 self.col = min(self.col + 1, self.cols - 1)
117
118 def _newline(self):
119 if self.row + 1 < self.rows:
120 self.row += 1
121
122 def _clamp(self, row, col):
123 self.row = max(0, min(row, self.rows - 1))
124 self.col = max(0, min(col, self.cols - 1))
125
126 def _blank(self, row, start, end):
127 for c in range(max(0, start), min(end, self.cols)):
128 self.grid[row][c] = " "
129
130 # ---- escapes ----
131
132 def _escape(self, s, i):
133 """Index just past the sequence starting at `i`, or None if it is not
134 all here yet."""
135 if i + 1 >= len(s):
136 return None
137 kind = s[i + 1]
138 if kind == "[":
139 j = i + 2
140 while j < len(s) and "\x30" <= s[j] <= "\x3f":
141 j += 1
142 while j < len(s) and "\x20" <= s[j] <= "\x2f":
143 j += 1
144 if j >= len(s):
145 return None
146 self._csi(s[i + 2 : j], s[j])
147 return j + 1
148 if kind in "]P^_":
149 # OSC, DCS, PM, APC: a string terminated by BEL or by ST.
150 j = i + 2
151 while j < len(s):
152 if s[j] == "\x07":
153 return j + 1
154 if s[j] == "\x1b":
155 if j + 1 >= len(s):
156 return None
157 if s[j + 1] == "\\":
158 return j + 2
159 j += 1
160 return None
161 if "\x20" <= kind <= "\x2f":
162 # A two-character intermediate, such as a character-set select.
163 return None if i + 2 >= len(s) else i + 3
164 return i + 2
165
166 def _csi(self, params, final):
167 # `?` and friends mark private modes: cursor visibility, the alternate
168 # screen, bracketed paste. None of them move the cursor or erase, so
169 # skipping them is not an approximation.
170 if params[:1] in ("?", "<", ">", "="):
171 return
172 nums = [int(p) if p.isdigit() else 0 for p in params.split(";")] if params else []
173
174 def arg(k, default=1):
175 # Zero means "the default" in every one of these, per ECMA-48.
176 return nums[k] if k < len(nums) and nums[k] else default
177
178 if final in "Hf":
179 self._clamp(arg(0) - 1, arg(1) - 1)
180 elif final == "A":
181 self._clamp(self.row - arg(0), self.col)
182 elif final == "B":
183 self._clamp(self.row + arg(0), self.col)
184 elif final == "C":
185 self._clamp(self.row, self.col + arg(0))
186 elif final == "D":
187 self._clamp(self.row, self.col - arg(0))
188 elif final == "G":
189 self._clamp(self.row, arg(0) - 1)
190 elif final == "d":
191 self._clamp(arg(0) - 1, self.col)
192 elif final == "J":
193 # Mode 0 is the default here, so `arg` is wrong for J and K.
194 mode = nums[0] if nums else 0
195 if mode == 0:
196 self._blank(self.row, self.col, self.cols)
197 for r in range(self.row + 1, self.rows):
198 self._blank(r, 0, self.cols)
199 elif mode == 1:
200 for r in range(0, self.row):
201 self._blank(r, 0, self.cols)
202 self._blank(self.row, 0, self.col + 1)
203 else:
204 for r in range(self.rows):
205 self._blank(r, 0, self.cols)
206 elif final == "K":
207 mode = nums[0] if nums else 0
208 if mode == 0:
209 self._blank(self.row, self.col, self.cols)
210 elif mode == 1:
211 self._blank(self.row, 0, self.col + 1)
212 else:
213 self._blank(self.row, 0, self.cols)
214 # SGR and everything else changes how the screen looks, not what it says.
215
216
217 class Session:
218 """A child under a pty of a known size, and the screen it is drawing."""
219
220 def __init__(self, argv, rows=ROWS, cols=COLS, log_path=None):
221 self.screen = Screen(rows, cols)
222 self.argv = argv
223 self.log = open(log_path, "wb") if log_path else None
224 self.pid, self.fd = pty.fork()
225 if self.pid == 0:
226 os.execvp(argv[0], argv)
227 os._exit(127)
228 fcntl.ioctl(self.fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
229 self.eof = False
230
231 def pump(self, seconds=0.2):
232 """Read for up to `seconds`, applying whatever arrives to the screen."""
233 end = time.time() + seconds
234 while time.time() < end:
235 r, _, _ = select.select([self.fd], [], [], max(0.0, end - time.time()))
236 if not r:
237 continue
238 try:
239 chunk = os.read(self.fd, 65536)
240 except OSError:
241 self.eof = True
242 return
243 if not chunk:
244 self.eof = True
245 return
246 self.screen.feed(chunk)
247 if self.log:
248 self.log.write(chunk)
249 self.log.flush()
250
251 def wait_for(self, pattern, seconds, poll=0.2):
252 """Block until the screen matches, and return the match.
253
254 A regex over the whole screen rather than over the stream: see the
255 module docstring for why the stream cannot answer this."""
256 end = time.time() + seconds
257 rx = re.compile(pattern)
258 while True:
259 hit = rx.search(self.screen.text())
260 if hit:
261 return hit
262 if time.time() >= end:
263 raise Timeout(pattern, seconds, self.screen.text())
264 if self.eof:
265 # One more look: the last frame before the child left may
266 # carry the thing being waited for.
267 if rx.search(self.screen.text()):
268 return rx.search(self.screen.text())
269 raise Timeout(pattern, seconds, self.screen.text() + "\n\n(the session ended)")
270 self.pump(poll)
271
272 def send(self, text, settle=0.4):
273 """Type, then let the frame that answers arrive before anyone reads."""
274 os.write(self.fd, text.encode())
275 self.pump(settle)
276
277 def close(self):
278 try:
279 os.kill(self.pid, signal.SIGKILL)
280 except ProcessLookupError:
281 pass
282 try:
283 os.waitpid(self.pid, 0)
284 except ChildProcessError:
285 pass
286 if self.log:
287 self.log.close()
288
289
290 class SocketSession(Session):
291 """The same screen, over a unix socket instead of a pty.
292
293 The guest's serial console is a chardev socket, and on the installer medium
294 GRUB draws its menu there and `alloy-debug-shell@ttyS0` puts a root bash
295 there. Both are things a scenario has to read and answer, and neither is
296 reachable through a pty.
297
298 Subclasses Session for its `wait_for` and `screen`, and replaces only the
299 transport: there is no child process here, so nothing to fork, size or
300 reap. The terminal size is the guest's `stty`, not an ioctl from this end.
301 """
302
303 def __init__(self, path, rows=ROWS, cols=COLS, log_path=None, connect_timeout=60.0):
304 self.screen = Screen(rows, cols)
305 self.rows, self.cols = rows, cols
306 self.log = open(log_path, "wb") if log_path else None
307 self.eof = False
308 end = time.time() + connect_timeout
309 while True:
310 try:
311 self.sock = socket.socket(socket.AF_UNIX)
312 self.sock.connect(path)
313 break
314 except (FileNotFoundError, ConnectionRefusedError):
315 if time.time() >= end:
316 raise
317 time.sleep(0.2)
318 self.fd = self.sock.fileno()
319
320 def send(self, text, settle=0.4):
321 self.sock.sendall(text.encode())
322 self.pump(settle)
323
324 def close(self):
325 try:
326 self.sock.close()
327 except OSError:
328 pass
329 if self.log:
330 self.log.close()
331
332
333 def self_test():
334 """Check the parser, because a wrong parse still produces a screen.
335
336 Same reason check-installed.sh and check-rust-stage.sh carry one: the part
337 that fails invisibly is the part that turns input into a verdict, and a
338 verdict is what this file produces. Every case below is a shape crossterm
339 actually emits.
340 """
341 fails = []
342
343 def check(label, got, want):
344 if got != want:
345 fails.append("%s\n got: %r\n want: %r" % (label, got, want))
346
347 # The case this file exists for, and the one a stream search cannot answer:
348 # a redraw that rewrites one cell. Search the stream for "step 2 of 6" and
349 # it is not there; read the screen and it is.
350 s = Screen(3, 20)
351 s.feed(b"\x1b[2J\x1b[1;1Hstep 1 of 6\x1b[1;6H2")
352 check("a one-cell redraw", s.text().splitlines()[0], "step 2 of 6")
353
354 # SGR carries how the screen looks, not what it says, and must leave
355 # nothing behind. This is the failure that fills a grid with `[38;5;`.
356 s = Screen(1, 30)
357 s.feed(b"\x1b[1;38;5;203mred\x1b[0m ok")
358 check("SGR is skipped", s.text(), "red ok")
359
360 # Erase to end of line, which is how a footer's hints are replaced when a
361 # step changes the shorter set for the longer one.
362 s = Screen(1, 20)
363 s.feed(b"\x1b[1;1Hesc back\x1b[1;1Hr reboot\x1b[K")
364 check("erase to end of line", s.text(), "r reboot")
365
366 # And erase-all, which is the first thing a full repaint does.
367 s = Screen(2, 10)
368 s.feed(b"\x1b[1;1Hgone\x1b[2;1Halso\x1b[2J\x1b[1;1Hnew")
369 check("erase display", s.text(), "new\n")
370
371 # An escape split across two reads. A parser that consumed the partial
372 # sequence would print `[2;` into the grid and lose the move.
373 s = Screen(2, 10)
374 s.feed(b"a\x1b[2")
375 s.feed(b";1Hb")
376 check("a split escape", s.text(), "a\nb")
377
378 # A UTF-8 character split across two reads. The list marker is three bytes,
379 # and a per-chunk decode would put a replacement character in the gutter
380 # and make every row read as unselected.
381 s = Screen(1, 10)
382 s.feed(b"\xe2\x96")
383 s.feed(b"\xb6 vda")
384 check("a split character", s.text(), "\u25b6 vda")
385
386 # OSC, which crossterm emits to set the window title. Terminated by BEL
387 # here and by ST in the next case; a parser that knew only one would print
388 # the other's payload.
389 s = Screen(1, 20)
390 s.feed(b"\x1b]0;alloy\x07ok")
391 check("OSC ended by BEL", s.text(), "ok")
392 s = Screen(1, 20)
393 s.feed(b"\x1b]0;alloy\x1b\\ok")
394 check("OSC ended by ST", s.text(), "ok")
395
396 # Relative moves and the column/row absolutes.
397 s = Screen(2, 10)
398 s.feed(b"\x1b[1;1Habc\x1b[1D\x1b[1BX\x1b[1;1H\x1b[5GY")
399 # `abc` leaves the cursor at column 3; one left is 2, one down writes the X
400 # there, and `5G` is column 4 counting from one.
401 check("relative moves", s.text(), "abc Y\n X")
402
403 # A pane the way one really arrives: a border, a gutter marker, and rows.
404 # The assertion is the question install_drive.py asks of the disk step.
405 s = Screen(4, 40)
406 s.feed(
407 b"\x1b[2J"
408 b"\x1b[1;1H\xe2\x94\x8c\x1b[1;2H select a disk"
409 b"\x1b[2;1H\xe2\x94\x82\x1b[2;3H vda 40.0 GiB"
410 b"\x1b[3;1H\xe2\x94\x82\x1b[3;3H sdb 931.5 GiB"
411 # ...and then the selection moves down, which rewrites two gutters.
412 b"\x1b[3;3H\xe2\x96\xb6"
413 )
414 row = re.search("\u25b6" + r"\s+(\S+)", s.text())
415 check("the selected row is readable", row and row.group(1), "sdb")
416
417 if fails:
418 for f in fails:
419 print("self-test: " + f, file=sys.stderr)
420 print("self-test: %d failed" % len(fails), file=sys.stderr)
421 return 1
422 print("self-test: the parser is right about all ten shapes")
423 return 0
424
425
426 def main():
427 """Ad-hoc use: run a command, wait for a pattern, print the screen."""
428 if len(sys.argv) > 1 and sys.argv[1] == "--self-test":
429 raise SystemExit(self_test())
430 if len(sys.argv) < 3:
431 raise SystemExit("usage: tui.py <seconds> <pattern> <command> [args...]\n"
432 " an empty pattern just reads for <seconds>")
433 seconds, pattern, argv = float(sys.argv[1]), sys.argv[2], sys.argv[3:]
434 s = Session(argv)
435 try:
436 if pattern:
437 s.wait_for(pattern, seconds)
438 else:
439 s.pump(seconds)
440 print(s.screen.text())
441 except Timeout as e:
442 print(e, file=sys.stderr)
443 raise SystemExit(1)
444 finally:
445 s.close()
446
447
448 if __name__ == "__main__":
449 main()
450