-
Notifications
You must be signed in to change notification settings - Fork 3
/
client.go
155 lines (129 loc) · 3.25 KB
/
client.go
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
package ssh
// This package manages SSH connections to remote hosts and implements cmd.Executor.
// Based on a fork of github.com/melbahja/goph with context cancellation support (github.com/antonsergeyev/goph).
import (
"bytes"
"context"
"fmt"
"github.com/melbahja/goph"
"github.com/pkg/errors"
"go.uber.org/zap"
"golang.org/x/crypto/ssh"
"os/exec"
"sync"
"time"
)
type Options struct {
Port uint
Username string
Password string
KeyFile string `yaml:"keyFile"`
KeyPassword string `yaml:"keyPassword"`
Debug bool
}
type Client struct {
options Options
executors sync.Map
logger *zap.SugaredLogger
auth goph.Auth
}
// HostExecutor implements a shell command executor on a remote machine over SSH
type HostExecutor struct {
host string
debug bool
sshConn *goph.Client
logger *zap.SugaredLogger
}
func NewClient(opts Options, logger *zap.SugaredLogger) (client *Client, err error) {
if opts.Port == 0 {
opts.Port = 22
}
client = &Client{
options: opts,
logger: logger,
executors: sync.Map{},
}
if len(opts.KeyFile) > 0 {
client.auth, err = goph.Key(opts.KeyFile, opts.KeyPassword)
if err != nil {
return nil, errors.Wrapf(
err,
"could not read SSH key from %s",
opts.KeyFile,
)
}
} else {
client.auth = goph.Password(opts.Password)
}
return client, nil
}
// GetByHost returns a shell command executor on a remote host.
// Attempts to only create SSH connections once, keeping them in a `sync.Map`.
func (c *Client) GetByHost(host string) (*HostExecutor, error) {
_, ok := c.executors.Load(host)
if !ok {
logger := c.logger.With("host", host)
logger.Debugw(
"creating SSH connection",
"host",
host,
)
sshConn, err := goph.NewConn(&goph.Config{
Auth: c.auth,
User: c.options.Username,
Addr: host,
Port: c.options.Port,
Timeout: time.Second * 2,
Callback: ssh.InsecureIgnoreHostKey(),
})
if err != nil {
return nil, errors.Wrapf(
err,
"could not create SSH connection to %s as %s",
host,
c.options.Username,
)
}
c.executors.Store(host, &HostExecutor{
host: host,
debug: c.options.Debug,
sshConn: sshConn,
logger: logger,
})
}
result, _ := c.executors.Load(host)
return result.(*HostExecutor), nil
}
func (h *HostExecutor) Execute(ctx context.Context, cmd *exec.Cmd) ([]byte, error) {
timeStarted := time.Now()
if h.debug {
fmt.Printf("\n---[SSH] executing command at %s:---\n%s\n", h.host, cmd.String())
}
sshCmd, err := h.sshConn.CommandContext(ctx, cmd.String())
if err != nil {
return nil, err
}
output, err := sshCmd.CombinedOutput()
if h.debug {
fmt.Printf(
"\n---[SSH] command at %s done in %s, output:---\n%s\n",
h.host,
time.Now().Sub(timeStarted).String(),
string(output),
)
}
return output, err
}
func (h *HostExecutor) Run(ctx context.Context, cmd *exec.Cmd) error {
_, err := h.Execute(ctx, cmd)
return err
}
func (h *HostExecutor) ReadFile(ctx context.Context, path string) ([]byte, error) {
var data bytes.Buffer
err := h.sshConn.ReadFile(path, &data)
return data.Bytes(), err
}
func (h *HostExecutor) WriteFile(ctx context.Context, path string, data []byte) error {
source := bytes.NewReader(data)
return h.sshConn.WriteFile(path, source)
}