Update app.py

app.py

@@ -1,12 +1,10 @@
 """
-THE Z AI — Computer Mode Server
-================================
-سيرفر التحكم بالكمبيوتر عن بعد للذكاء الاصطناعي
-يوفر:
-  - WebSocket للتحكم اللحظي (الماوس، لوحة المفاتيح، الترمينال)
-  - تصوير سكرين شوت وإرساله للعميل
-  - بث شاشة مباشر (base64 frames)
-  - تنفيذ أوامر bash وقراءة الناتج
+THE Z AI — Computer Mode Server (FIXED v2)
+==========================================
+إصلاحات:
+  - التقاط الشاشة الحقيقي عبر scrot (يعمل مع Xvfb)
+  - فتح التطبيقات مع انتظار حقيقي
+  - تسجيل دقيق لما يحدث فعلاً
 """
 
 import asyncio
@@ -15,14 +13,9 @@ import io
 import json
 import os
 import subprocess
-import threading
 import time
 import traceback
-from pathlib import Path
 
-import pyautogui
-import pyperclip
-from PIL import Image, ImageGrab
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
@@ -31,11 +24,8 @@ import uvicorn
 # ─── ENV ────────────────────────────────────────────
 DISPLAY = os.environ.get("DISPLAY", ":1")
 os.environ["DISPLAY"] = DISPLAY
-pyautogui.FAILSAFE = False
-pyautogui.PAUSE = 0.05
-
-app = FastAPI(title="Z-Computer-Mode API")
 
+app = FastAPI(title="Z-Computer-Mode API v2")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -44,20 +34,67 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ─── Active Connections ──────────────────────────────
 active_connections: list[WebSocket] = []
 stream_active = False
-stream_quality = 60   # JPEG quality
-stream_fps    = 3     # frames per second (low to save bandwidth)
-stream_scale  = 0.5   # downscale factor
+stream_fps = 3
+stream_quality = 60
+stream_scale = 0.5
+
+
+# ─── Screen Capture via scrot (works with Xvfb) ────
+
+def capture_screen(scale=0.5, quality=60) -> str:
+    """Capture real screen using scrot → base64 JPEG"""
+    try:
+        tmp = f"/tmp/zscreen_{int(time.time()*1000)}.png"
+        result = subprocess.run(
+            ["scrot", "-q", "90", tmp],
+            env={**os.environ, "DISPLAY": DISPLAY},
+            timeout=8, capture_output=True
+        )
+        if result.returncode != 0 or not os.path.exists(tmp):
+            # fallback: xwd → PIL
+            return capture_screen_xwd(scale, quality)
+
+        from PIL import Image
+        img = Image.open(tmp)
+        os.unlink(tmp)
+
+        if scale < 1.0:
+            w = int(img.width * scale)
+            h = int(img.height * scale)
+            img = img.resize((w, h), Image.LANCZOS)
 
+        buf = io.BytesIO()
+        img.convert("RGB").save(buf, format="JPEG", quality=quality, optimize=True)
+        return base64.b64encode(buf.getvalue()).decode()
+    except Exception as e:
+        print(f"[capture_screen] error: {e}")
+        return ""
 
-# ─── Utility ─────────────────────────────────────────
 
-def capture_screen(scale=stream_scale, quality=stream_quality) -> str:
-    """Capture screen → base64 JPEG string"""
+def capture_screen_xwd(scale=0.5, quality=60) -> str:
+    """Fallback: use xwd + PIL"""
     try:
-        img = ImageGrab.grab()
+        tmp_xwd = f"/tmp/zscreen_{int(time.time()*1000)}.xwd"
+        tmp_png = tmp_xwd.replace(".xwd", ".png")
+        subprocess.run(
+            ["xwd", "-root", "-silent", "-out", tmp_xwd],
+            env={**os.environ, "DISPLAY": DISPLAY},
+            timeout=8, capture_output=True
+        )
+        subprocess.run(
+            ["convert", tmp_xwd, tmp_png],
+            timeout=8, capture_output=True
+        )
+        from PIL import Image
+        if not os.path.exists(tmp_png):
+            return ""
+        img = Image.open(tmp_png)
+        for f in [tmp_xwd, tmp_png]:
+            try: os.unlink(f)
+            except: pass
+
         if scale < 1.0:
             w = int(img.width * scale)
             h = int(img.height * scale)
