/
pulley.go
136 lines (113 loc) · 3.17 KB
/
pulley.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
// Package pulley is a wrapper around the golang.org/x/crypto/ssh package
// providing a suckless experience.
package pulley
import (
"fmt"
"io/ioutil"
"os"
"path/filepath"
"golang.org/x/crypto/ssh"
)
// Client is the client for the ssh connection and the primary way to interact
// with pulley.
type Client struct {
HostName string
Port string
User string
config *ssh.ClientConfig
connection *ssh.Client
}
// New creates a default Client. Modify this client via it's public members,
// HostName, Port, and User. These default to localhost, 22, and the current
// user running the process.
func New(user string) *Client {
return &Client{
HostName: "localhost",
Port: "22",
User: user,
config: &ssh.ClientConfig{
User: user,
},
}
}
// Session will return a new session for the current connection or an error if
// there was one. You only need to use this if you're doing something advanced.
func (s *Client) Session() (*ssh.Session, error) {
return s.connection.NewSession()
}
// Connect will connect the client to it's hostname and port, if LoadKey or
// LoadDefaultKey have not been called yet this will call LoadDefaultKey
func (s *Client) Connect() error {
var err error
if s.config.Auth == nil {
kerr := s.LoadDefaultKey()
if kerr != nil {
return kerr
}
}
s.connection, err = ssh.Dial("tcp",
fmt.Sprintf("%s:%s", s.HostName, s.Port),
s.config)
return err
}
// LoadDefaultKey will load the ssh key at $HOME/.ssh/id_rsa
func (s *Client) LoadDefaultKey() error {
keyFile := filepath.Join(os.Getenv("HOME"), ".ssh", "id_rsa")
key, err := ioutil.ReadFile(keyFile)
if err != nil {
return err
}
return s.LoadKey(key)
}
// LoadKey will load the given key
func (s *Client) LoadKey(key []byte) error {
parsedKey, err := ssh.ParsePrivateKey(key)
if err != nil {
return err
}
s.config.Auth = append(s.config.Auth, ssh.PublicKeys(parsedKey))
return nil
}
// Exec runs the command on the server that's connected to by this client, if
// It will handle sessions automatically and return a pulley.Result
func (s *Client) Exec(cmd string) Result {
sess, serr := s.Session()
if serr != nil {
return Result{err: serr}
}
defer sess.Close()
var r Result
r.Output, r.err = sess.Output(cmd)
return r
}
// ExecErr is the same as exec however the result's output will have both
// stdout and stderr.
func (s *Client) ExecErr(cmd string) Result {
sess, serr := s.Session()
if serr != nil {
return Result{err: serr}
}
defer sess.Close()
var r Result
r.Output, r.err = sess.CombinedOutput(cmd)
return r
}
// ExecAsync is the same as exec however will execute in a go routine and takes
// a channel which it will send the result over.
func (s *Client) ExecAsync(cmd string, rc chan Result) {
go func() {
rc <- s.Exec(cmd)
}()
}
// ExecAsyncErr is the same as ExecAsync however the result's output will have
// both stderr and stdout.
func (s *Client) ExecAsyncErr(cmd string, rc chan Result) {
go func() {
rc <- s.ExecErr(cmd)
}()
}
// Ugly will return the underlying ssh.Client and ssh.ClientConfig in case you
// need those structs directly.
func (s *Client) Ugly() (*ssh.Client, *ssh.ClientConfig) {
return s.connection, s.config
}