后台屏幕截图

发布时间 2023-08-02 15:17:31作者: Lyxin_c
import os
from ctypes import windll

import win32gui
import win32ui
from PIL import Image

import const
import logs
import utils

hwnd_title = dict()


# hwnd :    handle to a window

def get_all_hwnd(hwnd, mouse):
    if win32gui.IsWindow(hwnd) and win32gui.IsWindowEnabled(hwnd) and win32gui.IsWindowVisible(hwnd):
        hwnd_title.update({hwnd: win32gui.GetWindowText(hwnd)})


def screenshot(startTitle="", pic_name=None):
    # 遍历所有进程,获取到 窗口句柄、title
    win32gui.EnumWindows(get_all_hwnd, 0)

    for h, t in hwnd_title.items():
        if str(t).startswith(startTitle):
            mhxy_title = t

            # 获取窗口句柄
            hwnd = win32gui.FindWindow(None, mhxy_title)
            if hwnd is None:
                logs.debug("获取窗口句柄失败")
                continue

            # 获取窗口的大小
            window_rect = win32gui.GetWindowRect(hwnd)
            if window_rect is None:
                logs.debug("获取窗口大小失败")
                continue

            left, top, right, bottom = window_rect
            width = right - left
            height = bottom - top

            # 获取进程的 DC
            hdc = win32gui.GetWindowDC(hwnd)
            # 创建一个虚拟句柄,用于绘制截图
            mfc_dc = win32ui.CreateDCFromHandle(hdc)
            save_dc = mfc_dc.CreateCompatibleDC()

            # 创建一个图片对象,保存截图
            save_bitmap = win32ui.CreateBitmap()
            save_bitmap.CreateCompatibleBitmap(mfc_dc, width, height)
            save_dc.SelectObject(save_bitmap)

            # 复制窗口到虚拟句柄
            result = windll.user32.PrintWindow(hwnd, save_dc.GetSafeHdc(), 1)

            # Convert the raw bitmap data to an image
            bmpinfo = save_bitmap.GetInfo()
            bmpstr = save_bitmap.GetBitmapBits(True)
            image = Image.frombuffer('RGB', (bmpinfo['bmWidth'], bmpinfo['bmHeight']), bmpstr, 'raw', 'BGRX', 0, 1)

            # 清理句柄对象
            save_dc.DeleteDC()
            mfc_dc.DeleteDC()
            win32gui.ReleaseDC(hwnd, hdc)

            if result == 1:
                if pic_name is None:
                    pic_name = os.path.join(const.CACHE, str(utils.generate_tid()) + ".png")
                image.save(pic_name)
                return pic_name
            else:
                logs.debug("Failed to capture window content.")

    return None