bunny-lab/Borealis-Github-Replica
1
0
Fork 0
mirror of https://github.com/bunny-lab-io/Borealis.git synced 2025-12-19 02:35:48 -07:00
Code Issues Packages Projects Releases Wiki Activity

Overhaul of VPN Codebase

Browse Source
This commit is contained in:
Nicole Rappe
2025-12-18 01:35:03 -07:00
parent 2f81061a1b
commit 6ceb59f717
56 changed files with 1786 additions and 4778 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/__init__.py
View File

@@ -1,3 +0,0 @@
"""Namespace package for reverse tunnel domain handlers (Engine side)."""
__all__ = ["remote_interactive_shell", "remote_management", "remote_video"]

78
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_interactive_shell/Protocols/Bash.py
View File

@@ -1,78 +0,0 @@
"""Placeholder Bash channel server (Engine side)."""
from __future__ import annotations
from collections import deque
class BashChannelServer:
"""Stub Bash handler until the agent-side channel is implemented."""
protocol_name = "bash"
def __init__(self, bridge, service, *, channel_id: int = 1, frame_cls=None, close_frame_fn=None):
self.bridge = bridge
self.service = service
self.channel_id = channel_id
self.logger = service.logger.getChild(f"bash.{bridge.lease.tunnel_id}")
self._open_sent = False
self._ack_received = False
self._closed = False
self._output = deque()
self._frame_cls = frame_cls
self._close_frame_fn = close_frame_fn
def handle_agent_frame(self, frame) -> None:
# No-op placeholder; output collection for future Bash support.
try:
if frame.msg_type == 0x04: # MSG_CHANNEL_ACK
self._ack_received = True
except Exception:
return
def open_channel(self, *, cols: int = 120, rows: int = 32) -> None:
self._open_sent = True
self.logger.info(
"bash channel placeholder open sent tunnel_id=%s channel_id=%s cols=%s rows=%s",
self.bridge.lease.tunnel_id,
self.channel_id,
cols,
rows,
)
def send_input(self, data: str) -> None:
# Placeholder: no agent-side Bash yet.
self.logger.info("bash placeholder send_input ignored tunnel_id=%s", self.bridge.lease.tunnel_id)
def send_resize(self, cols: int, rows: int) -> None:
# Placeholder: not implemented.
return
def close(self, code: int = 6, reason: str = "operator_close") -> None:
if self._closed:
return
self._closed = True
if callable(self._close_frame_fn):
try:
frame = self._close_frame_fn(self.channel_id, code, reason)
self.bridge.operator_to_agent(frame)
except Exception:
pass
def drain_output(self):
items = []
while self._output:
items.append(self._output.popleft())
return items
def status(self):
return {
"channel_id": self.channel_id,
"open_sent": self._open_sent,
"ack": self._ack_received,
"closed": self._closed,
"close_reason": None,
"close_code": None,
}
__all__ = ["BashChannelServer"]

139
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_interactive_shell/Protocols/Powershell.py
View File

