| 1 | import base64 |
| 2 | import ctypes |
| 3 | import hashlib |
| 4 | import json |
| 5 | import os |
| 6 | import struct |
| 7 | import subprocess |
| 8 | import sys |
| 9 | import tempfile |
| 10 | import threading |
| 11 | import time |
| 12 | import zlib |
| 13 | from ctypes import wintypes |
| 14 | from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler |
| 15 | |
| 16 | BASE = os.path.dirname(os.path.abspath(__file__)) |
| 17 | with open(__file__, 'rb') as source: |
| 18 | AGENT_SHA256 = hashlib.sha256(source.read()).hexdigest() |
| 19 | LOCK = threading.Lock() |
| 20 | |
| 21 | user32 = ctypes.windll.user32 |
| 22 | gdi32 = ctypes.windll.gdi32 |
| 23 | |
| 24 | SRCCOPY = 0x00CC0020 |
| 25 | CAPTUREBLT = 0x40000000 |
| 26 | BI_RGB = 0 |
| 27 | DIB_RGB_COLORS = 0 |
| 28 | |
| 29 | # restype must be set on x64: the default c_int return truncates a 64-bit HANDLE. |
| 30 | user32.GetDC.restype = wintypes.HDC |
| 31 | user32.GetDC.argtypes = [wintypes.HWND] |
| 32 | user32.ReleaseDC.restype = ctypes.c_int |
| 33 | user32.ReleaseDC.argtypes = [wintypes.HWND, wintypes.HDC] |
| 34 | user32.GetSystemMetrics.restype = ctypes.c_int |
| 35 | user32.GetSystemMetrics.argtypes = [ctypes.c_int] |
| 36 | gdi32.CreateCompatibleDC.restype = wintypes.HDC |
| 37 | gdi32.CreateCompatibleDC.argtypes = [wintypes.HDC] |
| 38 | gdi32.CreateCompatibleBitmap.restype = wintypes.HBITMAP |
| 39 | gdi32.CreateCompatibleBitmap.argtypes = [wintypes.HDC, ctypes.c_int, ctypes.c_int] |
| 40 | gdi32.SelectObject.restype = wintypes.HGDIOBJ |
| 41 | gdi32.SelectObject.argtypes = [wintypes.HDC, wintypes.HGDIOBJ] |
| 42 | gdi32.BitBlt.restype = wintypes.BOOL |
| 43 | gdi32.BitBlt.argtypes = [wintypes.HDC, ctypes.c_int, ctypes.c_int, ctypes.c_int, |
| 44 | ctypes.c_int, wintypes.HDC, ctypes.c_int, ctypes.c_int, |
| 45 | wintypes.DWORD] |
| 46 | gdi32.GetDIBits.restype = ctypes.c_int |
| 47 | gdi32.GetDIBits.argtypes = [wintypes.HDC, wintypes.HBITMAP, wintypes.UINT, |
| 48 | wintypes.UINT, ctypes.c_void_p, ctypes.c_void_p, |
| 49 | wintypes.UINT] |
| 50 | gdi32.DeleteObject.restype = wintypes.BOOL |
| 51 | gdi32.DeleteObject.argtypes = [wintypes.HGDIOBJ] |
| 52 | gdi32.DeleteDC.restype = wintypes.BOOL |
| 53 | gdi32.DeleteDC.argtypes = [wintypes.HDC] |
| 54 | |
| 55 | WNDENUMPROC = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.HWND, wintypes.LPARAM) |
| 56 | user32.GetForegroundWindow.restype = wintypes.HWND |
| 57 | user32.GetForegroundWindow.argtypes = [] |
| 58 | user32.GetWindowTextW.restype = ctypes.c_int |
| 59 | user32.GetWindowTextW.argtypes = [wintypes.HWND, wintypes.LPWSTR, ctypes.c_int] |
| 60 | user32.GetClassNameW.restype = ctypes.c_int |
| 61 | user32.GetClassNameW.argtypes = [wintypes.HWND, wintypes.LPWSTR, ctypes.c_int] |
| 62 | user32.GetClientRect.restype = wintypes.BOOL |
| 63 | user32.GetClientRect.argtypes = [wintypes.HWND, ctypes.POINTER(wintypes.RECT)] |
| 64 | user32.EnumChildWindows.restype = wintypes.BOOL |
| 65 | user32.EnumChildWindows.argtypes = [wintypes.HWND, WNDENUMPROC, wintypes.LPARAM] |
| 66 | |
| 67 | |
| 68 | class BITMAPINFOHEADER(ctypes.Structure): |
| 69 | _fields_ = [("biSize", wintypes.DWORD), |
| 70 | ("biWidth", wintypes.LONG), |
| 71 | ("biHeight", wintypes.LONG), |
| 72 | ("biPlanes", wintypes.WORD), |
| 73 | ("biBitCount", wintypes.WORD), |
| 74 | ("biCompression", wintypes.DWORD), |
| 75 | ("biSizeImage", wintypes.DWORD), |
| 76 | ("biXPelsPerMeter", wintypes.LONG), |
| 77 | ("biYPelsPerMeter", wintypes.LONG), |
| 78 | ("biClrUsed", wintypes.DWORD), |
| 79 | ("biClrImportant", wintypes.DWORD)] |
| 80 | |
| 81 | |
| 82 | PREAMBLE = ( |
| 83 | "#Requires AutoHotkey v2.0\n" |
| 84 | "#SingleInstance Off\n" |
| 85 | "#Warn All, Off\n" |
| 86 | 'FileEncoding "UTF-8-RAW"\n' |
| 87 | 'SendMode "Input"\n' |
| 88 | "SetWorkingDir A_ScriptDir\n" |
| 89 | 'CoordMode "Mouse", "Screen"\n' |
| 90 | 'CoordMode "Pixel", "Screen"\n' |
| 91 | 'CoordMode "ToolTip", "Screen"\n' |
| 92 | "SetTitleMatchMode 2\n" |
| 93 | "DetectHiddenWindows true\n" |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | def ahk_path(): |
| 98 | u64 = os.path.join(BASE, "vendor", "ahk", "AutoHotkey64.exe") |
| 99 | u32 = os.path.join(BASE, "vendor", "ahk", "AutoHotkey32.exe") |
| 100 | return u64 if os.path.exists(u64) else u32 |
| 101 | |
| 102 | |
| 103 | def _chunk(tag, data): |
| 104 | return (struct.pack(">I", len(data)) + tag + data + |
| 105 | struct.pack(">I", zlib.crc32(tag + data) & 0xffffffff)) |
| 106 | |
| 107 | |
| 108 | def _to_png(bgra, w, h): |
| 109 | arr = bytearray(bgra) |
| 110 | r = arr[2::4] |
| 111 | arr[2::4] = arr[0::4] |
| 112 | arr[0::4] = r |
| 113 | del arr[3::4] |
| 114 | stride = w * 3 |
| 115 | mv = memoryview(arr) |
| 116 | raw = bytearray() |
| 117 | for y in range(h): |
| 118 | raw.append(0) |
| 119 | raw += mv[y * stride:(y + 1) * stride] |
| 120 | ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0) |
| 121 | return (b"\x89PNG\r\n\x1a\n" + _chunk(b"IHDR", ihdr) + |
| 122 | _chunk(b"IDAT", zlib.compress(bytes(raw))) + _chunk(b"IEND", b"")) |
| 123 | |
| 124 | |
| 125 | def capture(): |
| 126 | w = user32.GetSystemMetrics(0) |
| 127 | h = user32.GetSystemMetrics(1) |
| 128 | hdc = user32.GetDC(None) |
| 129 | mem = None |
| 130 | hbm = None |
| 131 | try: |
| 132 | mem = gdi32.CreateCompatibleDC(hdc) |
| 133 | hbm = gdi32.CreateCompatibleBitmap(hdc, w, h) |
| 134 | old = gdi32.SelectObject(mem, hbm) |
| 135 | gdi32.BitBlt(mem, 0, 0, w, h, hdc, 0, 0, SRCCOPY | CAPTUREBLT) |
| 136 | gdi32.SelectObject(mem, old) |
| 137 | bmi = BITMAPINFOHEADER() |
| 138 | bmi.biSize = ctypes.sizeof(BITMAPINFOHEADER) |
| 139 | bmi.biWidth = w |
| 140 | bmi.biHeight = -h # negative => top-down rows, matches screenshot orientation |
| 141 | bmi.biPlanes = 1 |
| 142 | bmi.biBitCount = 32 |
| 143 | bmi.biCompression = BI_RGB |
| 144 | buf = ctypes.create_string_buffer(w * h * 4) |
| 145 | if gdi32.GetDIBits(mem, hbm, 0, h, buf, ctypes.byref(bmi), |
| 146 | DIB_RGB_COLORS) == 0: |
| 147 | raise RuntimeError("GetDIBits failed") |
| 148 | data = buf.raw |
| 149 | finally: |
| 150 | if hbm: |
| 151 | gdi32.DeleteObject(hbm) |
| 152 | if mem: |
| 153 | gdi32.DeleteDC(mem) |
| 154 | user32.ReleaseDC(None, hdc) |
| 155 | return _to_png(data, w, h), w, h |
| 156 | |
| 157 | |
| 158 | def _win_text(fn, hwnd, n=512): |
| 159 | buf = ctypes.create_unicode_buffer(n) |
| 160 | fn(hwnd, buf, n) |
| 161 | return buf.value |
| 162 | |
| 163 | |
| 164 | def active_window(): |
| 165 | hwnd = user32.GetForegroundWindow() |
| 166 | if not hwnd: |
| 167 | return {"title": "", "class": "", "dialog": False} |
| 168 | cls = _win_text(user32.GetClassNameW, hwnd) |
| 169 | return {"title": _win_text(user32.GetWindowTextW, hwnd), |
| 170 | "class": cls, "dialog": cls == "#32770"} |
| 171 | |
| 172 | |
| 173 | def control_tree(): |
| 174 | hwnd = user32.GetForegroundWindow() |
| 175 | rows = [] |
| 176 | |
| 177 | def add(h): |
| 178 | rect = wintypes.RECT() |
| 179 | user32.GetClientRect(h, ctypes.byref(rect)) |
| 180 | rows.append((_win_text(user32.GetClassNameW, h), |
| 181 | _win_text(user32.GetWindowTextW, h), |
| 182 | rect.left, rect.top, |
| 183 | rect.right - rect.left, rect.bottom - rect.top)) |
| 184 | |
| 185 | if hwnd: |
| 186 | add(hwnd) |
| 187 | |
| 188 | def cb(child, _lparam): |
| 189 | add(child) |
| 190 | return True |
| 191 | user32.EnumChildWindows(hwnd, WNDENUMPROC(cb), 0) |
| 192 | lines = ["%-24s %-28s %s" % ("class", "text", "l,t,w,h")] |
| 193 | for cls, text, l, t, w, h in rows: |
| 194 | lines.append("%-24s %-28s %d,%d,%d,%d" % (cls, text, l, t, w, h)) |
| 195 | return "\n".join(lines) |
| 196 | |
| 197 | |
| 198 | def _run(argv, timeout_ms, encoding, on_timeout=None): |
| 199 | def dec(b): |
| 200 | return b.decode(encoding, "replace") if b else "" |
| 201 | |
| 202 | p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE) |
| 203 | try: |
| 204 | out, err = p.communicate(timeout=timeout_ms / 1000.0) |
| 205 | return p.returncode, dec(out), dec(err), None |
| 206 | except subprocess.TimeoutExpired: |
| 207 | if on_timeout: |
| 208 | on_timeout() |
| 209 | subprocess.call(["taskkill", "/F", "/T", "/PID", str(p.pid)]) |
| 210 | try: |
| 211 | # A process the script launched can outlive taskkill still holding the |
| 212 | # inherited stdout pipe; an unbounded wait here wedges the agent's lock. |
| 213 | out, err = p.communicate(timeout=5) |
| 214 | except subprocess.TimeoutExpired: |
| 215 | out = err = b"" |
| 216 | for pipe in (p.stdout, p.stderr): |
| 217 | if pipe: |
| 218 | pipe.close() |
| 219 | msg = "timed out after %d ms, process killed" % timeout_ms |
| 220 | return None, dec(out), dec(err), msg |
| 221 | |
| 222 | |
| 223 | def run_ahk(script, timeout_ms, on_timeout=None): |
| 224 | fd, path = tempfile.mkstemp(suffix=".ahk") |
| 225 | os.close(fd) |
| 226 | # AutoHotkey (Unicode-only) needs a UTF-8 BOM to read the file as UTF-8. |
| 227 | with open(path, "wb") as f: |
| 228 | f.write(b"\xef\xbb\xbf" + (PREAMBLE + script).encode("utf-8")) |
| 229 | try: |
| 230 | # /ErrorStdOut sends syntax errors to stderr instead of a blocking modal. |
| 231 | code, out, err, error = _run([ahk_path(), "/ErrorStdOut", path], timeout_ms, |
| 232 | "utf-8", on_timeout) |
| 233 | finally: |
| 234 | os.remove(path) |
| 235 | if error: |
| 236 | error = "script " + error |
| 237 | return code, out, err, error |
| 238 | |
| 239 | |
| 240 | def _safe_capture(error): |
| 241 | try: |
| 242 | png, w, h = capture() |
| 243 | return base64.b64encode(png).decode("ascii"), w, h, error |
| 244 | except Exception as e: |
| 245 | note = "screenshot failed: " + repr(e) |
| 246 | return "", 0, 0, (error + "; " + note if error else note) |
| 247 | |
| 248 | |
| 249 | LOG_PATH = os.path.join(BASE, "agent.log") |
| 250 | LOG_LOCK = threading.Lock() |
| 251 | |
| 252 | |
| 253 | def log(text, stamp=True): |
| 254 | """Console and agent.log, so a session can be reconstructed after the fact.""" |
| 255 | line = time.strftime("%Y-%m-%d %H:%M:%S ") + text if stamp else text |
| 256 | with LOG_LOCK: |
| 257 | print(line, end="" if line.endswith("\n") else "\n", flush=True) |
| 258 | try: |
| 259 | with open(LOG_PATH, "a", encoding="utf-8") as f: |
| 260 | f.write(line if line.endswith("\n") else line + "\n") |
| 261 | except OSError: |
| 262 | pass # a log that cannot be written must not take the agent down |
| 263 | |
| 264 | |
| 265 | class Handler(BaseHTTPRequestHandler): |
| 266 | def log_message(self, *a): |
| 267 | pass |
| 268 | |
| 269 | def _send(self, code, obj): |
| 270 | body = json.dumps(obj).encode("utf-8") |
| 271 | self.send_response(code) |
| 272 | self.send_header("Content-Type", "application/json") |
| 273 | self.send_header("Content-Length", str(len(body))) |
| 274 | self.end_headers() |
| 275 | self.wfile.write(body) |
| 276 | |
| 277 | def _log(self, start, code, body=None): |
| 278 | log("%s %s %dms exit=%s" % (self.command, self.path, |
| 279 | int((time.time() - start) * 1000), code)) |
| 280 | if body: |
| 281 | log("".join(" | %s\n" % l for l in body.splitlines()), stamp=False) |
| 282 | |
| 283 | def _auth_ok(self): |
| 284 | token = os.environ.get("WIN7_TOKEN") |
| 285 | return not token or self.headers.get("X-Win7-Token") == token |
| 286 | |
| 287 | def _body(self): |
| 288 | n = int(self.headers.get("Content-Length") or 0) |
| 289 | raw = self.rfile.read(n) if n else b"" |
| 290 | return json.loads(raw.decode("utf-8")) if raw else {} |
| 291 | |
| 292 | def do_GET(self): |
| 293 | start = time.time() |
| 294 | if not self._auth_ok(): |
| 295 | self._send(401, {"error": "bad token"}) |
| 296 | elif self.path == "/health": |
| 297 | self._send(200, {"ok": True, "w": user32.GetSystemMetrics(0), |
| 298 | "h": user32.GetSystemMetrics(1), |
| 299 | "hostname": os.environ.get("COMPUTERNAME", ""), |
| 300 | "ahk": ahk_path(), "python": sys.executable, |
| 301 | "agent_sha256": AGENT_SHA256}) |
| 302 | else: |
| 303 | self._send(404, {"error": "not found"}) |
| 304 | self._log(start, None) |
| 305 | |
| 306 | def do_POST(self): |
| 307 | start = time.time() |
| 308 | code = None |
| 309 | if not self._auth_ok(): |
| 310 | self._send(401, {"error": "bad token"}) |
| 311 | self._log(start, None) |
| 312 | return |
| 313 | try: |
| 314 | data = self._body() |
| 315 | except Exception: |
| 316 | self._send(400, {"error": "bad json"}) |
| 317 | self._log(start, None) |
| 318 | return |
| 319 | body = {"/exec": data.get("script"), "/cmd": data.get("command"), |
| 320 | "/spawn": data.get("command"), |
| 321 | "/put": data.get("path"), "/get": data.get("path")}.get(self.path) |
| 322 | try: |
| 323 | if self.path == "/exec": |
| 324 | resp = self._exec(data) |
| 325 | code = resp["exit"] |
| 326 | elif self.path == "/cmd": |
| 327 | resp = self._cmd(data) |
| 328 | code = resp["exit"] |
| 329 | elif self.path == "/spawn": |
| 330 | resp = self._spawn(data) |
| 331 | code = 0 |
| 332 | elif self.path == "/put": |
| 333 | resp = self._put(data) |
| 334 | elif self.path == "/get": |
| 335 | resp = self._get(data) |
| 336 | elif self.path == "/shot": |
| 337 | resp = self._shot() |
| 338 | elif self.path == "/ui": |
| 339 | resp = self._ui() |
| 340 | else: |
| 341 | self._send(404, {"error": "not found"}) |
| 342 | self._log(start, None) |
| 343 | return |
| 344 | self._send(200, resp) |
| 345 | except Exception as e: |
| 346 | self._send(200, {"error": repr(e), "exit": None, "stdout": "", |
| 347 | "stderr": "", "png_b64": "", "w": 0, "h": 0}) |
| 348 | self._log(start, code, body) |
| 349 | |
| 350 | def _exec(self, data): |
| 351 | pre = [] |
| 352 | |
| 353 | def grab(): # capture before taskkill so the blocker is still on screen |
| 354 | pre.append((_safe_capture(None), active_window())) |
| 355 | |
| 356 | with LOCK: |
| 357 | code, out, err, error = run_ahk(data.get("script", ""), |
| 358 | data.get("timeout_ms", 60000), grab) |
| 359 | if pre: |
| 360 | (b64, w, h, note), win = pre[0] |
| 361 | else: |
| 362 | time.sleep(data.get("shot_delay_ms", 500) / 1000.0) |
| 363 | b64, w, h, note = _safe_capture(None) |
| 364 | win = active_window() |
| 365 | if note: |
| 366 | error = error + "; " + note if error else note |
| 367 | return {"exit": code, "stdout": out, "stderr": err, "png_b64": b64, |
| 368 | "w": w, "h": h, "error": error, "win": win} |
| 369 | |
| 370 | def _put(self, data): |
| 371 | path = data.get("path", "") |
| 372 | blob = base64.b64decode(data.get("b64", "")) |
| 373 | parent = os.path.dirname(path) |
| 374 | if parent and not os.path.isdir(parent): |
| 375 | os.makedirs(parent) |
| 376 | with open(path, "wb") as f: |
| 377 | f.write(blob) |
| 378 | return {"bytes": len(blob), "path": os.path.abspath(path), "error": None} |
| 379 | |
| 380 | def _get(self, data): |
| 381 | with open(data.get("path", ""), "rb") as f: |
| 382 | blob = f.read() |
| 383 | return {"b64": base64.b64encode(blob).decode("ascii"), "bytes": len(blob), |
| 384 | "error": None} |
| 385 | |
| 386 | def _cmd(self, data): |
| 387 | with LOCK: |
| 388 | # A raw command-line string, not a list: on Windows list2cmdline would |
| 389 | # backslash-escape the embedded quotes before cmd.exe ever sees them. |
| 390 | code, out, err, error = _run("cmd.exe /c " + data.get("command", ""), |
| 391 | data.get("timeout_ms", 60000), "oem") |
| 392 | return {"exit": code, "stdout": out, "stderr": err, "error": error} |
| 393 | |
| 394 | def _spawn(self, data): |
| 395 | flags = subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.DETACHED_PROCESS |
| 396 | process = subprocess.Popen(data.get("command", ""), stdin=subprocess.DEVNULL, |
| 397 | stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, |
| 398 | close_fds=True, creationflags=flags) |
| 399 | return {"pid": process.pid, "error": None} |
| 400 | |
| 401 | def _shot(self): |
| 402 | b64, w, h, error = _safe_capture(None) |
| 403 | win = active_window() |
| 404 | return {"png_b64": b64, "w": w, "h": h, "error": error, "win": win} |
| 405 | |
| 406 | def _ui(self): |
| 407 | win = active_window() |
| 408 | controls = control_tree() |
| 409 | return {"win": win, "controls": controls, "error": None} |
| 410 | |
| 411 | |
| 412 | def keep_awake(): |
| 413 | """Hold the display and system awake for as long as this process lives. |
| 414 | |
| 415 | Screen blanking alone is harmless -- BitBlt still captures the composited |
| 416 | desktop -- but sleep kills the agent outright, and a locked session moves |
| 417 | the desktop to Winlogon where neither AutoHotkey nor the capture can reach. |
| 418 | """ |
| 419 | ES_CONTINUOUS, ES_SYSTEM_REQUIRED, ES_DISPLAY_REQUIRED = 0x80000000, 0x1, 0x2 |
| 420 | kernel32 = ctypes.windll.kernel32 |
| 421 | kernel32.SetThreadExecutionState.restype = wintypes.DWORD |
| 422 | kernel32.SetThreadExecutionState.argtypes = [wintypes.DWORD] |
| 423 | return kernel32.SetThreadExecutionState( |
| 424 | ES_CONTINUOUS | ES_SYSTEM_REQUIRED | ES_DISPLAY_REQUIRED) |
| 425 | |
| 426 | |
| 427 | def main(): |
| 428 | user32.SetProcessDPIAware() |
| 429 | awake = keep_awake() |
| 430 | port = int(os.environ.get("WIN7_PORT", "8777")) |
| 431 | srv = ThreadingHTTPServer(("0.0.0.0", port), Handler) |
| 432 | log("win7-agent listening on 0.0.0.0:%d ahk=%s log=%s%s" |
| 433 | % (port, ahk_path(), LOG_PATH, |
| 434 | "" if awake else " (WARNING: could not inhibit sleep)")) |
| 435 | srv.serve_forever() |
| 436 | |
| 437 | |
| 438 | if __name__ == "__main__": |
| 439 | main() |