GO语言开发实践-ssh批量远程执行命令

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

package gossh

import (
    "fmt"
    "io/ioutil"
    "time"

    "github.com/pkg/errors"
    "golang.org/x/crypto/ssh"
)

type Args struct {
    Config  *ssh.ClientConfig
    Host    string
    Command string
    Port    int
    Timeout int
}

// 初始化配置
func NewConfig(keyFile, user string, timeout time.Duration) (config *ssh.ClientConfig, err error) {
    // 读取 rsa 私钥
    key, err := ioutil.ReadFile(keyFile)
    if err != nil {
        err = fmt.Errorf("unable to read private key: %v", err)
        return
    }

    // 解析私钥是否正确
    signer, err := ssh.ParsePrivateKey(key)
    if err != nil {
        err = fmt.Errorf("unable to parse private key: %v", err)
        return
    }

    // 初始化 ssh 配置
    config = &ssh.ClientConfig{
        User: user,
        Auth: []ssh.AuthMethod{
            ssh.PublicKeys(signer),
        },
        HostKeyCallback: ssh.InsecureIgnoreHostKey(),
        Timeout:         timeout,
    }
    return
}

// 根据配置执行命令
func Run(config *ssh.ClientConfig, host, command string, port int) ([]byte, error) {
    // 生成客户端
    client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", host, port), config)
    if err != nil {
        msg := fmt.Sprintf("unable to connect: %s error %v", host, err)
        return nil, errors.Wrap(err, msg)
    }
    defer client.Close()

    // 创建会话窗口
    session, err := client.NewSession()
    if err != nil {
        msg := fmt.Sprintf("ssh new session error %v", err)
        return nil, errors.Wrap(err, msg)
    }
    defer session.Close()

    // 标准输出和标准错误输出一起输出
    outPut, err := session.CombinedOutput(command)
    if err != nil {
        msg := fmt.Sprintf("run command %s on host %s error %v", command, host, err)
        return nil, errors.Wrap(err, msg)
    }
    return outPut, nil
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

package gossh

import (
    "log"
    "reflect"
    "testing"
    "time"

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


func TestRun(t *testing.T) {
    type args struct {
        config  *ssh.ClientConfig
        host    string
        command string
        port    int
    }
    config, err := NewConfig("/root/.ssh/id_rsa", "devops", time.Duration(100))
    if err != nil {
        log.Println("获取key失败")
        return
    }
    // 创建 ssh 主机列表
    tests := []struct {
        name       string
        args       args
        wantOutPut []byte
        wantErr    bool
    }{
        {
            name: "test1",
            args: args{
                config:  config,
                host:    "127.0.0.1",
                command: "echo 'hello'",
                port:    22,
            },
            wantOutPut: []byte("hello"),
            wantErr: false,
        },
    }
    // 从主机列表循环拿出主机执行命令
    for _, tt := range tests {
        t.Run(tt.name, func(t *testing.T) {
            gotOutPut, err := Run(tt.args.config, tt.args.host, tt.args.command, tt.args.port)
            if (err != nil) != tt.wantErr {
                t.Errorf("Run() error = %v, wantErr %v", err, tt.wantErr)
                return
            }
            // 如果标准输出的值不是期望值，则输出错误日志
            if !reflect.DeepEqual(gotOutPut, tt.wantOutPut) {
                t.Errorf("Run() = %s, want %s", gotOutPut, tt.wantOutPut)
            }
            // 输出正确期望值的日志
            t.Logf("Run() = %s", gotOutPut)
        })
    }
}

1
2
3
4
5
6
7

=== RUN   TestRun
=== RUN   TestRun/test1
    gossh_test.go:57: Run() = hello
--- PASS: TestRun (0.35s)
    --- PASS: TestRun/test1 (0.35s)
PASS
ok      pkg/gossh     0.357s