From 52e40c375331f78970f8727087b571620664b35b Mon Sep 17 00:00:00 2001 From: Nicole Rappe Date: Mon, 1 Dec 2025 02:54:46 -0700 Subject: [PATCH] Fixed Agent Role Implementation Errors / Bugs --- .../services/WebSocket/Agent/ReverseTunnel.py | 9 +++- .../WebSocket/Agent/ReverseTunnel/__init__.py | 2 - .../Powershell.py | 44 +++++++++++-------- .../Agent/ReverseTunnelProtocols/__init__.py | 6 +++ 4 files changed, 38 insertions(+), 23 deletions(-) delete mode 100644 Data/Engine/services/WebSocket/Agent/ReverseTunnel/__init__.py rename Data/Engine/services/WebSocket/Agent/{ReverseTunnel => ReverseTunnelProtocols}/Powershell.py (77%) create mode 100644 Data/Engine/services/WebSocket/Agent/ReverseTunnelProtocols/__init__.py diff --git a/Data/Engine/services/WebSocket/Agent/ReverseTunnel.py b/Data/Engine/services/WebSocket/Agent/ReverseTunnel.py index 5bda734f..6beecfc2 100644 --- a/Data/Engine/services/WebSocket/Agent/ReverseTunnel.py +++ b/Data/Engine/services/WebSocket/Agent/ReverseTunnel.py @@ -29,7 +29,7 @@ from typing import Callable, Deque, Dict, Iterable, List, Optional, Tuple from collections import deque from threading import Thread -from .ReverseTunnel.Powershell import PowershellChannelServer +from .ReverseTunnelProtocols import PowershellChannelServer try: # websockets is added to engine requirements import websockets @@ -1000,7 +1000,12 @@ class ReverseTunnelService: if lease is None or (lease.domain or "").lower() != "ps": return None bridge = self.ensure_bridge(lease) - server = PowershellChannelServer(bridge=bridge, service=self) + server = PowershellChannelServer( + bridge=bridge, + service=self, + frame_cls=TunnelFrame, + close_frame_fn=close_frame, + ) self._ps_servers[tunnel_id] = server return server diff --git a/Data/Engine/services/WebSocket/Agent/ReverseTunnel/__init__.py b/Data/Engine/services/WebSocket/Agent/ReverseTunnel/__init__.py deleted file mode 100644 index db1c9337..00000000 --- a/Data/Engine/services/WebSocket/Agent/ReverseTunnel/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -"""Protocol-specific helpers for Reverse Tunnel (Engine side).""" - diff --git a/Data/Engine/services/WebSocket/Agent/ReverseTunnel/Powershell.py b/Data/Engine/services/WebSocket/Agent/ReverseTunnelProtocols/Powershell.py similarity index 77% rename from Data/Engine/services/WebSocket/Agent/ReverseTunnel/Powershell.py rename to Data/Engine/services/WebSocket/Agent/ReverseTunnelProtocols/Powershell.py index 72d2b82d..21275905 100644 --- a/Data/Engine/services/WebSocket/Agent/ReverseTunnel/Powershell.py +++ b/Data/Engine/services/WebSocket/Agent/ReverseTunnelProtocols/Powershell.py @@ -5,24 +5,21 @@ import json from collections import deque from typing import Any, Deque, Dict, List, Optional -from ..ReverseTunnel import ( - CLOSE_AGENT_SHUTDOWN, - CLOSE_OK, - CLOSE_PROTOCOL_ERROR, - MSG_CHANNEL_ACK, - MSG_CHANNEL_OPEN, - MSG_CLOSE, - MSG_CONTROL, - MSG_DATA, - TunnelFrame, - close_frame, -) +# 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): + 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 @@ -33,9 +30,11 @@ class PowershellChannelServer: 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: TunnelFrame) -> None: + def handle_agent_frame(self, frame) -> None: if frame.channel_id != self.channel_id: return if frame.msg_type == MSG_CHANNEL_ACK: @@ -73,7 +72,7 @@ class PowershellChannelServer: {"protocol": "ps", "metadata": {"cols": cols, "rows": rows}}, separators=(",", ":"), ).encode("utf-8") - frame = TunnelFrame(msg_type=MSG_CHANNEL_OPEN, channel_id=self.channel_id, payload=payload) + 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( @@ -88,21 +87,29 @@ class PowershellChannelServer: if self._closed: return payload = data.encode("utf-8", errors="replace") - frame = TunnelFrame(msg_type=MSG_DATA, channel_id=self.channel_id, payload=payload) + 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 = TunnelFrame(msg_type=MSG_CONTROL, channel_id=self.channel_id, payload=payload) + 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 - self.bridge.operator_to_agent(close_frame(self.channel_id, code, reason)) + 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]: @@ -127,4 +134,3 @@ class PowershellChannelServer: "close_reason": self._close_reason, "close_code": self._close_code, } - diff --git a/Data/Engine/services/WebSocket/Agent/ReverseTunnelProtocols/__init__.py b/Data/Engine/services/WebSocket/Agent/ReverseTunnelProtocols/__init__.py new file mode 100644 index 00000000..dfd26238 --- /dev/null +++ b/Data/Engine/services/WebSocket/Agent/ReverseTunnelProtocols/__init__.py @@ -0,0 +1,6 @@ +"""Protocol-specific helpers for Reverse Tunnel (Engine side).""" + +from .Powershell import PowershellChannelServer + +__all__ = ["PowershellChannelServer"] +