@@ -1,139 +0,0 @@
"""Engine-side PowerShell tunnel channel helper (remote interactive shell domain)."""
from __future__ import annotations
import json
from collections import deque
from typing import Any, Deque, Dict, List, Optional
# Mirror framing constants to avoid circular imports.
MSG_CHANNEL_OPEN = 0x03
MSG_CHANNEL_ACK = 0x04
MSG_DATA = 0x05
MSG_CONTROL = 0x09
MSG_CLOSE = 0x08
CLOSE_OK = 0
CLOSE_PROTOCOL_ERROR = 3
CLOSE_AGENT_SHUTDOWN = 6
class PowershellChannelServer:
"""Coordinate PowerShell channel frames over a TunnelBridge."""
def __init__(self, bridge, service, *, channel_id: int = 1, frame_cls=None, close_frame_fn=None):
self.bridge = bridge
self.service = service
self.channel_id = channel_id
self.logger = service.logger.getChild(f"ps.{bridge.lease.tunnel_id}")
self._open_sent = False
self._ack_received = False
self._closed = False
self._output: Deque[str] = deque()
self._close_reason: Optional[str] = None
self._close_code: Optional[int] = None
self._frame_cls = frame_cls
self._close_frame_fn = close_frame_fn
# ------------------------------------------------------------------ Agent frame handling
def handle_agent_frame(self, frame) -> None:
if frame.channel_id != self.channel_id:
return
if frame.msg_type == MSG_CHANNEL_ACK:
self._ack_received = True
self.logger.info("ps channel acked tunnel_id=%s", self.bridge.lease.tunnel_id)
return
if frame.msg_type == MSG_DATA:
try:
text = frame.payload.decode("utf-8", errors="replace")
except Exception:
text = ""
if text:
self._append_output(text)
return
if frame.msg_type == MSG_CLOSE:
try:
payload = json.loads(frame.payload.decode("utf-8"))
except Exception:
payload = {}
self._closed = True
self._close_code = payload.get("code") if isinstance(payload, dict) else None
self._close_reason = payload.get("reason") if isinstance(payload, dict) else None
self.logger.info(
"ps channel closed tunnel_id=%s code=%s reason=%s",
self.bridge.lease.tunnel_id,
self._close_code,
self._close_reason or "-",
)
# ------------------------------------------------------------------ Operator actions
def open_channel(self, *, cols: int = 120, rows: int = 32) -> None:
if self._open_sent:
return
payload = json.dumps(
{"protocol": "ps", "metadata": {"cols": cols, "rows": rows}},
separators=(",", ":"),
).encode("utf-8")
frame = self._frame_cls(msg_type=MSG_CHANNEL_OPEN, channel_id=self.channel_id, payload=payload)
self.bridge.operator_to_agent(frame)
self._open_sent = True
self.logger.info(
"ps channel open sent tunnel_id=%s channel_id=%s cols=%s rows=%s",
self.bridge.lease.tunnel_id,
self.channel_id,
cols,
rows,
)
def send_input(self, data: str) -> None:
if self._closed:
return
payload = data.encode("utf-8", errors="replace")
frame = self._frame_cls(msg_type=MSG_DATA, channel_id=self.channel_id, payload=payload)
self.bridge.operator_to_agent(frame)
def send_resize(self, cols: int, rows: int) -> None:
if self._closed:
return
payload = json.dumps({"cols": cols, "rows": rows}, separators=(",", ":")).encode("utf-8")
frame = self._frame_cls(msg_type=MSG_CONTROL, channel_id=self.channel_id, payload=payload)
self.bridge.operator_to_agent(frame)
def close(self, code: int = CLOSE_AGENT_SHUTDOWN, reason: str = "operator_close") -> None:
if self._closed:
return
self._closed = True
if callable(self._close_frame_fn):
frame = self._close_frame_fn(self.channel_id, code, reason)
else:
frame = self._frame_cls(
msg_type=MSG_CLOSE,
channel_id=self.channel_id,
payload=json.dumps({"code": code, "reason": reason}, separators=(",", ":")).encode("utf-8"),
)
self.bridge.operator_to_agent(frame)
# ------------------------------------------------------------------ Output polling
def drain_output(self) -> List[str]:
items: List[str] = []
while self._output:
items.append(self._output.popleft())
return items
def _append_output(self, text: str) -> None:
self._output.append(text)
# Cap buffer to avoid unbounded memory growth.
while len(self._output) > 500:
self._output.popleft()
# ------------------------------------------------------------------ Status helpers
def status(self) -> Dict[str, Any]:
return {
"channel_id": self.channel_id,
"open_sent": self._open_sent,
"ack": self._ack_received,
"closed": self._closed,
"close_reason": self._close_reason,
"close_code": self._close_code,
}
__all__ = ["PowershellChannelServer"]

6
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_interactive_shell/Protocols/__init__.py
View File

@@ -1,6 +0,0 @@
"""Protocol handlers for remote interactive shell tunnels (Engine side)."""
from .Powershell import PowershellChannelServer
from .Bash import BashChannelServer
__all__ = ["PowershellChannelServer", "BashChannelServer"]

3
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_interactive_shell/__init__.py
View File

@@ -1,3 +0,0 @@
"""Domain handlers for remote interactive shells (PowerShell/Bash)."""
__all__ = ["Protocols"]

