mirror of
https://github.com/bunny-lab-io/Borealis.git
synced 2026-02-04 12:20:31 -07:00
280 lines
9.5 KiB
Python
280 lines
9.5 KiB
Python
# ======================================================
|
|
# 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
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any, Callable, 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
|
|
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:
|
|
t = threading.Thread(target=self._read_loop, daemon=True)
|
|
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:
|
|
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.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()
|
|
except Exception:
|
|
pass
|
|
|
|
def send(self, payload: str) -> None:
|
|
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:
|
|
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, 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)
|
|
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)
|
|
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
|
|
except Exception as exc:
|
|
last_error = exc
|
|
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,
|
|
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)
|
|
|
|
def close(self, sid: str) -> None:
|
|
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()
|