network-topology-discovery/internal/ssh/client.go

package sshclient

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"time"

	"golang.org/x/crypto/ssh"
)

// Client SSH客户端
type Client struct {
	client           *ssh.Client
	timeout          time.Duration
	host             string
	port             int
	username         string
	password         string
	keyFile          string
	insecureCiphers  bool
}

// Config SSH客户端配置
type Config struct {
	Host              string
	Port              int
	Username          string
	Password          string
	KeyFile           string
	Timeout           time.Duration
	InsecureCiphers   bool // 启用不安全的加密算法(用于兼容老旧设备)
}

// NewClient 创建新的SSH客户端
func NewClient(config Config) *Client {
	if config.Port == 0 {
		config.Port = 22
	}
	if config.Timeout == 0 {
		config.Timeout = 10 * time.Second
	}
	return &Client{
		host:             config.Host,
		port:             config.Port,
		username:         config.Username,
		password:         config.Password,
		keyFile:          config.KeyFile,
		timeout:          config.Timeout,
		insecureCiphers:  config.InsecureCiphers,
	}
}

// Connect 连接到SSH服务器
func (c *Client) Connect() error {
	config := &ssh.ClientConfig{
		User: c.username,
		Auth: []ssh.AuthMethod{},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         c.timeout,
	}

	// 添加密码认证
	if c.password != "" {
		config.Auth = append(config.Auth, ssh.Password(c.password))
	}

	// 添加密钥认证
	if c.keyFile != "" {
		key, err := loadPrivateKey(c.keyFile)
		if err != nil {
			return fmt.Errorf("failed to load private key: %w", err)
		}
		config.Auth = append(config.Auth, ssh.PublicKeys(key))
	}

	// 如果启用不安全加密算法,添加旧版算法支持(用于兼容老旧设备)
	if c.insecureCiphers {
		config.Ciphers = []string{
			"aes128-ctr", "aes192-ctr", "aes256-ctr",
			"aes128-gcm@openssh.com", "aes256-gcm@openssh.com",
			"chacha20-poly1305@openssh.com",
			"aes128-cbc", "aes256-cbc", // 旧版CBC算法
		}
		config.KeyExchanges = []string{
			"curve25519-sha256", "curve25519-sha256@libssh.org",
			"ecdh-sha2-nistp256", "ecdh-sha2-nistp384", "ecdh-sha2-nistp521",
			"diffie-hellman-group14-sha256", "diffie-hellman-group16-sha512",
			"diffie-hellman-group14-sha1", "diffie-hellman-group1-sha1", // 旧版KEX算法
		}
		config.MACs = []string{
			"hmac-sha2-256-etm@openssh.com", "hmac-sha2-512-etm@openssh.com",
			"hmac-sha2-256", "hmac-sha2-512",
			"hmac-sha1", "hmac-sha1-96", // 旧版MAC算法
		}
	}

	// 连接
	addr := fmt.Sprintf("%s:%d", c.host, c.port)
	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return fmt.Errorf("failed to connect to %s: %w", addr, err)
	}

	c.client = client
	return nil
}

// Close 关闭SSH连接
func (c *Client) Close() error {
	if c.client != nil {
		return c.client.Close()
	}
	return nil
}

// ExecuteCommand 执行命令并返回输出
func (c *Client) ExecuteCommand(command string) (string, error) {
	if c.client == nil {
		return "", fmt.Errorf("not connected")
	}

	// 重试机制,处理会话创建失败的情况
	var session *ssh.Session
	var err error
	maxRetries := 3
	
	for i := 0; i < maxRetries; i++ {
		session, err = c.client.NewSession()
		if err == nil {
			break
		}
		
		// 如果是会话拒绝错误,等待后重试
		if i < maxRetries-1 {
			time.Sleep(time.Duration(i+1) * 500 * time.Millisecond)
		}
	}
	
	if err != nil {
		return "", fmt.Errorf("failed to create session after %d retries: %w", maxRetries, err)
	}
	defer session.Close()

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	session.Stderr = &stdoutBuf

	err = session.Run(command)
	if err != nil {
		return stdoutBuf.String(), fmt.Errorf("command execution failed: %w", err)
	}

	return stdoutBuf.String(), nil
}

// ExecuteCommands 执行多个命令
func (c *Client) ExecuteCommands(commands []string) ([]string, error) {
	results := make([]string, 0, len(commands))
	for _, cmd := range commands {
		result, err := c.ExecuteCommand(cmd)
		if err != nil {
			return results, fmt.Errorf("failed to execute command '%s': %w", cmd, err)
		}
		results = append(results, result)
	}
	return results, nil
}

// CheckSSH 检查主机是否开启SSH
func CheckSSH(host string, port int, timeout time.Duration) bool {
	if port == 0 {
		port = 22
	}
	if timeout == 0 {
		timeout = 2 * time.Second
	}

	addr := fmt.Sprintf("%s:%d", host, port)
	conn, err := net.DialTimeout("tcp", addr, timeout)
	if err != nil {
		return false
	}
	defer conn.Close()
	return true
}

// loadPrivateKey 加载私钥文件
func loadPrivateKey(keyFile string) (ssh.Signer, error) {
	keyData, err := os.ReadFile(keyFile)
	if err != nil {
		return nil, fmt.Errorf("failed to read key file: %w", err)
	}

	signer, err := ssh.ParsePrivateKey(keyData)
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %w", err)
	}

	return signer, nil
}

// Ping 检查主机是否可达 (使用ICMP)
func Ping(host string, timeout time.Duration) bool {
	// 简单的TCP ping,实际项目可以使用专门的ICMP库
	ports := []int{22, 80, 443, 3389}
	for _, port := range ports {
		addr := fmt.Sprintf("%s:%d", host, port)
		conn, err := net.DialTimeout("tcp", addr, timeout)
		if err == nil {
			conn.Close()
			return true
		}
	}
	return false
}