| 1 |
|
| 2 |
|
| 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 |
|
| 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 |
|
| 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 |
|