73
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_management/Protocols/SSH.py
View File

@@ -1,73 +0,0 @@
"""Placeholder SSH channel server (Engine side)."""
from __future__ import annotations
from collections import deque
class SSHChannelServer:
protocol_name = "ssh"
def __init__(self, bridge, service, *, channel_id: int = 1, frame_cls=None, close_frame_fn=None):
self.bridge = bridge
self.service = service
self.channel_id = channel_id
self.logger = service.logger.getChild(f"ssh.{bridge.lease.tunnel_id}")
self._open_sent = False
self._ack_received = False
self._closed = False
self._output = deque()
self._frame_cls = frame_cls
self._close_frame_fn = close_frame_fn
def handle_agent_frame(self, frame) -> None:
try:
if frame.msg_type == 0x04: # MSG_CHANNEL_ACK
self._ack_received = True
except Exception:
return
def open_channel(self, *, cols: int = 120, rows: int = 32) -> None:
self._open_sent = True
self.logger.info(
"ssh channel placeholder open sent tunnel_id=%s channel_id=%s cols=%s rows=%s",
self.bridge.lease.tunnel_id,
self.channel_id,
cols,
rows,
)
def send_input(self, data: str) -> None:
self.logger.info("ssh placeholder send_input ignored tunnel_id=%s", self.bridge.lease.tunnel_id)
def send_resize(self, cols: int, rows: int) -> None:
return
def close(self, code: int = 6, reason: str = "operator_close") -> None:
if self._closed:
return
self._closed = True
if callable(self._close_frame_fn):
try:
frame = self._close_frame_fn(self.channel_id, code, reason)
self.bridge.operator_to_agent(frame)
except Exception:
pass
def drain_output(self):
items = []
while self._output:
items.append(self._output.popleft())
return items
def status(self):
return {
"channel_id": self.channel_id,
"open_sent": self._open_sent,
"ack": self._ack_received,
"closed": self._closed,
"close_reason": None,
"close_code": None,
}
__all__ = ["SSHChannelServer"]

73
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_management/Protocols/WinRM.py
View File

@@ -1,73 +0,0 @@
"""Placeholder WinRM channel server (Engine side)."""
from __future__ import annotations
from collections import deque
class WinRMChannelServer:
protocol_name = "winrm"
def __init__(self, bridge, service, *, channel_id: int = 1, frame_cls=None, close_frame_fn=None):
self.bridge = bridge
self.service = service
self.channel_id = channel_id
self.logger = service.logger.getChild(f"winrm.{bridge.lease.tunnel_id}")
self._open_sent = False
self._ack_received = False
self._closed = False
self._output = deque()
self._frame_cls = frame_cls
self._close_frame_fn = close_frame_fn
def handle_agent_frame(self, frame) -> None:
try:
if frame.msg_type == 0x04: # MSG_CHANNEL_ACK
self._ack_received = True
except Exception:
return
def open_channel(self, *, cols: int = 120, rows: int = 32) -> None:
self._open_sent = True
self.logger.info(
"winrm channel placeholder open sent tunnel_id=%s channel_id=%s cols=%s rows=%s",
self.bridge.lease.tunnel_id,
self.channel_id,
cols,
rows,
)
def send_input(self, data: str) -> None:
self.logger.info("winrm placeholder send_input ignored tunnel_id=%s", self.bridge.lease.tunnel_id)
def send_resize(self, cols: int, rows: int) -> None:
return
def close(self, code: int = 6, reason: str = "operator_close") -> None:
if self._closed:
return
self._closed = True
if callable(self._close_frame_fn):
try:
frame = self._close_frame_fn(self.channel_id, code, reason)
self.bridge.operator_to_agent(frame)
except Exception:
pass
def drain_output(self):
items = []
while self._output:
items.append(self._output.popleft())
return items
def status(self):
return {
"channel_id": self.channel_id,
"open_sent": self._open_sent,
"ack": self._ack_received,
"closed": self._closed,
"close_reason": None,
"close_code": None,
}
__all__ = ["WinRMChannelServer"]

6
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_management/Protocols/__init__.py
View File

