More VPN Tunnel Changes

Data/Engine/services/WebSocket/__init__.py

@@ -103,7 +103,7 @@ def register_realtime(socket_server: SocketIO, context: EngineContext) -> None:
     adapters = EngineRealtimeAdapters(context)
     logger = context.logger.getChild("realtime.quick_jobs")
     agent_logger = context.logger.getChild("realtime.agents")
-    shell_bridge = VpnShellBridge(socket_server, context)
+    shell_bridge = VpnShellBridge(socket_server, context, adapters.service_log)
     agent_registry = AgentSocketRegistry(socket_server, agent_logger)
 
     def _emit_agent_event(agent_id: str, event: str, payload: Any) -> bool:
@@ -148,6 +148,24 @@ def register_realtime(socket_server: SocketIO, context: EngineContext) -> None:
         setattr(context, "vpn_tunnel_service", service)
         return service
 
+    def _tunnel_log(message: str, *, level: str = "INFO") -> None:
+        try:
+            adapters.service_log("VPN_Tunnel/tunnel", message, level=level)
+        except Exception:
+            agent_logger.debug("vpn_tunnel service log write failed", exc_info=True)
+
+    def _shell_log(message: str, *, level: str = "INFO") -> None:
+        try:
+            adapters.service_log("VPN_Tunnel/remote_shell", message, level=level)
+        except Exception:
+            agent_logger.debug("vpn_shell service log write failed", exc_info=True)
+
+    def _remote_addr() -> str:
+        forwarded = (request.headers.get("X-Forwarded-For") or "").strip()
+        if forwarded:
+            return forwarded.split(",")[0].strip()
+        return (request.remote_addr or "").strip()
+
     @socket_server.on("quick_job_result")
     def _handle_quick_job_result(data: Any) -> None:
         if not isinstance(data, dict):
@@ -317,18 +335,59 @@ def register_realtime(socket_server: SocketIO, context: EngineContext) -> None:
         elif isinstance(data, str):
             agent_id = data.strip()
         if not agent_id:
+            _shell_log(
+                "vpn_shell_open_missing sid={0} remote={1}".format(
+                    request.sid,
+                    _remote_addr() or "-",
+                ),
+                level="WARNING",
+            )
             return {"error": "agent_id_required"}
 
+        _shell_log(
+            "vpn_shell_open_request agent_id={0} sid={1} remote={2}".format(
+                agent_id,
+                request.sid,
+                _remote_addr() or "-",
+            )
+        )
         service = _get_tunnel_service()
         if service is None:
+            _shell_log(
+                "vpn_shell_open_failed agent_id={0} sid={1} reason=vpn_service_unavailable".format(
+                    agent_id,
+                    request.sid,
+                ),
+                level="WARNING",
+            )
             return {"error": "vpn_service_unavailable"}
         if not service.status(agent_id):
+            _shell_log(
+                "vpn_shell_open_failed agent_id={0} sid={1} reason=tunnel_down".format(
+                    agent_id,
+                    request.sid,
+                ),
+                level="WARNING",
+            )
             return {"error": "tunnel_down"}
 
         session = shell_bridge.open_session(request.sid, agent_id)
         if session is None:
+            _shell_log(
+                "vpn_shell_open_failed agent_id={0} sid={1} reason=shell_connect_failed".format(
+                    agent_id,
+                    request.sid,
+                ),
+                level="WARNING",
+            )
             return {"error": "shell_connect_failed"}
         service.bump_activity(agent_id)
+        _shell_log(
+            "vpn_shell_open_success agent_id={0} sid={1}".format(
+                agent_id,
+                request.sid,
+            )
+        )
         return {"status": "ok"}
 
     @socket_server.on("connect_agent")
@@ -341,16 +400,38 @@ def register_realtime(socket_server: SocketIO, context: EngineContext) -> None:
         elif isinstance(data, str):
             agent_id = data.strip()
         if not agent_id:
+            _tunnel_log(
+                "vpn_agent_socket_missing sid={0} remote={1}".format(
+                    request.sid,
+                    _remote_addr() or "-",
+                ),
+                level="WARNING",
+            )
             return {"error": "agent_id_required"}
 
         agent_registry.register(agent_id, request.sid)
         agent_logger.info("Agent socket registered agent_id=%s service_mode=%s sid=%s", agent_id, service_mode, request.sid)
+        _tunnel_log(
+            "vpn_agent_socket_register agent_id={0} service_mode={1} sid={2} remote={3}".format(
+                agent_id,
+                service_mode or "-",
+                request.sid,
+                _remote_addr() or "-",
+            )
+        )
 
         service = _get_tunnel_service()
         if service:
             payload = service.session_payload(agent_id, include_token=True)
             if payload:
