Update app.py

app.py

@@ -12,6 +12,7 @@ THE Z AI — Computer Mode Server v7 — SMART CONTROL
 
 import asyncio
 import base64
+import contextvars
 import io
 import json
 import os
@@ -25,23 +26,159 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 import uvicorn
 
-DISPLAY = os.environ.get("DISPLAY", ":1")
-os.environ["DISPLAY"] = DISPLAY
+# ─── DISPLAY الأساسي (احتياطي فقط — لا يُستخدم فعلياً بعد تفعيل العزل الكامل) ──
+_BASE_DISPLAY = os.environ.get("DISPLAY", ":1")
+os.environ["DISPLAY"] = _BASE_DISPLAY
+
+# ─── عزل كامل بين المستخدمين: كل اتصال WebSocket يحصل على Xvfb + Firefox خاصين به ──
+# يُخزَّن DISPLAY الحالي في contextvar بحيث كل الدوال (xdo, capture_screen_with_grid, ...)
+# التي تُستدعى من داخل معالجة رسائل هذا الاتصال تقرأ تلقائياً display الخاص بجلسته،
+# دون الحاجة لتمرير المعامل يدوياً عبر كل سلسلة الاستدعاءات.
+_current_display: contextvars.ContextVar[str] = contextvars.ContextVar(
+    "_current_display", default=_BASE_DISPLAY
+)
+
+def DISPLAY_OF() -> str:
+    """يُعيد DISPLAY الخاص بالجلسة الحالية (الاتصال الذي تجري معالجته الآن)."""
+    return _current_display.get()
+
+# نطاق أرقام الشاشات الافتراضية الديناميكية المخصصة لكل جلسة — تبدأ من 100 لتجنب التعارض
+# مع أي DISPLAY أساسي ثابت (:0, :1, :99) قد يكون مستخدَماً خارجياً.
+_next_display_num = 100
+_display_lock = asyncio.Lock()
+# session_id (id(ws)) → {"display": ":N", "xvfb_proc": Popen, "firefox_proc": Popen, "num": N}
+_active_sessions: dict[int, dict] = {}
+_freed_display_nums: list[int] = []  # أرقام شاشات تحررت يمكن إعادة استخدامها
+
+async def _allocate_display_num() -> int:
+    """يحجز رقم شاشة فريد لجلسة جديدة — يعيد استخدام أرقام محرَّرة من جلسات منتهية أولاً."""
+    global _next_display_num
+    async with _display_lock:
+        if _freed_display_nums:
+            return _freed_display_nums.pop()
+        n = _next_display_num
+        _next_display_num += 1
+        return n
+
+async def _start_session_display(num: int) -> tuple[str, "subprocess.Popen | None"]:
+    """يُشغّل Xvfb جديداً تماماً على الرقم المحدد ويُعيد (display_string, process)."""
+    disp = f":{num}"
+
+    def _spawn():
+        try:
+            return subprocess.Popen(
+                ["Xvfb", disp, "-screen", "0", "1920x1080x24", "-nolisten", "tcp"],
+                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+            )
+        except FileNotFoundError:
+            return None  # Xvfb غير مثبت في هذه البيئة
+
+    proc = await asyncio.to_thread(_spawn)
+    if proc is None:
+        # Xvfb غير متاح — رجوع آمن لـ DISPLAY الأساسي (بلا عزل، لكن بلا انهيار)
+        print("[session] ⚠️ Xvfb binary not found — falling back to shared base display (no isolation)")
+        return _BASE_DISPLAY, None
+    # انتظار قصير حتى يصبح X server جاهزاً فعلياً قبل استخدامه
+    def _check_ready():
+        try:
+            return subprocess.run(
+                ["xdpyinfo", "-display", disp], capture_output=True, timeout=2
+            ).returncode
+        except FileNotFoundError:
+            return None  # xdpyinfo غير مثبت — لا يمكن التحقق، سننتظر مهلة ثابتة فقط
+    for _ in range(30):
+        await asyncio.sleep(0.1)
+        rc = await asyncio.to_thread(_check_ready)
+        if rc is None:
+            await asyncio.sleep(0.5)  # مهلة احتياطية ثابتة بدل التحقق الدقيق
+            break
+        if rc == 0:
+            break
+    return disp, proc
+
+async def _start_session_browser(disp: str) -> "subprocess.Popen | None":
+    """يُشغّل متصفح Firefox خاص بهذه الجلسة على الشاشة الافتراضية الخاصة بها."""
+    env = {**os.environ, "DISPLAY": disp}
+
+    def _spawn():
+        try:
+            return subprocess.Popen(
+                [BROWSER, "--new-window", "--no-remote", "about:blank"],
+                env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+            )
+        except Exception as e:
+            print(f"[session] ⚠️ failed to launch browser on {disp}: {e}")
+            return None
+    proc = await asyncio.to_thread(_spawn)
+    return proc
+
+async def create_session_for(ws: "WebSocket") -> str:
+    """
+    يُنشئ بيئة كمبيوتر افتراضي معزولة تماماً (Xvfb + متصفح خاصين) لهذا الاتصال،
+    ويُعيد سلسلة DISPLAY الخاصة بها (مثل ':103'). يُ��تدعى عند فتح أي اتصال WebSocket جديد.
+    """
+    num = await _allocate_display_num()
+    disp, xvfb_proc = await _start_session_display(num)
+    if xvfb_proc is None:
+        # رجعنا لـ DISPLAY الأساسي المشترك (Xvfb غير متاح) — لا حاجة لمتصفح جديد منفصل،
+        # المتصفح الأساسي (إن وُجد خارجياً) يكفي، ولا داعي لتتبع جلسة فعلية لهذا الاتصال.
+        _freed_display_nums.append(num)
+        return disp
+    firefox_proc = await _start_session_browser(disp)
+    _active_sessions[id(ws)] = {
+        "display": disp, "num": num,
+        "xvfb_proc": xvfb_proc, "firefox_proc": firefox_proc,
+    }
+    print(f"[session] ✅ created isolated session display={disp} for ws={id(ws)} (total active: {len(_active_sessions)})")
+    return disp
+
+async def destroy_session_for(ws: "WebSocket"):
+    """يُغلق ويُحرّر بيئة الكمبيوتر الافتراضي الخاصة بهذا الاتصال بالكامل عند انقطاعه."""
+    sess = _active_sessions.pop(id(ws), None)
+    if not sess:
+        return
+
+    def _kill_all():
+        for key in ("firefox_proc", "xvfb_proc"):
+            p = sess.get(key)
+            if p:
+                try:
+                    p.terminate()
+                    try:
+                        p.wait(timeout=3)
+                    except Exception:
+                        p.kill()
+                except Exception:
+                    pass
+        # تأكيد إضافي: قتل أي عملية متبقية مرتبطة بهذا DISPLAY بالتحديد فقط
+        disp = sess.get("display", "")
+        if disp:
+            try:
+                subprocess.run(
+                    ["pkill", "-9", "-f", f"DISPLAY={disp}"],
+                    capture_output=True, timeout=3
+                )
+            except Exception:
+                pass
+
+    await asyncio.to_thread(_kill_all)
+    _freed_display_nums.append(sess["num"])
+    print(f"[session] 🗑️ destroyed session display={sess['display']} for ws={id(ws)} (remaining active: {len(_active_sessions)})")
 
 # ─── تشخيص أدوات الـ capture عند البداية ─────────────
 def _check_capture_tools():