@@ -1,6 +0,0 @@
"""Protocol handlers for remote management tunnels (Engine side)."""
from .SSH import SSHChannelServer
from .WinRM import WinRMChannelServer
__all__ = ["SSHChannelServer", "WinRMChannelServer"]

3
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_management/__init__.py
View File

@@ -1,3 +0,0 @@
"""Domain handlers for remote management tunnels (SSH/WinRM)."""
__all__ = ["Protocols"]

73
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_video/Protocols/RDP.py
View File

@@ -1,73 +0,0 @@
"""Placeholder RDP channel server (Engine side)."""
from __future__ import annotations
from collections import deque
class RDPChannelServer:
protocol_name = "rdp"
def __init__(self, bridge, service, *, channel_id: int = 1, frame_cls=None, close_frame_fn=None):
self.bridge = bridge
self.service = service
self.channel_id = channel_id
self.logger = service.logger.getChild(f"rdp.{bridge.lease.tunnel_id}")
self._open_sent = False
self._ack_received = False
self._closed = False
self._output = deque()
self._frame_cls = frame_cls
self._close_frame_fn = close_frame_fn
def handle_agent_frame(self, frame) -> None:
try:
if frame.msg_type == 0x04: # MSG_CHANNEL_ACK
self._ack_received = True
except Exception:
return
def open_channel(self, *, cols: int = 120, rows: int = 32) -> None:
self._open_sent = True
self.logger.info(
"rdp channel placeholder open sent tunnel_id=%s channel_id=%s cols=%s rows=%s",
self.bridge.lease.tunnel_id,
self.channel_id,
cols,
rows,
)
def send_input(self, data: str) -> None:
self.logger.info("rdp placeholder send_input ignored tunnel_id=%s", self.bridge.lease.tunnel_id)
def send_resize(self, cols: int, rows: int) -> None:
return
def close(self, code: int = 6, reason: str = "operator_close") -> None:
if self._closed:
return
self._closed = True
if callable(self._close_frame_fn):
try:
frame = self._close_frame_fn(self.channel_id, code, reason)
self.bridge.operator_to_agent(frame)
except Exception:
pass
def drain_output(self):
items = []
while self._output:
items.append(self._output.popleft())
return items
def status(self):
return {
"channel_id": self.channel_id,
"open_sent": self._open_sent,
"ack": self._ack_received,
"closed": self._closed,
"close_reason": None,
"close_code": None,
}
__all__ = ["RDPChannelServer"]

73
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_video/Protocols/VNC.py
View File

@@ -1,73 +0,0 @@
"""Placeholder VNC channel server (Engine side)."""
from __future__ import annotations
from collections import deque
class VNCChannelServer:
protocol_name = "vnc"
def __init__(self, bridge, service, *, channel_id: int = 1, frame_cls=None, close_frame_fn=None):
self.bridge = bridge
self.service = service
self.channel_id = channel_id
self.logger = service.logger.getChild(f"vnc.{bridge.lease.tunnel_id}")
self._open_sent = False
self._ack_received = False
self._closed = False
self._output = deque()
self._frame_cls = frame_cls
self._close_frame_fn = close_frame_fn
def handle_agent_frame(self, frame) -> None:
try:
if frame.msg_type == 0x04: # MSG_CHANNEL_ACK
self._ack_received = True
except Exception:
return
def open_channel(self, *, cols: int = 120, rows: int = 32) -> None:
self._open_sent = True
self.logger.info(
"vnc channel placeholder open sent tunnel_id=%s channel_id=%s cols=%s rows=%s",
self.bridge.lease.tunnel_id,
self.channel_id,
cols,
rows,
)
def send_input(self, data: str) -> None:
self.logger.info("vnc placeholder send_input ignored tunnel_id=%s", self.bridge.lease.tunnel_id)
def send_resize(self, cols: int, rows: int) -> None:
return
def close(self, code: int = 6, reason: str = "operator_close") -> None:
if self._closed:
return
self._closed = True
if callable(self._close_frame_fn):
try:
frame = self._close_frame_fn(self.channel_id, code, reason)
self.bridge.operator_to_agent(frame)
except Exception:
pass
def drain_output(self):
items = []
while self._output:
items.append(self._output.popleft())
return items
def status(self):
return {
"channel_id": self.channel_id,
"open_sent": self._open_sent,
"ack": self._ack_received,
"closed": self._closed,
"close_reason": None,
"close_code": None,
}
__all__ = ["VNCChannelServer"]

