mirror of
https://github.com/bunny-lab-io/Borealis.git
synced 2026-02-06 08:50:31 -07:00
Removed RDP in favor of VNC / Made WireGuard Tunnel Persistent
This commit is contained in:
@@ -1,77 +0,0 @@
|
||||
# ======================================================
|
||||
# Data\Agent\Roles\role_RDP.py
|
||||
# Description: Optional RDP readiness helper for Borealis (Windows-only).
|
||||
#
|
||||
# API Endpoints (if applicable): None
|
||||
# ======================================================
|
||||
|
||||
"""RDP readiness helper role (no-op unless enabled via environment flags)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
ROLE_NAME = "RDP"
|
||||
ROLE_CONTEXTS = ["system"]
|
||||
|
||||
|
||||
def _log_path() -> Path:
|
||||
root = Path(__file__).resolve().parents[2] / "Logs"
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
return root / "rdp.log"
|
||||
|
||||
|
||||
def _write_log(message: str) -> None:
|
||||
ts = time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime())
|
||||
try:
|
||||
_log_path().open("a", encoding="utf-8").write(f"[{ts}] [rdp-role] {message}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _bool_env(name: str) -> bool:
|
||||
value = os.environ.get(name)
|
||||
if value is None:
|
||||
return False
|
||||
return str(value).strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _enable_rdp_windows() -> None:
|
||||
command = (
|
||||
"Set-ItemProperty -Path 'HKLM:\\System\\CurrentControlSet\\Control\\Terminal Server' "
|
||||
"-Name fDenyTSConnections -Value 0; "
|
||||
"Set-Service -Name TermService -StartupType Automatic; "
|
||||
"Start-Service -Name TermService; "
|
||||
"Enable-NetFirewallRule -DisplayGroup 'Remote Desktop'"
|
||||
)
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["powershell.exe", "-NoProfile", "-Command", command],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
_write_log(f"RDP enable failed: {result.stderr.strip()}")
|
||||
else:
|
||||
_write_log("RDP enable applied (registry/service/firewall).")
|
||||
except Exception as exc:
|
||||
_write_log(f"RDP enable failed: {exc}")
|
||||
|
||||
|
||||
class Role:
|
||||
def __init__(self, ctx) -> None:
|
||||
self.ctx = ctx
|
||||
auto_enable = _bool_env("BOREALIS_RDP_AUTO_ENABLE")
|
||||
_write_log(f"RDP role loaded auto_enable={auto_enable}")
|
||||
if auto_enable and os.name == "nt":
|
||||
_enable_rdp_windows()
|
||||
|
||||
def register_events(self) -> None:
|
||||
return
|
||||
|
||||
def stop_all(self) -> None:
|
||||
return
|
||||
468
Data/Agent/Roles/role_VNC.py
Normal file
468
Data/Agent/Roles/role_VNC.py
Normal file
@@ -0,0 +1,468 @@
|
||||
# ======================================================
|
||||
# Data/Agent/Roles/role_VNC.py
|
||||
# Description: On-demand UltraVNC server lifecycle over WireGuard.
|
||||
#
|
||||
# API Endpoints (if applicable): None
|
||||
# ======================================================
|
||||
|
||||
"""UltraVNC role (Windows) for on-demand VNC sessions over WireGuard."""
|
||||
from __future__ import annotations
|
||||
|
||||
import ipaddress
|
||||
import os
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
ROLE_NAME = "VNC"
|
||||
ROLE_CONTEXTS = ["system"]
|
||||
|
||||
VNC_FIREWALL_RULE_NAME = "Borealis - VNC - UltraVNC"
|
||||
DEFAULT_VNC_PORT = 5900
|
||||
ULTRAVNC_SERVICE_NAME = os.environ.get("BOREALIS_ULTRAVNC_SERVICE") or "uvnc_service"
|
||||
|
||||
|
||||
def _log_path() -> Path:
|
||||
root = Path(__file__).resolve().parents[2] / "Logs" / "VPN_Tunnel"
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
return root / "vnc.log"
|
||||
|
||||
|
||||
def _write_log(message: str) -> None:
|
||||
ts = time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime())
|
||||
try:
|
||||
_log_path().open("a", encoding="utf-8").write(f"[{ts}] [vnc] {message}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _find_project_root() -> Optional[Path]:
|
||||
override = os.environ.get("BOREALIS_ROOT") or os.environ.get("BOREALIS_PROJECT_ROOT")
|
||||
if override:
|
||||
try:
|
||||
override_path = Path(override).expanduser().resolve()
|
||||
if override_path.is_dir():
|
||||
return override_path
|
||||
except Exception:
|
||||
pass
|
||||
current = Path(__file__).resolve()
|
||||
for parent in (current, *current.parents):
|
||||
try:
|
||||
if (parent / "Borealis.ps1").is_file() or (parent / "users.json").is_file():
|
||||
return parent
|
||||
except Exception:
|
||||
continue
|
||||
try:
|
||||
return current.parents[3]
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_vnc_root() -> Optional[Path]:
|
||||
root = _find_project_root()
|
||||
if not root:
|
||||
return None
|
||||
candidate = root / "Dependencies" / "UltraVNC_Server"
|
||||
if candidate.is_dir():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_vnc_exe() -> Optional[str]:
|
||||
override = os.environ.get("BOREALIS_VNC_SERVER_BIN")
|
||||
if override:
|
||||
try:
|
||||
if Path(override).is_file():
|
||||
return str(Path(override))
|
||||
except Exception:
|
||||
pass
|
||||
vnc_root = _resolve_vnc_root()
|
||||
if vnc_root:
|
||||
preferred = [
|
||||
vnc_root / "payload" / "x64" / "winvnc.exe",
|
||||
vnc_root / "payload" / "x64" / "winvnc64.exe",
|
||||
vnc_root / "winvnc64.exe",
|
||||
vnc_root / "winvnc.exe",
|
||||
vnc_root / "payload" / "x86" / "winvnc.exe",
|
||||
]
|
||||
for candidate in preferred:
|
||||
if candidate.is_file():
|
||||
return str(candidate)
|
||||
try:
|
||||
for candidate in vnc_root.rglob("winvnc*.exe"):
|
||||
if candidate.is_file():
|
||||
return str(candidate)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_vnc_port(value: Any = None) -> int:
|
||||
raw = value if value is not None else os.environ.get("BOREALIS_VNC_PORT")
|
||||
try:
|
||||
port = int(raw) if raw is not None else DEFAULT_VNC_PORT
|
||||
except Exception:
|
||||
port = DEFAULT_VNC_PORT
|
||||
if port < 1 or port > 65535:
|
||||
return DEFAULT_VNC_PORT
|
||||
return port
|
||||
|
||||
|
||||
def _resolve_vnc_password_tool(root: Optional[Path]) -> Optional[str]:
|
||||
override = os.environ.get("BOREALIS_VNC_PASSWORD_TOOL")
|
||||
if override:
|
||||
try:
|
||||
if Path(override).is_file():
|
||||
return str(Path(override))
|
||||
except Exception:
|
||||
pass
|
||||
candidates: list[Path] = []
|
||||
if root:
|
||||
candidates.extend(
|
||||
[
|
||||
root / "createpassword.exe",
|
||||
root / "tools" / "createpassword.exe",
|
||||
root / "createpassword64.exe",
|
||||
root / "tools" / "createpassword64.exe",
|
||||
root / "CreatePassword.exe",
|
||||
]
|
||||
)
|
||||
vnc_root = _resolve_vnc_root()
|
||||
if vnc_root and vnc_root != root:
|
||||
candidates.extend(
|
||||
[
|
||||
vnc_root / "createpassword.exe",
|
||||
vnc_root / "tools" / "createpassword.exe",
|
||||
vnc_root / "createpassword64.exe",
|
||||
vnc_root / "tools" / "createpassword64.exe",
|
||||
]
|
||||
)
|
||||
for candidate in candidates:
|
||||
if candidate.is_file():
|
||||
return str(candidate)
|
||||
try:
|
||||
for candidate in root.rglob("createpassword.exe"):
|
||||
if candidate.is_file():
|
||||
return str(candidate)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _ensure_ultravnc_ini(config_dir: Path, port: int) -> Optional[Path]:
|
||||
try:
|
||||
config_dir.mkdir(parents=True, exist_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
ini_path = config_dir / "ultravnc.ini"
|
||||
try:
|
||||
lines = [
|
||||
"UseRegistry=0",
|
||||
"AuthRequired=1",
|
||||
"MSLogonRequired=0",
|
||||
"NewMSLogon=0",
|
||||
f"PortNumber={port}",
|
||||
"AutoPortSelect=0",
|
||||
"SocketConnect=1",
|
||||
"HTTPConnect=0",
|
||||
"AllowShutdown=0",
|
||||
"DisableTrayIcon=1",
|
||||
"EnableFileTransfer=0",
|
||||
]
|
||||
ini_path.write_text("\n".join(lines), encoding="utf-8")
|
||||
return ini_path
|
||||
except Exception as exc:
|
||||
_write_log(f"Failed to write ultravnc.ini: {exc}")
|
||||
return None
|
||||
|
||||
|
||||
def _parse_allowed_ips(value: Any) -> Optional[str]:
|
||||
if isinstance(value, list):
|
||||
if not value:
|
||||
return None
|
||||
return str(value[0])
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value.strip()
|
||||
return None
|
||||
|
||||
|
||||
class VncManager:
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._last_port: Optional[int] = None
|
||||
self._last_password: Optional[str] = None
|
||||
self._vnc_exe = _resolve_vnc_exe()
|
||||
self._password_tool: Optional[str] = None
|
||||
|
||||
def _service_state(self) -> Optional[str]:
|
||||
if os.name != "nt":
|
||||
return None
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["sc.exe", "query", ULTRAVNC_SERVICE_NAME],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
return None
|
||||
for line in (result.stdout or "").splitlines():
|
||||
if "STATE" in line:
|
||||
parts = line.strip().split()
|
||||
if parts:
|
||||
return parts[-1].upper()
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
def _restart_service(self) -> None:
|
||||
if os.name != "nt":
|
||||
return
|
||||
state = self._service_state()
|
||||
if state != "RUNNING":
|
||||
return
|
||||
try:
|
||||
subprocess.run(["sc.exe", "stop", ULTRAVNC_SERVICE_NAME], capture_output=True, text=True, check=False)
|
||||
time.sleep(1)
|
||||
subprocess.run(["sc.exe", "start", ULTRAVNC_SERVICE_NAME], capture_output=True, text=True, check=False)
|
||||
except Exception as exc:
|
||||
_write_log(f"Failed to restart UltraVNC service: {exc}")
|
||||
|
||||
def _ensure_service_running(self) -> bool:
|
||||
if os.name != "nt":
|
||||
return False
|
||||
state = self._service_state()
|
||||
if state == "RUNNING":
|
||||
return True
|
||||
if not self._vnc_exe:
|
||||
self._vnc_exe = _resolve_vnc_exe()
|
||||
if not self._vnc_exe:
|
||||
return False
|
||||
try:
|
||||
if state is None:
|
||||
subprocess.run([self._vnc_exe, "-install"], capture_output=True, text=True, check=False)
|
||||
subprocess.run(
|
||||
["sc.exe", "config", ULTRAVNC_SERVICE_NAME, "start=", "auto"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
subprocess.run(["sc.exe", "start", ULTRAVNC_SERVICE_NAME], capture_output=True, text=True, check=False)
|
||||
except Exception as exc:
|
||||
_write_log(f"Failed to ensure UltraVNC service running: {exc}")
|
||||
return False
|
||||
return self._service_state() == "RUNNING"
|
||||
|
||||
def _normalize_firewall_remote(self, allowed_ips: Optional[str]) -> Optional[str]:
|
||||
if not allowed_ips:
|
||||
return None
|
||||
try:
|
||||
network = ipaddress.ip_network(str(allowed_ips).strip(), strict=False)
|
||||
except Exception:
|
||||
_write_log(f"Refusing to apply VNC firewall rule; invalid allowed_ips={allowed_ips}.")
|
||||
return None
|
||||
if network.prefixlen != 32:
|
||||
_write_log(f"Refusing to apply VNC firewall rule; allowed_ips not /32: {network}.")
|
||||
return None
|
||||
return str(network)
|
||||
|
||||
def _ensure_firewall(self, allowed_ips: Optional[str], port: int) -> None:
|
||||
if os.name != "nt":
|
||||
return
|
||||
remote = self._normalize_firewall_remote(allowed_ips)
|
||||
if not remote:
|
||||
return
|
||||
rule_name = VNC_FIREWALL_RULE_NAME.replace("'", "''")
|
||||
command = (
|
||||
"Remove-NetFirewallRule -DisplayName '{name}' -ErrorAction SilentlyContinue; "
|
||||
"New-NetFirewallRule -DisplayName '{name}' -Direction Inbound -Action Allow "
|
||||
"-Protocol TCP -LocalPort {port} -RemoteAddress {remote} -Profile Any"
|
||||
).format(name=rule_name, port=port, remote=remote)
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["powershell.exe", "-NoProfile", "-Command", command],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
_write_log(f"Failed to ensure VNC firewall rule: {result.stderr.strip()}")
|
||||
else:
|
||||
_write_log(f"Ensured VNC firewall rule for {remote} on port {port}.")
|
||||
except Exception as exc:
|
||||
_write_log(f"Failed to ensure VNC firewall rule: {exc}")
|
||||
|
||||
def _remove_firewall(self) -> None:
|
||||
if os.name != "nt":
|
||||
return
|
||||
rule_name = VNC_FIREWALL_RULE_NAME.replace("'", "''")
|
||||
command = "Remove-NetFirewallRule -DisplayName '{name}' -ErrorAction SilentlyContinue".format(
|
||||
name=rule_name
|
||||
)
|
||||
try:
|
||||
subprocess.run(
|
||||
["powershell.exe", "-NoProfile", "-Command", command],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _apply_password(self, config_dir: Path, password: str) -> Optional[str]:
|
||||
if not password:
|
||||
_write_log("VNC password missing; refusing to start without auth.")
|
||||
return None
|
||||
trimmed = str(password)[:8]
|
||||
if trimmed != password:
|
||||
_write_log("VNC password trimmed to 8 characters for UltraVNC compatibility.")
|
||||
if not self._password_tool:
|
||||
self._password_tool = _resolve_vnc_password_tool(config_dir)
|
||||
if not self._password_tool:
|
||||
_write_log("VNC password tool not found; expected createpassword.exe under Dependencies/UltraVNC_Server.")
|
||||
return None
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[self._password_tool, "-secure", trimmed],
|
||||
cwd=str(config_dir),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
_write_log(f"Failed to apply VNC password: {result.stderr.strip()}")
|
||||
return None
|
||||
return trimmed
|
||||
except Exception as exc:
|
||||
_write_log(f"Failed to apply VNC password: {exc}")
|
||||
return None
|
||||
|
||||
def start(
|
||||
self,
|
||||
*,
|
||||
port: Optional[int],
|
||||
allowed_ips: Optional[str],
|
||||
password: Optional[str],
|
||||
reason: str = "start",
|
||||
) -> None:
|
||||
with self._lock:
|
||||
port_value = _resolve_vnc_port(port)
|
||||
self._ensure_firewall(allowed_ips, port_value)
|
||||
|
||||
if not self._vnc_exe:
|
||||
self._vnc_exe = _resolve_vnc_exe()
|
||||
if not self._vnc_exe:
|
||||
_write_log("UltraVNC server binary not found; expected under Dependencies/UltraVNC_Server.")
|
||||
return
|
||||
|
||||
exe_path = Path(self._vnc_exe)
|
||||
config_dir = exe_path.parent
|
||||
ini_path = _ensure_ultravnc_ini(config_dir, port_value)
|
||||
if not ini_path:
|
||||
return
|
||||
applied_password = self._apply_password(config_dir, password or "")
|
||||
if not applied_password:
|
||||
return
|
||||
|
||||
if not self._ensure_service_running():
|
||||
_write_log("Failed to start UltraVNC service.")
|
||||
return
|
||||
|
||||
if self._last_port != port_value or self._last_password != applied_password:
|
||||
self._restart_service()
|
||||
self._last_port = port_value
|
||||
self._last_password = applied_password
|
||||
_write_log(f"VNC service running port={port_value} reason={reason}.")
|
||||
|
||||
def stop(self, *, reason: str = "stop") -> None:
|
||||
with self._lock:
|
||||
self._remove_firewall()
|
||||
self._last_port = None
|
||||
_write_log(f"VNC firewall closed reason={reason}.")
|
||||
|
||||
|
||||
class Role:
|
||||
def __init__(self, ctx) -> None:
|
||||
self.ctx = ctx
|
||||
self.vnc = VncManager()
|
||||
self._last_allowed_ips: Optional[str] = None
|
||||
hooks = getattr(ctx, "hooks", {}) or {}
|
||||
self._log_hook = hooks.get("log_agent")
|
||||
try:
|
||||
self.vnc.stop(reason="agent_startup")
|
||||
except Exception:
|
||||
self._log("Failed to preflight VNC cleanup.", error=True)
|
||||
try:
|
||||
self.vnc._ensure_service_running()
|
||||
except Exception:
|
||||
self._log("Failed to ensure UltraVNC service running.", error=True)
|
||||
|
||||
def _log(self, message: str, *, error: bool = False) -> None:
|
||||
if callable(self._log_hook):
|
||||
try:
|
||||
self._log_hook(message, fname="VPN_Tunnel/vnc.log")
|
||||
if error:
|
||||
self._log_hook(message, fname="agent.error.log")
|
||||
except Exception:
|
||||
pass
|
||||
_write_log(message)
|
||||
|
||||
def register_events(self) -> None:
|
||||
sio = self.ctx.sio
|
||||
|
||||
@sio.on("vpn_tunnel_start")
|
||||
async def _vpn_tunnel_start(payload):
|
||||
if isinstance(payload, dict):
|
||||
target_agent = payload.get("agent_id")
|
||||
if target_agent and str(target_agent).strip() != str(self.ctx.agent_id).strip():
|
||||
return
|
||||
allowed_ips = payload.get("allowed_ips")
|
||||
self._last_allowed_ips = _parse_allowed_ips(allowed_ips)
|
||||
|
||||
@sio.on("vpn_tunnel_stop")
|
||||
async def _vpn_tunnel_stop(payload):
|
||||
reason = "server_stop"
|
||||
if isinstance(payload, dict):
|
||||
target_agent = payload.get("agent_id")
|
||||
if target_agent and str(target_agent).strip() != str(self.ctx.agent_id).strip():
|
||||
return
|
||||
reason = payload.get("reason") or reason
|
||||
self._log(f"VNC stop requested (reason={reason}).")
|
||||
self.vnc.stop(reason=str(reason))
|
||||
|
||||
@sio.on("vnc_start")
|
||||
async def _vnc_start(payload):
|
||||
if isinstance(payload, dict):
|
||||
target_agent = payload.get("agent_id")
|
||||
if target_agent and str(target_agent).strip() != str(self.ctx.agent_id).strip():
|
||||
return
|
||||
port = payload.get("port")
|
||||
allowed_ips = payload.get("allowed_ips") or self._last_allowed_ips
|
||||
password = payload.get("password") or ""
|
||||
reason = payload.get("reason") or "vnc_session_start"
|
||||
else:
|
||||
port = None
|
||||
allowed_ips = self._last_allowed_ips
|
||||
password = ""
|
||||
reason = "vnc_session_start"
|
||||
self._log(f"VNC start request received (reason={reason}).")
|
||||
self.vnc.start(port=port, allowed_ips=allowed_ips, password=password, reason=str(reason))
|
||||
|
||||
@sio.on("vnc_stop")
|
||||
async def _vnc_stop(payload):
|
||||
reason = "vnc_session_end"
|
||||
if isinstance(payload, dict):
|
||||
target_agent = payload.get("agent_id")
|
||||
if target_agent and str(target_agent).strip() != str(self.ctx.agent_id).strip():
|
||||
return
|
||||
reason = payload.get("reason") or reason
|
||||
self._log(f"VNC stop requested (reason={reason}).")
|
||||
self.vnc.stop(reason=str(reason))
|
||||
|
||||
def stop_all(self) -> None:
|
||||
try:
|
||||
self.vnc.stop(reason="agent_shutdown")
|
||||
except Exception:
|
||||
self._log("Failed to stop VNC during shutdown.", error=True)
|
||||
@@ -8,10 +8,10 @@
|
||||
"""WireGuard client role (Windows) for reverse VPN tunnels.
|
||||
|
||||
This role prepares the WireGuard client config, manages a single active
|
||||
session, enforces idle teardown, and logs lifecycle events to
|
||||
Agent/Logs/VPN_Tunnel/tunnel.log. It binds to Engine Socket.IO events
|
||||
(`vpn_tunnel_start`, `vpn_tunnel_stop`, `vpn_tunnel_activity`) to start/stop
|
||||
the client session with the issued config/token.
|
||||
session, and keeps the tunnel online while the agent service runs. It logs
|
||||
lifecycle events to Agent/Logs/VPN_Tunnel/tunnel.log. It responds to Engine
|
||||
Socket.IO events (`vpn_tunnel_start`, `vpn_tunnel_activity`) and periodically
|
||||
ensures the persistent session via `/api/agent/vpn/ensure`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -55,6 +55,24 @@ TUNNEL_IDLE_ADDRESS = "169.254.255.254/32"
|
||||
FIREWALL_RULE_NAME = "Borealis - WireGuard - Shell"
|
||||
|
||||
|
||||
def _env_int(name: str, default: int, *, min_value: int = 1, max_value: int = 3600) -> int:
|
||||
raw = os.environ.get(name)
|
||||
try:
|
||||
value = int(raw) if raw is not None else default
|
||||
except Exception:
|
||||
value = default
|
||||
if value < min_value:
|
||||
return min_value
|
||||
if value > max_value:
|
||||
return max_value
|
||||
return value
|
||||
|
||||
|
||||
KEEPALIVE_SECONDS = _env_int("BOREALIS_WIREGUARD_KEEPALIVE_SECONDS", 30, min_value=10, max_value=600)
|
||||
ENSURE_INITIAL_DELAY_SECONDS = _env_int("BOREALIS_WIREGUARD_ENSURE_DELAY", 10, min_value=0, max_value=300)
|
||||
ENSURE_INTERVAL_SECONDS = _env_int("BOREALIS_WIREGUARD_ENSURE_INTERVAL", 60, min_value=15, max_value=3600)
|
||||
|
||||
|
||||
def _log_path() -> Path:
|
||||
root = Path(__file__).resolve().parents[2] / "Logs" / "VPN_Tunnel"
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
@@ -69,6 +87,8 @@ def _write_log(message: str) -> None:
|
||||
pass
|
||||
|
||||
|
||||
|
||||
|
||||
def _encode_key(raw: bytes) -> str:
|
||||
return base64.b64encode(raw).decode("ascii").strip()
|
||||
|
||||
@@ -125,6 +145,7 @@ class SessionConfig:
|
||||
self,
|
||||
*,
|
||||
token: Dict[str, Any],
|
||||
tunnel_id: str,
|
||||
virtual_ip: str,
|
||||
allowed_ips: str,
|
||||
endpoint: str,
|
||||
@@ -136,6 +157,7 @@ class SessionConfig:
|
||||
client_public_key: Optional[str] = None,
|
||||
) -> None:
|
||||
self.token = token
|
||||
self.tunnel_id = tunnel_id
|
||||
self.virtual_ip = virtual_ip
|
||||
self.allowed_ips = allowed_ips
|
||||
self.endpoint = endpoint
|
||||
@@ -199,7 +221,7 @@ class WireGuardClient:
|
||||
return True
|
||||
except Exception as exc:
|
||||
_write_log(f"WireGuard service registry check failed: {exc}")
|
||||
return False
|
||||
return False
|
||||
|
||||
def _service_image_path(self) -> Optional[str]:
|
||||
if winreg is None:
|
||||
@@ -223,6 +245,19 @@ class WireGuardClient:
|
||||
return Path(match.group(1))
|
||||
return None
|
||||
|
||||
def _service_state(self) -> Optional[str]:
|
||||
code, out, err = self._run(["sc.exe", "query", self._service_id()])
|
||||
if code != 0:
|
||||
return None
|
||||
text = out or err or ""
|
||||
for line in text.splitlines():
|
||||
if "STATE" not in line:
|
||||
continue
|
||||
match = re.search(r"STATE\s*:\s*\d+\s+(\w+)", line)
|
||||
if match:
|
||||
return match.group(1).upper()
|
||||
return None
|
||||
|
||||
def _wireguard_config_path(self) -> Path:
|
||||
settings_dir = self.temp_root.parent / "Settings" / "WireGuard"
|
||||
candidates = [
|
||||
@@ -457,7 +492,7 @@ class WireGuardClient:
|
||||
f"PublicKey = {session.server_public_key}",
|
||||
f"AllowedIPs = {session.allowed_ips}",
|
||||
f"Endpoint = {session.endpoint}",
|
||||
"PersistentKeepalive = 20",
|
||||
f"PersistentKeepalive = {KEEPALIVE_SECONDS}",
|
||||
]
|
||||
if session.preshared_key:
|
||||
lines.append(f"PresharedKey = {session.preshared_key}")
|
||||
@@ -478,11 +513,45 @@ class WireGuardClient:
|
||||
t.start()
|
||||
self._idle_thread = t
|
||||
|
||||
def _stop_session_locked(self, reason: str = "stop", ignore_missing: bool = False) -> None:
|
||||
self._remove_shell_firewall()
|
||||
if not self._service_exists():
|
||||
if not ignore_missing:
|
||||
_write_log("WireGuard tunnel service not found when stopping session.")
|
||||
self.session = None
|
||||
self.idle_deadline = None
|
||||
self._stop_event.set()
|
||||
return
|
||||
|
||||
idle_config = self._render_idle_config()
|
||||
wrote_idle = self._write_config(idle_config)
|
||||
service_config_path = self._service_config_path()
|
||||
if service_config_path and service_config_path != self.conf_path:
|
||||
wrote_idle = self._write_config_to(service_config_path, idle_config) or wrote_idle
|
||||
if wrote_idle:
|
||||
self._restart_service()
|
||||
self._ensure_adapter_name()
|
||||
self._ensure_service_display_name()
|
||||
_write_log(f"WireGuard client session stopped (reason={reason}).")
|
||||
elif not ignore_missing:
|
||||
_write_log("Failed to write idle WireGuard config.")
|
||||
self.session = None
|
||||
self.idle_deadline = None
|
||||
self._stop_event.set()
|
||||
|
||||
def start_session(self, session: SessionConfig, *, signing_client: Optional[Any] = None) -> None:
|
||||
with self._session_lock:
|
||||
if self.session:
|
||||
_write_log("Rejecting start_session: existing session already active.")
|
||||
return
|
||||
if self.session.tunnel_id == session.tunnel_id:
|
||||
_write_log("WireGuard session already active; reusing existing session.")
|
||||
self.bump_activity()
|
||||
return
|
||||
_write_log(
|
||||
"WireGuard session replace: existing_tunnel_id={0} new_tunnel_id={1}".format(
|
||||
self.session.tunnel_id, session.tunnel_id
|
||||
)
|
||||
)
|
||||
self._stop_session_locked(reason="session_replace", ignore_missing=True)
|
||||
|
||||
try:
|
||||
self._validate_token(session.token, signing_client=signing_client)
|
||||
@@ -519,41 +588,15 @@ class WireGuardClient:
|
||||
self._ensure_shell_firewall(session.allowed_ips)
|
||||
|
||||
self.session = session
|
||||
self.idle_deadline = time.time() + max(60, session.idle_seconds)
|
||||
_write_log("WireGuard client session started; idle timer armed.")
|
||||
self._start_idle_monitor()
|
||||
self.idle_deadline = None
|
||||
_write_log("WireGuard client session started (persistent mode).")
|
||||
|
||||
def stop_session(self, reason: str = "stop", ignore_missing: bool = False) -> None:
|
||||
with self._session_lock:
|
||||
self._remove_shell_firewall()
|
||||
if not self._service_exists():
|
||||
if not ignore_missing:
|
||||
_write_log("WireGuard tunnel service not found when stopping session.")
|
||||
self.session = None
|
||||
self.idle_deadline = None
|
||||
self._stop_event.set()
|
||||
return
|
||||
|
||||
idle_config = self._render_idle_config()
|
||||
wrote_idle = self._write_config(idle_config)
|
||||
service_config_path = self._service_config_path()
|
||||
if service_config_path and service_config_path != self.conf_path:
|
||||
wrote_idle = self._write_config_to(service_config_path, idle_config) or wrote_idle
|
||||
if wrote_idle:
|
||||
self._restart_service()
|
||||
self._ensure_adapter_name()
|
||||
self._ensure_service_display_name()
|
||||
_write_log(f"WireGuard client session stopped (reason={reason}).")
|
||||
elif not ignore_missing:
|
||||
_write_log("Failed to write idle WireGuard config.")
|
||||
self.session = None
|
||||
self.idle_deadline = None
|
||||
self._stop_event.set()
|
||||
self._stop_session_locked(reason=reason, ignore_missing=ignore_missing)
|
||||
|
||||
def bump_activity(self) -> None:
|
||||
if self.session and self.idle_deadline:
|
||||
self.idle_deadline = time.time() + max(60, self.session.idle_seconds)
|
||||
_write_log("WireGuard client activity bump; idle timer reset.")
|
||||
return
|
||||
|
||||
|
||||
_client: Optional[WireGuardClient] = None
|
||||
@@ -686,10 +729,9 @@ class Role:
|
||||
self._log_hook = hooks.get("log_agent")
|
||||
self._http_client_factory = hooks.get("http_client")
|
||||
self._get_server_url = hooks.get("get_server_url")
|
||||
try:
|
||||
self.client.stop_session(reason="agent_startup", ignore_missing=True)
|
||||
except Exception:
|
||||
self._log("Failed to preflight WireGuard session cleanup.", error=True)
|
||||
self._ensure_stop = threading.Event()
|
||||
self._ensure_thread = threading.Thread(target=self._ensure_loop, daemon=True)
|
||||
self._ensure_thread.start()
|
||||
|
||||
def _log(self, message: str, *, error: bool = False) -> None:
|
||||
if callable(self._log_hook):
|
||||
@@ -751,6 +793,11 @@ class Role:
|
||||
self._log("WireGuard start missing token payload.", error=True)
|
||||
return None
|
||||
|
||||
tunnel_id = payload.get("tunnel_id") or token.get("tunnel_id")
|
||||
if not tunnel_id:
|
||||
self._log("WireGuard start missing tunnel_id.", error=True)
|
||||
return None
|
||||
|
||||
virtual_ip = payload.get("virtual_ip") or payload.get("client_virtual_ip")
|
||||
endpoint = payload.get("endpoint") or payload.get("server_endpoint")
|
||||
endpoint = self._resolve_endpoint(endpoint, token)
|
||||
@@ -780,6 +827,7 @@ class Role:
|
||||
|
||||
return SessionConfig(
|
||||
token=token,
|
||||
tunnel_id=str(tunnel_id),
|
||||
virtual_ip=str(virtual_ip),
|
||||
allowed_ips=str(allowed_ips),
|
||||
endpoint=str(endpoint),
|
||||
@@ -791,6 +839,51 @@ class Role:
|
||||
client_public_key=payload.get("client_public_key"),
|
||||
)
|
||||
|
||||
def _request_persistent_session(self) -> Optional[Dict[str, Any]]:
|
||||
client = self._http_client()
|
||||
if client is None:
|
||||
return None
|
||||
try:
|
||||
payload = client.post_json(
|
||||
"/api/agent/vpn/ensure",
|
||||
{"agent_id": self.ctx.agent_id, "reason": "agent_boot"},
|
||||
require_auth=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
self._log(f"WireGuard ensure request failed: {exc}", error=True)
|
||||
return None
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
return payload
|
||||
|
||||
def _ensure_loop(self) -> None:
|
||||
if ENSURE_INITIAL_DELAY_SECONDS:
|
||||
time.sleep(ENSURE_INITIAL_DELAY_SECONDS)
|
||||
while not self._ensure_stop.is_set():
|
||||
payload = self._request_persistent_session()
|
||||
if payload:
|
||||
incoming_tunnel = str(payload.get("tunnel_id") or "")
|
||||
current_tunnel = ""
|
||||
if self.client.session is not None:
|
||||
try:
|
||||
current_tunnel = str(self.client.session.tunnel_id)
|
||||
except Exception:
|
||||
current_tunnel = ""
|
||||
state = None
|
||||
try:
|
||||
state = self.client._service_state()
|
||||
except Exception:
|
||||
state = None
|
||||
service_ready = state in ("RUNNING", "START_PENDING")
|
||||
if incoming_tunnel and incoming_tunnel == current_tunnel and service_ready:
|
||||
self._ensure_stop.wait(ENSURE_INTERVAL_SECONDS)
|
||||
continue
|
||||
session = self._build_session(payload)
|
||||
if session:
|
||||
self._log("WireGuard persistent session ensure received.")
|
||||
self.client.start_session(session, signing_client=self._http_client())
|
||||
self._ensure_stop.wait(ENSURE_INTERVAL_SECONDS)
|
||||
|
||||
def register_events(self) -> None:
|
||||
sio = self.ctx.sio
|
||||
|
||||
@@ -810,14 +903,18 @@ class Role:
|
||||
if target_agent and str(target_agent).strip() != str(self.ctx.agent_id).strip():
|
||||
return
|
||||
reason = payload.get("reason") or reason
|
||||
self._log(f"WireGuard stop requested (reason={reason}).")
|
||||
self.client.stop_session(reason=str(reason))
|
||||
self._log(f"WireGuard stop requested (reason={reason}); persistent tunnels ignore stop.")
|
||||
|
||||
|
||||
@sio.on("vpn_tunnel_activity")
|
||||
async def _vpn_tunnel_activity(payload):
|
||||
self.client.bump_activity()
|
||||
|
||||
def stop_all(self) -> None:
|
||||
try:
|
||||
self._ensure_stop.set()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
self.client.stop_session(reason="agent_shutdown")
|
||||
except Exception:
|
||||
|
||||
Reference in New Issue
Block a user