MyAvaloniaDemo/AuroraDesk.Infrastructure/Services/SshSessionService.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Reactive.Linq;
using System.Reactive.Subjects;
using System.Threading;
using System.Threading.Tasks;
using AuroraDesk.Core.Entities;
using AuroraDesk.Core.Interfaces;
using Microsoft.Extensions.Logging;

namespace AuroraDesk.Infrastructure.Services;

/// <summary>
/// 基于系统 ssh 可执行文件的多会话管理服务，实现实时消息收发
/// </summary>
public sealed class SshSessionService : ISshSessionService
{
    private readonly ConcurrentDictionary<Guid, SessionContext> _sessions = new();
    private readonly ILogger<SshSessionService>? _logger;
    private readonly TimeSpan _shutdownGracePeriod = TimeSpan.FromSeconds(2);

    public SshSessionService(ILogger<SshSessionService>? logger = null)
    {
        _logger = logger;
    }

    public async Task<SshSessionInfo> StartSessionAsync(SshSessionOptions options, CancellationToken cancellationToken = default)
    {
        if (options is null)
        {
            throw new ArgumentNullException(nameof(options));
        }

        ValidateOptions(options);

        var sessionId = Guid.NewGuid();
        var displayName = string.IsNullOrWhiteSpace(options.DisplayName)
            ? $"{options.UserName}@{options.Host}:{options.Port}"
            : options.DisplayName;

        var info = new SshSessionInfo(sessionId, displayName, options.Host, options.Port, options.UserName, SshSessionStatus.Connecting, DateTime.UtcNow);
        var context = new SessionContext(info, options);

        if (!_sessions.TryAdd(sessionId, context))
        {
            throw new InvalidOperationException("无法注册新的 SSH 会话。");
        }

        try
        {
            await StartProcessAsync(context, cancellationToken).ConfigureAwait(false);
            context.UpdateStatus(SshSessionStatus.Connected);
            _logger?.LogInformation("SSH 会话 {SessionId} 已连接到 {User}@{Host}:{Port}", sessionId, options.UserName, options.Host, options.Port);
            return context.Info;
        }
        catch (Exception ex)
        {
            context.UpdateStatus(SshSessionStatus.Error);
            context.PublishSystemMessage($"启动会话失败: {ex.Message}", isError: true);
            _sessions.TryRemove(sessionId, out _);
            context.Dispose();
            _logger?.LogError(ex, "启动 SSH 会话失败: {User}@{Host}:{Port}", options.UserName, options.Host, options.Port);
            throw;
        }
    }

    public async Task StopSessionAsync(Guid sessionId)
    {
        if (!_sessions.TryRemove(sessionId, out var context))
        {
            return;
        }

        context.UpdateStatus(SshSessionStatus.Disconnected);
        context.PublishSystemMessage("会话已手动关闭。");

        try
        {
            if (context.Process is { HasExited: false })
            {
                await context.SendLock.WaitAsync().ConfigureAwait(false);
                try
                {
                    await context.Process.StandardInput.WriteLineAsync("exit").ConfigureAwait(false);
                    await context.Process.StandardInput.FlushAsync().ConfigureAwait(false);
                }
                finally
                {
                    context.SendLock.Release();
                }

                if (!context.Process.WaitForExit((int)_shutdownGracePeriod.TotalMilliseconds))
                {
                    context.Process.Kill(entireProcessTree: true);
                }
            }
        }
        catch (Exception ex)
        {
            _logger?.LogWarning(ex, "关闭 SSH 会话时出现异常: {SessionId}", sessionId);
        }
        finally
        {
            context.Complete();
            context.Dispose();
        }
    }

    public async Task<bool> SendAsync(Guid sessionId, string payload, CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(payload))
        {
            return false;
        }

        if (!_sessions.TryGetValue(sessionId, out var context))
        {
            return false;
        }

        if (context.StatusSubject.Value != SshSessionStatus.Connected || context.Process is null || context.Process.HasExited)
        {
            context.PublishSystemMessage("会话未连接，无法发送数据。", isError: true);
            return false;
        }