-                agent_registry.emit(agent_id, "vpn_tunnel_start", payload)
+                if agent_registry.emit(agent_id, "vpn_tunnel_start", payload):
+                    _tunnel_log(
+                        "vpn_agent_socket_emit_start agent_id={0} tunnel_id={1} sid={2}".format(
+                            agent_id,
+                            payload.get("tunnel_id", "-"),
+                            request.sid,
+                        )
+                    )
 
         return {"status": "ok"}
 
@@ -363,11 +444,28 @@ def register_realtime(socket_server: SocketIO, context: EngineContext) -> None:
             payload = data
         if payload is None:
             return {"error": "payload_required"}
+        try:
+            payload_len = len(str(payload))
+        except Exception:
+            payload_len = 0
+        _shell_log(
+            "vpn_shell_send_request sid={0} bytes={1} remote={2}".format(
+                request.sid,
+                payload_len,
+                _remote_addr() or "-",
+            )
+        )
         shell_bridge.send(request.sid, str(payload))
         return {"status": "ok"}
 
     @socket_server.on("vpn_shell_close")
     def _vpn_shell_close(data: Any = None) -> Dict[str, Any]:
+        _shell_log(
+            "vpn_shell_close_request sid={0} remote={1}".format(
+                request.sid,
+                _remote_addr() or "-",
+            )
+        )
         shell_bridge.close(request.sid)
         return {"status": "ok"}
 
@@ -376,4 +474,18 @@ def register_realtime(socket_server: SocketIO, context: EngineContext) -> None:
         agent_id = agent_registry.unregister(request.sid)
         if agent_id:
             agent_logger.info("Agent socket disconnected agent_id=%s sid=%s", agent_id, request.sid)
+            _tunnel_log(
+                "vpn_agent_socket_disconnect agent_id={0} sid={1}".format(
+                    agent_id,
+                    request.sid,
+                )
+            )
+        else:
+            _shell_log(
+                "vpn_shell_client_disconnect sid={0} remote={1}".format(
+                    request.sid,
+                    _remote_addr() or "-",
+                ),
+                level="WARNING",
+            )
         shell_bridge.close(request.sid)

Data/Engine/services/WebSocket/vpn_shell.py

@@ -15,7 +15,7 @@ import socket
 import threading
 import time
 from dataclasses import dataclass
-from typing import Any, Dict, Optional
+from typing import Any, Callable, Dict, Optional
 
 
 def _b64encode(data: bytes) -> str:
@@ -32,6 +32,11 @@ class ShellSession:
     agent_id: str
     socketio: Any
     tcp: socket.socket
+    service_log: Optional[Callable[[str, str, Optional[str]], None]] = None
+    output_lines: int = 0
+    output_bytes: int = 0
+    input_messages: int = 0
+    input_bytes: int = 0
     _reader: Optional[threading.Thread] = None
 
     def start_reader(self) -> None:
@@ -39,15 +44,31 @@ class ShellSession:
         t.start()
         self._reader = t
 
+    def _service_log_event(self, message: str, *, level: str = "INFO") -> None:
+        if not callable(self.service_log):
+            return
+        try:
+            self.service_log("VPN_Tunnel/remote_shell", message, level=level)
+        except Exception:
+            pass
+
     def _read_loop(self) -> None:
         buffer = b""
+        reason = "remote_closed"
+        error_detail = ""
         try:
             while True:
                 try:
                     data = self.tcp.recv(4096)
                 except (socket.timeout, TimeoutError):
+                    # No data ready; keep the session alive.
+                    continue
+                except Exception as exc:
+                    reason = "read_error"
+                    error_detail = f"{type(exc).__name__}:{exc}"
                     break
                 if not data:
+                    reason = "remote_closed"
                     break
                 buffer += data
                 while b"\n" in buffer:
@@ -64,8 +85,37 @@ class ShellSession:
                             decoded = _b64decode(str(payload)).decode("utf-8", errors="replace")
                         except Exception:
                             decoded = ""
+                        self.output_lines += 1
+                        self.output_bytes += len(line)
                         self.socketio.emit("vpn_shell_output", {"data": decoded}, to=self.sid)
         finally:
+            if reason == "read_error":
+                self._service_log_event(
+                    "vpn_shell_read_error agent_id={0} sid={1} reason={2} error={3}".format(
+                        self.agent_id,
+                        self.sid,
+                        reason,
+                        error_detail or "-",
+                    ),
+                    level="WARNING",
+                )
+            self._service_log_event(
+                "vpn_shell_closed agent_id={0} sid={1} reason={2}".format(
+                    self.agent_id,
+                    self.sid,
+                    reason,
+                )
+            )
+            self._service_log_event(
+                "vpn_shell_output_summary agent_id={0} sid={1} lines={2} bytes={3} inputs={4} input_bytes={5}".format(
+                    self.agent_id,
+                    self.sid,
+                    self.output_lines,
+                    self.output_bytes,
+                    self.input_messages,
+                    self.input_bytes,
+                )
+            )
             self.socketio.emit("vpn_shell_closed", {"agent_id": self.agent_id}, to=self.sid)
             try:
                 self.tcp.close()
