From 9e86f160e7d88660d6a8a81efdb56f98e751dbf4 Mon Sep 17 00:00:00 2001
From: FredWuCZ <fredwucz@outlook.com>
Date: Fri, 18 Oct 2024 18:44:53 +0800
Subject: [PATCH] Capture cursor on Windows

---
 desktop_env/server/main.py | 62 +++++++++++++++++++++++++++-----------
 1 file changed, 44 insertions(+), 18 deletions(-)

diff --git a/desktop_env/server/main.py b/desktop_env/server/main.py
index 7ba8800..f2bf187 100644
--- a/desktop_env/server/main.py
+++ b/desktop_env/server/main.py
@@ -14,7 +14,7 @@ import lxml.etree
 import pyautogui
 import requests
 import re
-from PIL import Image
+from PIL import Image, ImageGrab
 from Xlib import display, X
 from flask import Flask, request, jsonify, send_file, abort  # , send_from_directory
 from lxml.etree import _Element
@@ -35,6 +35,7 @@ elif platform_name == "Windows":
     from pywinauto import Desktop
     from pywinauto.base_wrapper import BaseWrapper
     import pywinauto.application
+    import win32ui, win32gui
 
     Accessible = Any
 
@@ -88,7 +89,7 @@ def execute_command():
     # Execute the command without any safety checks.
     try:
         result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=shell, text=True,
-                                timeout=120)
+                                timeout=120, creationflags=subprocess.CREATE_NO_WINDOW)
         return jsonify({
             'status': 'success',
             'output': result.stdout,
@@ -150,22 +151,47 @@ def capture_screen_with_cursor():
 
     # fixme: This is a temporary fix for the cursor not being captured on Windows and Linux
     if user_platform == "Windows":
-        def _download_image(url, path):
-            response = requests.get(url)
-            with open(path, 'wb') as file:
-                file.write(response.content)
-
-        cursor_path = os.path.join("screenshots", "cursor.png")
-        if not os.path.exists(cursor_path):
-            cursor_url = "https://vip.helloimg.com/images/2023/12/02/oQPzmt.png"
-            _download_image(cursor_url, cursor_path)
-        screenshot = pyautogui.screenshot()
-        cursor_x, cursor_y = pyautogui.position()
-        cursor = Image.open(cursor_path)
-        # make the cursor smaller
-        cursor = cursor.resize((int(cursor.width / 1.5), int(cursor.height / 1.5)))
-        screenshot.paste(cursor, (cursor_x, cursor_y), cursor)
-        screenshot.save(file_path)
+        def get_cursor():
+            hcursor = win32gui.GetCursorInfo()[1]
+            hdc = win32ui.CreateDCFromHandle(win32gui.GetDC(0))
+            hbmp = win32ui.CreateBitmap()
+            hbmp.CreateCompatibleBitmap(hdc, 36, 36)
+            hdc = hdc.CreateCompatibleDC()
+            hdc.SelectObject(hbmp)
+            hdc.DrawIcon((0,0), hcursor)
+ 
+            bmpinfo = hbmp.GetInfo()
+            bmpstr = hbmp.GetBitmapBits(True)
+            cursor = Image.frombuffer('RGB', (bmpinfo['bmWidth'], bmpinfo['bmHeight']), bmpstr, 'raw', 'BGRX', 0, 1).convert("RGBA")
+ 
+            win32gui.DestroyIcon(hcursor)    
+            win32gui.DeleteObject(hbmp.GetHandle())
+            hdc.DeleteDC()
+ 
+            pixdata = cursor.load()
+ 
+            width, height = cursor.size
+            for y in range(height):
+                for x in range(width):
+                    if pixdata[x, y] == (0, 0, 0, 255):
+                        pixdata[x, y] = (0, 0, 0, 0)
+ 
+            hotspot = win32gui.GetIconInfo(hcursor)[1:3]
+ 
+            return (cursor, hotspot)
+ 
+        cursor, (hotspotx, hotspoty) = get_cursor()
+ 
+        ratio = ctypes.windll.shcore.GetScaleFactorForDevice(0) / 100
+ 
+        img = ImageGrab.grab(bbox=None, include_layered_windows=True)
+ 
+        pos_win = win32gui.GetCursorPos()
+        pos = (round(pos_win[0]*ratio - hotspotx), round(pos_win[1]*ratio - hotspoty))
+ 
+        img.paste(cursor, pos, cursor)
+ 
+        img.save(file_path)
     elif user_platform == "Linux":
         cursor_obj = Xcursor()
         imgarray = cursor_obj.getCursorImageArrayFast()