Skip to main content

max / alloy

4.3 KB · 145 lines History Blame Raw
1 #!/usr/bin/env python3
2 """Drive the qemu monitor: sendkey a string, and screendump to PNG.
3
4 Usage:
5 vm.py type "some text" type printable ASCII into the guest
6 vm.py key ret tab spc send named keys
7 vm.py shot name screendump to name.png
8 """
9 import socket, sys, time, zlib, struct, os, re
10
11 SCRATCH = os.environ.get(
12 "VM_STATE", os.path.join(os.path.dirname(os.path.abspath(__file__)), "state")
13 )
14 SOCK = os.path.join(SCRATCH, "monitor.sock")
15
16 NAMED = {
17 " ": "spc", "-": "minus", "=": "equal", "[": "bracket_left",
18 "]": "bracket_right", ";": "semicolon", "'": "apostrophe",
19 "`": "grave_accent", "\\": "backslash", ",": "comma", ".": "dot",
20 "/": "slash",
21 }
22 SHIFTED = {
23 "!": "1", "@": "2", "#": "3", "$": "4", "%": "5", "^": "6", "&": "7",
24 "*": "8", "(": "9", ")": "0", "_": "minus", "+": "equal",
25 "{": "bracket_left", "}": "bracket_right", ":": "semicolon",
26 '"': "apostrophe", "~": "grave_accent", "|": "backslash", "<": "comma",
27 ">": "dot", "?": "slash",
28 }
29
30
31 class Monitor:
32 def __init__(self):
33 self.s = socket.socket(socket.AF_UNIX)
34 self.s.connect(SOCK)
35 time.sleep(0.2)
36 self.drain()
37
38 def drain(self):
39 self.s.setblocking(False)
40 out = b""
41 try:
42 while True:
43 chunk = self.s.recv(65536)
44 if not chunk:
45 break
46 out += chunk
47 except BlockingIOError:
48 pass
49 self.s.setblocking(True)
50 return out.decode(errors="replace")
51
52 def cmd(self, line, wait=0.05):
53 self.s.sendall((line + "\n").encode())
54 time.sleep(wait)
55 return self.drain()
56
57 def sendkey(self, k):
58 self.cmd("sendkey " + k, wait=0.04)
59
60 def type(self, text):
61 for c in text:
62 if c.isalnum():
63 self.sendkey(c if not c.isupper() else "shift-" + c.lower())
64 elif c in NAMED:
65 self.sendkey(NAMED[c])
66 elif c in SHIFTED:
67 self.sendkey("shift-" + SHIFTED[c])
68 elif c == "\n":
69 self.sendkey("ret")
70 else:
71 raise SystemExit("no key mapping for %r" % c)
72
73
74 def ppm_to_png(ppm_path, png_path):
75 with open(ppm_path, "rb") as f:
76 data = f.read()
77 # P6 header: magic, width height, maxval, each possibly comment-separated
78 fields, idx = [], 2
79 while len(fields) < 3:
80 while data[idx:idx + 1].isspace():
81 idx += 1
82 if data[idx:idx + 1] == b"#":
83 while data[idx:idx + 1] != b"\n":
84 idx += 1
85 continue
86 start = idx
87 while not data[idx:idx + 1].isspace():
88 idx += 1
89 fields.append(int(data[start:idx]))
90 idx += 1
91 w, h, _maxval = fields
92 pixels = data[idx:]
93 raw = b"".join(b"\x00" + pixels[y * w * 3:(y + 1) * w * 3] for y in range(h))
94
95 def chunk(tag, payload):
96 return (struct.pack(">I", len(payload)) + tag + payload
97 + struct.pack(">I", zlib.crc32(tag + payload) & 0xFFFFFFFF))
98
99 png = (b"\x89PNG\r\n\x1a\n"
100 + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
101 + chunk(b"IDAT", zlib.compress(raw, 6))
102 + chunk(b"IEND", b""))
103 with open(png_path, "wb") as f:
104 f.write(png)
105 return w, h
106
107
108 def main():
109 if len(sys.argv) < 2:
110 raise SystemExit(__doc__)
111 verb = sys.argv[1]
112 if verb == "shot":
113 name = sys.argv[2]
114 ppm = os.path.join(SCRATCH, name + ".ppm")
115 png = os.path.join(SCRATCH, name + ".png")
116 m = Monitor()
117 m.cmd("screendump " + ppm, wait=1.0)
118 for _ in range(20):
119 if os.path.exists(ppm) and os.path.getsize(ppm) > 1000:
120 break
121 time.sleep(0.3)
122 print(ppm_to_png(ppm, png), png)
123 return
124 m = Monitor()
125 if verb == "type":
126 m.type(sys.argv[2])
127 elif verb == "key":
128 for k in sys.argv[2:]:
129 m.sendkey(k)
130 elif verb == "hold":
131 # Keep the GRUB countdown from expiring: one connection, many keys.
132 end = time.time() + float(sys.argv[2])
133 while time.time() < end:
134 m.sendkey("down")
135 m.sendkey("up")
136 time.sleep(0.2)
137 elif verb == "cmd":
138 print(m.cmd(" ".join(sys.argv[2:]), wait=0.5))
139 else:
140 raise SystemExit(__doc__)
141
142
143 if __name__ == "__main__":
144 main()
145