@@ -66,52 +103,62 @@ def capture_screen(scale=stream_scale, quality=stream_quality) -> str:
         img.convert("RGB").save(buf, format="JPEG", quality=quality, optimize=True)
         return base64.b64encode(buf.getvalue()).decode()
     except Exception as e:
+        print(f"[capture_screen_xwd] error: {e}")
         return ""
 
 
+# ─── Mouse/Keyboard via xdotool ─────────────────────
+
+def xdo(args: list, timeout=10) -> dict:
+    result = subprocess.run(
+        ["xdotool"] + args,
+        env={**os.environ, "DISPLAY": DISPLAY},
+        timeout=timeout, capture_output=True, text=True
+    )
+    return {"rc": result.returncode, "out": result.stdout, "err": result.stderr}
+
+
 def run_command(cmd: str, timeout: int = 30) -> dict:
-    """Run a shell command and return stdout/stderr"""
     try:
         result = subprocess.run(
             cmd, shell=True, capture_output=True, text=True,
             timeout=timeout, env={**os.environ, "DISPLAY": DISPLAY}
         )
         return {
-            "stdout": result.stdout[-8000:] if len(result.stdout) > 8000 else result.stdout,
-            "stderr": result.stderr[-4000:] if len(result.stderr) > 4000 else result.stderr,
+            "stdout": result.stdout[-8000:],
+            "stderr": result.stderr[-4000:],
             "returncode": result.returncode,
         }
     except subprocess.TimeoutExpired:
-        return {"stdout": "", "stderr": f"Command timed out after {timeout}s", "returncode": -1}
+        return {"stdout": "", "stderr": f"Timeout after {timeout}s", "returncode": -1}
     except Exception as e:
         return {"stdout": "", "stderr": str(e), "returncode": -1}
 
 
 async def broadcast(msg: dict):
-    """Send JSON message to all connected WebSocket clients"""
     txt = json.dumps(msg)
-    disconnected = []
+    dead = []
     for ws in active_connections:
         try:
             await ws.send_text(txt)
-        except Exception:
-            disconnected.append(ws)
-    for ws in disconnected:
+        except:
+            dead.append(ws)
+    for ws in dead:
         if ws in active_connections:
             active_connections.remove(ws)
 
 
-# ─── Screen Streaming Loop ───────────────────────────
+# ─── Stream Loop ─────────────────────────────────────
 
 async def screen_stream_loop():
     global stream_active
-    interval = 1.0 / stream_fps
+    interval = 1.0 / max(1, stream_fps)
     while stream_active and active_connections:
         try:
-            frame = capture_screen()
+            frame = capture_screen(scale=stream_scale, quality=stream_quality)
             if frame:
                 await broadcast({"type": "frame", "data": frame})
-        except Exception:
+        except:
             pass
         await asyncio.sleep(interval)
     stream_active = False
@@ -121,107 +168,120 @@ async def screen_stream_loop():
 
 async def handle_action(ws: WebSocket, msg: dict):
     action = msg.get("action", "")
-    data   = msg.get("data", {})
+    data = msg.get("data", {})
+
+    async def send(obj):
+        await ws.send_text(json.dumps(obj))
 
-    # ── SCREENSHOT ──────────────────────────────────
+    async def auto_screenshot(delay=0.8, sc=0.7, q=70):
+        await asyncio.sleep(delay)
+        frame = capture_screen(scale=sc, quality=q)
+        if frame:
+            await send({"type": "screenshot", "data": frame,
+                        "ts": int(time.time() * 1000), "auto": True})
+
+    # ── SCREENSHOT ──
     if action == "screenshot":
         frame = capture_screen(scale=0.8, quality=75)
-        await ws.send_text(json.dumps({
-            "type": "screenshot",
-            "data": frame,
-            "ts": int(time.time() * 1000)
-        }))
+        await send({"type": "screenshot", "data": frame, "ts": int(time.time() * 1000)})
 
-    # ── TERMINAL COMMAND ─────────────────────────────
+    # ── TERMINAL ──
     elif action == "terminal":
         cmd = data.get("cmd", "")
         if cmd:
