feat(compute): add SSH console type support

This commit is contained in:
Cristi 2026-04-30 15:52:47 +03:00
parent f0ef65b8e8
commit d64eca0418
18 changed files with 421 additions and 74 deletions

View File

@ -23,12 +23,14 @@ import tempfile
import psutil
import platform
import re
import asyncssh
from fastapi import WebSocketDisconnect
from gns3server.utils.interfaces import interfaces
from gns3server.compute.compute_error import ComputeError
from ..compute.port_manager import PortManager
from ..utils.asyncio import wait_run_in_executor, locking
from ..utils.asyncio.ssh_server import AsyncioSSHServer
from ..utils.asyncio.telnet_server import AsyncioTelnetServer
from gns3server.compute.ubridge.hypervisor import Hypervisor
from gns3server.compute.ubridge.ubridge_error import UbridgeError
@ -54,8 +56,8 @@ class BaseNode:
:param aux: auxiliary console TCP port
:param aux_type: auxiliary console type
:param linked_clone: The node base image is duplicate/overlay (Each node data are independent)
:param wrap_console: The console is wrapped using AsyncioTelnetServer
:param wrap_aux: The auxiliary console is wrapped using AsyncioTelnetServer
:param wrap_console: The console is wrapped using a proxy transport server
:param wrap_aux: The auxiliary console is wrapped using a proxy transport server
"""
def __init__(
@ -409,10 +411,10 @@ class BaseNode:
return vnc_console_start_port_range, vnc_console_end_port_range
async def _wrap_telnet_proxy(self, internal_port, external_port):
async def _wrap_console_proxy(self, internal_port, external_port, console_type):
"""
Start a telnet proxy for the console allowing multiple telnet clients
to be connected at the same time
Start a console proxy allowing multiple external clients to be
connected at the same time.
"""
remaining_trial = 60
@ -420,7 +422,7 @@ class BaseNode:
try:
(self._wrap_console_reader, self._wrap_console_writer) = await asyncio.open_connection(
host="127.0.0.1",
port=self._internal_console_port
port=internal_port
)
break
except (OSError, ConnectionRefusedError) as e:
@ -428,40 +430,46 @@ class BaseNode:
raise e
await asyncio.sleep(0.1)
remaining_trial -= 1
await AsyncioTelnetServer.write_client_intro(self._wrap_console_writer, echo=True)
server = AsyncioTelnetServer(
reader=self._wrap_console_reader,
writer=self._wrap_console_writer,
binary=True,
echo=True
)
if console_type == "telnet":
await AsyncioTelnetServer.write_client_intro(self._wrap_console_writer, echo=True)
server = AsyncioTelnetServer(
reader=self._wrap_console_reader,
writer=self._wrap_console_writer,
binary=True,
echo=True
)
elif console_type == "ssh":
server = AsyncioSSHServer(reader=self._wrap_console_reader, writer=self._wrap_console_writer)
else:
raise NodeError(f"Console wrapper does not support type {console_type}")
# warning: this will raise OSError exception if there is a problem...
telnet_server = await server.start(self._manager.port_manager.console_host, external_port)
self._wrapper_telnet_servers.append(telnet_server)
proxy_server = await server.start(self._manager.port_manager.console_host, external_port)
self._wrapper_telnet_servers.append(proxy_server)
async def start_wrap_console(self):
"""
Start a Telnet proxy servers for the console and auxiliary console allowing multiple telnet clients
Start console proxy servers for the console and auxiliary console allowing multiple clients
to be connected at the same time
"""
if self._wrap_console and self._console_type == "telnet":
await self._wrap_telnet_proxy(self._internal_console_port, self.console)
if self._wrap_console and self._console_type in ("telnet", "ssh"):
await self._wrap_console_proxy(self._internal_console_port, self.console, self._console_type)
log.info(
f"New Telnet proxy server for console started "
f"New {self._console_type.upper()} proxy server for console started "
f"(internal port = {self._internal_console_port}, external port = {self.console})"
)
if self._wrap_aux and self._aux_type == "telnet":
await self._wrap_telnet_proxy(self._internal_aux_port, self.aux)
if self._wrap_aux and self._aux_type in ("telnet", "ssh"):
await self._wrap_console_proxy(self._internal_aux_port, self.aux, self._aux_type)
log.info(
f"New Telnet proxy server for auxiliary console started "
f"New {self._aux_type.upper()} proxy server for auxiliary console started "
f"(internal port = {self._internal_aux_port}, external port = {self.aux})"
)
async def stop_wrap_console(self):
"""
Stops the telnet proxy servers.
Stops the console proxy servers.
"""
if self._wrap_console_writer:
@ -474,7 +482,7 @@ class BaseNode:
async def reset_wrap_console(self):
"""
Reset the wrap console (restarts the Telnet proxy)
Reset the wrap console (restarts the console proxy)
"""
await self.stop_wrap_console()
@ -493,27 +501,66 @@ class BaseNode:
)
if self.status != "started":
raise NodeError(f"Node {self.name} is not started")
await websocket.close(code=1000)
log.warning(f"Cannot open console WebSocket: node {self.name} is not started")
return
if self._console_type != "telnet":
raise NodeError(f"Node {self.name} console type is not telnet")
if self._console_type not in ("telnet", "ssh"):
await websocket.close(code=1000)
log.warning(
f"Cannot open console WebSocket: node {self.name} console type '{self._console_type}' "
f"is not supported"
)
return
telnet_reader = None
telnet_writer = None
ssh_connection = None
ssh_process = None
try:
host = self._manager.port_manager.console_host
port = self.console
(telnet_reader, telnet_writer) = await asyncio.open_connection(host, port)
log.info(f"Connected to local Telnet server {host}:{port}")
except ConnectionError as e:
raise NodeError(f"Cannot connect to node {self.name} telnet server: {e}")
if self._console_type == "ssh":
# For SSH consoles (wrapped or not), connect to the external SSH port via an
# SSH client. The AsyncioSSHServer on that port handles multi-client broadcasting
# to/from the node's process streams. Connecting directly to _internal_console_port
# would conflict with the exclusive connection already held by the SSH proxy.
port = self.console
ssh_connection = await asyncssh.connect(
host,
port=port,
username="gns3",
known_hosts=None,
encoding=None,
)
ssh_process = await ssh_connection.create_process(encoding=None, term_type="xterm")
telnet_reader = ssh_process.stdout
telnet_writer = ssh_process.stdin
else:
port = self.console
(telnet_reader, telnet_writer) = await asyncio.open_connection(host, port)
log.info(f"Connected to local console stream {host}:{port} (console type={self._console_type})")
except (ConnectionError, OSError, asyncssh.Error) as e:
await websocket.close(code=1000)
log.warning(f"Cannot connect to node {self.name} console server: {e}")
return
async def ws_forward(telnet_writer):
try:
while True:
data = await websocket.receive_text()
if data:
telnet_writer.write(data.encode())
await telnet_writer.drain()
msg = await websocket.receive()
if msg["type"] == "websocket.disconnect":
break
if "text" in msg and msg["text"]:
data = msg["text"].encode()
elif "bytes" in msg and msg["bytes"]:
data = msg["bytes"]
else:
continue
telnet_writer.write(data)
await telnet_writer.drain()
except WebSocketDisconnect:
log.info(
f"Client {websocket.client.host}:{websocket.client.port} has disconnected from compute"
@ -541,6 +588,17 @@ class BaseNode:
for task in pending:
task.cancel()
if ssh_connection:
if ssh_process:
ssh_process.close()
ssh_connection.close()
await ssh_connection.wait_closed()
if telnet_writer and hasattr(telnet_writer, "close"):
telnet_writer.close()
if hasattr(telnet_writer, "wait_closed"):
await telnet_writer.wait_closed()
async def start_vnc_websocket_console(self, websocket):
"""
Connect to VNC console using WebSocket.
@ -555,10 +613,14 @@ class BaseNode:
if self.status != "started":
await websocket.close(code=1000)
raise NodeError(f"Node {self.name} is not started")
log.warning(f"Cannot open VNC WebSocket: node {self.name} is not started")
return
if self._console_type != "vnc":
await websocket.close(code=1000)
raise NodeError(f"Node {self.name} console type is not vnc")
log.warning(
f"Cannot open VNC WebSocket: node {self.name} console type '{self._console_type}' is not vnc"
)
return
try:
vnc_reader, vnc_writer = await asyncio.open_connection(
@ -568,7 +630,8 @@ class BaseNode:
log.info(f"Connected to VNC server {self._manager.port_manager.console_host}:{self.console}")
except ConnectionError as e:
await websocket.close(code=1000)
raise NodeError(f"Cannot connect to node {self.name} VNC server: {e}")
log.warning(f"Cannot connect to node {self.name} VNC server: {e}")
return
async def ws_forward(vnc_writer):
# Browser → VNC: Forward binary WebSocket data to VNC server

View File

@ -28,6 +28,7 @@ import subprocess
import os
import re
from gns3server.utils.asyncio.ssh_server import AsyncioSSHServer
from gns3server.utils.asyncio.telnet_server import AsyncioTelnetServer
from gns3server.utils.asyncio.raw_command_server import AsyncioRawCommandServer
from gns3server.utils.asyncio import wait_for_file_creation
@ -667,7 +668,7 @@ class DockerVM(BaseNode):
log.error(line)
raise DockerError(logdata)
if self.console_type == "telnet":
if self.console_type in ("telnet", "ssh"):
await self._start_console()
elif self.console_type == "http" or self.console_type == "https":
await self._start_http()
@ -702,14 +703,19 @@ class DockerVM(BaseNode):
)
except OSError as e:
raise DockerError(f"Could not start auxiliary console process: {e}")
server = AsyncioTelnetServer(reader=process.stdout, writer=process.stdin, binary=True, echo=True)
if self.aux_type == "telnet":
server = AsyncioTelnetServer(reader=process.stdout, writer=process.stdin, binary=True, echo=True)
transport = "Telnet"
else:
server = AsyncioSSHServer(reader=process.stdout, writer=process.stdin)
transport = "SSH"
try:
self._telnet_servers.append(await server.start(self._manager.port_manager.console_host, self.aux))
except OSError as e:
raise DockerError(
f"Could not start Telnet server on socket {self._manager.port_manager.console_host}:{self.aux}: {e}"
f"Could not start {transport} server on socket {self._manager.port_manager.console_host}:{self.aux}: {e}"
)
log.debug(f"Docker container '{self.name}' started listen for auxiliary telnet on {self.aux}")
log.debug(f"Docker container '{self.name}' started listen for auxiliary {self.aux_type} on {self.aux}")
async def _fix_permissions(self):
"""
@ -872,7 +878,7 @@ class DockerVM(BaseNode):
async def _start_console(self):
"""
Starts streaming the console via telnet
Starts streaming the console via telnet or ssh
"""
class InputStream:
@ -889,18 +895,23 @@ class DockerVM(BaseNode):
output_stream = asyncio.StreamReader()
input_stream = InputStream()
telnet = AsyncioTelnetServer(
reader=output_stream,
writer=input_stream,
echo=True,
naws=True,
window_size_changed_callback=self._window_size_changed_callback,
)
if self.console_type == "telnet":
telnet = AsyncioTelnetServer(
reader=output_stream,
writer=input_stream,
echo=True,
naws=True,
window_size_changed_callback=self._window_size_changed_callback,
)
transport = "Telnet"
else:
telnet = AsyncioSSHServer(reader=output_stream, writer=input_stream)
transport = "SSH"
try:
self._telnet_servers.append(await telnet.start(self._manager.port_manager.console_host, self.console))
except OSError as e:
raise DockerError(
f"Could not start Telnet server on socket {self._manager.port_manager.console_host}:{self.console}: {e}"
f"Could not start {transport} server on socket {self._manager.port_manager.console_host}:{self.console}: {e}"
)
self._console_websocket = await self.manager.websocket_query(
@ -936,6 +947,9 @@ class DockerVM(BaseNode):
Reset the console.
"""
if self.console_type not in ("telnet", "ssh"):
return
if self._console_websocket:
await self._console_websocket.close()
await self._clean_servers()

