Overhaul of Reverse Tunnel Code

This commit is contained in:
2025-12-06 20:07:08 -07:00
parent 737bf1faef
commit 178257c588
42 changed files with 1240 additions and 357 deletions

View File

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

View File

@@ -0,0 +1,49 @@
"""Placeholder Bash channel (Agent side)."""
from __future__ import annotations
import asyncio
from typing import Any, Dict, Optional
MSG_DATA = 0x05
MSG_CONTROL = 0x09
MSG_CLOSE = 0x08
CLOSE_PROTOCOL_ERROR = 3
CLOSE_AGENT_SHUTDOWN = 6
class BashChannel:
"""Stub Bash handler that immediately reports unsupported."""
def __init__(self, role, tunnel, channel_id: int, metadata: Optional[Dict[str, Any]]):
self.role = role
self.tunnel = tunnel
self.channel_id = channel_id
self.metadata = metadata or {}
self.loop = getattr(role, "loop", None) or asyncio.get_event_loop()
self._closed = False
async def start(self) -> None:
# Until Bash support is implemented, close the channel to free resources.
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="bash_unsupported")
async def on_frame(self, frame) -> None:
if self._closed:
return
if frame.msg_type in (MSG_DATA, MSG_CONTROL):
# Ignore payloads but acknowledge by stopping the channel to avoid leaks.
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="bash_unsupported")
elif frame.msg_type == MSG_CLOSE:
await self.stop(code=CLOSE_AGENT_SHUTDOWN, reason="operator_close")
async def stop(self, code: int = CLOSE_PROTOCOL_ERROR, reason: str = "") -> None:
if self._closed:
return
self._closed = True
try:
await self.role._send_frame(self.tunnel, self.role.close_frame(self.channel_id, code, reason or "bash_closed"))
except Exception:
pass
self.role._log(f"reverse_tunnel bash channel stopped channel={self.channel_id} reason={reason or 'closed'}")
__all__ = ["BashChannel"]

View File

@@ -0,0 +1,35 @@
"""Expose the PowerShell channel under the domain path, with file-based import fallback."""
from __future__ import annotations
import importlib.util
from pathlib import Path
powershell_module = None
# Attempt package-relative import first
try: # pragma: no cover - best effort
from ....ReverseTunnel import tunnel_Powershell as powershell_module # type: ignore
except Exception:
powershell_module = None
# Fallback: load directly from file path to survive non-package runtimes
if powershell_module is None:
try:
base = Path(__file__).resolve().parents[3] / "ReverseTunnel" / "tunnel_Powershell.py"
spec = importlib.util.spec_from_file_location("tunnel_Powershell", base)
if spec and spec.loader:
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module) # type: ignore
powershell_module = module
except Exception:
powershell_module = None
if powershell_module and hasattr(powershell_module, "PowershellChannel"):
PowershellChannel = powershell_module.PowershellChannel # type: ignore
else: # pragma: no cover - safety guard
class PowershellChannel: # type: ignore
def __init__(self, *args, **kwargs):
raise ImportError("PowerShell channel unavailable")
__all__ = ["PowershellChannel"]

View File

@@ -0,0 +1,6 @@
"""Protocol handlers for interactive shell tunnels (Agent side)."""
from .Powershell import PowershellChannel
from .Bash import BashChannel
__all__ = ["PowershellChannel", "BashChannel"]

View File

@@ -0,0 +1,3 @@
"""Interactive shell domain (PowerShell/Bash) handlers."""
__all__ = ["tunnel", "Protocols"]

View File

@@ -0,0 +1,5 @@
"""Placeholder module for remote interactive shell tunnel domain (Agent side)."""
DOMAIN_NAME = "remote-interactive-shell"
__all__ = ["DOMAIN_NAME"]

View File