-            result = run_command(cmd, timeout=data.get("timeout", 30))
-            await ws.send_text(json.dumps({
+            res = run_command(cmd, timeout=data.get("timeout", 30))
+            await send({
                 "type": "terminal_result",
                 "cmd": cmd,
-                "stdout": result["stdout"],
-                "stderr": result["stderr"],
-                "returncode": result["returncode"],
-            }))
-            # Auto-screenshot after terminal command
-            await asyncio.sleep(0.5)
-            frame = capture_screen(scale=0.7, quality=70)
-            if frame:
-                await ws.send_text(json.dumps({
-                    "type": "screenshot",
-                    "data": frame,
-                    "ts": int(time.time() * 1000),
-                    "auto": True
-                }))
-
-    # ── MOUSE MOVE ───────────────────────────────────
+                "stdout": res["stdout"],
+                "stderr": res["stderr"],
+                "returncode": res["returncode"],
+            })
+            await auto_screenshot(0.6)
+
+    # ── MOUSE MOVE ──
     elif action == "mouse_move":
         x, y = int(data.get("x", 0)), int(data.get("y", 0))
-        pyautogui.moveTo(x, y, duration=0.1)
-        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_move"}))
+        xdo(["mousemove", str(x), str(y)])
+        await send({"type": "ack", "action": "mouse_move"})
 
-    # ���─ MOUSE CLICK ──────────────────────────────────
+    # ── MOUSE CLICK ──
     elif action == "mouse_click":
-        x      = int(data.get("x", 0))
-        y      = int(data.get("y", 0))
+        x, y = int(data.get("x", 0)), int(data.get("y", 0))
         button = data.get("button", "left")
+        btn_num = {"left": "1", "middle": "2", "right": "3"}.get(button, "1")
         double = data.get("double", False)
-        pyautogui.moveTo(x, y, duration=0.08)
+        xdo(["mousemove", str(x), str(y)])
+        await asyncio.sleep(0.05)
         if double:
-            pyautogui.doubleClick(x, y, button=button)
+            xdo(["click", "--repeat", "2", "--delay", "100", btn_num])
         else:
-            pyautogui.click(x, y, button=button)
-        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_click"}))
+            xdo(["click", btn_num])
+        await send({"type": "ack", "action": "mouse_click"})
+        await auto_screenshot(0.5)
 
-    # ── MOUSE DRAG ───────────────────────────────────
+    # ── MOUSE DRAG ──
     elif action == "mouse_drag":
         x1, y1 = int(data.get("x1", 0)), int(data.get("y1", 0))
         x2, y2 = int(data.get("x2", 0)), int(data.get("y2", 0))
-        pyautogui.moveTo(x1, y1)
-        pyautogui.dragTo(x2, y2, duration=0.3, button="left")
-        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_drag"}))
-
-    # ── KEYBOARD TYPE ────────────────────────────────
+        xdo(["mousemove", str(x1), str(y1)])
+        xdo(["mousedown", "1"])
+        await asyncio.sleep(0.1)
+        xdo(["mousemove", str(x2), str(y2)])
+        await asyncio.sleep(0.1)
+        xdo(["mouseup", "1"])
+        await send({"type": "ack", "action": "mouse_drag"})
+
+    # ── KEYBOARD TYPE ──
     elif action == "keyboard_type":
         text = data.get("text", "")
-        pyautogui.typewrite(text, interval=0.03)
-        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_type"}))
+        if text:
+            # Use xdotool type for proper Unicode support
+            xdo(["type", "--clearmodifiers", "--delay", "30", text])
+        await send({"type": "ack", "action": "keyboard_type"})
 
-    # ── KEYBOARD HOTKEY ──────────────────────────────
+    # ── KEYBOARD HOTKEY ──
     elif action == "keyboard_hotkey":
         keys = data.get("keys", [])
         if keys:
-            pyautogui.hotkey(*keys)
-        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_hotkey"}))
+            combo = "+".join(keys)
+            xdo(["key", "--clearmodifiers", combo])
+        await send({"type": "ack", "action": "keyboard_hotkey"})
 
