LokiStack/backend/services/ssh.go

package services

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"lokistack/config"
	"lokistack/models"
	"lokistack/pkg/crypto"

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

type SSHClient struct {
	client *ssh.Client
}

type HostInfo struct {
	Hostname string `json:"hostname"`
	OS       string `json:"os"`
	Version  string `json:"version"`
	Arch     string `json:"arch"`
	Kernel   string `json:"kernel"`
	CPU      string `json:"cpu"`
	Memory   string `json:"memory"`
	Disk     string `json:"disk"`
	Uptime   string `json:"uptime"`

	// 资源使用率（百分比）
	CPUUsage  float64 `json:"cpuUsage"`
	MemUsage  float64 `json:"memUsage"`
	DiskUsage float64 `json:"diskUsage"`
	// 资源已用/总量
	MemTotal  string `json:"memTotal"`
	MemUsed   string `json:"memUsed"`
	DiskTotal string `json:"diskTotal"`
	DiskUsed  string `json:"diskUsed"`
}

// SSHAuthConfig 重导出models.SSHAuthConfig以兼容现有调用
// Deprecated: 直接使用 models.SSHAuthConfig
type SSHAuthConfig = models.SSHAuthConfig

// validatePort 校验SSH端口范围
func validatePort(port int) int {
	if port <= 0 {
		return 22
	}
	if port > 65535 {
		return 22
	}
	return port
}

// hostKeyPath 返回主机密钥存储路径
func hostKeyPath(host string) string {
	cfg := config.GetConfig()
	// 对主机名进行安全编码作为文件名
	safeHost := base64.URLEncoding.EncodeToString([]byte(host))
	return filepath.Join(cfg.HostKeyDir, safeHost+".pub")
}

// fingerprint 计算SSH主机密钥的SHA256指纹
func fingerprint(key ssh.PublicKey) string {
	hash := sha256.Sum256(key.Marshal())
	return base64.StdEncoding.EncodeToString(hash[:])
}

// TrustOnFirstUse 实现TOFU(Trust On First Use)主机密钥验证
// 首次连接时记录主机密钥，后续连接时校验
func TrustOnFirstUse(host string) ssh.HostKeyCallback {
	return func(hostname string, remote net.Addr, key ssh.PublicKey) error {
		fp := fingerprint(key)
		hostKeyFile := hostKeyPath(host)

		storedKey, err := os.ReadFile(hostKeyFile)
		if err != nil {
			// 首次连接：保存主机密钥
			if os.IsNotExist(err) {
				if writeErr := os.WriteFile(hostKeyFile, ssh.MarshalAuthorizedKey(key), 0600); writeErr != nil {
					return fmt.Errorf("failed to save host key: %w", writeErr)
				}
				// 记录首次信任的主机密钥指纹（安全审计日志）
				fmt.Fprintf(os.Stderr, "[SSH-TOFU] First connection to %s, host key fingerprint saved: SHA256:%s\n", host, fp)
				return nil
			}
			return fmt.Errorf("failed to read host key file: %w", err)
		}

		// 解析存储的密钥并与当前密钥比较
		parsedKey, _, _, _, err := ssh.ParseAuthorizedKey(storedKey)
		if err != nil {
			return fmt.Errorf("failed to parse stored host key: %w", err)
		}

		if !bytes.Equal(parsedKey.Marshal(), key.Marshal()) {
			return fmt.Errorf("HOST KEY MISMATCH for %s! Expected SHA256:%s, got SHA256:%s. Possible MITM attack!",
				host, fingerprint(parsedKey), fp)
		}

		return nil
	}
}

// InsecureIgnoreHostKey 保留但不导出，用于向后兼容的内部测试场景
// 生产环境必须使用TrustOnFirstUse
func InsecureIgnoreHostKey() ssh.HostKeyCallback {
	return ssh.InsecureIgnoreHostKey()
}