View File

@ -80,7 +80,9 @@ class Router(BaseNode):
raise DynamipsError(f"{name} is an invalid name to create a Dynamips node")
super().__init__(
name, node_id, project, manager, console=console, console_type=console_type, aux=aux, aux_type=aux_type
name, node_id, project, manager, console=console, console_type=console_type, aux=aux, aux_type=aux_type,
wrap_console=(console_type == "ssh"),
wrap_aux=(aux_type == "ssh"),
)
self._working_directory = os.path.join(
@ -248,7 +250,10 @@ class Router(BaseNode):
)
if self._console is not None:
await self._hypervisor.send(f'vm set_con_tcp_port "{self._name}" {self._console}')
# For SSH console, tell Dynamips to listen on the internal port so that
# the AsyncioSSHServer proxy can wrap it on the external console port.
con_port = self._internal_console_port if self._wrap_console and self._internal_console_port else self._console
await self._hypervisor.send(f'vm set_con_tcp_port "{self._name}" {con_port}')
if self.aux is not None:
await self._hypervisor.send(f'vm set_aux_tcp_port "{self._name}" {self.aux}')
@ -327,6 +332,12 @@ class Router(BaseNode):
self._memory_watcher = FileWatcher(self._memory_files(), self._memory_changed, strategy="hash", delay=30)
monitor_process(self._hypervisor.process, self._termination_callback)
if self._console_type == "ssh":
try:
await self.start_wrap_console()
except OSError as e:
raise DynamipsError(f"Could not start SSH Dynamips console {e}")
async def _termination_callback(self, returncode):
"""
Called when the process has stopped.
@ -362,6 +373,7 @@ class Router(BaseNode):
self._memory_watcher.close()
self._memory_watcher = None
await self.save_configs()
await self.stop_wrap_console()
async def reload(self):
"""
@ -1017,8 +1029,21 @@ class Router(BaseNode):
self.console_type = console_type
if self._console and console_type == "telnet":
await self._hypervisor.send(f'vm set_con_tcp_port "{self._name}" {self._console}')
if self._console:
if console_type == "ssh":
# Switching to SSH: ensure an internal port is allocated and redirect Dynamips to it.
self._wrap_console = True
if self._internal_console_port is None:
self._internal_console_port = self._manager.port_manager.get_free_tcp_port(self._project)
await self._hypervisor.send(f'vm set_con_tcp_port "{self._name}" {self._internal_console_port}')
elif console_type == "telnet":
# Switching to telnet: stop SSH wrapper if running and redirect Dynamips to external port.
await self.stop_wrap_console()
self._wrap_console = False
if self._internal_console_port is not None:
self._manager.port_manager.release_tcp_port(self._internal_console_port, self._project)
self._internal_console_port = None
await self._hypervisor.send(f'vm set_con_tcp_port "{self._name}" {self._console}')
async def set_aux(self, aux):
"""
@ -1047,8 +1072,20 @@ class Router(BaseNode):
self.aux_type = aux_type
if self._aux and aux_type == "telnet":
await self._hypervisor.send(f'vm set_aux_tcp_port "{self._name}" {self._aux}')
if self._aux:
if aux_type == "ssh":
# Switching to SSH: ensure an internal aux port is allocated and redirect Dynamips to it.
self._wrap_aux = True
if self._internal_aux_port is None:
self._internal_aux_port = self._manager.port_manager.get_free_tcp_port(self._project)
await self._hypervisor.send(f'vm set_aux_tcp_port "{self._name}" {self._internal_aux_port}')
elif aux_type == "telnet":
# Switching to telnet: clear SSH wrapper state and redirect Dynamips to external port.
self._wrap_aux = False
if self._internal_aux_port is not None:
self._manager.port_manager.release_tcp_port(self._internal_aux_port, self._project)
self._internal_aux_port = None
await self._hypervisor.send(f'vm set_aux_tcp_port "{self._name}" {self._aux}')
async def reset_console(self):
"""

