Refactor with pyautogui

desktop_env/server/main.py

@@ -0,0 +1,64 @@
+import os
+import platform
+import subprocess
+
+import Xlib.display
+import pyautogui
+from PIL import ImageGrab
+from flask import Flask, request, jsonify, send_file
+
+app = Flask(__name__)
+
+pyautogui.PAUSE = 0
+pyautogui.DARWIN_CATCH_UP_TIME = 0
+
+
+@app.route('/execute', methods=['POST'])
+def execute_command():
+    data = request.json
+    # The 'command' key in the JSON request should contain the command to be executed.
+    command = data.get('command', '')
+
+    # Execute the command without any safety checks.
+    try:
+        result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        return jsonify({
+            'status': 'success',
+            'output': result.stdout,
+            'error': result.stderr
+        })
+    except Exception as e:
+        return jsonify({
+            'status': 'error',
+            'message': str(e)
+        }), 500
+
+
+@app.route('/screenshot', methods=['GET'])
+def capture_screen_with_cursor():
+    file_path = os.path.join("screenshots", "screenshot.png")
+    user_platform = platform.system()
+
+    # Ensure the screenshots directory exists
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+    if user_platform == "Windows":
+        screenshot = pyautogui.screenshot()
+        screenshot.save(file_path)
+    elif user_platform == "Linux":
+        # Use xlib to prevent scrot dependency for Linux
+        screen = Xlib.display.Display().screen()
+        size = screen.width_in_pixels, screen.height_in_pixels
+        screenshot = ImageGrab.grab(bbox=(0, 0, size[0], size[1]))
+        screenshot.save(file_path)
+    elif user_platform == "Darwin":  # (Mac OS)
+        # Use the screencapture utility to capture the screen with the cursor
+        subprocess.run(["screencapture", "-C", file_path])
+    else:
+        print(f"The platform you're using ({user_platform}) is not currently supported")
+
+    return send_file(file_path, mimetype='image/png')
+
+
+if __name__ == '__main__':
+    app.run(debug=True, host="0.0.0.0")