@@ -0,0 +1,47 @@
"""Placeholder SSH channel (Agent side)."""
from __future__ import annotations
import asyncio
from typing import Any, Dict, Optional
MSG_DATA = 0x05
MSG_CONTROL = 0x09
MSG_CLOSE = 0x08
CLOSE_PROTOCOL_ERROR = 3
CLOSE_AGENT_SHUTDOWN = 6
class SSHChannel:
"""Stub SSH handler that marks the channel unsupported for now."""
def __init__(self, role, tunnel, channel_id: int, metadata: Optional[Dict[str, Any]]):
self.role = role
self.tunnel = tunnel
self.channel_id = channel_id
self.metadata = metadata or {}
self.loop = getattr(role, "loop", None) or asyncio.get_event_loop()
self._closed = False
async def start(self) -> None:
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="ssh_unsupported")
async def on_frame(self, frame) -> None:
if self._closed:
return
if frame.msg_type in (MSG_DATA, MSG_CONTROL):
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="ssh_unsupported")
elif frame.msg_type == MSG_CLOSE:
await self.stop(code=CLOSE_AGENT_SHUTDOWN, reason="operator_close")
async def stop(self, code: int = CLOSE_PROTOCOL_ERROR, reason: str = "") -> None:
if self._closed:
return
self._closed = True
try:
await self.role._send_frame(self.tunnel, self.role.close_frame(self.channel_id, code, reason or "ssh_closed"))
except Exception:
pass
self.role._log(f"reverse_tunnel ssh channel stopped channel={self.channel_id} reason={reason or 'closed'}")
__all__ = ["SSHChannel"]

View File

@@ -0,0 +1,47 @@
"""Placeholder WinRM channel (Agent side)."""
from __future__ import annotations
import asyncio
from typing import Any, Dict, Optional
MSG_DATA = 0x05
MSG_CONTROL = 0x09
MSG_CLOSE = 0x08
CLOSE_PROTOCOL_ERROR = 3
CLOSE_AGENT_SHUTDOWN = 6
class WinRMChannel:
"""Stub WinRM handler that marks the channel unsupported for now."""
def __init__(self, role, tunnel, channel_id: int, metadata: Optional[Dict[str, Any]]):
self.role = role
self.tunnel = tunnel
self.channel_id = channel_id
self.metadata = metadata or {}
self.loop = getattr(role, "loop", None) or asyncio.get_event_loop()
self._closed = False
async def start(self) -> None:
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="winrm_unsupported")
async def on_frame(self, frame) -> None:
if self._closed:
return
if frame.msg_type in (MSG_DATA, MSG_CONTROL):
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="winrm_unsupported")
elif frame.msg_type == MSG_CLOSE:
await self.stop(code=CLOSE_AGENT_SHUTDOWN, reason="operator_close")
async def stop(self, code: int = CLOSE_PROTOCOL_ERROR, reason: str = "") -> None:
if self._closed:
return
self._closed = True
try:
await self.role._send_frame(self.tunnel, self.role.close_frame(self.channel_id, code, reason or "winrm_closed"))
except Exception:
pass
self.role._log(f"reverse_tunnel winrm channel stopped channel={self.channel_id} reason={reason or 'closed'}")
__all__ = ["WinRMChannel"]

View File

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

View File

@@ -0,0 +1,3 @@
"""Remote management domain (SSH/WinRM) handlers."""
__all__ = ["tunnel", "Protocols"]

View File

@@ -0,0 +1,5 @@
"""Placeholder module for remote management domain (Agent side)."""
DOMAIN_NAME = "remote-management"
__all__ = ["DOMAIN_NAME"]

View File

@@ -0,0 +1,47 @@
"""Placeholder RDP channel (Agent side)."""
from __future__ import annotations
import asyncio
from typing import Any, Dict, Optional
MSG_DATA = 0x05
MSG_CONTROL = 0x09
MSG_CLOSE = 0x08
CLOSE_PROTOCOL_ERROR = 3
CLOSE_AGENT_SHUTDOWN = 6
class RDPChannel:
"""Stub RDP handler that marks the channel unsupported for now."""
def __init__(self, role, tunnel, channel_id: int, metadata: Optional[Dict[str, Any]]):
self.role = role
self.tunnel = tunnel
self.channel_id = channel_id
self.metadata = metadata or {}
self.loop = getattr(role, "loop", None) or asyncio.get_event_loop()
self._closed = False
async def start(self) -> None:
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="rdp_unsupported")
async def on_frame(self, frame) -> None:
if self._closed:
return
if frame.msg_type in (MSG_DATA, MSG_CONTROL):
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="rdp_unsupported")
elif frame.msg_type == MSG_CLOSE:
await self.stop(code=CLOSE_AGENT_SHUTDOWN, reason="operator_close")
async def stop(self, code: int = CLOSE_PROTOCOL_ERROR, reason: str = "") -> None:
if self._closed:
return
self._closed = True
try:
await self.role._send_frame(self.tunnel, self.role.close_frame(self.channel_id, code, reason or "rdp_closed"))
except Exception:
pass
self.role._log(f"reverse_tunnel rdp channel stopped channel={self.channel_id} reason={reason or 'closed'}")
__all__ = ["RDPChannel"]

