forked from coreos/fleet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ssh.go
152 lines (124 loc) · 3.19 KB
/
ssh.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
package ssh
import (
"bufio"
"errors"
"net"
"os"
"time"
gossh "github.com/coreos/fleet/third_party/code.google.com/p/go.crypto/ssh"
"github.com/coreos/fleet/third_party/code.google.com/p/go.crypto/ssh/terminal"
)
func Execute(client *gossh.ClientConn, cmd string) (*bufio.Reader, error) {
session, err := client.NewSession()
if err != nil {
return nil, err
}
stdout, _ := session.StdoutPipe()
bstdout := bufio.NewReader(stdout)
session.Start(cmd)
go session.Wait()
return bstdout, nil
}
func Shell(client *gossh.ClientConn) error {
session, err := client.NewSession()
if err != nil {
return err
}
defer session.Close()
modes := gossh.TerminalModes{
gossh.ECHO: 1, // enable echoing
gossh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
gossh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
}
fd := int(os.Stdin.Fd())
oldState, err := terminal.MakeRaw(fd)
defer terminal.Restore(fd, oldState)
termWidth, termHeight, err := terminal.GetSize(fd)
if err != nil {
return err
}
session.Stdout = os.Stdout
session.Stderr = os.Stderr
session.Stdin = os.Stdin
if err := session.RequestPty("xterm-256color", termHeight, termWidth, modes); err != nil {
return err
}
if err = session.Shell(); err != nil {
return err
}
session.Wait()
return nil
}
func sshClientConfig(user string) (*gossh.ClientConfig, error) {
sock := os.Getenv("SSH_AUTH_SOCK")
if sock == "" {
return nil, errors.New("SSH_AUTH_SOCK environment variable is not set. Verify ssh-agent is running. See https://github.com/coreos/fleet/blob/master/Documentation/remote-access.md for help.")
}
agent, err := net.Dial("unix", sock)
if err != nil {
return nil, err
}
cfg := gossh.ClientConfig{
User: user,
Auth: []gossh.ClientAuth{
gossh.ClientAuthAgent(gossh.NewAgentClient(agent)),
},
}
return &cfg, nil
}
func NewSSHClient(user, addr string) (*gossh.ClientConn, error) {
clientConfig, err := sshClientConfig(user)
if err != nil {
return nil, err
}
var client *gossh.ClientConn
dialFunc := func(echan chan error) {
var err error
client, err = gossh.Dial("tcp", addr, clientConfig)
echan <- err
}
err = timeoutSSHDial(dialFunc)
return client, err
}
func NewTunnelledSSHClient(user, tunaddr, tgtaddr string) (*gossh.ClientConn, error) {
clientConfig, err := sshClientConfig(user)
if err != nil {
return nil, err
}
var tunnelClient *gossh.ClientConn
dialFunc := func(echan chan error) {
var err error
tunnelClient, err = gossh.Dial("tcp", tunaddr, clientConfig)
echan <- err
}
err = timeoutSSHDial(dialFunc)
if err != nil {
return nil, err
}
var targetConn net.Conn
dialFunc = func(echan chan error) {
var err error
targetConn, err = tunnelClient.Dial("tcp", tgtaddr)
echan <- err
}
err = timeoutSSHDial(dialFunc)
if err != nil {
return nil, err
}
targetClient, err := gossh.Client(targetConn, clientConfig)
if err != nil {
return nil, err
}
return targetClient, nil
}
func timeoutSSHDial(dial func(chan error)) error {
var err error
echan := make(chan error)
go dial(echan)
select {
case <-time.After(time.Duration(time.Second * 10)):
return errors.New("Timed out while initiating SSH connection")
case err = <-echan:
return err
}
}