lan-manager/server/services/ssh.go

package services

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"lan-manager/server/config"
	"lan-manager/server/models"
	"strconv"
	"strings"
	"time"

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

// SSH command timeout for each command execution
const sshCommandTimeout = 8 * time.Second

func GetSSHInfo(ip string, port int, user, pass string) (*models.SSHInfoResult, error) {
	if port <= 0 {
		port = 22
	}
	cfg := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(pass),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         time.Duration(config.Load().SSHTimeout) * time.Second,
	}

	client, err := ssh.Dial("tcp", ip+":"+strconv.Itoa(port), cfg)
	if err != nil {
		return nil, err
	}
	defer client.Close()

	result := &models.SSHInfoResult{}

	// Hostname
	if out, err := runSSHCommand(client, "hostname"); err == nil {
		result.Hostname = strings.TrimSpace(out)
	}

	// OS Version: source /etc/os-release is the most compatible way across Linux distros
	if out, err := runSSHCommand(client, `. /etc/os-release 2>/dev/null && printf '%s\n' "$PRETTY_NAME"`); err == nil && out != "" {
		result.OsVersion = strings.TrimSpace(out)
	} else if out, err := runSSHCommand(client, "uname -sr"); err == nil {
		result.OsVersion = strings.TrimSpace(out)
	}

	// CPU: read /proc/stat twice with 1s sleep in Go to avoid complex shell quoting
	if cpuUsage, cores, err := getCPUInfo(client); err == nil {
		result.CPU = cpuUsage + "%"
		if cores != "" {
			result.CPU += " (" + cores + " cores)"
		}
		result.RawCPUInfo = result.CPU
	}

	// Memory: use /proc/meminfo instead of free for better compatibility
	if out, err := runSSHCommand(client, "awk '/MemTotal/{t=$2} /MemAvailable/{a=$2} END{u=t-a; if(t>0) printf \"%.1f/%.1f (%.0f%%)\", u/1024/1024, t/1024/1024, u*100/t}' /proc/meminfo"); err == nil && out != "" {
		result.Memory = strings.TrimSpace(out)
		result.RawMemoryInfo = result.Memory
	}

	// Disk: keep real block devices only, exclude all virtual filesystems
	if out, err := runSSHCommand(client, "df -hP | awk 'NR>1 && ($1 ~ /^\\/dev/ || $1 ~ /^(LABEL=|UUID=)/) {printf \"%s %s/%s (%s)  \", $6,$3,$2,$5}'"); err == nil && out != "" {
		result.Disk = strings.TrimSpace(out)
		result.RawDiskInfo = result.Disk
	}

	// Uptime
	if out, err := runSSHCommand(client, "uptime -p 2>/dev/null || uptime | awk -F',' '{print $1}'"); err == nil {
		result.Uptime = strings.TrimSpace(out)
	}

	// Listen ports: try ss first, fallback to netstat, fallback to /proc/net/tcp
	ports, _ := getListenPorts(client)
	result.ListenPorts = ports

	return result, nil
}

func runSSHCommand(client *ssh.Client, cmd string) (string, error) {
	session, err := client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()

	var b bytes.Buffer
	session.Stdout = &b

	// Use context to enforce command timeout and prevent hanging
	ctx, cancel := context.WithTimeout(context.Background(), sshCommandTimeout)
	defer cancel()

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

	select {
	case err := <-done:
		return b.String(), err
	case <-ctx.Done():
		// Force close session to unblock Run()
		session.Signal(ssh.SIGKILL)
		session.Close()
		return "", fmt.Errorf("ssh command timed out after %v: %s", sshCommandTimeout, cmd)
	}
}

func getCPUInfo(client *ssh.Client) (usage string, cores string, err error) {
	readStat := func() (idle, total float64, e error) {
		out, e := runSSHCommand(client, "awk '/^cpu /{print $2,$3,$4,$5}' /proc/stat")
		if e != nil {
			return 0, 0, e
		}
		parts := strings.Fields(strings.TrimSpace(out))
		if len(parts) < 4 {
			return 0, 0, fmt.Errorf("unexpected /proc/stat format")
		}
		vals := make([]float64, 4)
		for i := 0; i < 4; i++ {
			v, e := strconv.ParseFloat(parts[i], 64)
			if e != nil {
				return 0, 0, e
			}
			vals[i] = v
		}
		user, nice, system, idleTick := vals[0], vals[1], vals[2], vals[3]
		idle = idleTick
		total = user + nice + system + idleTick
		return idle, total, nil
	}

	idle1, total1, err := readStat()
	if err != nil {
		return "", "", err
	}
	time.Sleep(1 * time.Second)
	idle2, total2, err := readStat()
	if err != nil {
		return "", "", err
	}

	dTotal := total2 - total1
	dIdle := idle2 - idle1
	if dTotal > 0 {
		pct := 100.0 * (dTotal - dIdle) / dTotal
		if pct < 0 {
			pct = 0
		}
		if pct > 100 {
			pct = 100
		}
		usage = fmt.Sprintf("%.1f", pct)
	} else {
		usage = "0.0"
	}

	if cout, e := runSSHCommand(client, "nproc"); e == nil {
		cores = strings.TrimSpace(cout)
	}
	return usage, cores, nil
}

func getListenPorts(client *ssh.Client) ([]int, error) {
	var out string
	var err error

	// try ss without -p (no need for process info, more compatible)
	out, err = runSSHCommand(client, "ss -tln 2>/dev/null | awk 'NR>1 {print $4}' | sed 's/.*://' | sort -n | uniq")
	if err != nil || strings.TrimSpace(out) == "" {
		// fallback to netstat
		out, err = runSSHCommand(client, "netstat -tln 2>/dev/null | awk 'NR>2 {print $4}' | sed 's/.*://' | sort -n | uniq")
	}
	if err != nil || strings.TrimSpace(out) == "" {
		// fallback to /proc/net/tcp
		out, err = runSSHCommand(client, "awk 'NR>1 {printf \"%04X\\n\", strtonum(\"0x\"$2)}' /proc/net/tcp")
	}
	if err != nil {
		return nil, err
	}

	ports := []int{}
	scanner := bufio.NewScanner(strings.NewReader(out))
	for scanner.Scan() {
		p := strings.TrimSpace(scanner.Text())
		if p == "" {
			continue
		}
		// handle both decimal ports and hex ports from /proc/net/tcp
		var pi int
		var e error
		if len(p) == 4 && isHex(p) {
			var h uint64
			h, e = strconv.ParseUint(p, 16, 32)
			pi = int(h)
		} else {
			pi, e = strconv.Atoi(p)
		}
		if e == nil && pi > 0 {
			ports = append(ports, pi)
		}
	}
	return ports, nil
}

func isHex(s string) bool {
	for _, c := range s {
		if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F')) {
			return false
		}
	}
	return true
}