Update app.py for persistence

app.py

@@ -6,14 +6,23 @@ import fcntl
 import termios
 import json
 import subprocess
+import uuid
+from pathlib import Path
 
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 
 app = FastAPI()
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
+PERSISTENT_DIR = Path("/persistent")
+SESSION_DIR = PERSISTENT_DIR / "sessions"
+PKGS_FILE = PERSISTENT_DIR / "installed_packages.txt"
+SESSION_STATE = PERSISTENT_DIR / "session_state.json"
+
+SESSION_DIR.mkdir(parents=True, exist_ok=True)
+
 @app.get("/")
 async def index():
     return FileResponse("static/index.html")
@@ -22,74 +31,144 @@ async def index():
 async def get_pass():
     return {"pass": os.environ.get("TERMINAL_PASS", "")}
 
+@app.get("/api/sessions")
+async def list_sessions():
+    sessions = []
+    for f in SESSION_DIR.glob("*.json"):
+        try:
+            data = json.loads(f.read_text())
+            sid = data.get("id", f.stem)
+            alive = subprocess.run(
+                ["screen", "-ls", f"term_{sid}"],
+                capture_output=True
+            ).returncode == 0
+            data["alive"] = alive
+            sessions.append(data)
+        except:
+            pass
+    return JSONResponse(sessions)
+
 @app.websocket("/ws")
 async def terminal(websocket: WebSocket):
     await websocket.accept()
 
-    master_fd, slave_fd = pty.openpty()
+    session_id = None
+    read_task = None
 
-    proc = await asyncio.create_subprocess_exec(
-        "/bin/bash",
-        stdin=slave_fd,
-        stdout=slave_fd,
-        stderr=slave_fd,
-        env={
-            **os.environ,
-            "TERM": "xterm-256color",
-            "HOME": "/persistent",
-            "USER": "ubuntu"
-        }
-    )
+    try:
+        while True:
+            message = await websocket.receive_text()
+            msg = json.loads(message)
 
-    os.close(slave_fd)
+            if msg.get("type") == "attach":
+                session_id = msg.get("session_id")
+                if not session_id or not (SESSION_DIR / f"{session_id}.json").exists():
+                    session_id = str(uuid.uuid4())[:8]
+                    (SESSION_DIR / f"{session_id}.json").write_text(json.dumps({
+                        "id": session_id,
+                        "created": str(subprocess.getoutput("date -Iseconds")),
+                    }))
 
-    loop = asyncio.get_event_loop()
+                await websocket.send_text(json.dumps({
+                    "type": "session_info",
+                    "session_id": session_id
+                }))
 
-    async def read_pty():
-        while True:
-            try:
-                data = await loop.run_in_executor(
-                    None,
-                    lambda: os.read(master_fd, 4096)
+                alive = subprocess.run(
+                    ["screen", "-ls", f"term_{session_id}"],
+                    capture_output=True
+                ).returncode == 0
+
+                if not alive:
+                    subprocess.Popen(
+                        ["screen", "-dmS", f"term_{session_id}", "-L", "-Logfile",
+                         f"/persistent/sessions/{session_id}.log", "/bin/bash", "--login"],
+                        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+                    )
+                    await asyncio.sleep(0.3)
+
+                    state_file = PERSISTENT_DIR / "session_state.json"
+                    if state_file.exists():
+                        try:
+                            state = json.loads(state_file.read_text())
+                            cwd = state.get("cwd", "/persistent")
+                            subprocess.run(
+                                ["screen", "-S", f"term_{session_id}", "-X", "stuff",
+                                 f"cd {cwd}\nclear\n"],
+                                capture_output=True, timeout=3
+                            )
+                        except:
+                            pass
+
+                subprocess.run(
+                    ["screen", "-S", f"term_{session_id}", "-X", "stuff", " \n"],
+                    capture_output=True, timeout=3
+                )
+
+                if read_task:
+                    read_task.cancel()
+                read_task = asyncio.create_task(read_screen(websocket, session_id))
+                continue
+
+            if msg.get("type") == "input" and session_id:
+                data = msg["data"]
+                subprocess.run(
+                    ["screen", "-S", f"term_{session_id}", "-X", "stuff", data],
+                    capture_output=True, timeout=5
+                )
+
+            elif msg.get("type") == "resize" and session_id:
+                cols = msg.get("cols", 200)
+                rows = msg.get("rows", 50)
+                subprocess.run(
+                    ["screen", "-S", f"term_{session_id}", "-X", "width", "cols", str(cols)],
+                    capture_output=True, timeout=3
+                )
+                subprocess.run(
+                    ["screen", "-S", f"term_{session_id}", "-X", "width", "lines", str(rows)],
+                    capture_output=True, timeout=3
+                )
+
+            elif msg.get("type") == "key" and session_id:
+                key = msg.get("key", "")
+                subprocess.run(
+                    ["screen", "-S", f"term_{session_id}", "-X", "stuff", key],
+                    capture_output=True, timeout=3
                 )
-                if data:
-                    try:
-                        text = data.decode("utf-8")
-                    except UnicodeDecodeError:
-                        text = data.decode("utf-8", errors="ignore")
-                    await websocket.send_text(text)
-            except OSError:
-                break
 
-    read_task = asyncio.create_task(read_pty())
+    except (WebSocketDisconnect, Exception):
+        pass
+    finally:
+        if read_task:
+            read_task.cancel()
 
+async def read_screen(websocket: WebSocket, session_id: str):
+    last_content = ""
     try:
         while True:
-            message = await websocket.receive_text()
+            await asyncio.sleep(0.05)
             try:
-                msg = json.loads(message)
-                if msg.get("type") == "input":
-                    os.write(master_fd, msg["data"].encode("utf-8"))
-                elif msg.get("type") == "resize":
-                    cols = msg.get("cols", 80)
-                    rows = msg.get("rows", 24)
-                    fcntl.ioctl(
-                        master_fd,
-                        termios.TIOCSWINSZ,
-                        struct.pack("HHHH", rows, cols, 0, 0)
-                    )
-            except (json.JSONDecodeError, OSError):
+                result = subprocess.run(
+                    ["screen", "-S", f"term_{session_id}", "-X", "hardcopy", "-h", "/tmp/screen_out.txt"],
+                    capture_output=True, timeout=3
+                )
+                if result.returncode == 0:
+                    try:
+                        content = Path("/tmp/screen_out.txt").read_text(errors="ignore")
+                        if content != last_content:
+                            if len(content) > len(last_content):
+                                new_data = content[len(last_content):]
+                            else:
+                                new_data = content
+                            last_content = content
+                            if new_data:
+                                await websocket.send_text(json.dumps({
+                                    "type": "output",
+                                    "data": new_data
+                                }))
+                    except:
+                        pass
+            except:
                 break
-    except (WebSocketDisconnect, Exception):
+    except asyncio.CancelledError:
         pass
-    finally:
-        read_task.cancel()
-        try:
-            proc.terminate()
-            await asyncio.wait_for(proc.wait(), timeout=2)
-        except Exception:
-            proc.kill()
-        try:
-            os.close(master_fd)
-        except:
-            pass