-    # ── KEYBOARD PRESS ───────────────────────────────
+    # ── KEYBOARD PRESS ──
     elif action == "keyboard_press":
         key = data.get("key", "")
         if key:
-            pyautogui.press(key)
-        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_press"}))
+            xdo(["key", "--clearmodifiers", key])
+        await send({"type": "ack", "action": "keyboard_press"})
 
-    # ── CLIPBOARD WRITE ──────────────────────────────
+    # ── CLIPBOARD ──
     elif action == "clipboard_write":
         text = data.get("text", "")
-        pyperclip.copy(text)
-        await ws.send_text(json.dumps({"type": "ack", "action": "clipboard_write"}))
+        proc = subprocess.Popen(
+            ["xclip", "-selection", "clipboard"],
+            stdin=subprocess.PIPE,
+            env={**os.environ, "DISPLAY": DISPLAY}
+        )
+        proc.communicate(text.encode())
+        await send({"type": "ack", "action": "clipboard_write"})
 
-    # ── CLIPBOARD READ ───────────────────────────────
     elif action == "clipboard_read":
-        text = pyperclip.paste()
-        await ws.send_text(json.dumps({"type": "clipboard_content", "text": text}))
+        res = run_command("xclip -selection clipboard -o", timeout=5)
+        await send({"type": "clipboard_content", "text": res["stdout"]})
 
-    # ── SCROLL ───────────────────────────────────────
+    # ── SCROLL ──
     elif action == "scroll":
-        x      = int(data.get("x", 0))
-        y      = int(data.get("y", 0))
+        x = int(data.get("x", 0))
+        y = int(data.get("y", 0))
         clicks = int(data.get("clicks", 3))
-        pyautogui.scroll(clicks, x=x, y=y)
-        await ws.send_text(json.dumps({"type": "ack", "action": "scroll"}))
+        btn = "4" if clicks > 0 else "5"
+        xdo(["mousemove", str(x), str(y)])
+        for _ in range(abs(clicks)):
+            xdo(["click", btn])
+        await send({"type": "ack", "action": "scroll"})
 
-    # ── OPEN APP ─────────────────────────────────────
+    # ── OPEN APP ──
     elif action == "open_app":
         app_cmd = data.get("cmd", "")
         if app_cmd:
@@ -230,49 +290,51 @@ async def handle_action(ws: WebSocket, msg: dict):
                 env={**os.environ, "DISPLAY": DISPLAY},
                 stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
             )
-        await ws.send_text(json.dumps({"type": "ack", "action": "open_app"}))
-        # Auto-screenshot after app opens (wait for it to load)
-        await asyncio.sleep(3)
-        frame = capture_screen(scale=0.7, quality=70)
-        if frame:
-            await ws.send_text(json.dumps({
-                "type": "screenshot",
-                "data": frame,
-                "ts": int(time.time() * 1000),
-                "auto": True
-            }))
-
-    # ── START STREAM ─────────────────────────────────
+            await send({"type": "ack", "action": "open_app", "cmd": app_cmd})
+            # Wait longer for app to open then take screenshot
+            await asyncio.sleep(4)
+            frame = capture_screen(scale=0.75, quality=72)
+            if frame:
+                await send({"type": "screenshot", "data": frame,
+                            "ts": int(time.time() * 1000), "auto": True, "label": f"بعد فتح: {app_cmd}"})
+        else:
+            await send({"type": "ack", "action": "open_app"})
+
+    # ── STREAM ──
     elif action == "start_stream":
         global stream_active, stream_fps, stream_quality, stream_scale
-        stream_fps     = int(data.get("fps", 3))
+        stream_fps = int(data.get("fps", 3))
         stream_quality = int(data.get("quality", 60))
-        stream_scale   = float(data.get("scale", 0.5))
+        stream_scale = float(data.get("scale", 0.5))
         if not stream_active:
             stream_active = True
             asyncio.create_task(screen_stream_loop())
-        await ws.send_text(json.dumps({"type": "ack", "action": "start_stream"}))
+        await send({"type": "ack", "action": "start_stream"})
 
-    # ── STOP STREAM ──────────────────────────────────
     elif action == "stop_stream":
         stream_active = False