View File

@@ -0,0 +1,47 @@
"""Placeholder VNC channel (Agent side)."""
from __future__ import annotations
import asyncio
from typing import Any, Dict, Optional
MSG_DATA = 0x05
MSG_CONTROL = 0x09
MSG_CLOSE = 0x08
CLOSE_PROTOCOL_ERROR = 3
CLOSE_AGENT_SHUTDOWN = 6
class VNCChannel:
"""Stub VNC handler that marks the channel unsupported for now."""
def __init__(self, role, tunnel, channel_id: int, metadata: Optional[Dict[str, Any]]):
self.role = role
self.tunnel = tunnel
self.channel_id = channel_id
self.metadata = metadata or {}
self.loop = getattr(role, "loop", None) or asyncio.get_event_loop()
self._closed = False
async def start(self) -> None:
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="vnc_unsupported")
async def on_frame(self, frame) -> None:
if self._closed:
return
if frame.msg_type in (MSG_DATA, MSG_CONTROL):
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="vnc_unsupported")
elif frame.msg_type == MSG_CLOSE:
await self.stop(code=CLOSE_AGENT_SHUTDOWN, reason="operator_close")
async def stop(self, code: int = CLOSE_PROTOCOL_ERROR, reason: str = "") -> None:
if self._closed:
return
self._closed = True
try:
await self.role._send_frame(self.tunnel, self.role.close_frame(self.channel_id, code, reason or "vnc_closed"))
except Exception:
pass
self.role._log(f"reverse_tunnel vnc channel stopped channel={self.channel_id} reason={reason or 'closed'}")
__all__ = ["VNCChannel"]

View File

@@ -0,0 +1,47 @@
"""Placeholder WebRTC channel (Agent side)."""
from __future__ import annotations
import asyncio
from typing import Any, Dict, Optional
MSG_DATA = 0x05
MSG_CONTROL = 0x09
MSG_CLOSE = 0x08
CLOSE_PROTOCOL_ERROR = 3
CLOSE_AGENT_SHUTDOWN = 6
class WebRTCChannel:
"""Stub WebRTC handler that marks the channel unsupported for now."""
def __init__(self, role, tunnel, channel_id: int, metadata: Optional[Dict[str, Any]]):
self.role = role
self.tunnel = tunnel
self.channel_id = channel_id
self.metadata = metadata or {}
self.loop = getattr(role, "loop", None) or asyncio.get_event_loop()
self._closed = False
async def start(self) -> None:
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="webrtc_unsupported")
async def on_frame(self, frame) -> None:
if self._closed:
return
if frame.msg_type in (MSG_DATA, MSG_CONTROL):
await self.stop(code=CLOSE_PROTOCOL_ERROR, reason="webrtc_unsupported")
elif frame.msg_type == MSG_CLOSE:
await self.stop(code=CLOSE_AGENT_SHUTDOWN, reason="operator_close")
async def stop(self, code: int = CLOSE_PROTOCOL_ERROR, reason: str = "") -> None:
if self._closed:
return
self._closed = True
try:
await self.role._send_frame(self.tunnel, self.role.close_frame(self.channel_id, code, reason or "webrtc_closed"))
except Exception:
pass
self.role._log(f"reverse_tunnel webrtc channel stopped channel={self.channel_id} reason={reason or 'closed'}")
__all__ = ["WebRTCChannel"]

View File