@@ -73,8 +123,21 @@ class ShellSession:
                 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")
+        payload_bytes = payload.encode("utf-8")
+        data = json.dumps({"type": "stdin", "data": _b64encode(payload_bytes)})
+        self.input_messages += 1
+        self.input_bytes += len(payload_bytes)
+        try:
+            self.tcp.sendall(data.encode("utf-8") + b"\n")
+        except Exception as exc:
+            self._service_log_event(
+                "vpn_shell_send_failed agent_id={0} sid={1} error={2}".format(
+                    self.agent_id,
+                    self.sid,
+                    f"{type(exc).__name__}:{exc}",
+                ),
+                level="WARNING",
+            )
 
     def close(self) -> None:
         try:
@@ -89,11 +152,20 @@ class ShellSession:
 
 
 class VpnShellBridge:
-    def __init__(self, socketio, context) -> None:
+    def __init__(self, socketio, context, service_log=None) -> None:
         self.socketio = socketio
         self.context = context
         self._sessions: Dict[str, ShellSession] = {}
         self.logger = context.logger.getChild("vpn_shell")
+        self.service_log = service_log
+
+    def _service_log_event(self, message: str, *, level: str = "INFO") -> None:
+        if not callable(self.service_log):
+            return
+        try:
+            self.service_log("VPN_Tunnel/remote_shell", message, level=level)
+        except Exception:
+            self.logger.debug("vpn_shell service log write failed", exc_info=True)
 
     def open_session(self, sid: str, agent_id: str) -> Optional[ShellSession]:
         service = getattr(self.context, "vpn_tunnel_service", None)
@@ -107,6 +179,15 @@ class VpnShellBridge:
         tcp = None
         last_error: Optional[Exception] = None
         for attempt in range(3):
+            self._service_log_event(
+                "vpn_shell_connect_attempt agent_id={0} sid={1} host={2} port={3} attempt={4}".format(
+                    agent_id,
+                    sid,
+                    host,
+                    port,
+                    attempt + 1,
+                )
+            )
             try:
                 tcp = socket.create_connection((host, port), timeout=5)
                 break
@@ -115,26 +196,72 @@ class VpnShellBridge:
                 if attempt == 0:
                     try:
                         service.request_agent_start(agent_id)
+                        self._service_log_event(
+                            "vpn_shell_agent_start_emit agent_id={0} sid={1}".format(agent_id, sid)
+                        )
                     except Exception:
                         self.logger.debug("Failed to re-emit vpn_tunnel_start for agent=%s", agent_id, exc_info=True)
+                        self._service_log_event(
+                            "vpn_shell_agent_start_failed agent_id={0} sid={1}".format(agent_id, sid),
+                            level="WARNING",
+                        )
                 time.sleep(1)
         if tcp is None:
+            self._service_log_event(
+                "vpn_shell_connect_failed agent_id={0} sid={1} host={2} port={3} error={4}".format(
+                    agent_id,
+                    sid,
+                    host,
+                    port,
+                    str(last_error) if last_error else "-",
+                ),
+                level="WARNING",
+            )
             self.logger.warning("Failed to connect vpn shell to %s:%s", host, port, exc_info=last_error)
             return None
-        session = ShellSession(sid=sid, agent_id=agent_id, socketio=self.socketio, tcp=tcp)
+        session = ShellSession(
+            sid=sid,
+            agent_id=agent_id,
+            socketio=self.socketio,
+            tcp=tcp,
+            service_log=self.service_log,
+        )
         try:
             session.tcp.settimeout(15)
         except Exception:
             pass
         self._sessions[sid] = session
+        self._service_log_event(
+            "vpn_shell_connect_success agent_id={0} sid={1} host={2} port={3}".format(
+                agent_id,
+                sid,
+                host,
+                port,
+            )
+        )
         session.start_reader()
         return session
 
     def send(self, sid: str, payload: str) -> None:
         session = self._sessions.get(sid)
         if not session:
+            self._service_log_event(
+                "vpn_shell_send_missing sid={0}".format(sid or "-"),
+                level="WARNING",
+            )
             return
         session.send(payload)
+        try:
+            payload_len = len(str(payload))
+        except Exception:
+            payload_len = 0
+        self._service_log_event(
+            "vpn_shell_send agent_id={0} sid={1} bytes={2}".format(
+                session.agent_id,
+                sid,
+                payload_len,
+            )
+        )
         service = getattr(self.context, "vpn_tunnel_service", None)
         if service:
             service.bump_activity(session.agent_id)
@@ -143,4 +270,10 @@ class VpnShellBridge:
         session = self._sessions.pop(sid, None)
         if not session:
             return
+        self._service_log_event(
+            "vpn_shell_close_request agent_id={0} sid={1}".format(
+                session.agent_id,
+                sid,
+            )
+        )
         session.close()