147 lines
3.7 KiB
Go
147 lines
3.7 KiB
Go
package sshclient
|
||
|
||
import (
|
||
"fmt"
|
||
"net"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
|
||
"golang.org/x/crypto/ssh"
|
||
)
|
||
|
||
// Config SSH 连接配置（不含连接状态）
|
||
type Config struct {
|
||
User string
|
||
Password string
|
||
KeyFile string
|
||
Port int
|
||
Timeout time.Duration
|
||
SudoPwd string // 非 root 用户执行特权命令(sudo)时的密码
|
||
}
|
||
|
||
// NewConfig 创建配置
|
||
func NewConfig(user, password, keyFile string, port int, timeout time.Duration) *Config {
|
||
return &Config{
|
||
User: user,
|
||
Password: password,
|
||
KeyFile: keyFile,
|
||
Port: port,
|
||
Timeout: timeout,
|
||
}
|
||
}
|
||
|
||
// Connect 建立 SSH 连接，由调用方负责 Close
|
||
func Connect(ip string, cfg *Config) (*ssh.Client, error) {
|
||
sshCfg, err := cfg.buildSSHConfig()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
addr := net.JoinHostPort(ip, fmt.Sprintf("%d", cfg.Port))
|
||
client, err := ssh.Dial("tcp", addr, sshCfg)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("SSH 连接 %s 失败: %w", addr, err)
|
||
}
|
||
return client, nil
|
||
}
|
||
|
||
// RunCommand 在已有 SSH 连接上执行命令，返回输出（实现 stress.Executor 接口）
|
||
func RunCommand(client *ssh.Client, command string) (string, error) {
|
||
session, err := client.NewSession()
|
||
if err != nil {
|
||
return "", fmt.Errorf("创建会话失败: %w", err)
|
||
}
|
||
defer session.Close()
|
||
|
||
out, err := session.CombinedOutput(command)
|
||
if err != nil {
|
||
return cleanOutput(out), fmt.Errorf("命令执行失败: %w", err)
|
||
}
|
||
return cleanOutput(out), nil
|
||
}
|
||
|
||
// isNoiseLine 判断某行是否为登录横幅/欢迎语等噪声（fnOS 的 admin 用户会把
|
||
// "Could not chdir to home directory" 之类写入 stdout，污染命令解析）。
|
||
func isNoiseLine(line string) bool {
|
||
s := strings.TrimSpace(line)
|
||
if s == "" {
|
||
return false
|
||
}
|
||
prefixes := []string{
|
||
"Could not chdir to",
|
||
"Last login:",
|
||
"Welcome to",
|
||
"/etc/motd",
|
||
"Permission denied",
|
||
"bash: line 1:",
|
||
}
|
||
for _, p := range prefixes {
|
||
if strings.HasPrefix(s, p) {
|
||
return true
|
||
}
|
||
}
|
||
// 某些横幅是 "Could: command not found" 这类由前面噪声行产生，
|
||
// 形如 "xxx: command not found" 且前半段是噪声关键词的也过滤
|
||
if strings.Contains(s, ": command not found") {
|
||
return true
|
||
}
|
||
return false
|
||
}
|
||
|
||
// cleanOutput 过滤掉登录横幅等噪声行，仅保留真实命令输出
|
||
func cleanOutput(raw []byte) string {
|
||
lines := strings.Split(string(raw), "\n")
|
||
var kept []string
|
||
for _, ln := range lines {
|
||
if isNoiseLine(ln) {
|
||
continue
|
||
}
|
||
kept = append(kept, ln)
|
||
}
|
||
return strings.Join(kept, "\n")
|
||
}
|
||
|
||
// EnsureHome 确保远端登录用户的 HOME 目录存在，消除 fnOS "Could not chdir to
|
||
// home directory" 登录横幅（该横幅会写入 stdout 污染命令输出）。
|
||
func EnsureHome(client *ssh.Client) {
|
||
out, err := RunCommand(client, "echo $HOME")
|
||
home := strings.TrimSpace(out)
|
||
if err != nil || home == "" || home == "/" {
|
||
return
|
||
}
|
||
RunCommand(client, fmt.Sprintf("mkdir -p %s", home))
|
||
}
|
||
|
||
// buildSSHConfig 构建认证配置
|
||
func (c *Config) buildSSHConfig() (*ssh.ClientConfig, error) {
|
||
var authMethods []ssh.AuthMethod
|
||
|
||
if c.Password != "" {
|
||
authMethods = append(authMethods, ssh.Password(c.Password))
|
||
}
|
||
|
||
if c.KeyFile != "" {
|
||
key, err := os.ReadFile(c.KeyFile)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("读取密钥文件失败: %w", err)
|
||
}
|
||
signer, err := ssh.ParsePrivateKey(key)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("解析私钥失败: %w", err)
|
||
}
|
||
authMethods = append(authMethods, ssh.PublicKeys(signer))
|
||
}
|
||
|
||
if len(authMethods) == 0 {
|
||
return nil, fmt.Errorf("未提供认证方式（密码或密钥）")
|
||
}
|
||
|
||
return &ssh.ClientConfig{
|
||
User: c.User,
|
||
Auth: authMethods,
|
||
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
|
||
Timeout: c.Timeout,
|
||
}, nil
|
||
}
|