// buildAuthMethods 构建SSH认证方法列表（优先密钥认证）
func buildAuthMethods(auth *models.SSHAuthConfig) []ssh.AuthMethod {
	methods := []ssh.AuthMethod{}

	// 1. 优先尝试私钥认证
	if auth.PrivateKey != "" {
		privateKey := auth.PrivateKey
		// 如果是加密存储的密钥，先解密
		if strings.HasPrefix(privateKey, "ENC:") {
			privateKey = crypto.DecryptField(privateKey)
		}

		var signer ssh.Signer
		var err error
		keyPhrase := auth.KeyPhrase
		if strings.HasPrefix(keyPhrase, "ENC:") {
			keyPhrase = crypto.DecryptField(keyPhrase)
		}

		if keyPhrase != "" {
			signer, err = ssh.ParsePrivateKeyWithPassphrase([]byte(privateKey), []byte(keyPhrase))
		} else {
			signer, err = ssh.ParsePrivateKey([]byte(privateKey))
		}

		if err == nil {
			methods = append(methods, ssh.PublicKeys(signer))
		}
	}

	// 2. 密码认证作为后备
	if auth.Password != "" {
		password := auth.Password
		if strings.HasPrefix(password, "ENC:") {
			password = crypto.DecryptField(password)
		}
		methods = append(methods, ssh.Password(password))
	}

	return methods
}

// HostSSHAuth 从 Host 模型构建完整 SSH 认证配置（密钥优先、密码后备）。
// 凭证保持 ENC: 原样传入，由 buildAuthMethods 统一解密。
func HostSSHAuth(host *models.Host) *models.SSHAuthConfig {
	return &models.SSHAuthConfig{
		Host:       host.HostIP,
		Port:       host.Port,
		User:       host.User,
		Password:   host.Password,
		PrivateKey: host.PrivateKey,
		KeyPhrase:  host.KeyPhrase,
	}
}

// NewSSHClient 创建SSH客户端（支持密钥+密码认证，TOFU主机密钥验证）
func NewSSHClient(host string, port int, user, password string) (*SSHClient, error) {
	auth := &SSHAuthConfig{
		Host:     host,
		Port:     port,
		User:     user,
		Password: password,
	}
	return NewSSHClientWithAuth(auth)
}

// NewSSHClientWithAuth 使用完整认证配置创建SSH客户端
func NewSSHClientWithAuth(auth *models.SSHAuthConfig) (*SSHClient, error) {
	port := validatePort(auth.Port)
	if auth.User == "" {
		return nil, errors.New("SSH user is required")
	}

	authMethods := buildAuthMethods(auth)
	if len(authMethods) == 0 {
		return nil, errors.New("no authentication method available")
	}

	cfg := &ssh.ClientConfig{
		User:            auth.User,
		HostKeyCallback: TrustOnFirstUse(auth.Host),
		Timeout:         10 * time.Second,
		Auth:            authMethods,
	}

	addr := fmt.Sprintf("%s:%d", auth.Host, port)
	client, err := ssh.Dial("tcp", addr, cfg)
	if err != nil {
		return nil, errors.New(FormatSSHError(err))
	}

	return &SSHClient{client: client}, nil
}

// NewSSHClientUnsafe 创建不验证主机密钥的SSH客户端（仅用于内部测试）
// Deprecated: 生产环境请使用NewSSHClient或NewSSHClientWithAuth
func NewSSHClientUnsafe(host string, port int, user, password string) (*SSHClient, error) {
	port = validatePort(port)
	cfg := &ssh.ClientConfig{
		User:            user,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
		Auth:            []ssh.AuthMethod{ssh.Password(password)},
	}
	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", host, port), cfg)
	if err != nil {
		return nil, errors.New(FormatSSHError(err))
	}
	return &SSHClient{client: client}, nil
}

func (c *SSHClient) Execute(ctx context.Context, command string) (string, string, error) {
	session, err := c.client.NewSession()
	if err != nil {
		return "", "", err
	}
	defer session.Close()

	var stdout, stderr bytes.Buffer
	session.Stdout = &stdout
	session.Stderr = &stderr

	done := make(chan error, 1)
	go func() { done <- session.Run(command) }()

	select {
	case err := <-done:
		return stdout.String(), stderr.String(), err
	case <-ctx.Done():
		session.Signal(ssh.SIGKILL)
		return "", "", ctx.Err()
	}
}

