Skip to main content

max / alloy

1.3 KB · 50 lines History Blame Raw
1 #!/usr/bin/env python3
2 """Run ssh under a real pty of a known size and capture what it draws.
3
4 ssh_pty.py <keyfile> <user> [seconds] [keys-to-send]
5
6 Keys are sent after 4s, then output is captured until the timeout.
7 """
8 import os, pty, sys, time, select, struct, fcntl, termios, signal
9
10 key, user = sys.argv[1], sys.argv[2]
11 seconds = float(sys.argv[3]) if len(sys.argv) > 3 else 12.0
12 to_send = sys.argv[4] if len(sys.argv) > 4 else ""
13
14 pid, fd = pty.fork()
15 if pid == 0:
16 os.execvp("ssh", [
17 "ssh", "-tt",
18 "-o", "StrictHostKeyChecking=no",
19 "-o", "UserKnownHostsFile=/dev/null",
20 "-o", "LogLevel=ERROR",
21 "-p", "2222", "-i", key, f"{user}@127.0.0.1",
22 ])
23
24 fcntl.ioctl(fd, termios.TIOCSWINSZ, struct.pack("HHHH", 40, 120, 0, 0))
25
26 out = b""
27 deadline = time.time() + seconds
28 sent = False
29 while time.time() < deadline:
30 if not sent and time.time() > deadline - seconds + 5:
31 if to_send:
32 os.write(fd, to_send.encode())
33 sent = True
34 r, _, _ = select.select([fd], [], [], 0.5)
35 if r:
36 try:
37 chunk = os.read(fd, 65536)
38 except OSError:
39 break
40 if not chunk:
41 break
42 out += chunk
43
44 try:
45 os.kill(pid, signal.SIGKILL)
46 except ProcessLookupError:
47 pass
48 os.waitpid(pid, 0)
49 sys.stdout.buffer.write(out)
50