Overhaul of VPN Codebase

Data/Engine/services/WebSocket/vpn_shell.py

@@ -0,0 +1,127 @@
+# ======================================================
+# Data\Engine\services\WebSocket\vpn_shell.py
+# Description: Socket.IO handlers bridging UI shell to agent TCP server over WireGuard.
+#
+# API Endpoints (if applicable): None
+# ======================================================
+
+"""WireGuard VPN PowerShell bridge (Engine side)."""
+
+from __future__ import annotations
+
+import base64
+import json
+import socket
+import threading
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+
+def _b64encode(data: bytes) -> str:
+    return base64.b64encode(data).decode("ascii").strip()
+
+
+def _b64decode(value: str) -> bytes:
+    return base64.b64decode(value.encode("ascii"))
+
+
+@dataclass
+class ShellSession:
+    sid: str
+    agent_id: str
+    socketio: Any
+    tcp: socket.socket
+    _reader: Optional[threading.Thread] = None
+
+    def start_reader(self) -> None:
+        t = threading.Thread(target=self._read_loop, daemon=True)
+        t.start()
+        self._reader = t
+
+    def _read_loop(self) -> None:
+        buffer = b""
+        try:
+            while True:
+                data = self.tcp.recv(4096)
+                if not data:
+                    break
+                buffer += data
+                while b"\n" in buffer:
+                    line, buffer = buffer.split(b"\n", 1)
+                    if not line:
+                        continue
+                    try:
+                        msg = json.loads(line.decode("utf-8"))
+                    except Exception:
+                        continue
+                    if msg.get("type") == "stdout":
+                        payload = msg.get("data") or ""
+                        try:
+                            decoded = _b64decode(str(payload)).decode("utf-8", errors="replace")
+                        except Exception:
+                            decoded = ""
+                        self.socketio.emit("vpn_shell_output", {"data": decoded}, to=self.sid)
+        finally:
+            self.socketio.emit("vpn_shell_closed", {"agent_id": self.agent_id}, to=self.sid)
+            try:
+                self.tcp.close()
+            except Exception:
+                pass
+
+    def send(self, payload: str) -> None:
+        data = json.dumps({"type": "stdin", "data": _b64encode(payload.encode("utf-8"))})
+        self.tcp.sendall(data.encode("utf-8") + b"\n")
+
+    def close(self) -> None:
+        try:
+            data = json.dumps({"type": "close"})
+            self.tcp.sendall(data.encode("utf-8") + b"\n")
+        except Exception:
+            pass
+        try:
+            self.tcp.close()
+        except Exception:
+            pass
+
+
+class VpnShellBridge:
+    def __init__(self, socketio, context) -> None:
+        self.socketio = socketio
+        self.context = context
+        self._sessions: Dict[str, ShellSession] = {}
+        self.logger = context.logger.getChild("vpn_shell")
+
+    def open_session(self, sid: str, agent_id: str) -> Optional[ShellSession]:
+        service = getattr(self.context, "vpn_tunnel_service", None)
+        if service is None:
+            return None
+        status = service.status(agent_id)
+        if not status:
+            return None
+        host = str(status.get("virtual_ip") or "").split("/")[0]
+        port = int(self.context.wireguard_shell_port)
+        try:
+            tcp = socket.create_connection((host, port), timeout=5)
+        except Exception:
+            self.logger.debug("Failed to connect vpn shell to %s:%s", host, port, exc_info=True)
+            return None
+        session = ShellSession(sid=sid, agent_id=agent_id, socketio=self.socketio, tcp=tcp)
+        self._sessions[sid] = session
+        session.start_reader()
+        return session
+
+    def send(self, sid: str, payload: str) -> None:
+        session = self._sessions.get(sid)
+        if not session:
+            return
+        session.send(payload)
+        service = getattr(self.context, "vpn_tunnel_service", None)
+        if service:
+            service.bump_activity(session.agent_id)
+
+    def close(self, sid: str) -> None:
+        session = self._sessions.pop(sid, None)
+        if not session:
+            return
+        session.close()
+