Golang ssh.Session类代码示例

func pipe(ch ssh.Channel, client *ssh.Client, session *ssh.Session, command string) (int, error) {
	targetStderr, err := session.StderrPipe()
	if err != nil {
		return -1, errors.New("fail to pipe stderr: " + err.Error())
	}
	targetStdout, err := session.StdoutPipe()
	if err != nil {
		return -1, errors.New("fail to pipe stdout: " + err.Error())
	}
	targetStdin, err := session.StdinPipe()
	if err != nil {
		return -1, errors.New("fail to pipe stdin: " + err.Error())
	}

	go io.Copy(targetStdin, ch)
	go io.Copy(ch.Stderr(), targetStderr)
	go io.Copy(ch, targetStdout)

	err = session.Start(command)
	if err != nil {
		ch.Write([]byte("Error when starting '" + command + "': " + err.Error()))
		ch.Close()
	}

	err = session.Wait()
	if err != nil {
		if err, ok := err.(*ssh.ExitError); ok {
			return err.ExitStatus(), nil
		} else {
			return -1, errors.New("failed to wait ssh command: " + err.Error())
		}
	}

	return 0, nil
}

func (n *SSHNode) getClientAndSession() (*ssh.Client, *ssh.Session, error) {
	var client *ssh.Client
	var s *ssh.Session
	var err error

	// Retry few times if ssh connection fails
	for i := 0; i < MaxSSHRetries; i++ {
		client, err = n.dial()
		if err != nil {
			time.Sleep(SSHRetryDelay)
			continue
		}

		s, err = client.NewSession()
		if err != nil {
			client.Close()
			time.Sleep(SSHRetryDelay)
			continue
		}
		modes := ssh.TerminalModes{
			ssh.ECHO:          0,
			ssh.TTY_OP_ISPEED: 14400,
			ssh.TTY_OP_OSPEED: 14400,
		}
		// Request pseudo terminal
		if err := s.RequestPty("xterm", 40, 80, modes); err != nil {
			return nil, nil, fmt.Errorf("failed to get pseudo-terminal: %v", err)
		}

		return client, s, nil
	}

	return nil, nil, err
}

/*
	Expected to be invoked as a goroutine which runs in parallel to sending the ssh command to the
	far side. This function reads from the input buffer reader br and writes to the target stripping
	blank and comment lines as it goes.
*/
func send_script(sess *ssh.Session, argv0 string, env_file string, br *bufio.Reader) {

	target, err := sess.StdinPipe() // we create the pipe here so that we can close here
	if err != nil {
		fmt.Fprintf(os.Stderr, "unable to create stdin for session: %s\n", err)
		return
	}
	defer target.Close()

	if argv0 != "" {
		target.Write([]byte("ARGV0=\"" + argv0 + "\"\n")) // $0 isn't valid using this, so simulate $0 with argv0
	}

	if env_file != "" { // must push out the environment first
		env_file, err = find_file(env_file) // find it in the path if not a qualified name
		if err == nil {
			ef, err := os.Open(env_file)

			if err != nil {
				fmt.Fprintf(os.Stderr, "ssh_broker: could not open environment file: %s: %s\n", env_file, err)
			} else {
				ebr := bufio.NewReader(ef) // get a buffered reader for the file
				send_file(ebr, target)
				ef.Close()
			}
		} else {
			fmt.Fprintf(os.Stderr, "ssh_broker: could not find  environment file: %s: %s\n", env_file, err)
		}
	}

	send_file(br, target)
}