// GetHostInfo 采集主机静态信息与资源使用率
func (c *SSHClient) GetHostInfo(ctx context.Context) (*HostInfo, error) {
	info := &HostInfo{}

	// 静态信息（逐条，best-effort）
	if out, _, err := c.Execute(ctx, "hostname"); err == nil {
		info.Hostname = strings.TrimSpace(out)
	}
	if out, _, err := c.Execute(ctx, `cat /etc/os-release | grep '^ID=' | cut -d= -f2 | tr -d '"'`); err == nil {
		info.OS = strings.TrimSpace(out)
	}
	if out, _, err := c.Execute(ctx, `cat /etc/os-release | grep 'VERSION_ID=' | cut -d= -f2 | tr -d '"'`); err == nil {
		info.Version = strings.TrimSpace(out)
	}
	if out, _, err := c.Execute(ctx, "uname -m"); err == nil {
		info.Arch = strings.TrimSpace(out)
	}
	if out, _, err := c.Execute(ctx, "uname -r"); err == nil {
		info.Kernel = strings.TrimSpace(out)
	}
	if out, _, err := c.Execute(ctx, "nproc"); err == nil {
		info.CPU = strings.TrimSpace(out) + " 核"
	}
	if out, _, err := c.Execute(ctx, "free -h | awk '/^Mem:/ {print $2}'"); err == nil {
		info.Memory = strings.TrimSpace(out)
	}
	if out, _, err := c.Execute(ctx, "df -h / | awk 'NR==2 {print $2}'"); err == nil {
		info.Disk = strings.TrimSpace(out)
	}
	if out, _, err := c.Execute(ctx, "uptime -p 2>/dev/null || uptime"); err == nil {
		info.Uptime = strings.TrimSpace(out)
	}

	// 资源使用率：一次合并命令（减少 SSH 往返）
	// 输出形如：CPUUSAGE:12.3 / MEM:16384:8192 / DISK:51200:25600:50%
	metricScript := `top -bn1 2>/dev/null | awk '/Cpu\(s\)/{print "CPUUSAGE:"100-$8; exit}'
free -m | awk '/^Mem:/{print "MEM:"$2":"$3}'
df -m / | awk 'NR==2{print "DISK:"$2":"$3":"$5}'`
	if out, _, err := c.Execute(ctx, metricScript); err == nil {
		for _, line := range strings.Split(out, "\n") {
			line = strings.TrimSpace(line)
			switch {
			case strings.HasPrefix(line, "CPUUSAGE:"):
				info.CPUUsage, _ = strconv.ParseFloat(strings.TrimPrefix(line, "CPUUSAGE:"), 64)
			case strings.HasPrefix(line, "MEM:"):
				parts := strings.Split(strings.TrimPrefix(line, "MEM:"), ":")
				if len(parts) == 2 {
					total, _ := strconv.Atoi(parts[0])
					used, _ := strconv.Atoi(parts[1])
					info.MemTotal = fmt.Sprintf("%d MB", total)
					info.MemUsed = fmt.Sprintf("%d MB", used)
					if total > 0 {
						info.MemUsage = float64(used) / float64(total) * 100
					}
				}
			case strings.HasPrefix(line, "DISK:"):
				parts := strings.Split(strings.TrimPrefix(line, "DISK:"), ":")
				if len(parts) >= 3 {
					info.DiskTotal = parts[0] + " MB"
					info.DiskUsed = parts[1] + " MB"
					pct := strings.TrimSuffix(strings.TrimSpace(parts[2]), "%")
					info.DiskUsage, _ = strconv.ParseFloat(pct, 64)
				}
			}
		}
	}

	return info, nil
}

func (c *SSHClient) Close() error {
	if c.client != nil {
		return c.client.Close()
	}
	return nil
}

// TestSSHConnection 测试SSH连接（使用TOFU模式）
func TestSSHConnection(host string, port int, user, password string) error {
	client, err := NewSSHClient(host, port, user, password)
	if err != nil {
		return err
	}
	defer client.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, _, err = client.Execute(ctx, "echo 'connected'")
	return err
}

// TestSSHConnectionWithAuth 使用完整认证配置测试SSH连接
func TestSSHConnectionWithAuth(auth *models.SSHAuthConfig) error {
	client, err := NewSSHClientWithAuth(auth)
	if err != nil {
		return err
	}
	defer client.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, _, err = client.Execute(ctx, "echo 'connected'")
	return err
}

// GetHostKeyFingerprint 获取已保存的主机密钥指纹
func GetHostKeyFingerprint(host string) string {
	hostKeyFile := hostKeyPath(host)
	storedKey, err := os.ReadFile(hostKeyFile)
	if err != nil {
		return ""
	}
	parsedKey, _, _, _, err := ssh.ParseAuthorizedKey(storedKey)
	if err != nil {
		return ""
	}
	return "SHA256:" + fingerprint(parsedKey)
}