View File

@ -41,6 +41,7 @@ from .utils.iou_import import nvram_import
from .utils.iou_export import nvram_export
from gns3server.compute.ubridge.ubridge_error import UbridgeError
from gns3server.utils.file_watcher import FileWatcher
from gns3server.utils.asyncio.ssh_server import AsyncioSSHServer
from gns3server.utils.asyncio.telnet_server import AsyncioTelnetServer
from gns3server.utils.hostname import is_ios_hostname_valid
from gns3server.utils.asyncio import locking
@ -643,19 +644,25 @@ class IOUVM(BaseNode):
async def start_console(self):
"""
Start the Telnet server to provide console access.
Start the console server to provide console access.
"""
if self.console and self.console_type == "telnet":
server = AsyncioTelnetServer(
reader=self._iou_process.stdout, writer=self._iou_process.stdin, binary=True, echo=True
)
if self.console and self.console_type in ("telnet", "ssh"):
if self.console_type == "telnet":
server = AsyncioTelnetServer(
reader=self._iou_process.stdout, writer=self._iou_process.stdin, binary=True, echo=True
)
error_prefix = "Telnet"
else:
server = AsyncioSSHServer(reader=self._iou_process.stdout, writer=self._iou_process.stdin)
error_prefix = "SSH"
try:
self._telnet_server = await server.start(self._manager.port_manager.console_host, self.console)
except OSError as e:
await self.stop()
raise IOUError(
"Could not start Telnet server on socket {}:{}: {}".format(
"Could not start {} server on socket {}:{}: {}".format(
error_prefix,
self._manager.port_manager.console_host, self.console, e
)
)

View File

@ -1922,7 +1922,7 @@ class QemuVM(BaseNode):
def _console_options(self):
if self._console_type == "telnet" and self._wrap_console:
if self._console_type in ("telnet", "ssh") and self._wrap_console:
return self._serial_options(self._internal_console_port, self.console)
elif self._console_type == "vnc":
return self._vnc_options(self.console)
@ -1938,7 +1938,7 @@ class QemuVM(BaseNode):
if self._aux_type != "none" and self._aux_type == self._console_type:
raise QemuError(f"Auxiliary console type {self._aux_type} cannot be the same as console type")
if self._aux_type == "telnet" and self._wrap_aux:
if self._aux_type in ("telnet", "ssh") and self._wrap_aux:
return self._serial_options(self._internal_aux_port, self.aux)
elif self._aux_type == "vnc":
return self._vnc_options(self.aux)
@ -2663,7 +2663,7 @@ class QemuVM(BaseNode):
await self._clear_save_vm_stated()
else:
command.extend(await self._saved_state_option())
if self._console_type == "telnet":
if self._console_type in ("telnet", "ssh"):
command.extend(await self._disable_graphics())
if self._tpm:
command.extend(self._tpm_options())

View File

@ -30,6 +30,7 @@ import tempfile
import xml.etree.ElementTree as ET
from gns3server.utils import parse_version
from gns3server.utils.asyncio.ssh_server import AsyncioSSHServer
from gns3server.utils.asyncio.telnet_server import AsyncioTelnetServer
from gns3server.utils.asyncio.serial import asyncio_open_serial
from gns3server.utils.asyncio import locking
@ -1011,20 +1012,25 @@ class VirtualBoxVM(BaseNode):
Starts remote console support for this VM.
"""
if self.console and self.console_type == "telnet":
if self.console and self.console_type in ("telnet", "ssh"):
pipe_name = self._get_pipe_name()
try:
self._remote_pipe = await asyncio_open_serial(pipe_name)
except OSError as e:
raise VirtualBoxError(f"Could not open serial pipe '{pipe_name}': {e}")
server = AsyncioTelnetServer(reader=self._remote_pipe, writer=self._remote_pipe, binary=True, echo=True)
if self.console_type == "telnet":
server = AsyncioTelnetServer(reader=self._remote_pipe, writer=self._remote_pipe, binary=True, echo=True)
transport = "Telnet"
else:
server = AsyncioSSHServer(reader=self._remote_pipe, writer=self._remote_pipe)
transport = "SSH"
try:
self._telnet_server = await server.start(self._manager.port_manager.console_host, self.console)
except OSError as e:
self.project.emit(
"log.warning",
{
"message": f"Could not start Telnet server on socket {self._manager.port_manager.console_host}:{self.console}: {e}"
"message": f"Could not start {transport} server on socket {self._manager.port_manager.console_host}:{self.console}: {e}"
},
)

View File

@ -24,6 +24,7 @@ import asyncio
import tempfile
import platform
from gns3server.utils.asyncio.ssh_server import AsyncioSSHServer
from gns3server.utils.asyncio.telnet_server import AsyncioTelnetServer
from gns3server.utils.asyncio.serial import asyncio_open_serial
from gns3server.utils import parse_version
@ -912,20 +913,25 @@ class VMwareVM(BaseNode):
Starts remote console support for this VM.
"""
if self.console and self.console_type == "telnet":
if self.console and self.console_type in ("telnet", "ssh"):
pipe_name = self._get_pipe_name()
try:
self._remote_pipe = await asyncio_open_serial(self._get_pipe_name())
except OSError as e:
raise VMwareError(f"Could not open serial pipe '{pipe_name}': {e}")
server = AsyncioTelnetServer(reader=self._remote_pipe, writer=self._remote_pipe, binary=True, echo=True)
if self.console_type == "telnet":
server = AsyncioTelnetServer(reader=self._remote_pipe, writer=self._remote_pipe, binary=True, echo=True)
transport = "Telnet"
else:
server = AsyncioSSHServer(reader=self._remote_pipe, writer=self._remote_pipe)
transport = "SSH"
try:
self._telnet_server = await server.start(self._manager.port_manager.console_host, self.console)
except OSError as e:
self.project.emit(
"log.warning",
{
"message": f"Could not start Telnet server on socket {self._manager.port_manager.console_host}:{self.console}: {e}"
"message": f"Could not start {transport} server on socket {self._manager.port_manager.console_host}:{self.console}: {e}"
},
)

View File

@ -614,7 +614,7 @@ class Node:
Reset the console
"""
if self._console and self._console_type == "telnet":
if self._console and self._console_type in ("telnet", "ssh"):
try:
await self.post("/console/reset", timeout=240)
except asyncio.TimeoutError:

View File

@ -55,6 +55,7 @@ class ConsoleType(str, Enum):
vnc = "vnc"
telnet = "telnet"
ssh = "ssh"
http = "http"
https = "https"
spice = "spice"
@ -68,4 +69,5 @@ class AuxType(str, Enum):
"""
telnet = "telnet"
ssh = "ssh"
none = "none"

View File

@ -88,6 +88,7 @@ class UDPPort(BaseModel):
class CloudConsoleType(str, Enum):
telnet = "telnet"
ssh = "ssh"
vnc = "vnc"
spice = "spice"
http = "http"

View File

@ -87,6 +87,7 @@ class DynamipsConsoleType(str, Enum):
"""
telnet = "telnet"
ssh = "ssh"
none = "none"

View File

@ -61,6 +61,7 @@ class QemuConsoleType(str, Enum):
vnc = "vnc"
telnet = "telnet"
ssh = "ssh"
spice = "spice"
spice_agent = "spice+agent"
none = "none"

View File

@ -28,6 +28,7 @@ class VirtualBoxConsoleType(str, Enum):
"""
telnet = "telnet"
ssh = "ssh"
none = "none"

View File

@ -28,6 +28,7 @@ class VMwareConsoleType(str, Enum):
"""
telnet = "telnet"
ssh = "ssh"
none = "none"

View File

@ -28,6 +28,7 @@ class ConsoleType(str, Enum):
"""
telnet = "telnet"
ssh = "ssh"
none = "none"

View File

@ -162,6 +162,7 @@ class QemuConsoleType(str, Enum):
"""Qemu console type enum"""
telnet = 'telnet'
ssh = 'ssh'
vnc = 'vnc'
spice = 'spice'
spice_agent = 'spice+agent'
@ -279,6 +280,7 @@ class DockerConsoleType(str, Enum):
"""Docker console type enum"""
telnet = 'telnet'
ssh = 'ssh'
vnc = 'vnc'
http = 'http'
https = 'https'

View File

@ -0,0 +1,203 @@
#
# Copyright (C) 2026 GNS3 Technologies Inc.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import asyncio
import contextlib
import logging
import socket
import asyncssh
log = logging.getLogger(__name__)
READ_SIZE = 1024
BROADCAST_DRAIN_TIMEOUT = 10
class _NoAuthSSHServer(asyncssh.SSHServer):
"""Allow console transport without interactive SSH authentication prompts."""
def begin_auth(self, username):
return False
class _ManagedSSHListener:
"""Compatibility wrapper that owns AsyncioSSHServer shutdown."""
def __init__(self, ssh_server, listener):
self._ssh_server = ssh_server
self._listener = listener
self._close_task = None
def close(self):
if self._close_task is None:
self._close_task = asyncio.create_task(self._ssh_server.close())
async def wait_closed(self):
self.close()
with contextlib.suppress(asyncio.CancelledError):
await self._close_task
def __getattr__(self, attribute):
return getattr(self._listener, attribute)
class AsyncioSSHServer:
def __init__(self, reader=None, writer=None):
self._reader = reader
self._writer = writer
self._sessions = {}
self._sessions_lock = asyncio.Lock()
self._writer_lock = asyncio.Lock()
self._close_lock = asyncio.Lock()
self._broadcast_task = None
self._server = None
self._server_handle = None
self._host_key = asyncssh.generate_private_key("ssh-rsa")
async def start(self, host, port):
if self._server is not None:
raise RuntimeError("AsyncioSSHServer is already started")
self._server = await asyncssh.listen(
host=host,
port=port,
server_factory=_NoAuthSSHServer,
server_host_keys=[self._host_key],
process_factory=self._run_client_session,
encoding=None,
reuse_address=True,
)
self._server_handle = _ManagedSSHListener(self, self._server)
if self._reader is not None and self._broadcast_task is None:
self._broadcast_task = asyncio.create_task(self._broadcast_from_upstream())
return self._server_handle
async def close(self):
async with self._close_lock:
if self._broadcast_task is not None:
broadcast_task = self._broadcast_task
self._broadcast_task = None
broadcast_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await broadcast_task
# Always disconnect all active SSH client sessions so that
# server.wait_closed() does not block indefinitely waiting for
# them to finish on their own.
await self._disconnect_all_clients()
if self._server is not None:
self._server.close()
await self._server.wait_closed()
self._server = None
self._server_handle = None
async def _run_client_session(self, process):
self._set_socket_options(process)
async with self._sessions_lock:
self._sessions[process] = process
try:
while True:
data = await process.stdin.read(READ_SIZE)
if not data:
break
if self._writer is not None:
async with self._writer_lock:
self._writer.write(data)
await self._writer.drain()
except asyncio.CancelledError:
raise
except (ConnectionError, OSError, asyncssh.Error):
pass
finally:
await self._disconnect_client(process)
async def _broadcast_from_upstream(self):
try:
while True:
data = await self._reader.read(READ_SIZE)
if not data:
break
for process in await self._get_sessions_snapshot():
try:
process.stdout.write(data)
await asyncio.wait_for(process.stdout.drain(), timeout=BROADCAST_DRAIN_TIMEOUT)
except (OSError, ConnectionError, asyncio.TimeoutError, asyncssh.Error):
await self._disconnect_client(process)
except asyncio.CancelledError:
raise
except (ConnectionError, OSError):
pass
finally:
await self._disconnect_all_clients()
async def _get_sessions_snapshot(self):
async with self._sessions_lock:
return list(self._sessions.keys())
async def _disconnect_all_clients(self):
async with self._sessions_lock:
sessions = list(self._sessions.keys())
for process in sessions:
await self._disconnect_client(process)
async def _disconnect_client(self, process):
async with self._sessions_lock:
self._sessions.pop(process, None)
with contextlib.suppress(Exception):
process.exit(0)
channel = process.get_extra_info("channel")
if channel is not None:
with contextlib.suppress(Exception):
channel.close()
wait_closed = getattr(channel, "wait_closed", None)
if callable(wait_closed):
with contextlib.suppress(Exception):
await wait_closed()
@staticmethod
def _set_socket_options(process):
sock = process.get_extra_info("socket")
if sock is None:
return
with contextlib.suppress(OSError):
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
try:
if hasattr(socket, "TCP_KEEPIDLE"):
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60)
elif hasattr(socket, "TCP_KEEPALIVE"):
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, 60)
else:
raise AttributeError("No TCP keepalive idle socket option is available")
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 10)
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 4)
except (AttributeError, OSError):
log.debug("Failed to tune TCP keepalive for SSH client; using OS defaults", exc_info=True)

View File

@ -26,6 +26,7 @@ truststore>=0.10.4; python_version >= '3.10'
# Shared dependencies (also used by AI Copilot)
telnetlib3==4.0.2
asyncssh>=2.21.0,<3
typing-extensions>=4.15.0
requests>=2.33.1
urllib3>=2.6.2