func (s *sshConn) runSudoCommand(sess *ssh.Session, stdin io.Writer, cmd string) error {
	// Wrap the output
	if s.DebugEnabled() {
		debugOutWriter := util.NewDebugLogWriter("SSH OUT:", s.Context)
		if sess.Stdout != nil {
			sess.Stdout = io.MultiWriter(sess.Stdout, debugOutWriter)
		} else {
			sess.Stdout = debugOutWriter
		}
		debugErrWriter := util.NewDebugLogWriter("SSH ERR:", s.Context)
		if sess.Stderr != nil {
			sess.Stderr = io.MultiWriter(sess.Stderr, debugErrWriter)
		} else {
			sess.Stderr = debugErrWriter
		}
	}
	// We need a checker to enter the password
	passwordTyper := util.NewExpectListener(stdin, shell.SudoPasswordPromptMatch, s.server.SSH.Pass+"\n")

	if sess.Stdout == nil {
		sess.Stdout = passwordTyper
	} else {
		sess.Stdout = io.MultiWriter(sess.Stdout, passwordTyper)
	}
	if sess.Stderr == nil {
		sess.Stderr = passwordTyper
	} else {
		sess.Stderr = io.MultiWriter(sess.Stderr, passwordTyper)
	}
	if err := sess.Run("sudo -S " + cmd); err != nil {
		return fmt.Errorf("Error running command %v: %v", cmd, err)
	}
	return nil
}

func (cmd *command) StartCommand() (error, *ssh.Session) {
	var (
		session *ssh.Session
		err     error
	)

	z := cmd.zh

	// open ssh link
	if z.client == nil {
		if err = z.dialSSH(); err != nil {
			return err, nil
		}
	}

	// establish ssh session
	if session, err = z.client.NewSession(); err != nil {
		return err, nil
	}

	// setup env, stdin, stdout, stderr
	if err = prepareCommand(session, cmd); err != nil {
		return err, nil
	}

	// start remote command
	err = session.Start(cmd.Path)
	if err == nil {
		return err, session
	}
	return err, nil
}

func (c *command) Run(arg ...string) ([][]string, error) {

	var err error
	var cmd waitable
	var session *ssh.Session

	joinedArgs := strings.Join(arg, " ")
	c.Path = c.Command + " " + joinedArgs
	c.Env = []string{"LC_CTYPE=C", "LANG=en_US.UTF-8"}
	id := uuid.New()
	logger.Log([]string{"ID:" + id, "START", c.Path})
	if c.zh.Local {
		lcmd := c.LocalPrepare(arg...)
		err = lcmd.Start()
		cmd = lcmd
	} else {
		err, session = c.StartCommand()
		if session != nil {
			defer func() {
				session.Close()
			}()
		}
		cmd = session
	}

	logger.Log([]string{"ID:" + id, "DONE"})

	if err != nil {
		return nil, &Error{
			Err:    err,
			Debug:  strings.Join([]string{c.Command, joinedArgs}, " "),
			Stderr: c.stderr.String(),
		}
	}

	if err = cmd.Wait(); err != nil {
		return nil, &Error{
			Err:    err,
			Stderr: c.stderr.String(),
			Debug:  strings.Join([]string{c.Command, joinedArgs}, " "),
		}
	}

	// assume if you passed in something for stdout, that you know what to do with it
	if c.Stdout != nil {
		return nil, nil
	}

	lines := strings.Split(c.stdout.String(), "\n")
	//last line is always blank
	lines = lines[0 : len(lines)-1]
	output := make([][]string, len(lines))

	for i, l := range lines {
		output[i] = strings.Fields(l)
	}
	return output, err

}

func execLine(session *ssh.Session, cmd string) {
	out, err := session.CombinedOutput(cmd)
	if err != nil {
		// panic(err)
		log.Fatal(err.Error())
	}
	fmt.Println(string(out))
}

//NewSSHClientSession creates a new ssh session instance
func NewSSHClientSession(s *ssh.Session, in io.Reader) *SSHClientSession {
	out := new(bytes.Buffer)
	err := new(bytes.Buffer)
	s.Stdin = in
	s.Stdout = out
	s.Stderr = err
	return &SSHClientSession{s, in, out, err}
}

