this repo has no description
feat(toolbox): add ssh connector primitives
+172
-1
+4
-1
toolbox/go.mod
+4
-1
toolbox/go.mod
+168
toolbox/internal/cluster/cluster.go
+168
toolbox/internal/cluster/cluster.go
···
1
+
package cluster
2
+
3
+
import (
4
+
"context"
5
+
"encoding/json"
6
+
"fmt"
7
+
"os"
8
+
"strings"
9
+
"time"
10
+
11
+
"golang.org/x/crypto/ssh"
12
+
"golang.org/x/crypto/ssh/knownhosts"
13
+
)
14
+
15
+
const (
16
+
defaultSSHPort = 22
17
+
defaultSSHTimeout = 10 * time.Second
18
+
)
19
+
20
+
type SSHConfig struct {
21
+
Host string
22
+
User string
23
+
KeyPath string
24
+
KnownHostsPath string
25
+
Timeout time.Duration
26
+
}
27
+
28
+
type Connector struct {
29
+
sshClient *ssh.Client
30
+
}
31
+
32
+
func Connect(cfg SSHConfig) (*Connector, error) {
33
+
if cfg.Timeout == 0 {
34
+
cfg.Timeout = defaultSSHTimeout
35
+
}
36
+
37
+
sshClient, err := dialSSH(cfg)
38
+
if err != nil {
39
+
return nil, fmt.Errorf("ssh connect: %w", err)
40
+
}
41
+
42
+
return &Connector{sshClient: sshClient}, nil
43
+
}
44
+
45
+
func (c *Connector) RunCommand(cmd string) ([]byte, error) {
46
+
return c.RunCommandContext(context.Background(), cmd)
47
+
}
48
+
49
+
func (c *Connector) RunCommandContext(ctx context.Context, cmd string) ([]byte, error) {
50
+
session, err := c.sshClient.NewSession()
51
+
if err != nil {
52
+
return nil, fmt.Errorf("create session: %w", err)
53
+
}
54
+
defer session.Close()
55
+
56
+
type commandResult struct {
57
+
output []byte
58
+
err error
59
+
}
60
+
61
+
resultCh := make(chan commandResult, 1)
62
+
go func() {
63
+
output, err := session.CombinedOutput(cmd)
64
+
resultCh <- commandResult{output: output, err: err}
65
+
}()
66
+
67
+
select {
68
+
case <-ctx.Done():
69
+
_ = session.Signal(ssh.SIGTERM)
70
+
return nil, fmt.Errorf("run command canceled: %w", ctx.Err())
71
+
case result := <-resultCh:
72
+
if result.err != nil {
73
+
return nil, fmt.Errorf("run command: %w (output: %s)", result.err, result.output)
74
+
}
75
+
return result.output, nil
76
+
}
77
+
}
78
+
79
+
func (c *Connector) Close() error {
80
+
if c.sshClient != nil {
81
+
if err := c.sshClient.Close(); err != nil {
82
+
return fmt.Errorf("close ssh: %w", err)
83
+
}
84
+
}
85
+
return nil
86
+
}
87
+
88
+
type HostInfo struct {
89
+
IPv6Address string `json:"ipv6_address"`
90
+
}
91
+
92
+
func LoadHost(hostsFile, hostname string) (string, error) {
93
+
data, err := os.ReadFile(hostsFile)
94
+
if err != nil {
95
+
return "", fmt.Errorf("read file: %w", err)
96
+
}
97
+
98
+
var hosts map[string]HostInfo
99
+
if err := json.Unmarshal(data, &hosts); err != nil {
100
+
return "", fmt.Errorf("parse JSON: %w", err)
101
+
}
102
+
103
+
hostInfo, ok := hosts[hostname]
104
+
if !ok {
105
+
available := make([]string, 0, len(hosts))
106
+
for name := range hosts {
107
+
available = append(available, name)
108
+
}
109
+
return "", fmt.Errorf("host %q not found (available: %s)", hostname, strings.Join(available, ", "))
110
+
}
111
+
112
+
if hostInfo.IPv6Address == "" {
113
+
return "", fmt.Errorf("host %q has no ipv6_address", hostname)
114
+
}
115
+
116
+
return hostInfo.IPv6Address, nil
117
+
}
118
+
119
+
func dialSSH(cfg SSHConfig) (*ssh.Client, error) {
120
+
keyData, err := os.ReadFile(cfg.KeyPath)
121
+
if err != nil {
122
+
return nil, fmt.Errorf("read key %s: %w", cfg.KeyPath, err)
123
+
}
124
+
125
+
signer, err := ssh.ParsePrivateKey(keyData)
126
+
if err != nil {
127
+
return nil, fmt.Errorf("parse key: %w", err)
128
+
}
129
+
130
+
hostKeyCallback, err := buildHostKeyCallback(cfg.KnownHostsPath)
131
+
if err != nil {
132
+
return nil, fmt.Errorf("build host key callback: %w", err)
133
+
}
134
+
135
+
config := &ssh.ClientConfig{
136
+
User: cfg.User,
137
+
Auth: []ssh.AuthMethod{
138
+
ssh.PublicKeys(signer),
139
+
},
140
+
HostKeyCallback: hostKeyCallback,
141
+
Timeout: cfg.Timeout,
142
+
}
143
+
144
+
addr := fmt.Sprintf("%s:%d", cfg.Host, defaultSSHPort)
145
+
if strings.Contains(cfg.Host, ":") {
146
+
addr = fmt.Sprintf("[%s]:%d", cfg.Host, defaultSSHPort)
147
+
}
148
+
149
+
client, err := ssh.Dial("tcp", addr, config)
150
+
if err != nil {
151
+
return nil, fmt.Errorf("dial %s: %w", addr, err)
152
+
}
153
+
154
+
return client, nil
155
+
}
156
+
157
+
func buildHostKeyCallback(path string) (ssh.HostKeyCallback, error) {
158
+
if path == "" {
159
+
return nil, fmt.Errorf("known_hosts path is required")
160
+
}
161
+
162
+
callback, err := knownhosts.New(path)
163
+
if err != nil {
164
+
return nil, fmt.Errorf("load known_hosts file %s: %w", path, err)
165
+
}
166
+
167
+
return callback, nil
168
+
}