More VPN Tunnel Changes

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()