73
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_video/Protocols/WebRTC.py
View File

@@ -1,73 +0,0 @@
"""Placeholder WebRTC channel server (Engine side)."""
from __future__ import annotations
from collections import deque
class WebRTCChannelServer:
protocol_name = "webrtc"
def __init__(self, bridge, service, *, channel_id: int = 1, frame_cls=None, close_frame_fn=None):
self.bridge = bridge
self.service = service
self.channel_id = channel_id
self.logger = service.logger.getChild(f"webrtc.{bridge.lease.tunnel_id}")
self._open_sent = False
self._ack_received = False
self._closed = False
self._output = deque()
self._frame_cls = frame_cls
self._close_frame_fn = close_frame_fn
def handle_agent_frame(self, frame) -> None:
try:
if frame.msg_type == 0x04: # MSG_CHANNEL_ACK
self._ack_received = True
except Exception:
return
def open_channel(self, *, cols: int = 120, rows: int = 32) -> None:
self._open_sent = True
self.logger.info(
"webrtc channel placeholder open sent tunnel_id=%s channel_id=%s cols=%s rows=%s",
self.bridge.lease.tunnel_id,
self.channel_id,
cols,
rows,
)
def send_input(self, data: str) -> None:
self.logger.info("webrtc placeholder send_input ignored tunnel_id=%s", self.bridge.lease.tunnel_id)
def send_resize(self, cols: int, rows: int) -> None:
return
def close(self, code: int = 6, reason: str = "operator_close") -> None:
if self._closed:
return
self._closed = True
if callable(self._close_frame_fn):
try:
frame = self._close_frame_fn(self.channel_id, code, reason)
self.bridge.operator_to_agent(frame)
except Exception:
pass
def drain_output(self):
items = []
while self._output:
items.append(self._output.popleft())
return items
def status(self):
return {
"channel_id": self.channel_id,
"open_sent": self._open_sent,
"ack": self._ack_received,
"closed": self._closed,
"close_reason": None,
"close_code": None,
}
__all__ = ["WebRTCChannelServer"]

7
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_video/Protocols/__init__.py
View File

@@ -1,7 +0,0 @@
"""Protocol handlers for remote video tunnels (Engine side)."""
from .WebRTC import WebRTCChannelServer
from .RDP import RDPChannelServer
from .VNC import VNCChannelServer
__all__ = ["WebRTCChannelServer", "RDPChannelServer", "VNCChannelServer"]

3
Data/Engine/services/WebSocket/Agent/Reverse_Tunnels/remote_video/__init__.py
View File

@@ -1,3 +0,0 @@
"""Domain handlers for remote video/desktop tunnels (RDP/VNC/WebRTC)."""
__all__ = ["Protocols"]

10
Data/Engine/services/WebSocket/Agent/__init__.py
View File

@@ -1,10 +0,0 @@
# ======================================================
# Data\Engine\services\WebSocket\Agent\__init__.py
# Description: Package marker for Agent-facing WebSocket services (reverse tunnel scaffolding).
#
# API Endpoints (if applicable): None
# ======================================================
"""Agent-facing WebSocket services for the Engine runtime."""
__all__ = []

1361
Data/Engine/services/WebSocket/Agent/reverse_tunnel_orchestrator.py
View File

File diff suppressed because it is too large Load Diff

318
Data/Engine/services/WebSocket/__init__.py
View File

