OTSSignsOrchestrator/OTSSignsOrchestrator.Desktop/Services/SshConnectionService.cs

using Microsoft.Extensions.Logging;
using OTSSignsOrchestrator.Core.Models.Entities;
using Renci.SshNet;

namespace OTSSignsOrchestrator.Desktop.Services;

/// <summary>
/// Manages SSH connections to remote Docker Swarm hosts.
/// Creates and caches SshClient instances with key or password authentication.
/// </summary>
public class SshConnectionService : IDisposable
{
    private readonly ILogger<SshConnectionService> _logger;
    private readonly Dictionary<Guid, SshClient> _clients = new();
    private readonly object _lock = new();

    public SshConnectionService(ILogger<SshConnectionService> logger)
    {
        _logger = logger;
    }

    /// <summary>
    /// Get or create a connected SshClient for a given SshHost.
    /// </summary>
    public SshClient GetClient(SshHost host)
    {
        lock (_lock)
        {
            if (_clients.TryGetValue(host.Id, out var existing) && existing.IsConnected)
                return existing;

            // Dispose old client if disconnected
            if (existing != null)
            {
                existing.Dispose();
                _clients.Remove(host.Id);
            }

            var client = CreateClient(host);
            client.Connect();
            _clients[host.Id] = client;

            _logger.LogInformation("SSH connected to {Host}:{Port} as {User}", host.Host, host.Port, host.Username);
            return client;
        }
    }

    /// <summary>
    /// Test the SSH connection to a host. Returns (success, message).
    /// </summary>
    public async Task<(bool Success, string Message)> TestConnectionAsync(SshHost host)
    {
        return await Task.Run(() =>
        {
            try
            {
                using var client = CreateClient(host);
                client.ConnectionInfo.Timeout = TimeSpan.FromSeconds(10);
                client.Connect();

                if (client.IsConnected)
                {
                    // Quick test: run a simple command
                    using var cmd = client.RunCommand("docker --version");
                    client.Disconnect();

                    if (cmd.ExitStatus == 0)
                        return (true, $"Connected. {cmd.Result.Trim()}");
                    else
                        return (true, $"Connected but docker not available: {cmd.Error}");
                }

                return (false, "Failed to connect.");
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "SSH connection test failed for {Host}:{Port}", host.Host, host.Port);
                return (false, $"Connection failed: {ex.Message}");
            }
        });
    }

    /// <summary>
    /// Run a command on the remote host and return (exitCode, stdout, stderr).
    /// </summary>
    public async Task<(int ExitCode, string Stdout, string Stderr)> RunCommandAsync(SshHost host, string command)
    {
        return await Task.Run(() =>
        {
            var client = GetClient(host);
            using var cmd = client.RunCommand(command);
            return (cmd.ExitStatus ?? -1, cmd.Result, cmd.Error);
        });
    }

    /// <summary>
    /// Run a command on the remote host with a timeout.
    /// Returns exit code -1 and an error message if the command times out.
    /// </summary>
    public async Task<(int ExitCode, string Stdout, string Stderr)> RunCommandAsync(SshHost host, string command, TimeSpan timeout)
    {
        return await Task.Run(() =>
        {
            var client = GetClient(host);
            using var cmd = client.CreateCommand(command);
            cmd.CommandTimeout = timeout;
            try
            {
                cmd.Execute();
                return (cmd.ExitStatus ?? -1, cmd.Result, cmd.Error);
            }
            catch (Renci.SshNet.Common.SshOperationTimeoutException)
            {
                _logger.LogWarning("SSH command timed out after {Timeout}s: {Command}",
                    timeout.TotalSeconds, command.Length > 120 ? command[..120] + "…" : command);
                return (-1, string.Empty, $"Command timed out after {timeout.TotalSeconds}s");
            }
        });
    }

    /// <summary>
    /// Run a command that requires stdin input (e.g., docker stack deploy --compose-file -).
    /// </summary>
    public async Task<(int ExitCode, string Stdout, string Stderr)> RunCommandWithStdinAsync(
        SshHost host, string command, string stdinContent)
    {
        return await Task.Run(() =>
        {
            var client = GetClient(host);

            // Use shell stream approach for stdin piping
            // We pipe via: echo '<content>' | <command>
            // But for large YAML, use a heredoc approach
            var safeContent = stdinContent.Replace("'", "'\\''");
            var fullCommand = $"printf '%s' '{safeContent}' | {command}";

            using var cmd = client.RunCommand(fullCommand);
            return (cmd.ExitStatus ?? -1, cmd.Result, cmd.Error);
        });
    }

    /// <summary>
    /// Disconnect and remove a cached client.
    /// </summary>
    public void Disconnect(Guid hostId)
    {
        lock (_lock)
        {
            if (_clients.TryGetValue(hostId, out var client))
            {
                client.Disconnect();
                client.Dispose();
                _clients.Remove(hostId);
                _logger.LogInformation("SSH disconnected from host {HostId}", hostId);
            }
        }
    }

    private SshClient CreateClient(SshHost host)
    {
        var authMethods = new List<AuthenticationMethod>();

        if (host.UseKeyAuth && !string.IsNullOrEmpty(host.PrivateKeyPath))
        {
            var keyFile = string.IsNullOrEmpty(host.KeyPassphrase)
                ? new PrivateKeyFile(host.PrivateKeyPath)
                : new PrivateKeyFile(host.PrivateKeyPath, host.KeyPassphrase);

            authMethods.Add(new PrivateKeyAuthenticationMethod(host.Username, keyFile));
        }

        if (!string.IsNullOrEmpty(host.Password))
        {
            authMethods.Add(new PasswordAuthenticationMethod(host.Username, host.Password));
        }

        if (authMethods.Count == 0)
        {
            // Fall back to default SSH agent / key in ~/.ssh/
            var defaultKeyPath = Path.Combine(
                Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), ".ssh", "id_rsa");

            if (File.Exists(defaultKeyPath))
            {
                authMethods.Add(new PrivateKeyAuthenticationMethod(host.Username, new PrivateKeyFile(defaultKeyPath)));
            }
            else
            {
                throw new InvalidOperationException(
                    $"No authentication method configured for SSH host '{host.Label}'. " +
                    "Provide a private key path or password.");
            }
        }

        var connInfo = new ConnectionInfo(host.Host, host.Port, host.Username, authMethods.ToArray());
        return new SshClient(connInfo);
    }

    /// <summary>
    /// Opens an SSH local port-forward from 127.0.0.1:&lt;auto&gt; → <paramref name="remoteHost"/>:<paramref name="remotePort"/>
    /// through the existing SSH connection for <paramref name="host"/>.
    /// The caller must dispose the returned <see cref="ForwardedPortLocal"/> to close the tunnel.
    /// </summary>
    public ForwardedPortLocal OpenForwardedPort(SshHost host, string remoteHost, uint remotePort)
    {
        var client = GetClient(host);
        // Port 0 lets the OS assign a free local port; SSH.NET updates BoundPort after Start().
        var tunnel = new ForwardedPortLocal("127.0.0.1", 0, remoteHost, remotePort);
        client.AddForwardedPort(tunnel);
        tunnel.Start();
        _logger.LogDebug("SSH tunnel opened: 127.0.0.1:{LocalPort} → {RemoteHost}:{RemotePort}",
            tunnel.BoundPort, remoteHost, remotePort);
        return tunnel;
    }

    public void Dispose()
    {
        lock (_lock)
        {
            foreach (var client in _clients.Values)
            {
                try { client.Disconnect(); } catch { }
                client.Dispose();
            }
            _clients.Clear();
        }
    }
}