// RequestAgentForwarding sets up agent forwarding for the session.
// ForwardToAgent or ForwardToRemote should be called to route
// the authentication requests.
func RequestAgentForwarding(session *ssh.Session) error {
	ok, err := session.SendRequest("[email protected]", true, nil)
	if err != nil {
		return err
	}
	if !ok {
		return errors.New("forwarding request denied")
	}
	return nil
}

func sshExecNative(c conf.ServerInfo, cmd string, sudo bool) (result execResult) {
	result.Servername = c.ServerName
	result.Host = c.Host
	result.Port = c.Port

	var client *ssh.Client
	var err error
	if client, err = sshConnect(c); err != nil {
		result.Error = err
		result.ExitStatus = 999
		return
	}
	defer client.Close()

	var session *ssh.Session
	if session, err = client.NewSession(); err != nil {
		result.Error = fmt.Errorf(
			"Failed to create a new session. servername: %s, err: %s",
			c.ServerName, err)
		result.ExitStatus = 999
		return
	}
	defer session.Close()

	// http://blog.ralch.com/tutorial/golang-ssh-connection/
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}
	if err = session.RequestPty("xterm", 400, 256, modes); err != nil {
		result.Error = fmt.Errorf(
			"Failed to request for pseudo terminal. servername: %s, err: %s",
			c.ServerName, err)
		result.ExitStatus = 999
		return
	}

	var stdoutBuf, stderrBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	session.Stderr = &stderrBuf

	cmd = decolateCmd(c, cmd, sudo)
	if err := session.Run(cmd); err != nil {
		if exitErr, ok := err.(*ssh.ExitError); ok {
			result.ExitStatus = exitErr.ExitStatus()
		} else {
			result.ExitStatus = 999
		}
	} else {
		result.ExitStatus = 0
	}

	result.Stdout = stdoutBuf.String()
	result.Stderr = stderrBuf.String()
	result.Cmd = strings.Replace(cmd, "\n", "", -1)
	return
}

func configureSessionEnv(session *ssh.Session, env map[string]string) error {
	fmt.Println("Setting environment variables")

	for key, value := range env {
		if err := session.Setenv(key, value); err != nil {
			fmt.Printf("Unable to set environment %s = %s\n", key, value)
			return err
		}
	}

	fmt.Println("Done setting environment variables")
	return nil
}

// RunCommand runs a shell command in a vagrant node and returns it's exit status
func (n *VagrantNode) RunCommand(cmd string) error {
	var (
		s   *ssh.Session
		err error
	)

	if s, err = n.client.NewSession(); err != nil {
		return err
	}
	defer s.Close()

	return s.Run(newCmdStrWithSource(cmd))
}

func (client *clientSSH) prepareCommand(session *ssh.Session, cmd *sshCommand) error {
	for _, env := range cmd.env {
		variable := strings.Split(env, "=")
		if len(variable) != 2 {
			continue
		}

		if err := session.Setenv(variable[0], variable[1]); err != nil {
			return err
		}
	}

	return nil
}

// RunCommandWithOutput runs a shell command in a vagrant node and returns it's
// exit status and output
func (n *VagrantNode) RunCommandWithOutput(cmd string) (string, error) {
	var (
		s   *ssh.Session
		err error
	)

	if s, err = n.client.NewSession(); err != nil {
		return "", err
	}
	defer s.Close()

	output, err := s.CombinedOutput(newCmdStrWithSource(cmd))
	return string(output), err
}

// NewClient creates a new SFTP client on top of an already created
// ssh.Session.
func NewClient(s *ssh.Session) (*Client, error) {
	stdin, err := s.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdout, err := s.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderr, err := s.StderrPipe()
	if err != nil {
		return nil, err
	}

	if err := s.RequestSubsystem("sftp"); err != nil {
		return nil, err
	}
	sftp := &Client{
		stdin:   stdin,
		stdout:  stdout,
		stderr:  stderr,
		chans:   &fxpChanList{},
		session: s,
	}

	if err := sftp.init(); err != nil {
		return nil, err
	}
	return sftp, nil
}