exec/ssh/ssh.go

package ssh

import (
	"fmt"
	"io"
	"os"
	"path"
	"strings"

	"golang.org/x/crypto/ssh"
	"sirherobrine23.com.br/go-bds/exec"
)

var _ exec.Proc = (*Session)(nil)

type Session struct {
	Client *ssh.Client

	Rootfs string

	SshExit *ssh.ExitError
	Session *ssh.Session
}

func NewSession(client *ssh.Client) (*Session, error) {
	session, err := client.NewSession()
	if err != nil {
		return nil, err
	}
	return &Session{Client: client, Session: session}, nil
}

func (session *Session) Close() error { return session.Session.Close() }
func (session *Session) Kill() error  { return session.Signal(os.Kill) }
func (session *Session) Signal(s os.Signal) error {
	return session.Session.Signal(ssh.Signal(s.String()))
}

func (session *Session) Wait() error { return session.Session.Wait() }
func (session *Session) ExitCode() (int, error) {
	err := session.Session.Wait()
	if err == nil {
		return 0, nil
	}

	if exitErr, ok := err.(*ssh.ExitError); ok {
		session.SshExit = exitErr
		return exitErr.ExitStatus(), nil
	}
	return -1, err
}

func (session *Session) AttachStdin(r io.Reader) error {
	if session.Session == nil {
		return io.ErrUnexpectedEOF
	}
	session.Session.Stdin = r
	return nil
}
func (session *Session) AttachStdout(w io.Writer) error {
	if session.Session == nil {
		return io.ErrUnexpectedEOF
	}
	session.Session.Stdout = w
	return nil
}

func (session *Session) AttachStderr(w io.Writer) error {
	if session.Session == nil {
		return io.ErrUnexpectedEOF
	}
	session.Session.Stderr = w
	return nil
}

func (session *Session) Start(command *exec.Exec) (err error) {
	if session.Session == nil {
		if session.Session, err = session.Client.NewSession(); err != nil {
			return
		}
	}

	for key, value := range command.Environment {
		session.Session.Setenv(key, value)
	}

	var commandString string

	if command.Cwd != "" {
		commandString = "cd " + command.Cwd + ";"
	}

	if session.Rootfs != "" {
		commandString = "cd " + session.Rootfs + ";"
		if command.Cwd != "" {
			commandString = "cd " + path.Join(session.Rootfs, command.Cwd) + ";"
		}
	}

	if command.Stdin != nil {
		session.Session.Stdin = command.Stdin
	}
	if command.Stdout != nil {
		session.Session.Stdout = command.Stdout
		session.Session.Stderr = command.Stdout
	}
	if command.Stderr != nil {
		session.Session.Stderr = command.Stderr
	}

	commandStr, args := command.Arguments[0], command.Arguments[1:]
	if err = session.Session.Start(fmt.Sprintf("%s%s %s; exit $?", commandString, commandStr, strings.Join(args, " "))); err != nil {
		return
	}

	if command.Context != nil {
		go func() {
			<-command.Context.Done()
			session.Session.Signal(ssh.SIGKILL)
		}()
	}

	return nil
}