-    tools = ["scrot", "import", "xwd", "convert", "ffmpeg", "xdotool"]
+    tools = ["scrot", "import", "xwd", "convert", "ffmpeg", "xdotool", "Xvfb", "xdpyinfo"]
     print("─── Capture Tools Check ───")
     for t in tools:
         r = subprocess.run(["which", t], capture_output=True, text=True)
         status = "✅" if r.returncode == 0 else "❌"
         print(f"  {status} {t}: {r.stdout.strip() or 'not found'}")
     try:
-        r2 = subprocess.run(["xdpyinfo", "-display", DISPLAY],
+        r2 = subprocess.run(["xdpyinfo", "-display", _BASE_DISPLAY],
                             capture_output=True, text=True, timeout=3)
-        print(f"  {'✅' if r2.returncode==0 else '❌'} DISPLAY={DISPLAY}: {'active' if r2.returncode==0 else 'not active'}")
+        print(f"  {'✅' if r2.returncode==0 else '❌'} DISPLAY={_BASE_DISPLAY}: {'active' if r2.returncode==0 else 'not active'}")
     except Exception as ex:
-        print(f"  ❌ DISPLAY={DISPLAY}: {ex}")
+        print(f"  ❌ DISPLAY={_BASE_DISPLAY}: {ex}")
     print("───────────────────────────")
 
 try:
@@ -184,9 +321,9 @@ def _is_empty_result(stdout: str) -> bool:
 # ─── Screen Capture ─────────────────────────────────
 
 def _get_screen_size() -> tuple[int, int]:
-    """يُعيد حجم الشاشة الحقيقي"""
+    """يُعيد حجم الشاشة الحقيقي (لشاشة الجلسة الحالية)"""
     r = subprocess.run(["xdotool", "getdisplaygeometry"],
-                       env={**os.environ, "DISPLAY": DISPLAY},
+                       env={**os.environ, "DISPLAY": DISPLAY_OF()},
                        capture_output=True, text=True, timeout=5)
     try:
         parts = r.stdout.strip().split()
@@ -196,9 +333,9 @@ def _get_screen_size() -> tuple[int, int]:
 
 
 def _get_mouse_pos() -> tuple[int, int]:
-    """يُعيد موقع الماوس الحالي"""
+    """يُعيد موقع الماوس الحالي (لشاشة الجلسة الحالية)"""
     r = subprocess.run(["xdotool", "getmouselocation"],
-                       env={**os.environ, "DISPLAY": DISPLAY},
+                       env={**os.environ, "DISPLAY": DISPLAY_OF()},
                        capture_output=True, text=True, timeout=5)
     try:
         mx = int(re.search(r"x:(\d+)", r.stdout).group(1))
@@ -220,7 +357,8 @@ def capture_screen_raw() -> tuple[object, int, int]:
     5. python-xlib / Xlib مباشرة
     """
     from PIL import Image
-    env = {**os.environ, "DISPLAY": DISPLAY}
+    disp = DISPLAY_OF()
+    env = {**os.environ, "DISPLAY": disp}
     tmp = f"/tmp/zs_{int(time.time()*1000)}"
     _tmp_files_to_clean = []  # كل الملفات المؤقتة التي أُنشئت في هذه المحاولة
 
@@ -314,7 +452,7 @@ def capture_screen_raw() -> tuple[object, int, int]:
                 "ffmpeg", "-y",
                 "-f", "x11grab",
                 "-video_size", f"{sw}x{sh}",
-                "-i", DISPLAY,
+                "-i", disp,
                 "-vframes", "1",
                 "-q:v", "2",
                 p
@@ -331,7 +469,7 @@ def capture_screen_raw() -> tuple[object, int, int]:
         try:
             from Xlib import display as Xdisplay
             from Xlib.ext.xtest import fake_input
-            xdisp = Xdisplay.Display(DISPLAY)
+            xdisp = Xdisplay.Display(disp)
             root = xdisp.screen().root
             geom = root.get_geometry()
             w, h = geom.width, geom.height
@@ -359,7 +497,7 @@ def capture_screen_raw() -> tuple[object, int, int]:
             img = Image.new("RGB", (sw or 1280, sh or 720), (20, 20, 30))
             draw = ImageDraw.Draw(img)
             draw.rectangle([(0, 0), (sw, 40)], fill=(30, 30, 50))
-            draw.text((10, 10), f"⚠️ Screenshot failed — DISPLAY={DISPLAY}", fill=(255, 100, 100))
+            draw.text((10, 10), f"⚠️ Screenshot failed — DISPLAY={disp}", fill=(255, 100, 100))
             draw.text((10, 50), "Methods tried: scrot, import, xwd, ffmpeg, xlib", fill=(150, 150, 150))
             return img, sw or 1280, sh or 720
         except:
@@ -511,7 +649,7 @@ def _sanitize_pkill_self_match(cmd: str) -> str:
 
 def run_raw_command(cmd: str, timeout: int = 60) -> dict:
     cmd = _sanitize_pkill_self_match(cmd)
-    env = {**os.environ, "DISPLAY": DISPLAY,
+    env = {**os.environ, "DISPLAY": DISPLAY_OF(),
            "PYTHONIOENCODING": "utf-8", "LANG": "en_US.UTF-8"}
     try:
         result = subprocess.run(cmd, shell=True, capture_output=True,
@@ -591,9 +729,10 @@ def run_command_smart(cmd: str, timeout: int = 60) -> dict:
 
 async def xdo(args: list, timeout=10) -> dict:
     """ينفّذ xdotool في thread منفصل حتى لا يُجمّد حلقة asyncio أثناء الانتظار."""
+    disp = DISPLAY_OF()
     def _run():
         r = subprocess.run(["xdotool"] + args,
-                           env={**os.environ, "DISPLAY": DISPLAY},
+                           env={**os.environ, "DISPLAY": disp},
                            timeout=timeout, capture_output=True, text=True)
         return {"rc": r.returncode, "out": r.stdout, "err": r.stderr}
     return await asyncio.to_thread(_run)
@@ -610,11 +749,12 @@ async def type_text_smart(text: str) -> dict:
     if has_arabic:
         # طريقة Clipboard لضمان كتابة العربية بشكل صحيح
         try:
+            disp = DISPLAY_OF()
             def _xclip_write():
                 proc = subprocess.Popen(
                     ["xclip", "-selection", "clipboard"],
                     stdin=subprocess.PIPE,
-                    env={**os.environ, "DISPLAY": DISPLAY}
+                    env={**os.environ, "DISPLAY": disp}
                 )
                 proc.communicate(text.encode("utf-8"))
             await asyncio.to_thread(_xclip_write)
@@ -641,7 +781,7 @@ def open_browser_smart(url: str = "") -> str:
     # جرب أولاً: BROWSER العادي
     cmd = f"{browser_cmd} --new-window '{url}' &"
     proc = subprocess.Popen(cmd, shell=True,
-                            env={**os.environ, "DISPLAY": DISPLAY},
+                            env={**os.environ, "DISPLAY": DISPLAY_OF()},
                             stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
     return f"{browser_cmd} → {url}"
 
@@ -919,11 +1059,12 @@ async def handle_action(ws: WebSocket, msg: dict):
     elif action == "clipboard_write":
         text = data.get("text", "")
         try:
+            disp = DISPLAY_OF()
             def _xclip_write2():
                 proc = subprocess.Popen(
                     ["xclip", "-selection", "clipboard"],
                     stdin=subprocess.PIPE,
-                    env={**os.environ, "DISPLAY": DISPLAY}
+                    env={**os.environ, "DISPLAY": disp}
                 )
                 proc.communicate(text.encode("utf-8"))
             await asyncio.to_thread(_xclip_write2)
@@ -960,7 +1101,7 @@ async def handle_action(ws: WebSocket, msg: dict):
                     fixed_cmd = app_cmd.replace("firefox-esr", BROWSER)
 
             subprocess.Popen(fixed_cmd, shell=True,
-                             env={**os.environ, "DISPLAY": DISPLAY},
+                             env={**os.environ, "DISPLAY": DISPLAY_OF()},
                              stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
             await send({"type": "ack", "action": "open_app", "cmd": fixed_cmd})
 
@@ -1016,11 +1157,12 @@ async def handle_action(ws: WebSocket, msg: dict):
     elif action == "paste":
         text = data.get("text", "")
         if text:
+            disp = DISPLAY_OF()
             def _xclip_paste():
                 proc = subprocess.Popen(
                     ["xclip", "-selection", "clipboard"],
                     stdin=subprocess.PIPE,
-                    env={**os.environ, "DISPLAY": DISPLAY}
+                    env={**os.environ, "DISPLAY": disp}
                 )
                 proc.communicate(text.encode("utf-8"))
             await asyncio.to_thread(_xclip_paste)
@@ -1043,13 +1185,21 @@ async def handle_action(ws: WebSocket, msg: dict):
 async def websocket_endpoint(ws: WebSocket):
     await ws.accept()
     active_connections.append(ws)
-    w, h = _get_screen_size()
+
+    # ── عزل كامل: إنشاء Xvfb + متصفح خاصين بهذا الاتصال فقط ──
+    # كل اتصال WebSocket يحصل على بيئة كمبيوتر افتراضي مستقلة تماماً عن أي اتصال آخر،
+    # بحيث لا يرى أي مستخدم أي أثر من نشاط مستخدم آخر مهما كانوا متصلين في نفس اللحظة.
+    session_display = await create_session_for(ws)
+    _current_display.set(session_display)  # يُربط بـ task هذا الاتصال — لا يؤثر على اتصالات أخرى
+
+    w, h = await asyncio.to_thread(_get_screen_size)
     await ws.send_text(json.dumps({
         "type": "connected",
         "screen_width": w,
         "screen_height": h,
         "browser": BROWSER,
-        "msg": f"Z Computer Mode v7 — Smart Control | Browser: {BROWSER} | Screen: {w}x{h}",
+        "display": session_display,
+        "msg": f"Z Computer Mode v7 — Smart Control | Browser: {BROWSER} | Screen: {w}x{h} | Isolated session: {session_display}",
         "commands": [
             "screenshot", "terminal {cmd}", "mouse_move {x,y}", "mouse_click {x,y,button}",
             "keyboard_type {text}", "keyboard_hotkey {keys}", "scroll {x,y,clicks}",
@@ -1057,7 +1207,7 @@ async def websocket_endpoint(ws: WebSocket):
             "Simple cmds: 'open firefox', 'mouse go x y', 'click x y', 'type TEXT', 'key enter'"
         ]
     }, ensure_ascii=False))
-    # أرسل screenshot أولية مع grid
+    # أرسل screenshot أولية مع grid (لشاشة هذه الجلسة المعزولة فقط)
     result = await asyncio.to_thread(capture_screen_with_grid, scale=0.65, quality=72)
     if result["data"]:
         await ws.send_text(json.dumps({
@@ -1082,6 +1232,8 @@ async def websocket_endpoint(ws: WebSocket):
     finally:
         if ws in active_connections:
             active_connections.remove(ws)
+        # ── تحرير بيئة هذه الجلسة بالكامل (Xvfb + المتصفح) فوراً عند انقطاع الاتصال ──
+        await destroy_session_for(ws)
 
 
 # ─── REST ────────────────────────────────────────────