import json import logging import random from typing import Any, Dict, Optional import requests from desktop_env.envs.actions import KEYBOARD_KEYS logger = logging.getLogger("desktopenv.pycontroller") class PythonController: def __init__(self, vm_ip: str, pkgs_prefix: str = "python -c \"import pyautogui; {command}\""): self.vm_ip = vm_ip self.http_server = f"http://{vm_ip}:5000" self.pkgs_prefix = pkgs_prefix # fixme: this is a hacky way to execute python commands. fix it and combine it with installation of packages def get_screenshot(self): """ Gets a screenshot from the server. With the cursor. """ response = requests.get(self.http_server + "/screenshot") if response.status_code == 200: return response.content else: logger.error("Failed to get screenshot. Status code: %d", response.status_code) return None def get_accessibility_tree(self) -> Optional[str]: response: requests.Response = requests.get(self.http_server + "/accessibility") if response.status_code == 200: return response.json()["AT"] else: logger.error("Failed to get accessibility tree. Status code: %d", response.status_code) return None def get_file(self, file_path: str): """ Gets a file from the server. """ response = requests.post(self.http_server + "/file", data={"file_path": file_path}) if response.status_code == 200: logger.info("File downloaded successfully") return response.content else: logger.error("Failed to get file. Status code: %d", response.status_code) return None def execute_python_command(self, command: str) -> None: """ Executes a python command on the server. It can be used to execute the pyautogui commands, or... any other python command. who knows? """ command = self.pkgs_prefix.format(command=command) payload = json.dumps({"command": command, "shell": True}) headers = { 'Content-Type': 'application/json' } try: response = requests.post(self.http_server + "/execute", headers=headers, data=payload) if response.status_code == 200: logger.info("Command executed successfully: %s", response.text) else: logger.error("Failed to execute command. Status code: %d", response.status_code) return response.json() except requests.exceptions.RequestException as e: logger.error("An error occurred while trying to execute the command: %s", e) def execute_action(self, action: Dict[str, Any]): """ Executes an action on the server computer. """ action_type = action["action_type"] parameters = action["parameters"] if "parameters" in action else {} move_mode = random.choice( ["pyautogui.easeInQuad", "pyautogui.easeOutQuad", "pyautogui.easeInOutQuad", "pyautogui.easeInBounce", "pyautogui.easeInElastic"]) duration = random.uniform(0.5, 1) if action_type == "MOVE_TO": if parameters == {} or None: self.execute_python_command("pyautogui.moveTo()") elif "x" in parameters and "y" in parameters: x = parameters["x"] y = parameters["y"] self.execute_python_command(f"pyautogui.moveTo({x}, {y}, {duration}, {move_mode})") else: raise Exception(f"Unknown parameters: {parameters}") elif action_type == "CLICK": if parameters == {} or None: self.execute_python_command("pyautogui.click()") elif "button" in parameters and "x" in parameters and "y" in parameters: button = parameters["button"] x = parameters["x"] y = parameters["y"] if "num_clicks" in parameters: num_clicks = parameters["num_clicks"] self.execute_python_command( f"pyautogui.click(button='{button}', x={x}, y={y}, clicks={num_clicks})") else: self.execute_python_command(f"pyautogui.click(button='{button}', x={x}, y={y})") elif "button" in parameters and "x" not in parameters and "y" not in parameters: button = parameters["button"] if "num_clicks" in parameters: num_clicks = parameters["num_clicks"] self.execute_python_command(f"pyautogui.click(button='{button}', clicks={num_clicks})") else: self.execute_python_command(f"pyautogui.click(button='{button}')") elif "button" not in parameters and "x" in parameters and "y" in parameters: x = parameters["x"] y = parameters["y"] if "num_clicks" in parameters: num_clicks = parameters["num_clicks"] self.execute_python_command(f"pyautogui.click(x={x}, y={y}, clicks={num_clicks})") else: self.execute_python_command(f"pyautogui.click(x={x}, y={y})") else: raise Exception(f"Unknown parameters: {parameters}") elif action_type == "MOUSE_DOWN": if parameters == {} or None: self.execute_python_command("pyautogui.mouseDown()") elif "button" in parameters: button = parameters["button"] self.execute_python_command(f"pyautogui.mouseDown(button='{button}')") else: raise Exception(f"Unknown parameters: {parameters}") elif action_type == "MOUSE_UP": if parameters == {} or None: self.execute_python_command("pyautogui.mouseUp()") elif "button" in parameters: button = parameters["button"] self.execute_python_command(f"pyautogui.mouseUp(button='{button}')") else: raise Exception(f"Unknown parameters: {parameters}") elif action_type == "RIGHT_CLICK": if parameters == {} or None: self.execute_python_command("pyautogui.rightClick()") elif "x" in parameters and "y" in parameters: x = parameters["x"] y = parameters["y"] self.execute_python_command(f"pyautogui.rightClick(x={x}, y={y})") else: raise Exception(f"Unknown parameters: {parameters}") elif action_type == "DOUBLE_CLICK": if parameters == {} or None: self.execute_python_command("pyautogui.doubleClick()") elif "x" in parameters and "y" in parameters: x = parameters["x"] y = parameters["y"] self.execute_python_command(f"pyautogui.doubleClick(x={x}, y={y})") else: raise Exception(f"Unknown parameters: {parameters}") elif action_type == "DRAG_TO": if "x" in parameters and "y" in parameters: x = parameters["x"] y = parameters["y"] self.execute_python_command( f"pyautogui.dragTo({x}, {y}, duration=1.0, button='left', mouseDownUp=True)") elif action_type == "SCROLL": # todo: check if it is related to the operating system, as https://github.com/TheDuckAI/DuckTrack/blob/main/ducktrack/playback.py pointed out if "dx" in parameters and "dy" in parameters: dx = parameters["dx"] dy = parameters["dy"] self.execute_python_command(f"pyautogui.hscroll({dx})") self.execute_python_command(f"pyautogui.vscroll({dy})") elif "dx" in parameters and "dy" not in parameters: dx = parameters["dx"] self.execute_python_command(f"pyautogui.hscroll({dx})") elif "dx" not in parameters and "dy" in parameters: dy = parameters["dy"] self.execute_python_command(f"pyautogui.vscroll({dy})") else: raise Exception(f"Unknown parameters: {parameters}") elif action_type == "TYPING": if "text" not in parameters: raise Exception(f"Unknown parameters: {parameters}") text = parameters["text"] self.execute_python_command(f"pyautogui.typewrite('{text}')") elif action_type == "PRESS": if "key" not in parameters: raise Exception(f"Unknown parameters: {parameters}") key = parameters["key"] if key.lower() not in KEYBOARD_KEYS: raise Exception(f"Key must be one of {KEYBOARD_KEYS}") self.execute_python_command(f"pyautogui.press('{key}')") elif action_type == "KEY_DOWN": if "key" not in parameters: raise Exception(f"Unknown parameters: {parameters}") key = parameters["key"] if key.lower() not in KEYBOARD_KEYS: raise Exception(f"Key must be one of {KEYBOARD_KEYS}") self.execute_python_command(f"pyautogui.keyDown('{key}')") elif action_type == "KEY_UP": if "key" not in parameters: raise Exception(f"Unknown parameters: {parameters}") key = parameters["key"] if key.lower() not in KEYBOARD_KEYS: raise Exception(f"Key must be one of {KEYBOARD_KEYS}") self.execute_python_command(f"pyautogui.keyUp('{key}')") elif action_type == "HOTKEY": if "keys" not in parameters: raise Exception(f"Unknown parameters: {parameters}") keys = parameters["keys"] if not isinstance(keys, list): raise Exception("Keys must be a list of keys") for key in keys: if key.lower() not in KEYBOARD_KEYS: raise Exception(f"Key must be one of {KEYBOARD_KEYS}") keys_para_rep = "', '".join(keys) self.execute_python_command(f"pyautogui.hotkey('{keys_para_rep}')") else: raise Exception(f"Unknown action type: {action_type}") # Additional info def get_vm_platform(self): """ Gets the size of the vm screen. """ return self.execute_python_command("import platform; print(platform.system())")['output'].strip() def get_vm_screen_size(self): """ Gets the size of the vm screen. """ response = requests.post(self.http_server + "/screen_size") if response.status_code == 200: return response.json() else: logger.error("Failed to get screen size. Status code: %d", response.status_code) return None def get_vm_window_size(self, app_class_name: str): """ Gets the size of the vm app window. """ response = requests.post(self.http_server + "/window_size", data={"app_class_name": app_class_name}) if response.status_code == 200: return response.json() else: logger.error("Failed to get window size. Status code: %d", response.status_code) return None def get_vm_wallpaper(self): """ Gets the wallpaper of the vm. """ response = requests.post(self.http_server + "/wallpaper") if response.status_code == 200: logger.info("Wallpaper downloaded successfully") return response.content else: logger.error("Failed to get wallpaper. Status code: %d", response.status_code) return None