github.com/billybanfield/evergreen@v0.0.0-20170525200750-eeee692790f7/remote/ssh.go (about)

     1  package remote
     2  
     3  import (
     4  	"io"
     5  	"time"
     6  
     7  	"github.com/mongodb/grip"
     8  	"github.com/pkg/errors"
     9  	"golang.org/x/crypto/ssh"
    10  )
    11  
    12  var (
    13  	ErrCmdTimedOut = errors.New("ssh command timed out")
    14  )
    15  
    16  // SSHCommand abstracts a single command to be run via ssh, on a remote machine.
    17  type SSHCommand struct {
    18  	// the command to be run on the remote machine
    19  	Command string
    20  
    21  	// the remote host to connect to
    22  	Host string
    23  
    24  	// the user to connect with
    25  	User string
    26  
    27  	// the location of the private key file (PEM-encoded) to be used
    28  	Keyfile string
    29  
    30  	// stdin for the remote command
    31  	Stdin io.Reader
    32  
    33  	// the threshold at which the command is considered to time out, and will be killed
    34  	Timeout time.Duration
    35  }
    36  
    37  // Run the command via ssh. Returns the combined stdout and stderr, as well as any
    38  // error that occurs.
    39  func (cmd *SSHCommand) Run() ([]byte, error) {
    40  
    41  	// configure appropriately
    42  	clientConfig, err := createClientConfig(cmd.User, cmd.Keyfile)
    43  	if err != nil {
    44  		return nil, errors.Wrap(err, "error configuring ssh")
    45  	}
    46  
    47  	// open a connection to the ssh server
    48  	conn, err := ssh.Dial("tcp", cmd.Host, clientConfig)
    49  	if err != nil {
    50  		return nil, errors.Wrapf(err, "error connecting to ssh server at `%v`", cmd.Host)
    51  	}
    52  
    53  	// initiate a session for running an ssh command
    54  	session, err := conn.NewSession()
    55  	if err != nil {
    56  		return nil, errors.Wrapf(err, "error creating an ssh session to `%v`", cmd.Host)
    57  	}
    58  	defer session.Close()
    59  
    60  	// set stdin appropriately
    61  	session.Stdin = cmd.Stdin
    62  
    63  	// terminal modes for the pty we'll be using
    64  	modes := ssh.TerminalModes{
    65  		ssh.ECHO:          0,     // disable echoing
    66  		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
    67  		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
    68  	}
    69  
    70  	// request a pseudo terminal
    71  	if err := session.RequestPty("xterm", 80, 40, modes); err != nil {
    72  		return nil, errors.Errorf("error requesting pty: %v", err)
    73  	}
    74  
    75  	// kick the ssh command off
    76  	errChan := make(chan error)
    77  	output := []byte{}
    78  	go func() {
    79  		output, err = session.CombinedOutput(cmd.Command)
    80  		errChan <- errors.WithStack(err)
    81  	}()
    82  
    83  	// wait for the command to finish, or time out
    84  	select {
    85  	case err := <-errChan:
    86  		return output, errors.WithStack(err)
    87  	case <-time.After(cmd.Timeout):
    88  		// command timed out; kill the remote process
    89  		grip.CatchError(errors.WithStack(session.Signal(ssh.SIGKILL)))
    90  		return nil, ErrCmdTimedOut
    91  	}
    92  
    93  }