2 Star 6 Fork 1

卧雪Sirk / gosshtool

Create your Gitee Account
Explore and code with more than 6 million developers,Free private repositories !:)
Sign up
Clone or download
client.go 4.51 KB
Copy Edit Web IDE Raw Blame History
scottkiss authored 2016-11-03 17:59 . support file upload
package gosshtool
import (
"bytes"
"golang.org/x/crypto/ssh"
"log"
"net"
"strings"
"sync"
"time"
)
const (
DEFAULT_CHUNK_SIZE = 65536
MIN_CHUNKS = 10
THROUGHPUT_SLEEP_INTERVAL = 100
MIN_THROUGHPUT = DEFAULT_CHUNK_SIZE * MIN_CHUNKS * (1000 / THROUGHPUT_SLEEP_INTERVAL)
)
var (
maxThroughputChan = make(chan bool, MIN_CHUNKS)
maxThroughput uint64
maxThroughputMutex sync.Mutex
)
type SSHClient struct {
SSHClientConfig
remoteConn *ssh.Client
isConnected bool
}
func (c *SSHClient) maxThroughputControl() {
for {
if c.MaxDataThroughput > 0 && c.MaxDataThroughput < MIN_THROUGHPUT {
log.Panicf("Minimal throughput is %d Bps", MIN_THROUGHPUT)
}
maxThroughputMutex.Lock()
throughput := c.MaxDataThroughput
maxThroughputMutex.Unlock()
chunks := throughput / DEFAULT_CHUNK_SIZE * THROUGHPUT_SLEEP_INTERVAL / 1000
if chunks < MIN_CHUNKS {
chunks = MIN_CHUNKS
}
for i := uint64(0); i < chunks; i++ {
maxThroughputChan <- true
}
if throughput > 0 {
time.Sleep(THROUGHPUT_SLEEP_INTERVAL * time.Millisecond)
}
}
}
func (c *SSHClient) Connect() (conn *ssh.Client, err error) {
if c.remoteConn != nil {
return
}
port := "22"
host := c.Host
hstr := strings.SplitN(host, ":", 2)
if len(hstr) == 2 {
host = hstr[0]
port = hstr[1]
}
config := makeConfig(c.User, c.Password, c.Privatekey)
if c.DialTimeoutSecond > 0 {
connNet, err := net.DialTimeout("tcp", host+":"+port, time.Duration(c.DialTimeoutSecond)*time.Second)
if err != nil {
return nil, err
}
sc, chans, reqs, err := ssh.NewClientConn(connNet, host+":"+port, config)
if err != nil {
return nil, err
}
conn = ssh.NewClient(sc, chans, reqs)
} else {
conn, err = ssh.Dial("tcp", host+":"+port, config)
if err != nil {
return
}
}
log.Println("dial ssh success")
c.remoteConn = conn
return
}
func (c *SSHClient) TransferData(target string, data []byte) (stdout, stderr string, err error) {
go c.maxThroughputControl()
if c.isConnected == false {
_, err = c.Connect()
if err != nil {
return
}
}
currentSession, err := NewSession(c.remoteConn, nil, 0)
if err != nil {
return
}
defer currentSession.Close()
cmd := "cat >'" + strings.Replace(target, "'", "'\\''", -1) + "'"
stdinPipe, err := currentSession.StdinPipe()
if err != nil {
return
}
var stdoutBuf bytes.Buffer
var stderrBuf bytes.Buffer
currentSession.Stdout = &stdoutBuf
currentSession.Stderr = &stderrBuf
err = currentSession.session.Start(cmd)
if err != nil {
return
}
for start, max := 0, len(data); start < max; start += DEFAULT_CHUNK_SIZE {
<-maxThroughputChan
end := start + DEFAULT_CHUNK_SIZE
if end > max {
end = max
}
_, err = stdinPipe.Write(data[start:end])
if err != nil {
return
}
}
err = stdinPipe.Close()
if err != nil {
return
}
err = currentSession.Wait()
stdout = stdoutBuf.String()
stderr = stderrBuf.String()
return
}
func (c *SSHClient) Cmd(cmd string, sn *SshSession, deadline *time.Time, idleTimeout int) (output, errput string, currentSession *SshSession, err error) {
if c.isConnected == false {
_, err = c.Connect()
if err != nil {
return
}
}
if sn == nil {
currentSession, err = NewSession(c.remoteConn, deadline, idleTimeout)
} else {
currentSession = sn
currentSession.SetDeadline(deadline)
}
if err != nil {
return
}
var stdoutBuf bytes.Buffer
var stderrBuf bytes.Buffer
currentSession.Stdout = &stdoutBuf
currentSession.Stderr = &stderrBuf
err = currentSession.Run(cmd)
defer currentSession.Close()
output = stdoutBuf.String()
errput = stderrBuf.String()
return
}
func (c *SSHClient) Pipe(rw ReadWriteCloser, pty *PtyInfo, deadline *time.Time, idleTimeout int) (currentSession *SshSession, err error) {
if c.isConnected == false {
_, err := c.Connect()
if err != nil {
return nil, err
}
}
currentSession, err = NewSession(c.remoteConn, deadline, idleTimeout)
if err != nil {
return
}
if err = currentSession.RequestPty(pty.Term, pty.H, pty.W, pty.Modes); err != nil {
return
}
wc, err := currentSession.StdinPipe()
if err != nil {
return
}
go CopyIOAndUpdateSessionDeadline(wc, rw, currentSession)
r, err := currentSession.StdoutPipe()
if err != nil {
return
}
go CopyIOAndUpdateSessionDeadline(rw, r, currentSession)
er, err := currentSession.StderrPipe()
if err != nil {
return
}
go CopyIOAndUpdateSessionDeadline(rw, er, currentSession)
err = currentSession.Shell()
if err != nil {
return
}
err = currentSession.Wait()
if err != nil {
return
}
defer currentSession.Close()
return
}

Comment ( 0 )

Sign in for post a comment

Go
1
https://git.oschina.net/scottkiss/gosshtool.git
git@git.oschina.net:scottkiss/gosshtool.git
scottkiss
gosshtool
gosshtool
master

Search