-        await ws.send_text(json.dumps({"type": "ack", "action": "stop_stream"}))
+        await send({"type": "ack", "action": "stop_stream"})
 
-    # ── GET SCREEN INFO ──────────────────────────────
+    # ── SCREEN INFO ──
     elif action == "screen_info":
+        res = run_command("xdotool getdisplaygeometry", timeout=5)
         try:
-            w, h = pyautogui.size()
-            mx, my = pyautogui.position()
+            parts = res["stdout"].strip().split()
+            w, h = int(parts[0]), int(parts[1])
         except:
-            w, h, mx, my = 1920, 1080, 0, 0
-        await ws.send_text(json.dumps({
-            "type": "screen_info",
-            "width": w, "height": h,
-            "mouse_x": mx, "mouse_y": my,
-        }))
+            w, h = 1920, 1080
+        res2 = run_command("xdotool getmouselocation", timeout=5)
+        try:
+            import re
+            mx = int(re.search(r"x:(\d+)", res2["stdout"]).group(1))
+            my = int(re.search(r"y:(\d+)", res2["stdout"]).group(1))
+        except:
+            mx, my = 0, 0
+        await send({"type": "screen_info", "width": w, "height": h,
+                    "mouse_x": mx, "mouse_y": my})
 
     else:
-        await ws.send_text(json.dumps({"type": "error", "msg": f"Unknown action: {action}"}))
+        await send({"type": "error", "msg": f"Unknown action: {action}"})
 
 
 # ─── WebSocket Endpoint ──────────────────────────────
@@ -281,17 +343,25 @@ async def handle_action(ws: WebSocket, msg: dict):
 async def websocket_endpoint(ws: WebSocket):
     await ws.accept()
     active_connections.append(ws)
-    # Send welcome
+    # Send welcome + first real screenshot
+    res = run_command("xdotool getdisplaygeometry", timeout=5)
     try:
-        w, h = pyautogui.size()
+        parts = res["stdout"].strip().split()
+        w, h = int(parts[0]), int(parts[1])
     except:
         w, h = 1920, 1080
     await ws.send_text(json.dumps({
         "type": "connected",
-        "screen_width": w,
-        "screen_height": h,
-        "msg": "Z Computer Mode — Connected"
+        "screen_width": w, "screen_height": h,
+        "msg": "Z Computer Mode v2 — Connected"
     }))
+    # Send initial screenshot immediately
+    frame = capture_screen(scale=0.75, quality=72)
+    if frame:
+        await ws.send_text(json.dumps({
+            "type": "screenshot", "data": frame,
+            "ts": int(time.time() * 1000), "label": "Initial screen"
+        }))
     try:
         while True:
             raw = await ws.receive_text()
@@ -300,13 +370,13 @@ async def websocket_endpoint(ws: WebSocket):
     except WebSocketDisconnect:
         pass
     except Exception as e:
-        pass
+        print(f"[ws] error: {e}")
     finally:
         if ws in active_connections:
             active_connections.remove(ws)
 
 
-# ─── REST Fallback ───────────────────────────────────
+# ─── REST ────────────────────────────────────────────
 
 @app.get("/screenshot")
 async def rest_screenshot():
@@ -316,18 +386,19 @@ async def rest_screenshot():
 
 @app.post("/terminal")
 async def rest_terminal(body: dict):
-    cmd    = body.get("cmd", "")
-    result = run_command(cmd, timeout=body.get("timeout", 30))
-    return JSONResponse(result)
+    cmd = body.get("cmd", "")
+    res = run_command(cmd, timeout=body.get("timeout", 30))
+    return JSONResponse(res)
 
 
 @app.get("/health")
 async def health():
-    return {"status": "ok", "display": DISPLAY}
-
+    # Check if display is really working
+    res = run_command("xdotool getdisplaygeometry", timeout=4)
+    display_ok = res["returncode"] == 0
+    return {"status": "ok", "display": DISPLAY, "display_working": display_ok}
 
-# ─── Entry ───────────────────────────────────────────
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
-    uvicorn.run("app:app", host="0.0.0.0", port=port, log_level="warning")
+    uvicorn.run("app:app", host="0.0.0.0", port=port, log_level="info")