@@ -0,0 +1,7 @@
"""Protocol handlers for remote video tunnels (Agent side)."""
from .WebRTC import WebRTCChannel
from .RDP import RDPChannel
from .VNC import VNCChannel
__all__ = ["WebRTCChannel", "RDPChannel", "VNCChannel"]

View File

@@ -0,0 +1,3 @@
"""Remote video/desktop domain (RDP/VNC/WebRTC) handlers."""
__all__ = ["tunnel", "Protocols"]

View File

@@ -0,0 +1,5 @@
"""Placeholder module for remote video domain (Agent side)."""
DOMAIN_NAME = "remote-video"
__all__ = ["DOMAIN_NAME"]

View File

@@ -14,25 +14,109 @@ import aiohttp
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ed25519
# Capture import errors for the PowerShell handler so we can report why it is missing.
# Capture import errors for protocol handlers so we can report why they are missing.
PS_IMPORT_ERROR: Optional[str] = None
BASH_IMPORT_ERROR: Optional[str] = None
tunnel_SSH = None
tunnel_WinRM = None
tunnel_VNC = None
tunnel_RDP = None
tunnel_WebRTC = None
tunnel_Powershell = None
tunnel_Bash = None
def _load_protocol_module(module_name: str, rel_parts: list[str]) -> tuple[Optional[object], Optional[str]]:
"""Load a protocol handler directly from a file path to survive non-package runtimes."""
base = Path(__file__).parent
path = base
for part in rel_parts:
path = path / part
if not path.exists():
return None, f"path_missing:{path}"
try:
spec = importlib.util.spec_from_file_location(module_name, path)
if not spec or not spec.loader:
return None, "spec_failed"
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module) # type: ignore
return module, None
except Exception as exc: # pragma: no cover - defensive
return None, repr(exc)
try:
from .ReverseTunnel import tunnel_Powershell # type: ignore
from .Reverse_Tunnels.remote_interactive_shell.Protocols import Powershell as tunnel_Powershell # type: ignore
except Exception as exc: # pragma: no cover - best-effort logging only
PS_IMPORT_ERROR = repr(exc)
# Try manual import from file to survive non-package execution.
try:
_ps_path = Path(__file__).parent / "ReverseTunnel" / "tunnel_Powershell.py"
if _ps_path.exists():
spec = importlib.util.spec_from_file_location("tunnel_Powershell", _ps_path)
if spec and spec.loader:
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module) # type: ignore
tunnel_Powershell = module
PS_IMPORT_ERROR = None
except Exception as exc2: # pragma: no cover - diagnostic only
PS_IMPORT_ERROR = f"{PS_IMPORT_ERROR} | fallback_load_failed={exc2!r}"
_module, _err = _load_protocol_module(
"tunnel_Powershell",
["Reverse_Tunnels", "remote_interactive_shell", "Protocols", "Powershell.py"],
)
if _module:
tunnel_Powershell = _module
PS_IMPORT_ERROR = None
else:
try:
from .ReverseTunnel import tunnel_Powershell # type: ignore # legacy fallback
PS_IMPORT_ERROR = None
except Exception as exc2: # pragma: no cover - diagnostic only
PS_IMPORT_ERROR = f"{PS_IMPORT_ERROR} | legacy_fallback={exc2!r} | file_load_failed={_err}"
try:
from .Reverse_Tunnels.remote_interactive_shell.Protocols import Bash as tunnel_Bash # type: ignore
except Exception as exc: # pragma: no cover - best-effort logging only
BASH_IMPORT_ERROR = repr(exc)
_module, _err = _load_protocol_module(
"tunnel_Bash",
["Reverse_Tunnels", "remote_interactive_shell", "Protocols", "Bash.py"],
)
if _module:
tunnel_Bash = _module
BASH_IMPORT_ERROR = None
else:
BASH_IMPORT_ERROR = f"{BASH_IMPORT_ERROR} | file_load_failed={_err}"
try:
from .Reverse_Tunnels.remote_management.Protocols import SSH as tunnel_SSH # type: ignore
except Exception:
_module, _err = _load_protocol_module(
"tunnel_SSH",
["Reverse_Tunnels", "remote_management", "Protocols", "SSH.py"],
)
tunnel_SSH = _module
try:
from .Reverse_Tunnels.remote_management.Protocols import WinRM as tunnel_WinRM # type: ignore
except Exception:
_module, _err = _load_protocol_module(
"tunnel_WinRM",
["Reverse_Tunnels", "remote_management", "Protocols", "WinRM.py"],
)
tunnel_WinRM = _module
try:
from .Reverse_Tunnels.remote_video.Protocols import VNC as tunnel_VNC # type: ignore
except Exception:
_module, _err = _load_protocol_module(
"tunnel_VNC",
["Reverse_Tunnels", "remote_video", "Protocols", "VNC.py"],
)
tunnel_VNC = _module
try:
from .Reverse_Tunnels.remote_video.Protocols import RDP as tunnel_RDP # type: ignore
except Exception:
_module, _err = _load_protocol_module(
"tunnel_RDP",
["Reverse_Tunnels", "remote_video", "Protocols", "RDP.py"],
)
tunnel_RDP = _module
try:
from .Reverse_Tunnels.remote_video.Protocols import WebRTC as tunnel_WebRTC # type: ignore
except Exception:
_module, _err = _load_protocol_module(
"tunnel_WebRTC",
["Reverse_Tunnels", "remote_video", "Protocols", "WebRTC.py"],
)
tunnel_WebRTC = _module
ROLE_NAME = "reverse_tunnel"
ROLE_CONTEXTS = ["interactive", "system"]
@@ -185,10 +269,14 @@ class Role:
self._active: Dict[str, ActiveTunnel] = {}
self._domain_claims: Dict[str, str] = {}
self._domain_limits: Dict[str, Optional[int]] = {
"ps": 1,
"remote-interactive-shell": 2,
"remote-management": 1,
"remote-video": 2,
# Legacy / protocol fallbacks
"ps": 2,
"rdp": 1,
"vnc": 1,
"webrtc": 1,
"webrtc": 2,
"ssh": None,
"winrm": None,
}
@@ -211,6 +299,28 @@ class Role:
)
except Exception as exc:
self._log(f"reverse_tunnel ps handler registration failed: {exc}", error=True)
try:
if tunnel_Bash and hasattr(tunnel_Bash, "BashChannel"):
self._protocol_handlers["bash"] = tunnel_Bash.BashChannel
module_path = getattr(tunnel_Bash, "__file__", None)
self._log(f"reverse_tunnel bash handler registered (BashChannel) module={module_path}")
elif BASH_IMPORT_ERROR:
self._log(f"reverse_tunnel bash handler NOT registered (missing module/class) import_error={BASH_IMPORT_ERROR}", error=True)
except Exception as exc:
self._log(f"reverse_tunnel bash handler registration failed: {exc}", error=True)
try:
if tunnel_SSH and hasattr(tunnel_SSH, "SSHChannel"):
self._protocol_handlers["ssh"] = tunnel_SSH.SSHChannel
if tunnel_WinRM and hasattr(tunnel_WinRM, "WinRMChannel"):
self._protocol_handlers["winrm"] = tunnel_WinRM.WinRMChannel
if tunnel_VNC and hasattr(tunnel_VNC, "VNCChannel"):
self._protocol_handlers["vnc"] = tunnel_VNC.VNCChannel
if tunnel_RDP and hasattr(tunnel_RDP, "RDPChannel"):
self._protocol_handlers["rdp"] = tunnel_RDP.RDPChannel
if tunnel_WebRTC and hasattr(tunnel_WebRTC, "WebRTCChannel"):
self._protocol_handlers["webrtc"] = tunnel_WebRTC.WebRTCChannel
except Exception as exc:
self._log(f"reverse_tunnel protocol handler registration failed: {exc}", error=True)
# ------------------------------------------------------------------ Logging
def _log(self, message: str, *, error: bool = False) -> None:
@@ -728,6 +838,7 @@ class Role:
try:
handler = handler_cls(self, tunnel, frame.channel_id, metadata)
except Exception:
self._log(f"reverse_tunnel channel handler fallback to BaseChannel protocol={protocol}", error=True)
handler = BaseChannel(self, tunnel, frame.channel_id, metadata)
tunnel.channels[frame.channel_id] = handler
await handler.start()