LokiStack/backend/services/ssh.go

package services

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"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"`
}

func NewSSHClient(host string, port int, user, password string) (*SSHClient, error) {
	if port == 0 {
		port = 22
	}
	config := &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), config)
	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
}

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
}