741 lines
31 KiB
Python
741 lines
31 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import shlex
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
import asyncssh
|
|
from fastapi import WebSocket
|
|
|
|
from app.core.db import utc_now
|
|
from app.models.connection import Connection
|
|
from app.services.session_updates import SessionBroadcaster
|
|
from app.services.types import PtyInfo, SessionBuffer
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
|
|
TARGET_PROBE_RTT_WINDOW = 8
|
|
TARGET_PROBE_TIMEOUT_SECONDS = 5.0
|
|
TARGET_PROBE_COMMAND = "printf '__wsg_probe__'"
|
|
|
|
|
|
@dataclass
|
|
class ManagedSession:
|
|
session_id: str
|
|
connection_id: int
|
|
user_id: int
|
|
client: asyncssh.SSHClientConnection
|
|
channel: asyncssh.SSHClientChannel | None
|
|
buffer: SessionBuffer
|
|
last_activity: datetime
|
|
websockets: set[WebSocket] = field(default_factory=set)
|
|
status: str = "active"
|
|
enhanced_enabled: bool = False
|
|
enhanced_fingerprint: str | None = None
|
|
tmux_binary_path: str | None = None
|
|
resize_seq: int = 0
|
|
target_profile: str = "unknown"
|
|
target_rtt_ms: int | None = None
|
|
target_avg_rtt_ms: int | None = None
|
|
target_jitter_ms: int = 0
|
|
target_probe_error_streak: int = 0
|
|
target_measured_at: datetime | None = None
|
|
target_rtt_samples: list[int] = field(default_factory=list)
|
|
|
|
def status_payload(self) -> str:
|
|
return json.dumps(
|
|
{
|
|
"id": self.session_id,
|
|
"status": self.status,
|
|
"last_activity": self.last_activity.isoformat(),
|
|
"target_profile": self.target_profile,
|
|
"target_rtt_ms": self.target_rtt_ms,
|
|
"target_avg_rtt_ms": self.target_avg_rtt_ms,
|
|
"target_jitter_ms": self.target_jitter_ms,
|
|
"target_probe_error_streak": self.target_probe_error_streak,
|
|
"target_measured_at": self.target_measured_at.isoformat() if self.target_measured_at else None,
|
|
}
|
|
)
|
|
|
|
async def send(self, data: str) -> None:
|
|
if not self.channel:
|
|
return
|
|
try:
|
|
self.channel.write(data)
|
|
except Exception:
|
|
self.status = "disconnected"
|
|
self.last_activity = utc_now()
|
|
raise
|
|
self.last_activity = utc_now()
|
|
|
|
async def _collect_tmux_metrics(self, phase: str, expected_rows: int, expected_cols: int) -> None:
|
|
if not (self.enhanced_enabled and self.enhanced_fingerprint and self.tmux_binary_path):
|
|
return
|
|
quoted_remote_binary = shlex.quote(self.tmux_binary_path)
|
|
quoted_fingerprint = shlex.quote(self.enhanced_fingerprint)
|
|
target_window = shlex.quote(f"{self.enhanced_fingerprint}:0")
|
|
target_pane = shlex.quote(f"{self.enhanced_fingerprint}:0.0")
|
|
try:
|
|
pane_res = await self.client.run(
|
|
f'{quoted_remote_binary} display-message -p -t {target_pane} "pane=#{{pane_height}}x#{{pane_width}} window=#{{window_height}}x#{{window_width}}"',
|
|
check=False,
|
|
)
|
|
window_mode_res = await self.client.run(
|
|
f"{quoted_remote_binary} show-window-options -t {target_window} -v window-size",
|
|
check=False,
|
|
)
|
|
clients_res = await self.client.run(
|
|
f'{quoted_remote_binary} list-clients -t {quoted_fingerprint} -F "#{{client_tty}} client=#{{client_height}}x#{{client_width}}"',
|
|
check=False,
|
|
)
|
|
LOGGER.info(
|
|
"tmux-metrics phase=%s session_id=%s fingerprint=%s expected=%sx%s pane='%s' window_mode='%s' clients='%s' exits(pane=%s,window=%s,clients=%s)",
|
|
phase,
|
|
self.session_id,
|
|
self.enhanced_fingerprint,
|
|
expected_rows,
|
|
expected_cols,
|
|
(pane_res.stdout or "").strip(),
|
|
(window_mode_res.stdout or "").strip(),
|
|
(clients_res.stdout or "").strip().replace("\n", " | "),
|
|
pane_res.exit_status,
|
|
window_mode_res.exit_status,
|
|
clients_res.exit_status,
|
|
)
|
|
except Exception as exc:
|
|
LOGGER.warning(
|
|
"tmux-metrics-failed phase=%s session_id=%s fingerprint=%s error=%s",
|
|
phase,
|
|
self.session_id,
|
|
self.enhanced_fingerprint,
|
|
exc,
|
|
)
|
|
|
|
async def resize(self, rows: int, cols: int) -> None:
|
|
if not self.channel:
|
|
return
|
|
self.resize_seq += 1
|
|
seq = self.resize_seq
|
|
pty_resize_ok = False
|
|
LOGGER.info(
|
|
"resize-recv session_id=%s seq=%s rows=%s cols=%s enhanced=%s",
|
|
self.session_id,
|
|
seq,
|
|
rows,
|
|
cols,
|
|
self.enhanced_enabled,
|
|
)
|
|
try:
|
|
# AsyncSSH expects width/height, not cols/rows keyword arguments.
|
|
self.channel.change_terminal_size(cols, rows)
|
|
pty_resize_ok = True
|
|
LOGGER.info(
|
|
"resize-pty-ok session_id=%s seq=%s rows=%s cols=%s",
|
|
self.session_id,
|
|
seq,
|
|
rows,
|
|
cols,
|
|
)
|
|
except Exception as exc:
|
|
LOGGER.warning(
|
|
"resize-pty-failed session_id=%s seq=%s rows=%s cols=%s error=%s",
|
|
self.session_id,
|
|
seq,
|
|
rows,
|
|
cols,
|
|
exc,
|
|
)
|
|
if self.enhanced_enabled and self.enhanced_fingerprint and self.tmux_binary_path:
|
|
quoted_remote_binary = shlex.quote(self.tmux_binary_path)
|
|
quoted_fingerprint = shlex.quote(self.enhanced_fingerprint)
|
|
target_window = shlex.quote(f"{self.enhanced_fingerprint}:0")
|
|
target_pane = shlex.quote(f"{self.enhanced_fingerprint}:0.0")
|
|
try:
|
|
res_mode = await self.client.run(
|
|
f"{quoted_remote_binary} set-window-option -t {target_window} window-size manual",
|
|
check=False,
|
|
)
|
|
res_window = await self.client.run(
|
|
f"{quoted_remote_binary} resize-window -t {target_window} -x {cols} -y {rows}",
|
|
check=False,
|
|
)
|
|
res_pane = await self.client.run(
|
|
f"{quoted_remote_binary} resize-pane -t {target_pane} -x {cols} -y {rows}",
|
|
check=False,
|
|
)
|
|
clients = await self.client.run(
|
|
f"{quoted_remote_binary} list-clients -t {quoted_fingerprint} -F '#{{client_tty}}\t#{{client_control_mode}}'",
|
|
check=False,
|
|
)
|
|
refreshed = 0
|
|
control_clients = 0
|
|
client_count = 0
|
|
if clients.exit_status == 0:
|
|
for line in (clients.stdout or "").splitlines():
|
|
raw_tty, _, raw_control_mode = line.partition("\t")
|
|
tty = raw_tty.strip()
|
|
if not tty:
|
|
continue
|
|
client_count += 1
|
|
control_mode = raw_control_mode.strip().lower()
|
|
if control_mode not in {"1", "on", "yes", "true"}:
|
|
continue
|
|
control_clients += 1
|
|
refresh_res = await self.client.run(
|
|
f"{quoted_remote_binary} refresh-client -t {shlex.quote(tty)} -C {cols}x{rows}",
|
|
check=False,
|
|
)
|
|
if refresh_res.exit_status == 0:
|
|
refreshed += 1
|
|
LOGGER.info(
|
|
"resize-tmux-ok session_id=%s seq=%s rows=%s cols=%s pty_resize_ok=%s exits(mode=%s,window=%s,pane=%s,clients=%s) clients(total=%s,control=%s,refreshed=%s)",
|
|
self.session_id,
|
|
seq,
|
|
rows,
|
|
cols,
|
|
pty_resize_ok,
|
|
res_mode.exit_status,
|
|
res_window.exit_status,
|
|
res_pane.exit_status,
|
|
clients.exit_status,
|
|
client_count,
|
|
control_clients,
|
|
refreshed,
|
|
)
|
|
except Exception as exc:
|
|
# 弱网或链路抖动时不向用户终端注入 tmux 控制命令，避免命令串污染可见输出。
|
|
LOGGER.warning(
|
|
"resize-tmux-fallback-skipped session_id=%s seq=%s rows=%s cols=%s error=%s",
|
|
self.session_id,
|
|
seq,
|
|
rows,
|
|
cols,
|
|
exc,
|
|
)
|
|
await self._collect_tmux_metrics(f"resize_seq_{seq}", rows, cols)
|
|
self.last_activity = utc_now()
|
|
|
|
async def close(self) -> None:
|
|
if self.channel:
|
|
self.channel.close()
|
|
self.client.close()
|
|
|
|
|
|
class SessionManager:
|
|
def __init__(
|
|
self,
|
|
keepalive_interval: int,
|
|
broadcaster: SessionBroadcaster | None = None,
|
|
known_hosts: str | None = None,
|
|
allow_unknown_hosts: bool = False,
|
|
auto_add_known_hosts: bool = True,
|
|
on_session_disconnect: Callable[[str], None] | None = None,
|
|
keepalive_binary_dir: str | None = None,
|
|
) -> None:
|
|
self._sessions: dict[str, ManagedSession] = {}
|
|
self._network_probe_tasks: dict[str, asyncio.Task[None]] = {}
|
|
self._keepalive_interval = keepalive_interval
|
|
self._known_hosts = known_hosts
|
|
self._allow_unknown_hosts = allow_unknown_hosts
|
|
self._auto_add_known_hosts = auto_add_known_hosts
|
|
self._known_hosts_lock = asyncio.Lock()
|
|
self._logger = logging.getLogger(__name__)
|
|
self._broadcaster = broadcaster
|
|
self._on_session_disconnect = on_session_disconnect
|
|
self._keepalive_binary_dir = Path(keepalive_binary_dir).resolve() if keepalive_binary_dir else None
|
|
|
|
def _schedule_tmux_probe(self, session: ManagedSession, rows: int, cols: int) -> None:
|
|
async def _probe() -> None:
|
|
for delay in (0.2, 1.0, 3.0, 8.0):
|
|
await asyncio.sleep(delay)
|
|
await session._collect_tmux_metrics(f"create_probe_{delay:.1f}s", rows, cols)
|
|
|
|
asyncio.create_task(_probe())
|
|
|
|
def _target_profile_interval_seconds(self, profile: str) -> float:
|
|
if profile == "poor":
|
|
return 12.0
|
|
if profile == "degraded":
|
|
return 7.0
|
|
return 4.0
|
|
|
|
def _compute_target_profile(self, avg_rtt_ms: int | None, jitter_ms: int, error_streak: int) -> str:
|
|
if error_streak >= 2:
|
|
return "poor"
|
|
if avg_rtt_ms is None:
|
|
if error_streak >= 1:
|
|
return "degraded"
|
|
return "unknown"
|
|
if avg_rtt_ms >= 400 or jitter_ms >= 200:
|
|
return "poor"
|
|
if error_streak >= 1 or avg_rtt_ms >= 150 or jitter_ms >= 80:
|
|
return "degraded"
|
|
return "good"
|
|
|
|
def _target_snapshot(self, session: ManagedSession) -> tuple[str, int | None, int | None, int, int]:
|
|
return (
|
|
session.target_profile,
|
|
session.target_rtt_ms,
|
|
session.target_avg_rtt_ms,
|
|
session.target_jitter_ms,
|
|
session.target_probe_error_streak,
|
|
)
|
|
|
|
def _apply_probe_success(self, session: ManagedSession, rtt_ms: int) -> bool:
|
|
before = self._target_snapshot(session)
|
|
session.target_rtt_samples = [*session.target_rtt_samples[-(TARGET_PROBE_RTT_WINDOW - 1) :], rtt_ms]
|
|
session.target_rtt_ms = rtt_ms
|
|
session.target_avg_rtt_ms = round(sum(session.target_rtt_samples) / len(session.target_rtt_samples))
|
|
session.target_jitter_ms = (
|
|
max(session.target_rtt_samples) - min(session.target_rtt_samples) if len(session.target_rtt_samples) >= 2 else 0
|
|
)
|
|
session.target_probe_error_streak = 0
|
|
session.target_measured_at = utc_now()
|
|
session.target_profile = self._compute_target_profile(
|
|
session.target_avg_rtt_ms,
|
|
session.target_jitter_ms,
|
|
session.target_probe_error_streak,
|
|
)
|
|
return self._target_snapshot(session) != before
|
|
|
|
def _apply_probe_failure(self, session: ManagedSession) -> bool:
|
|
before = self._target_snapshot(session)
|
|
session.target_rtt_ms = None
|
|
session.target_probe_error_streak = min(session.target_probe_error_streak + 1, 10)
|
|
session.target_measured_at = utc_now()
|
|
session.target_profile = self._compute_target_profile(
|
|
session.target_avg_rtt_ms,
|
|
session.target_jitter_ms,
|
|
session.target_probe_error_streak,
|
|
)
|
|
return self._target_snapshot(session) != before
|
|
|
|
async def _broadcast_session_status(self, session: ManagedSession) -> None:
|
|
if not self._broadcaster:
|
|
return
|
|
await self._broadcaster.broadcast(session.user_id, session.status_payload())
|
|
|
|
async def _probe_target_network_once(self, session: ManagedSession) -> None:
|
|
start = time.monotonic()
|
|
changed = False
|
|
try:
|
|
result = await asyncio.wait_for(
|
|
session.client.run(TARGET_PROBE_COMMAND, check=False),
|
|
timeout=TARGET_PROBE_TIMEOUT_SECONDS,
|
|
)
|
|
if result.exit_status != 0:
|
|
raise RuntimeError(f"probe exit status {result.exit_status}")
|
|
elapsed_ms = max(1, round((time.monotonic() - start) * 1000))
|
|
changed = self._apply_probe_success(session, elapsed_ms)
|
|
except Exception as exc:
|
|
changed = self._apply_probe_failure(session)
|
|
self._logger.debug(
|
|
"target network probe failed session_id=%s error=%s",
|
|
session.session_id,
|
|
exc,
|
|
)
|
|
|
|
if changed:
|
|
await self._broadcast_session_status(session)
|
|
|
|
async def _run_target_network_probe_loop(self, session_id: str) -> None:
|
|
try:
|
|
while True:
|
|
session = self._sessions.get(session_id)
|
|
if not session or session.status != "active" or not session.channel:
|
|
return
|
|
await self._probe_target_network_once(session)
|
|
await asyncio.sleep(self._target_profile_interval_seconds(session.target_profile))
|
|
except asyncio.CancelledError:
|
|
return
|
|
|
|
def _cancel_target_network_probe(self, session_id: str) -> None:
|
|
task = self._network_probe_tasks.pop(session_id, None)
|
|
if task and not task.done():
|
|
task.cancel()
|
|
|
|
def _ensure_target_network_probe(self, session: ManagedSession) -> None:
|
|
self._cancel_target_network_probe(session.session_id)
|
|
self._network_probe_tasks[session.session_id] = asyncio.create_task(
|
|
self._run_target_network_probe_loop(session.session_id)
|
|
)
|
|
|
|
def list_sessions(self, user_id: int) -> list[ManagedSession]:
|
|
return [session for session in self._sessions.values() if session.user_id == user_id]
|
|
|
|
def get_session(self, session_id: str) -> ManagedSession | None:
|
|
return self._sessions.get(session_id)
|
|
|
|
def _resolve_known_hosts_path(self) -> str:
|
|
known_hosts = self._known_hosts or "~/.ssh/known_hosts"
|
|
return str(Path(known_hosts).expanduser())
|
|
|
|
def _ensure_known_hosts_path(self, known_hosts_path: str, create_file: bool) -> None:
|
|
path = Path(known_hosts_path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
if create_file and not path.exists():
|
|
path.touch(mode=0o600, exist_ok=True)
|
|
|
|
def _build_known_hosts_entry(self, host: str, port: int, host_key: asyncssh.SSHKey) -> str:
|
|
label = host if port == 22 else f"[{host}]:{port}"
|
|
key_data = host_key.export_public_key("openssh").decode("utf-8").strip()
|
|
return f"{label} {key_data}\n"
|
|
|
|
async def _fetch_host_key(self, connection: Connection) -> asyncssh.SSHKey:
|
|
host_key = await asyncssh.get_server_host_key(connection.host, port=connection.port)
|
|
if host_key is None:
|
|
raise ValueError("Host key not provided by server")
|
|
return host_key
|
|
|
|
async def _open_client(
|
|
self,
|
|
connection: Connection,
|
|
auth_payload: dict[str, Any],
|
|
known_hosts: str | None,
|
|
) -> asyncssh.SSHClientConnection:
|
|
client_keys = None
|
|
if auth_payload.get("private_key"):
|
|
try:
|
|
passphrase = auth_payload.get("key_passphrase")
|
|
key = asyncssh.import_private_key(auth_payload["private_key"], passphrase=passphrase)
|
|
client_keys = [key]
|
|
except Exception:
|
|
raise ValueError("无效的私钥格式或密码错误")
|
|
|
|
return await asyncssh.connect(
|
|
connection.host,
|
|
port=connection.port,
|
|
username=connection.username,
|
|
password=auth_payload.get("password"),
|
|
client_keys=client_keys,
|
|
known_hosts=known_hosts,
|
|
keepalive_interval=self._keepalive_interval,
|
|
)
|
|
|
|
async def _connect_with_known_hosts(
|
|
self,
|
|
connection: Connection,
|
|
auth_payload: dict[str, Any],
|
|
) -> asyncssh.SSHClientConnection:
|
|
known_hosts_path = self._resolve_known_hosts_path()
|
|
self._ensure_known_hosts_path(known_hosts_path, create_file=self._auto_add_known_hosts)
|
|
return await self._open_client(connection, auth_payload, known_hosts_path)
|
|
|
|
async def _append_known_host(self, known_hosts_path: str, entry: str) -> None:
|
|
async with self._known_hosts_lock:
|
|
self._ensure_known_hosts_path(known_hosts_path, create_file=self._auto_add_known_hosts)
|
|
with open(known_hosts_path, "a", encoding="utf-8") as handle:
|
|
handle.write(entry)
|
|
|
|
async def _try_auto_add_known_host(self, connection: Connection) -> None:
|
|
known_hosts_path = self._resolve_known_hosts_path()
|
|
host_key = await self._fetch_host_key(connection)
|
|
entry = self._build_known_hosts_entry(connection.host, connection.port, host_key)
|
|
await self._append_known_host(known_hosts_path, entry)
|
|
fingerprint = host_key.get_fingerprint("sha256")
|
|
self._logger.info(
|
|
"Added SSH host key to known_hosts (host=%s port=%s fingerprint=%s)",
|
|
connection.host,
|
|
connection.port,
|
|
fingerprint,
|
|
)
|
|
|
|
async def connect_client(self, connection: Connection, auth_payload: dict[str, Any]) -> asyncssh.SSHClientConnection:
|
|
if self._allow_unknown_hosts:
|
|
return await self._open_client(connection, auth_payload, None)
|
|
try:
|
|
return await self._connect_with_known_hosts(connection, auth_payload)
|
|
except asyncssh.HostKeyNotVerifiable:
|
|
if not self._auto_add_known_hosts:
|
|
raise
|
|
await self._try_auto_add_known_host(connection)
|
|
return await self._connect_with_known_hosts(connection, auth_payload)
|
|
|
|
async def detect_remote_platform(self, connection: Connection, auth_payload: dict[str, Any]) -> tuple[str, str]:
|
|
client = await self.connect_client(connection, auth_payload)
|
|
try:
|
|
arch_result = await client.run("uname -m", check=False)
|
|
os_result = await client.run("uname -s", check=False)
|
|
arch = (arch_result.stdout or "").strip()
|
|
remote_os = (os_result.stdout or "").strip()
|
|
return arch, remote_os
|
|
finally:
|
|
client.close()
|
|
|
|
def resolve_keepalive_binary(self, arch: str, remote_os: str) -> tuple[Path, str] | None:
|
|
if not self._keepalive_binary_dir:
|
|
return None
|
|
|
|
normalized_arch = arch.strip().lower()
|
|
normalized_os = remote_os.strip().lower()
|
|
|
|
arch_alias = {
|
|
"arm64": "aarch64",
|
|
"x64": "x86_64",
|
|
"amd64": "x86_64",
|
|
"aarch64": "aarch64",
|
|
"x86_64": "x86_64",
|
|
}.get(normalized_arch, normalized_arch)
|
|
|
|
if normalized_os == "darwin":
|
|
filename = "keeplive.macos-arm64" if arch_alias == "aarch64" else "keeplive.macos-x86_64" if arch_alias == "x86_64" else None
|
|
elif normalized_os == "linux":
|
|
filename = "keeplive.aarch64" if arch_alias == "aarch64" else "keeplive.x86_64" if arch_alias == "x86_64" else None
|
|
else:
|
|
filename = None
|
|
|
|
if not filename:
|
|
self._logger.warning(
|
|
"resolve_keepalive_binary unsupported arch/os arch=%s os=%s normalized_arch=%s normalized_os=%s",
|
|
arch,
|
|
remote_os,
|
|
arch_alias,
|
|
normalized_os,
|
|
)
|
|
return None
|
|
|
|
binary_path = self._keepalive_binary_dir / filename
|
|
if not binary_path.exists() or not binary_path.is_file():
|
|
self._logger.warning(
|
|
"resolve_keepalive_binary missing file path=%s arch=%s os=%s",
|
|
binary_path,
|
|
arch,
|
|
remote_os,
|
|
)
|
|
return None
|
|
|
|
remote_path = f"/tmp/tmux.{arch_alias}"
|
|
self._logger.info(
|
|
"resolve_keepalive_binary matched local=%s remote=%s arch=%s os=%s",
|
|
binary_path,
|
|
remote_path,
|
|
arch,
|
|
remote_os,
|
|
)
|
|
return binary_path, remote_path
|
|
|
|
async def _upload_file_to_remote(self, client: asyncssh.SSHClientConnection, local_path: Path, remote_path: str) -> None:
|
|
async with client.start_sftp_client() as sftp:
|
|
await sftp.put(str(local_path), remote_path)
|
|
|
|
async def _ensure_remote_keepalive_binary(
|
|
self,
|
|
client: asyncssh.SSHClientConnection,
|
|
local_path: Path,
|
|
remote_path: str,
|
|
) -> None:
|
|
quoted_remote = shlex.quote(remote_path)
|
|
check_result = await client.run(f"test -x {quoted_remote}", check=False)
|
|
if check_result.exit_status == 0:
|
|
self._logger.info("reuse remote keepalive binary: %s", remote_path)
|
|
return
|
|
|
|
self._logger.info("upload remote keepalive binary local=%s remote=%s", local_path, remote_path)
|
|
await self._upload_file_to_remote(client, local_path, remote_path)
|
|
await client.run(f"chmod +x {quoted_remote}", check=False)
|
|
|
|
async def create_session(
|
|
self,
|
|
connection: Connection,
|
|
auth_payload: dict[str, Any],
|
|
pty: PtyInfo,
|
|
enhanced_enabled: bool = False,
|
|
enhanced_fingerprint: str | None = None,
|
|
tmux_binary_path: str | None = None,
|
|
session_id: str | None = None,
|
|
) -> ManagedSession:
|
|
client = await self.connect_client(connection, auth_payload)
|
|
|
|
resolved_session_id = session_id or uuid.uuid4().hex
|
|
session = ManagedSession(
|
|
session_id=resolved_session_id,
|
|
connection_id=connection.id,
|
|
user_id=connection.user_id,
|
|
client=client,
|
|
channel=None,
|
|
buffer=SessionBuffer(),
|
|
last_activity=utc_now(),
|
|
enhanced_enabled=enhanced_enabled,
|
|
enhanced_fingerprint=enhanced_fingerprint,
|
|
tmux_binary_path=tmux_binary_path,
|
|
)
|
|
|
|
def _start_channel() -> asyncssh.SSHClientSession:
|
|
manager = self
|
|
|
|
class TerminalSession(asyncssh.SSHClientSession):
|
|
def data_received(self, data: str, datatype: asyncssh.DataType | None = None) -> None:
|
|
session.buffer.append(data)
|
|
session.last_activity = utc_now()
|
|
if session.websockets:
|
|
for websocket in list(session.websockets):
|
|
asyncio.create_task(websocket.send_text(data))
|
|
|
|
def connection_lost(self, exc: Exception | None) -> None:
|
|
current = manager._sessions.get(session.session_id)
|
|
# Ignore stale callbacks from replaced/explicitly-closed sessions.
|
|
if current is not session:
|
|
LOGGER.info(
|
|
"connection_lost ignored for stale session session_id=%s",
|
|
session.session_id,
|
|
)
|
|
return
|
|
session.status = "disconnected"
|
|
session.last_activity = utc_now()
|
|
manager._sessions.pop(session.session_id, None)
|
|
manager._cancel_target_network_probe(session.session_id)
|
|
if manager._broadcaster:
|
|
asyncio.create_task(manager._broadcaster.broadcast(session.user_id, session.status_payload()))
|
|
if manager._on_session_disconnect:
|
|
manager._on_session_disconnect(session.session_id)
|
|
|
|
return TerminalSession()
|
|
|
|
channel, _ = await client.create_session(
|
|
session_factory=_start_channel,
|
|
term_type=pty.term,
|
|
term_size=(pty.cols, pty.rows),
|
|
request_pty=True,
|
|
)
|
|
session.channel = channel
|
|
|
|
if enhanced_enabled and enhanced_fingerprint and tmux_binary_path:
|
|
keepalive = self.resolve_keepalive_binary(connection.remote_arch or "", connection.remote_os or "")
|
|
if keepalive:
|
|
local_binary, _ = keepalive
|
|
await self._ensure_remote_keepalive_binary(client, local_binary, tmux_binary_path)
|
|
quoted_remote_binary = shlex.quote(tmux_binary_path)
|
|
quoted_fingerprint = shlex.quote(enhanced_fingerprint)
|
|
target_window = shlex.quote(f"{enhanced_fingerprint}:0")
|
|
target_pane = shlex.quote(f"{enhanced_fingerprint}:0.0")
|
|
res_new = await client.run(
|
|
f"{quoted_remote_binary} new-session -Ad -s {quoted_fingerprint} -x {pty.cols} -y {pty.rows}",
|
|
check=False,
|
|
)
|
|
res_status = await client.run(f"{quoted_remote_binary} set-option -t {quoted_fingerprint} status off", check=False)
|
|
res_destroy_unattached = await client.run(
|
|
f"{quoted_remote_binary} set-option -t {quoted_fingerprint} destroy-unattached off",
|
|
check=False,
|
|
)
|
|
res_mouse = await client.run(
|
|
f"{quoted_remote_binary} set-option -t {quoted_fingerprint} mouse on",
|
|
check=False,
|
|
)
|
|
res_border = await client.run(f"{quoted_remote_binary} set-window-option -t {target_window} pane-border-status off", check=False)
|
|
res_aggressive = await client.run(f"{quoted_remote_binary} set-window-option -t {target_window} aggressive-resize on", check=False)
|
|
res_mode = await client.run(f"{quoted_remote_binary} set-window-option -t {target_window} window-size manual", check=False)
|
|
res_window = await client.run(
|
|
f"{quoted_remote_binary} resize-window -t {target_window} -x {pty.cols} -y {pty.rows}",
|
|
check=False,
|
|
)
|
|
res_pane = await client.run(
|
|
f"{quoted_remote_binary} resize-pane -t {target_pane} -x {pty.cols} -y {pty.rows}",
|
|
check=False,
|
|
)
|
|
res_has = await client.run(
|
|
f"{quoted_remote_binary} has-session -t {quoted_fingerprint}",
|
|
check=False,
|
|
)
|
|
self._logger.info(
|
|
"tmux-create session_id=%s fingerprint=%s requested=%sx%s exits(new=%s,status=%s,destroy_unattached=%s,mouse=%s,border=%s,aggressive=%s,mode=%s,window=%s,pane=%s,has=%s)",
|
|
resolved_session_id,
|
|
enhanced_fingerprint,
|
|
pty.rows,
|
|
pty.cols,
|
|
res_new.exit_status,
|
|
res_status.exit_status,
|
|
res_destroy_unattached.exit_status,
|
|
res_mouse.exit_status,
|
|
res_border.exit_status,
|
|
res_aggressive.exit_status,
|
|
res_mode.exit_status,
|
|
res_window.exit_status,
|
|
res_pane.exit_status,
|
|
res_has.exit_status,
|
|
)
|
|
if res_has.exit_status != 0:
|
|
raise RuntimeError(
|
|
f"tmux session not available after create (fingerprint={enhanced_fingerprint}, has_exit={res_has.exit_status})"
|
|
)
|
|
await session._collect_tmux_metrics("create_before_attach", pty.rows, pty.cols)
|
|
# Use exec to avoid silently falling back to an interactive shell
|
|
# when attach fails.
|
|
channel.write(f"exec {quoted_remote_binary} attach -t {quoted_fingerprint}\n")
|
|
self._schedule_tmux_probe(session, pty.rows, pty.cols)
|
|
|
|
session.status = "active"
|
|
replaced_session = self._sessions.get(resolved_session_id)
|
|
self._sessions[resolved_session_id] = session
|
|
self._ensure_target_network_probe(session)
|
|
if replaced_session and replaced_session is not session:
|
|
if replaced_session.websockets:
|
|
session.websockets.update(replaced_session.websockets)
|
|
replaced_session.websockets.clear()
|
|
try:
|
|
await replaced_session.close()
|
|
except Exception as exc:
|
|
self._logger.warning(
|
|
"failed to close replaced session session_id=%s error=%s",
|
|
resolved_session_id,
|
|
exc,
|
|
)
|
|
return session
|
|
|
|
async def _terminate_enhanced_tmux(self, session: ManagedSession) -> None:
|
|
if not session.enhanced_enabled or not session.enhanced_fingerprint or not session.tmux_binary_path:
|
|
return
|
|
|
|
quoted_remote_binary = shlex.quote(session.tmux_binary_path)
|
|
quoted_fingerprint = shlex.quote(session.enhanced_fingerprint)
|
|
try:
|
|
await session.client.run(
|
|
f"{quoted_remote_binary} kill-session -t {quoted_fingerprint}",
|
|
check=False,
|
|
)
|
|
except Exception as exc:
|
|
self._logger.warning(
|
|
"failed to terminate enhanced tmux session session_id=%s fingerprint=%s error=%s",
|
|
session.session_id,
|
|
session.enhanced_fingerprint,
|
|
exc,
|
|
)
|
|
|
|
async def close_session(self, session_id: str, terminate_enhanced: bool = False) -> None:
|
|
self._cancel_target_network_probe(session_id)
|
|
session = self._sessions.pop(session_id, None)
|
|
if session:
|
|
if terminate_enhanced:
|
|
await self._terminate_enhanced_tmux(session)
|
|
await session.close()
|
|
|
|
async def open_enhanced_attach(
|
|
self,
|
|
connection: Connection,
|
|
auth_payload: dict[str, Any],
|
|
pty: PtyInfo,
|
|
tmux_binary_path: str,
|
|
fingerprint: str,
|
|
) -> ManagedSession:
|
|
return await self.create_session(
|
|
connection=connection,
|
|
auth_payload=auth_payload,
|
|
pty=pty,
|
|
enhanced_enabled=True,
|
|
enhanced_fingerprint=fingerprint,
|
|
tmux_binary_path=tmux_binary_path,
|
|
)
|
|
|
|
def serialize_pty(self, pty: PtyInfo) -> str:
|
|
return json.dumps({"term": pty.term, "rows": pty.rows, "cols": pty.cols})
|
|
|
|
def deserialize_pty(self, payload: str) -> PtyInfo:
|
|
data: dict[str, Any] = json.loads(payload)
|
|
return PtyInfo(term=data.get("term", "xterm-256color"), rows=int(data.get("rows", 24)), cols=int(data.get("cols", 80)))
|