ykxmz/_build/shoot_pc_main.py

# -*- coding: utf-8 -*-
r'''拍电脑端主窗口截图（官网上用）。

踩过的坑：
  传 --code 启动是**一次性命令**，设完码就退出，窗口根本不会出现。
  必须先单独跑一次 --code，再不带参数启动，才有主窗口。

窗口标题现在是 "ykxmz · 第一次设置"（设过码之后变成控制窗口）。
'''
from __future__ import annotations

import ctypes
import os
import subprocess
import sys
import time
from ctypes import wintypes
from pathlib import Path

ROOT = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote")
SHOTS = ROOT / "_build" / "site_shots"
EXE = ROOT / "dist_exe" / "ykxmz" / "ykxmz.exe"
DATA = ROOT / "_build" / "_shots_data"

user32 = ctypes.windll.user32
gdi32 = ctypes.windll.gdi32


def windows_of(pid: int) -> list:
    out = []

    @ctypes.WINFUNCTYPE(ctypes.c_bool, wintypes.HWND, wintypes.LPARAM)
    def cb(hwnd, _lp):
        wpid = wintypes.DWORD()
        user32.GetWindowThreadProcessId(hwnd, ctypes.byref(wpid))
        if wpid.value == pid:
            rect = wintypes.RECT()
            user32.GetWindowRect(hwnd, ctypes.byref(rect))
            if user32.IsWindowVisible(hwnd) and (rect.right - rect.left) > 200:
                out.append(hwnd)
        return True

    user32.EnumWindows(cb, 0)
    return out


def capture(hwnd: int, path: Path) -> bool:
    rect = wintypes.RECT()
    user32.GetWindowRect(hwnd, ctypes.byref(rect))
    width = rect.right - rect.left
    height = rect.bottom - rect.top
    if width <= 0 or height <= 0:
        return False

    hdc_window = user32.GetWindowDC(hwnd)
    hdc_mem = gdi32.CreateCompatibleDC(hdc_window)
    bitmap = gdi32.CreateCompatibleBitmap(hdc_window, width, height)
    gdi32.SelectObject(hdc_mem, bitmap)
    user32.PrintWindow(hwnd, hdc_mem, 2)     # 2 = PW_RENDERFULLCONTENT

    class BMIH(ctypes.Structure):
        _fields_ = [("biSize", wintypes.DWORD), ("biWidth", wintypes.LONG),
                    ("biHeight", wintypes.LONG), ("biPlanes", wintypes.WORD),
                    ("biBitCount", wintypes.WORD),
                    ("biCompression", wintypes.DWORD),
                    ("biSizeImage", wintypes.DWORD),
                    ("biXPelsPerMeter", wintypes.LONG),
                    ("biYPelsPerMeter", wintypes.LONG),
                    ("biClrUsed", wintypes.DWORD),
                    ("biClrImportant", wintypes.DWORD)]

    header = BMIH()
    header.biSize = ctypes.sizeof(BMIH)
    header.biWidth = width
    header.biHeight = -height
    header.biPlanes = 1
    header.biBitCount = 32

    buf = ctypes.create_string_buffer(width * height * 4)
    gdi32.GetDIBits(hdc_mem, bitmap, 0, height, buf, ctypes.byref(header), 0)

    from PIL import Image
    img = Image.frombuffer("RGBA", (width, height), buf, "raw", "BGRA", 0, 1)
    img.convert("RGB").save(path)

    gdi32.DeleteObject(bitmap)
    gdi32.DeleteDC(hdc_mem)
    user32.ReleaseDC(hwnd, hdc_window)
    return True


def main() -> int:
    SHOTS.mkdir(parents=True, exist_ok=True)
    DATA.mkdir(parents=True, exist_ok=True)
    env = dict(os.environ)
    env["YKXMZ_DATA"] = str(DATA)

    # 先清掉旧进程
    subprocess.run(["powershell", "-NoProfile", "-Command",
                    "Get-Process ykxmz -ErrorAction SilentlyContinue | "
                    "Stop-Process -Force"],
                   capture_output=True, timeout=60)
    time.sleep(2)

    print("  一、清掉旧数据，模拟全新安装（这样才会出设置窗口）")
    import shutil
    # 设过配对码之后被控端就完全无窗口运行了（这是设计：
    # "用户要的是全程无窗口"）。所以想拍界面必须清掉数据当首次运行。
    shutil.rmtree(DATA, ignore_errors=True)
    DATA.mkdir(parents=True, exist_ok=True)
    print("      清好了")

    print("  二、启动（首次运行 -> 出「第一次设置」窗口）")
    proc = subprocess.Popen([str(EXE), "--port", "8898"],
                            cwd=str(ROOT), env=env,
                            stdout=subprocess.DEVNULL,
                            stderr=subprocess.DEVNULL)
    time.sleep(14)

    handles = windows_of(proc.pid)
    print("      找到 %d 个可见窗口" % len(handles))
    if not handles:
        print("      [失败] 没有窗口")
        proc.terminate()
        return 1

    hwnd = handles[0]
    user32.SetForegroundWindow(hwnd)
    user32.MoveWindow(hwnd, 80, 80, 520, 600, True)
    time.sleep(2)

    out = SHOTS / "pc-main.png"
    if capture(hwnd, out):
        print("      ✓ %s  %.1f KB" % (out.name, out.stat().st_size / 1024))
        # 顺便看看标题
        length = user32.GetWindowTextLengthW(hwnd)
        buf = ctypes.create_unicode_buffer(length + 1)
        user32.GetWindowTextW(hwnd, buf, length + 1)
        print("      窗口标题: %r" % buf.value)

    proc.terminate()
    try:
        proc.wait(timeout=10)
    except subprocess.TimeoutExpired:
        proc.kill()
    subprocess.run(["powershell", "-NoProfile", "-Command",
                    "Get-Process ykxmz -ErrorAction SilentlyContinue | "
                    "Stop-Process -Force"],
                   capture_output=True, timeout=60)

    print()
    print("  产物:")
    for p in sorted(SHOTS.glob("*.png")):
        print("    %-24s %6.1f KB" % (p.name, p.stat().st_size / 1024))
    return 0


if __name__ == "__main__":
    sys.exit(main())