@@ -1,6 +1,6 @@
# ======================================================
# Data\Engine\services\WebSocket\__init__.py
# Description: Socket.IO handlers for Engine runtime quick job updates and realtime notifications.
# Description: Socket.IO handlers for Engine runtime quick job updates and VPN shell bridging.
#
# API Endpoints (if applicable): None
# ======================================================
@@ -8,24 +8,20 @@
"""WebSocket service registration for the Borealis Engine runtime."""
from __future__ import annotations
import base64
import sqlite3
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, Optional
from flask import session, request
from flask import request
from flask_socketio import SocketIO
from ...database import initialise_engine_database
from ...security import signing
from ...server import EngineContext
from .Agent.reverse_tunnel_orchestrator import (
ReverseTunnelService,
TunnelBridge,
decode_frame,
TunnelFrame,
)
from ..VPN import VpnTunnelService
from .vpn_shell import VpnShellBridge
def _now_ts() -> int:
@@ -70,20 +66,31 @@ class EngineRealtimeAdapters:
def register_realtime(socket_server: SocketIO, context: EngineContext) -> None:
"""Register Socket.IO event handlers for the Engine runtime."""
from ..API import _make_db_conn_factory, _make_service_logger # Local import to avoid circular import at module load
adapters = EngineRealtimeAdapters(context)
logger = context.logger.getChild("realtime.quick_jobs")
tunnel_service = getattr(context, "reverse_tunnel_service", None)
if tunnel_service is None:
tunnel_service = ReverseTunnelService(
context,
signer=None,
shell_bridge = VpnShellBridge(socket_server, context)
def _get_tunnel_service() -> Optional[VpnTunnelService]:
service = getattr(context, "vpn_tunnel_service", None)
if service is not None:
return service
manager = getattr(context, "wireguard_server_manager", None)
if manager is None:
return None
try:
signer = signing.load_signer()
except Exception:
signer = None
service = VpnTunnelService(
context=context,
wireguard_manager=manager,
db_conn_factory=adapters.db_conn_factory,
socketio=socket_server,
service_log=adapters.service_log,
signer=signer,
)
tunnel_service.start()
setattr(context, "reverse_tunnel_service", tunnel_service)
setattr(context, "vpn_tunnel_service", service)
return service
@socket_server.on("quick_job_result")
def _handle_quick_job_result(data: Any) -> None:
@@ -246,252 +253,45 @@ def register_realtime(socket_server: SocketIO, context: EngineContext) -> None:
exc,
)
@socket_server.on("tunnel_bridge_attach")
def _tunnel_bridge_attach(data: Any) -> Any:
"""Placeholder operator bridge attach handler (no data channel yet)."""
@socket_server.on("vpn_shell_open")
def _vpn_shell_open(data: Any) -> Dict[str, Any]:
agent_id = ""
if isinstance(data, dict):
agent_id = str(data.get("agent_id") or "").strip()
elif isinstance(data, str):
agent_id = data.strip()
if not agent_id:
return {"error": "agent_id_required"}
if not isinstance(data, dict):
return {"error": "invalid_payload"}
service = _get_tunnel_service()
if service is None:
return {"error": "vpn_service_unavailable"}
if not service.status(agent_id):
return {"error": "tunnel_down"}
tunnel_id = str(data.get("tunnel_id") or "").strip()
operator_id = str(data.get("operator_id") or "").strip() or None
if not tunnel_id:
return {"error": "tunnel_id_required"}
session = shell_bridge.open_session(request.sid, agent_id)
if session is None:
return {"error": "shell_connect_failed"}
service.bump_activity(agent_id)
return {"status": "ok"}
try:
tunnel_service.operator_attach(tunnel_id, operator_id)
except ValueError as exc:
return {"error": str(exc)}
except Exception as exc: # pragma: no cover - defensive guard
logger.debug("tunnel_bridge_attach failed tunnel_id=%s: %s", tunnel_id, exc, exc_info=True)
return {"error": "bridge_attach_failed"}
return {"status": "ok", "tunnel_id": tunnel_id, "operator_id": operator_id or "-"}
def _encode_frame(frame: TunnelFrame) -> str:
return base64.b64encode(frame.encode()).decode("ascii")
def _decode_frame_payload(raw: Any) -> TunnelFrame:
if isinstance(raw, str):
try:
raw_bytes = base64.b64decode(raw)
except Exception:
raise ValueError("invalid_frame")
elif isinstance(raw, (bytes, bytearray)):
raw_bytes = bytes(raw)
@socket_server.on("vpn_shell_send")
def _vpn_shell_send(data: Any) -> Dict[str, Any]:
payload = None
if isinstance(data, dict):
payload = data.get("data")
else:
raise ValueError("invalid_frame")
return decode_frame(raw_bytes)
@socket_server.on("tunnel_operator_send")
def _tunnel_operator_send(data: Any) -> Any:
"""Operator -> agent frame enqueue (placeholder queue)."""
if not isinstance(data, dict):
return {"error": "invalid_payload"}
tunnel_id = str(data.get("tunnel_id") or "").strip()
frame_raw = data.get("frame")
if not tunnel_id or frame_raw is None:
return {"error": "tunnel_id_and_frame_required"}
try:
frame = _decode_frame_payload(frame_raw)
except Exception as exc:
return {"error": str(exc)}
bridge: Optional[TunnelBridge] = tunnel_service.get_bridge(tunnel_id)
if bridge is None:
return {"error": "unknown_tunnel"}
bridge.operator_to_agent(frame)
return {"status": "ok"}
@socket_server.on("tunnel_operator_poll")
def _tunnel_operator_poll(data: Any) -> Any:
"""Operator polls queued frames from agent."""
tunnel_id = ""
if isinstance(data, dict):
tunnel_id = str(data.get("tunnel_id") or "").strip()
if not tunnel_id:
return {"error": "tunnel_id_required"}
bridge: Optional[TunnelBridge] = tunnel_service.get_bridge(tunnel_id)
if bridge is None:
return {"error": "unknown_tunnel"}
frames = []
while True:
frame = bridge.next_for_operator()
if frame is None:
break
frames.append(_encode_frame(frame))
return {"frames": frames}
# WebUI operator bridge namespace for browser clients
tunnel_namespace = "/tunnel"
_operator_sessions: Dict[str, str] = {}
def _current_operator() -> Optional[str]:
username = session.get("username")
if username:
return str(username)
auth_header = (request.headers.get("Authorization") or "").strip()
token = None
if auth_header.lower().startswith("bearer "):
token = auth_header.split(" ", 1)[1].strip()
if not token:
token = request.cookies.get("borealis_auth")
return token or None
@socket_server.on("join", namespace=tunnel_namespace)
def _ws_tunnel_join(data: Any) -> Any:
if not isinstance(data, dict):
return {"error": "invalid_payload"}
operator_id = _current_operator()
if not operator_id:
return {"error": "unauthorized"}
tunnel_id = str(data.get("tunnel_id") or "").strip()
if not tunnel_id:
return {"error": "tunnel_id_required"}
bridge = tunnel_service.get_bridge(tunnel_id)
if bridge is None:
return {"error": "unknown_tunnel"}
try:
tunnel_service.operator_attach(tunnel_id, operator_id)
except Exception as exc:
logger.debug("ws_tunnel_join failed tunnel_id=%s: %s", tunnel_id, exc, exc_info=True)
return {"error": "attach_failed"}
sid = request.sid
_operator_sessions[sid] = tunnel_id
return {"status": "ok", "tunnel_id": tunnel_id}
@socket_server.on("send", namespace=tunnel_namespace)
def _ws_tunnel_send(data: Any) -> Any:
sid = request.sid
tunnel_id = _operator_sessions.get(sid)
if not tunnel_id:
return {"error": "not_joined"}
if not isinstance(data, dict):
return {"error": "invalid_payload"}
frame_raw = data.get("frame")
if frame_raw is None:
return {"error": "frame_required"}
try:
frame = _decode_frame_payload(frame_raw)
except Exception:
return {"error": "invalid_frame"}
bridge = tunnel_service.get_bridge(tunnel_id)
if bridge is None:
return {"error": "unknown_tunnel"}
bridge.operator_to_agent(frame)
return {"status": "ok"}
@socket_server.on("poll", namespace=tunnel_namespace)
def _ws_tunnel_poll() -> Any:
sid = request.sid
tunnel_id = _operator_sessions.get(sid)
if not tunnel_id:
return {"error": "not_joined"}
bridge = tunnel_service.get_bridge(tunnel_id)
if bridge is None:
return {"error": "unknown_tunnel"}
frames = []
while True:
frame = bridge.next_for_operator()
if frame is None:
break
frames.append(_encode_frame(frame))
return {"frames": frames}
def _require_ps_server():
sid = request.sid
tunnel_id = _operator_sessions.get(sid)
if not tunnel_id:
return None, None, {"error": "not_joined"}
server = tunnel_service.ensure_protocol_server(tunnel_id)
if server is None or not hasattr(server, "open_channel"):
return None, tunnel_id, {"error": "ps_unsupported"}
return server, tunnel_id, None
@socket_server.on("ps_open", namespace=tunnel_namespace)
def _ws_ps_open(data: Any) -> Any:
server, tunnel_id, error = _require_ps_server()
if server is None:
return error
cols = 120
rows = 32
if isinstance(data, dict):
try:
cols = int(data.get("cols", cols))
rows = int(data.get("rows", rows))
except Exception:
pass
cols = max(20, min(cols, 300))
rows = max(10, min(rows, 200))
try:
server.open_channel(cols=cols, rows=rows)
except Exception as exc:
logger.debug("ps_open failed tunnel_id=%s: %s", tunnel_id, exc, exc_info=True)
return {"error": "ps_open_failed"}
return {"status": "ok", "tunnel_id": tunnel_id, "cols": cols, "rows": rows}
@socket_server.on("ps_send", namespace=tunnel_namespace)
def _ws_ps_send(data: Any) -> Any:
server, tunnel_id, error = _require_ps_server()
if server is None:
return error
if data is None:
payload = data
if payload is None:
return {"error": "payload_required"}
text = data
if isinstance(data, dict):
text = data.get("data")
if text is None:
return {"error": "payload_required"}
try:
server.send_input(str(text))
except Exception as exc:
logger.debug("ps_send failed tunnel_id=%s: %s", tunnel_id, exc, exc_info=True)
return {"error": "ps_send_failed"}
shell_bridge.send(request.sid, str(payload))
return {"status": "ok"}
@socket_server.on("ps_resize", namespace=tunnel_namespace)
def _ws_ps_resize(data: Any) -> Any:
server, tunnel_id, error = _require_ps_server()
if server is None:
return error
cols = None
rows = None
if isinstance(data, dict):
cols = data.get("cols")
rows = data.get("rows")
try:
cols_int = int(cols) if cols is not None else 120
rows_int = int(rows) if rows is not None else 32
cols_int = max(20, min(cols_int, 300))
rows_int = max(10, min(rows_int, 200))
server.send_resize(cols_int, rows_int)
return {"status": "ok", "cols": cols_int, "rows": rows_int}
except Exception as exc:
logger.debug("ps_resize failed tunnel_id=%s: %s", tunnel_id, exc, exc_info=True)
return {"error": "ps_resize_failed"}
@socket_server.on("vpn_shell_close")
def _vpn_shell_close() -> Dict[str, Any]:
shell_bridge.close(request.sid)
return {"status": "ok"}
@socket_server.on("ps_poll", namespace=tunnel_namespace)
def _ws_ps_poll(data: Any = None) -> Any: # data is ignored; socketio passes it even when unused
server, tunnel_id, error = _require_ps_server()
if server is None:
return error
try:
output = server.drain_output()
status = server.status()
return {"output": output, "status": status}
except Exception as exc:
logger.debug("ps_poll failed tunnel_id=%s: %s", tunnel_id, exc, exc_info=True)
return {"error": "ps_poll_failed"}
@socket_server.on("disconnect", namespace=tunnel_namespace)
def _ws_tunnel_disconnect():
sid = request.sid
tunnel_id = _operator_sessions.pop(sid, None)
if tunnel_id and tunnel_id not in _operator_sessions.values():
try:
tunnel_service.stop_tunnel(tunnel_id, reason="operator_socket_disconnect")
except Exception as exc:
logger.debug("ws_tunnel_disconnect stop_tunnel failed tunnel_id=%s: %s", tunnel_id, exc, exc_info=True)
@socket_server.on("disconnect")
def _ws_disconnect() -> None:
shell_bridge.close(request.sid)

127
Data/Engine/services/WebSocket/vpn_shell.py Normal file
View File

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