"""herethere.here.server"""
import asyncio
import inspect
import logging
import os
import threading
import traceback
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from functools import partial
from io import StringIO
from typing import Any
import asyncssh
from herethere.everywhere.code import runcode
from herethere.everywhere.commands import (
BACKGROUND_COMMAND,
BACKGROUND_EXECUTE_COMMAND,
BACKGROUND_VALUE_COMMAND,
CODE_COMMAND,
EXECUTE_COMMAND,
PING_COMMAND,
RECENT_LOGS_COMMAND,
SHELL_COMMAND,
VALUE_COMMAND,
WORKER_EXECUTE_COMMAND,
)
from herethere.everywhere.live import LiveError, execute_live
from herethere.everywhere.logging import logger
from herethere.everywhere.protocol import (
CapturedStreamEvent,
EventStream,
OrderedBoundedOutputCollector,
decode_request_object,
write_event,
)
from herethere.everywhere.recent_logs import (
DEFAULT_MAX_LOG_RECORDS,
RECENT_LOGS_PROTOCOL_VERSION,
RecentLogHandler,
create_recent_log_handler,
)
from herethere.everywhere.redirected_output import redirect_output
from herethere.everywhere.values import dumps_error, dumps_value
from herethere.here.config import ServerConfig
from herethere.here.shell import handle_shell_command
CONNECTION_CLOSE_TIMEOUT = 1.0
MAX_RECENT_LOGS_REQUEST_BYTES = 64 * 1024
RECENT_LOGS_READ_CHUNK_SIZE = 4096
class _DiscardOutput:
"""Text writer which accepts and immediately discards output."""
def write(self, data: str) -> int:
"""Report a successful write without retaining data."""
return len(data)
def flush(self) -> None:
"""Provide the standard text-writer flush interface."""
@dataclass(frozen=True)
class WorkerExecutionOutcome:
"""Buffered stream events and final status returned by worker execution."""
events: tuple[CapturedStreamEvent, ...]
error: LiveError | None
async def handle_ping_command(process: asyncssh.SSHServerProcess, namespace: dict): # pylint: disable=unused-argument
"""Handler for SSH command 'ping'."""
process.stdout.write("pong")
async def handle_code_command(process: asyncssh.SSHServerProcess, namespace: dict):
"""Handler for SSH command 'code': execute code in the main thread.
Blocks main thread execution.
"""
data = await process.stdin.read()
runcode(data, stdout=process.stdout, stderr=process.stderr, namespace=namespace)
async def handle_background_code_command(
process: asyncssh.SSHServerProcess, namespace: dict
):
"""Handler for SSH command 'background': execute code in a separate thread.
Do not blocks main thread execution.
"""
server: SSHServerHere = process.channel.get_connection().get_owner()
data = await process.stdin.read()
await server.run_in_executor(
runcode,
code=data,
stdout=process.stdout,
stderr=process.stderr,
namespace=namespace,
)
async def maybe_await(value):
"""Await value if it is awaitable."""
if inspect.isawaitable(value):
return await value
return value
async def handle_value_command(process: asyncssh.SSHServerProcess, namespace: dict):
"""Handler for SSH command 'value': evaluate and return one Python value."""
data = await process.stdin.read()
try:
with redirect_output(stdout=StringIO(), stderr=StringIO()):
result = eval(data, namespace) # pylint: disable=eval-used
result = await maybe_await(result)
event = dumps_value(result)
except Exception as exc: # pylint: disable=broad-exception-caught
event = dumps_error(exc, traceback.format_exc())
process.stdout.write(event)
process.stdout.write("\n")
await process.stdout.drain()
def evaluate_value_sync(expression: str, namespace: dict):
"""Evaluate an expression in a worker while suppressing incidental output."""
output = _DiscardOutput()
with redirect_output(stdout=output, stderr=output):
return eval(expression, namespace) # pylint: disable=eval-used
async def handle_background_value_command(
process: asyncssh.SSHServerProcess, namespace: dict
):
"""Evaluate and return one Python value from an executor worker."""
server: SSHServerHere = process.channel.get_connection().get_owner()
data = await process.stdin.read()
try:
result = await server.run_in_executor(
evaluate_value_sync,
expression=data,
namespace=namespace,
)
output = _DiscardOutput()
with redirect_output(stdout=output, stderr=output):
result = await maybe_await(result)
event = dumps_value(result)
except Exception as exc: # pylint: disable=broad-exception-caught
event = dumps_error(exc, traceback.format_exc())
process.stdout.write(event)
process.stdout.write("\n")
await process.stdout.drain()
def _write_execute_result(process: asyncssh.SSHServerProcess, error) -> None:
"""Write the final structured execution status event."""
event = {"type": "result", "ok": error is None}
if error is not None:
event["error"] = error.asdict()
write_event(process.stdout, event)
async def handle_execute_command(process: asyncssh.SSHServerProcess, namespace: dict):
"""Execute code and return JSON-lines output events plus a final status."""
data = await process.stdin.read()
error = execute_live(
data,
namespace,
stdout=EventStream(process.stdout, "stdout"),
stderr=EventStream(process.stdout, "stderr"),
)
_write_execute_result(process, error)
async def handle_background_execute_command(
process: asyncssh.SSHServerProcess, namespace: dict
):
"""Execute structured Python code silently in an executor worker."""
server: SSHServerHere = process.channel.get_connection().get_owner()
data = await process.stdin.read()
output = _DiscardOutput()
error = await server.run_in_executor(
execute_live,
code=data,
namespace=namespace,
stdout=output,
stderr=output,
)
_write_execute_result(process, error)
def execute_worker_live(code: str, namespace: dict) -> WorkerExecutionOutcome:
"""Execute code with bounded ordered output capture in an executor worker."""
collector = OrderedBoundedOutputCollector()
error = execute_live(
code,
namespace,
stdout=collector.stdout,
stderr=collector.stderr,
)
return WorkerExecutionOutcome(events=tuple(collector.events), error=error)
async def handle_worker_execute_command(
process: asyncssh.SSHServerProcess, namespace: dict
):
"""Execute code in a worker, then replay buffered events on the event loop."""
server: SSHServerHere = process.channel.get_connection().get_owner()
data = await process.stdin.read()
outcome = await server.run_in_executor(
execute_worker_live,
code=data,
namespace=namespace,
)
for event in outcome.events:
write_event(process.stdout, event.asdict())
_write_execute_result(process, outcome.error)
async def handle_recent_logs_command(
process: asyncssh.SSHServerProcess,
namespace: dict,
recent_logs: RecentLogHandler,
):
"""Return a finite snapshot of buffered Python log records."""
try:
request_text = await _read_recent_logs_request(process.stdin)
max_records = _decode_recent_logs_request(request_text)
except (TypeError, ValueError) as exc:
process.stderr.write(f"Invalid recent-logs request: {exc}")
return
write_event(process.stdout, recent_logs.snapshot(max_records).asdict())
async def _read_recent_logs_request(reader) -> str:
"""Read one bounded recent-logs control request through EOF."""
chunks = []
size = 0
while chunk := await reader.read(RECENT_LOGS_READ_CHUNK_SIZE):
size += len(chunk.encode("utf-8"))
if size > MAX_RECENT_LOGS_REQUEST_BYTES:
raise ValueError(
"request is too large "
f"({size} bytes > {MAX_RECENT_LOGS_REQUEST_BYTES} bytes)"
)
chunks.append(chunk)
return "".join(chunks)
def _decode_recent_logs_request(request_text: str) -> int | None:
"""Decode an optional recent-log snapshot request."""
if not request_text:
return None
request = decode_request_object(request_text)
if request.get("version") != RECENT_LOGS_PROTOCOL_VERSION:
raise ValueError(
f"request must use protocol version {RECENT_LOGS_PROTOCOL_VERSION}"
)
max_records = request.get("records")
if max_records is None:
return None
if (
not isinstance(max_records, int)
or isinstance(max_records, bool)
or not 1 <= max_records <= DEFAULT_MAX_LOG_RECORDS
):
raise ValueError(f"records must be in the range 1..{DEFAULT_MAX_LOG_RECORDS}")
return max_records
async def handle_client(
process: asyncssh.SSHServerProcess,
namespace: dict,
recent_logs: RecentLogHandler | None = None,
):
"""SSH requests handler."""
if namespace is None:
namespace = {}
channel = process.channel
stdin_channel = process.stdin.channel
# Configure terminal input only for PTY sessions, where AsyncSSH line
# editor echo and line-mode controls are meaningful.
if (
process.get_terminal_type()
and hasattr(channel, "set_echo")
and hasattr(stdin_channel, "set_line_mode")
):
channel.set_echo(False)
stdin_channel.set_line_mode(True)
try:
processors = {
PING_COMMAND: handle_ping_command,
CODE_COMMAND: handle_code_command,
BACKGROUND_COMMAND: handle_background_code_command,
BACKGROUND_EXECUTE_COMMAND: handle_background_execute_command,
BACKGROUND_VALUE_COMMAND: handle_background_value_command,
SHELL_COMMAND: handle_shell_command,
VALUE_COMMAND: handle_value_command,
EXECUTE_COMMAND: handle_execute_command,
WORKER_EXECUTE_COMMAND: handle_worker_execute_command,
}
if recent_logs is not None:
processors[RECENT_LOGS_COMMAND] = partial(
handle_recent_logs_command,
recent_logs=recent_logs,
)
processor = processors[process.command]
except KeyError:
logger.error("Unknown command: %s", process.command[:64])
process.stderr.write("Unknown command")
process.exit(0)
return
await processor(process, namespace=namespace)
await process.stdout.drain()
await process.stderr.drain()
process.exit(0)
class SFTPServerHere(asyncssh.SFTPServer):
"""SFTP session handler for a given root directory."""
def __init__(self, chan: asyncssh.SSHLineEditorChannel, chroot: str):
os.makedirs(chroot, exist_ok=True)
super().__init__(chan, chroot=chroot)
[docs]
class SSHServerHere(asyncssh.SSHServer):
"""SSH server protocol handler with `username` and `password` options."""
def __init__(
self,
username: str,
password: str,
executor: ThreadPoolExecutor,
connections: set[asyncssh.SSHServerConnection] | None = None,
):
self.passwords = {username: password}
self.executor = executor
self.connections = connections
self.conn: asyncssh.SSHServerConnection | None = None
def connection_made(self, conn: asyncssh.SSHServerConnection):
"""Called when a connection is opened successfully."""
self.conn = conn
if self.connections is not None:
self.connections.add(conn)
peername = conn.get_extra_info("peername")
peer = peername[0] if peername else "unknown"
logger.info("SSH connection received from %s.", peer)
def connection_lost(self, exc: Exception | None):
"""Called when a connection is closed."""
if self.connections is not None and self.conn is not None:
self.connections.discard(self.conn)
self.conn = None
if exc:
logger.info("SSH connection lost: %s.", exc)
else:
logger.info("SSH connection closed.")
def password_auth_supported(self) -> bool:
"""Password authentication is supported."""
return True
def begin_auth(self, username: str) -> bool:
"""Allow authentication for the client."""
return True
def validate_password(self, username: str, password: str) -> bool:
"""Return whether password is valid for this user."""
expected = self.passwords.get(username, None)
return expected is not None and password == expected
async def run_in_executor(self, func: Callable[..., Any], **kwargs: Any):
"""Run a callable in the thread-pool executor and return its result."""
return await asyncio.get_running_loop().run_in_executor(
self.executor, partial(func, **kwargs)
)
[docs]
class RunningServer:
"""Wrapper for a running SSH server instance."""
def __init__(
self,
server: asyncio.AbstractServer,
namespace,
executor: ThreadPoolExecutor,
connections: set[asyncssh.SSHServerConnection],
recent_logs: RecentLogHandler | None = None,
):
self.server = server
self.executor = executor
self.connections = connections
self.recent_logs = recent_logs
self._recent_logs_attached = recent_logs is not None
self.namespace = namespace
self.namespace["ssh_server_closed"] = threading.Event()
def __getattr__(self, attr):
return getattr(self.server, attr)
@staticmethod
def _log_wait_closed_results(tasks, action: str):
for task in tasks:
try:
task.result()
except asyncio.CancelledError:
logger.debug("SSH connection %s wait was cancelled.", action)
except Exception as exc: # pylint: disable=broad-exception-caught # noqa: BLE001
logger.debug("SSH connection %s finished with error: %r", action, exc)
async def _wait_for_connection_tasks(self, tasks, action: str, timeout: float):
if not tasks:
return set()
done, pending = await asyncio.wait(
tasks,
timeout=timeout,
)
self._log_wait_closed_results(done, action)
return pending
async def _close_connections(self, timeout: float):
connections = tuple(self.connections)
for conn in connections:
conn.close()
wait_closed_tasks = [
asyncio.create_task(
conn.wait_closed(),
name="SSH connection wait_closed",
)
for conn in connections
]
if wait_closed_tasks:
pending = await self._wait_for_connection_tasks(
wait_closed_tasks,
"close",
timeout,
)
for task, conn in zip(wait_closed_tasks, connections, strict=True):
if task in pending:
logger.debug("SSH connection did not close in time; aborting.")
conn.abort()
if pending:
still_pending = await self._wait_for_connection_tasks(
pending,
"abort",
timeout,
)
for task in still_pending:
logger.debug(
"SSH connection wait_closed still pending; cancelling task."
)
task.cancel()
if still_pending:
await asyncio.gather(*still_pending, return_exceptions=True)
[docs]
async def stop(self, timeout: float = CONNECTION_CLOSE_TIMEOUT):
"""Stop SSH server."""
self.namespace["ssh_server_closed"].set()
if self._recent_logs_attached and self.recent_logs is not None:
logging.getLogger().removeHandler(self.recent_logs)
self.recent_logs.close()
self._recent_logs_attached = False
self.server.close()
await self._close_connections(timeout)
self.executor.shutdown(wait=False)
try:
await asyncio.wait_for(
self.server.wait_closed(),
timeout=timeout,
)
except asyncio.TimeoutError:
logger.debug("SSH server wait_closed timed out.")
def generate_private_key(path: str):
"""Generate and save private key to a given location."""
asyncssh.generate_private_key("ssh-ed25519").write_private_key(path)
[docs]
async def start_server(
config: ServerConfig,
namespace: dict = None,
server_factory: type[SSHServerHere] = SSHServerHere,
) -> RunningServer:
"""Start SSH server.
:param config: server configuration options
:param namespace: dictionary in which Python code commands will be executed
:param server_factory: optional protocol handler class
"""
if not issubclass(server_factory, SSHServerHere):
raise TypeError("server_factory must be a SSHServerHere sublcass.")
if not os.path.exists(config.key_path):
logger.info("Generating new private key.")
generate_private_key(config.key_path)
if namespace is None:
namespace = {}
executor = ThreadPoolExecutor(
max_workers=64, thread_name_prefix="SSHServerHereThread"
)
connections: set[asyncssh.SSHServerConnection] = set()
recent_logs = create_recent_log_handler()
root_logger = logging.getLogger()
root_logger.addHandler(recent_logs)
logger.debug(
"start_server host=%s port=%s sftp_root=%s",
config.host,
config.port,
config.sftp_root,
)
try:
server = await asyncssh.create_server(
host=config.host,
port=config.port,
server_host_keys=[config.key_path],
server_factory=partial(
server_factory,
username=config.username,
password=config.password,
executor=executor,
connections=connections,
),
process_factory=partial(
handle_client,
namespace=namespace,
recent_logs=recent_logs,
),
sftp_factory=config.sftp_root
and partial(SFTPServerHere, chroot=config.sftp_root),
reuse_address=True,
)
except Exception:
root_logger.removeHandler(recent_logs)
recent_logs.close()
executor.shutdown(wait=False)
raise
return RunningServer(
server=server,
namespace=namespace,
executor=executor,
connections=connections,
recent_logs=recent_logs,
)