        await context.SendLock.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            await context.Process.StandardInput.WriteLineAsync(payload).ConfigureAwait(false);
            await context.Process.StandardInput.FlushAsync().ConfigureAwait(false);
            context.PublishMessage(payload, SshMessageDirection.Outgoing);
            return true;
        }
        catch (Exception ex)
        {
            context.PublishSystemMessage($"发送失败: {ex.Message}", isError: true);
            _logger?.LogError(ex, "SSH 会话 {SessionId} 发送数据失败。", sessionId);
            return false;
        }
        finally
        {
            context.SendLock.Release();
        }
    }

    public IReadOnlyCollection<SshSessionInfo> GetSessions()
    {
        return _sessions.Values.Select(context => context.Info).ToArray();
    }

    public IObservable<SshSessionStatus> ObserveStatus(Guid sessionId)
    {
        if (_sessions.TryGetValue(sessionId, out var context))
        {
            return context.StatusSubject.AsObservable();
        }

        throw new KeyNotFoundException($"未找到会话 {sessionId}");
    }

    public IObservable<SshMessage> ObserveMessages(Guid sessionId)
    {
        if (_sessions.TryGetValue(sessionId, out var context))
        {
            return context.MessageSubject.AsObservable();
        }

        throw new KeyNotFoundException($"未找到会话 {sessionId}");
    }

    public void Dispose()
    {
        foreach (var sessionId in _sessions.Keys.ToArray())
        {
            try
            {
                StopSessionAsync(sessionId).GetAwaiter().GetResult();
            }
            catch (Exception ex)
            {
                _logger?.LogWarning(ex, "释放 SshSessionService 时关闭会话 {SessionId} 失败。", sessionId);
            }
        }
    }

    private static void ValidateOptions(SshSessionOptions options)
    {
        if (string.IsNullOrWhiteSpace(options.Host))
        {
            throw new ArgumentException("Host 不能为空。", nameof(options));
        }

        if (options.Port is < 1 or > 65535)
        {
            throw new ArgumentException("端口必须在 1~65535 之间。", nameof(options));
        }

        if (string.IsNullOrWhiteSpace(options.UserName))
        {
            throw new ArgumentException("UserName 不能为空。", nameof(options));
        }

        if (string.IsNullOrWhiteSpace(options.Password) && string.IsNullOrWhiteSpace(options.PrivateKeyPath))
        {
            throw new ArgumentException("至少需要提供密码或私钥路径。", nameof(options));
        }
    }

    private async Task StartProcessAsync(SessionContext context, CancellationToken cancellationToken)
    {
        var options = context.Options;
        var startInfo = new ProcessStartInfo
        {
            FileName = string.IsNullOrWhiteSpace(options.SshExecutablePath) ? "ssh" : options.SshExecutablePath,
            RedirectStandardInput = true,
            RedirectStandardOutput = true,
            RedirectStandardError = true,
            UseShellExecute = false,
            CreateNoWindow = true
        };

        startInfo.ArgumentList.Add("-tt");

        if (options.AllowAnyHostKey)
        {
            startInfo.ArgumentList.Add("-o");
            startInfo.ArgumentList.Add("StrictHostKeyChecking=no");
            startInfo.ArgumentList.Add("-o");
            var nullDevice = OperatingSystem.IsWindows() ? "NUL" : "/dev/null";
            startInfo.ArgumentList.Add($"UserKnownHostsFile={nullDevice}");
        }

        startInfo.ArgumentList.Add("-p");
        startInfo.ArgumentList.Add(options.Port.ToString(CultureInfo.InvariantCulture));

        if (!string.IsNullOrWhiteSpace(options.PrivateKeyPath))
        {
            startInfo.ArgumentList.Add("-i");
            startInfo.ArgumentList.Add(options.PrivateKeyPath!);
        }

        startInfo.ArgumentList.Add($"{options.UserName}@{options.Host}");

        if (!string.IsNullOrWhiteSpace(options.PreferredShell))
        {
            startInfo.ArgumentList.Add(options.PreferredShell!);
        }

        var process = new Process
        {
            StartInfo = startInfo,
            EnableRaisingEvents = true
        };

        var started = process.Start();
        if (!started)
        {
            throw new InvalidOperationException("无法启动 ssh 进程。");
        }

        process.StandardInput.AutoFlush = true;
        context.AttachProcess(process);

        var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        context.AttachCancellation(linkedCts);

        _ = PumpStreamAsync(process, process.StandardOutput, context, SshMessageDirection.Incoming, linkedCts.Token, treatAsError: false);
        _ = PumpStreamAsync(process, process.StandardError, context, SshMessageDirection.System, linkedCts.Token, treatAsError: true);

        process.Exited += (_, _) =>
        {
            linkedCts.Cancel();

            var exitCode = process.ExitCode;
            var isError = exitCode != 0 && context.StatusSubject.Value != SshSessionStatus.Disconnected;
            var nextStatus = isError ? SshSessionStatus.Error : SshSessionStatus.Disconnected;

            context.UpdateStatus(nextStatus);
            context.PublishSystemMessage($"ssh 进程已退出 (代码 {exitCode}).", isError);

            _sessions.TryRemove(context.Info.SessionId, out _);
            context.Complete();
            context.Dispose();
        };

        await Task.CompletedTask.ConfigureAwait(false);
    }

    private Task PumpStreamAsync(
        Process process,
        StreamReader reader,
        SessionContext context,
        SshMessageDirection direction,
        CancellationToken cancellationToken,
        bool treatAsError)
    {
        return Task.Run(async () =>
        {
            try
            {
                while (!cancellationToken.IsCancellationRequested)
                {
                    var line = await reader.ReadLineAsync().ConfigureAwait(false);
                    if (line is null)
                    {
                        if (reader.EndOfStream)
                        {
                            break;
                        }

                        await Task.Delay(25, cancellationToken).ConfigureAwait(false);
                        continue;
                    }

                    if (direction == SshMessageDirection.System)
                    {
                        if (!context.PasswordSent
                            && !string.IsNullOrWhiteSpace(context.Options.Password)
                            && ContainsPrompt(line, "password"))
                        {
                            await SendSecretAsync(context, context.Options.Password!, cancellationToken).ConfigureAwait(false);
                            context.PasswordSent = true;
                            context.PublishSystemMessage("已自动响应密码提示。");
                            continue;
                        }

                        if (!context.PassphraseSent
                            && !string.IsNullOrWhiteSpace(context.Options.PrivateKeyPassphrase)
                            && ContainsPrompt(line, "enter passphrase"))
                        {
                            await SendSecretAsync(context, context.Options.PrivateKeyPassphrase!, cancellationToken).ConfigureAwait(false);
                            context.PassphraseSent = true;
                            context.PublishSystemMessage("已自动响应私钥口令提示。");
                            continue;
                        }

                        if (line.Contains("Permission denied", StringComparison.OrdinalIgnoreCase))
                        {
                            context.UpdateStatus(SshSessionStatus.Error);
                        }
                    }

                    context.PublishMessage(line, direction, treatAsError);
                }
            }
            catch (OperationCanceledException)
            {
                // ignored
            }
            catch (Exception ex)
            {
                _logger?.LogError(ex, "SSH 会话 {SessionId} 读取流失败。", context.Info.SessionId);
                context.PublishSystemMessage($"读取数据失败: {ex.Message}", isError: true);
                context.UpdateStatus(SshSessionStatus.Error);
            }
        }, cancellationToken);
    }

    private static bool ContainsPrompt(string line, string keyword)
    {
        return line.Contains(keyword, StringComparison.OrdinalIgnoreCase)
            || line.TrimEnd().EndsWith(keyword + ":", StringComparison.OrdinalIgnoreCase);
    }

    private static async Task SendSecretAsync(SessionContext context, string secret, CancellationToken cancellationToken)
    {
        if (context.Process is null || context.Process.HasExited)
        {
            return;
        }

        await context.SendLock.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            await context.Process.StandardInput.WriteLineAsync(secret).ConfigureAwait(false);
            await context.Process.StandardInput.FlushAsync().ConfigureAwait(false);
        }
        finally
        {
            context.SendLock.Release();
        }
    }

    private sealed class SessionContext : IDisposable
    {
        private readonly ReplaySubject<SshMessage> _messageSubject = new(bufferSize: 200);
        private readonly BehaviorSubject<SshSessionStatus> _statusSubject;
        private readonly object _disposeLock = new();
        private bool _disposed;

        public SessionContext(SshSessionInfo info, SshSessionOptions options)
        {
            Info = info;
            Options = options;
            _statusSubject = new BehaviorSubject<SshSessionStatus>(info.Status);
            SendLock = new SemaphoreSlim(1, 1);
        }

        public SshSessionInfo Info { get; private set; }
        public SshSessionOptions Options { get; }
        public Process? Process { get; private set; }
        public BehaviorSubject<SshSessionStatus> StatusSubject => _statusSubject;
        public ReplaySubject<SshMessage> MessageSubject => _messageSubject;
        public SemaphoreSlim SendLock { get; }
        public CancellationTokenSource? Cancellation { get; private set; }
        public bool PasswordSent { get; set; }
        public bool PassphraseSent { get; set; }

        public void AttachProcess(Process process)
        {
            Process = process ?? throw new ArgumentNullException(nameof(process));
        }

        public void AttachCancellation(CancellationTokenSource cts)
        {
            Cancellation = cts;
        }

        public void UpdateStatus(SshSessionStatus status)
        {
            Info = Info with { Status = status };
            _statusSubject.OnNext(status);
        }

        public void PublishMessage(string content, SshMessageDirection direction, bool isErrorChannel = false)
        {
            var message = new SshMessage(
                Info.SessionId,
                direction,
                content,
                DateTime.UtcNow,
                isErrorChannel && direction == SshMessageDirection.System && !IsNonFatalWarning(content));
            _messageSubject.OnNext(message);
        }

        public void PublishSystemMessage(string content, bool isError = false)
        {
            PublishMessage(content, SshMessageDirection.System, isError);
        }

        public void Complete()
        {
            if (!_messageSubject.IsDisposed)
            {
                _messageSubject.OnCompleted();
            }

            if (!_statusSubject.IsDisposed)
            {
                _statusSubject.OnCompleted();
            }
        }

        public void Dispose()
        {
            lock (_disposeLock)
            {
                if (_disposed)
                {
                    return;
                }

                _disposed = true;

                try
                {
                    Cancellation?.Cancel();
                }
                catch
                {
                    // ignored
                }
                finally
                {
                    Cancellation?.Dispose();
                }

                try
                {
                    Process?.Dispose();
                }
                catch
                {
                    // ignored
                }

                SendLock.Dispose();
                _messageSubject.Dispose();
                _statusSubject.Dispose();
            }
        }

        private static bool IsNonFatalWarning(string? content)
        {
            if (string.IsNullOrWhiteSpace(content))
            {
                return false;
            }

            var trimmed = content.TrimStart();
            return trimmed.StartsWith("Warning: Permanently added", StringComparison.OrdinalIgnoreCase);
        }
    }
}