pot/modules/runtime/context.go

package runtime

import (
	"fmt"
	"io"
	"net/url"
	"os"
	"os/user"
	"strconv"
	"strings"

	"github.com/google/uuid"
	"github.com/kevinburke/ssh_config"
	"github.com/melbahja/goph"
	log "github.com/sirupsen/logrus"
)

type Context struct {
	id      string
	rawHost string
	host    string
	port    int
	user    string
	pass    *string
	pkg     *Package
	client  *goph.Client
	output  io.Writer
}

func NewContext(pkg *Package, host string, output io.Writer) (*Context, error) {
	u, err := url.Parse("ssh://" + host)
	if err != nil {
		return nil, err
	}

	rawHost := host

	var port = 22
	if u.Port() != "" {
		port, _ = strconv.Atoi(u.Port())
	}
	host = u.Hostname()

	var hostUser string
	var pass *string
	if u.User != nil {
		hostUser = u.User.Username()
		p, hasPass := u.User.Password()
		if hasPass {
			pass = &p
		}
	}

	if hostUser == "" {
		currUser, err := user.Current()
		if err != nil {
			return nil, err
		}
		hostUser = currUser.Username
		pass = nil
	}

	return &Context{
		id:      uuid.New().String(),
		pkg:     pkg,
		rawHost: rawHost,
		host:    host,
		port:    port,
		user:    hostUser,
		pass:    pass,
		output:  output,
	}, nil
}

func (ctx *Context) IsLogin() bool {
	return ctx.client != nil
}

func (ctx *Context) Close() error {
	if ctx.client != nil {
		return ctx.client.Close()
	}
	return nil
}

func findPublicKeyPath(cfgPath, host string, port int, user, curUser string) (string, error) {
	f, err := os.Open(cfgPath)
	if err != nil {
		if os.IsNotExist(err) {
			return "", nil
		}
		return "", err
	}
	defer f.Close()
	cfg, err := ssh_config.Decode(f)
	if err != nil {
		return "", err
	}

	var publicKeyPath string
	for _, sec := range cfg.Hosts {
		var foundHost bool
		for _, h := range sec.Patterns {
			if h.String() == host {
				foundHost = true
				break
			}
		}
		if !foundHost {
			continue
		}

		var cfgPort int
		var cfgUser string
		for _, node := range sec.Nodes {
			fields := strings.Fields(node.String())
			if len(fields) < 1 {
				continue
			}
			switch fields[0] {
			case "Port":
				cfgPort, _ = strconv.Atoi(fields[1])
			case "IdentityFile":
				publicKeyPath = fields[1]
			case "User":
				cfgUser = fields[1]
			default:
			}
		}

		if (cfgPort > 0 && cfgPort != port) ||
			(cfgUser != "" && cfgUser != user) ||
			(cfgUser == "" && user != curUser) {
			publicKeyPath = ""
			continue
		}

		break
	}
	return publicKeyPath, nil
}

func (ctx *Context) loginWithPrivateKey() error {
	u, err := user.Current()
	if err != nil {
		return err
	}

	p, err := findPublicKeyPath(u.HomeDir+"/.ssh/config", ctx.host, ctx.port,
		ctx.user, u.Username)
	if err != nil {
		return err
	}

	if p == "" {
		p = u.HomeDir + "/.ssh/id_rsa"
	} else {
		p = strings.Replace(p, "~", u.HomeDir, -1)
	}

	// FIXME: don't support passphrase
	auth, err := goph.Key(p, "")
	if err != nil {
		return err
	}

	callback, err := goph.DefaultKnownHosts()
	if err != nil {
		return err
	}

	ctx.client, err = goph.NewConn(&goph.Config{
		User:     ctx.user,
		Addr:     ctx.host,
		Port:     uint(ctx.port),
		Auth:     auth,
		Timeout:  goph.DefaultTimeout,
		Callback: callback,
	})
	return err
}

func (ctx *Context) loginWithPass() error {
	callback, err := goph.DefaultKnownHosts()
	if err != nil {
		return err
	}

	ctx.client, err = goph.NewConn(&goph.Config{
		User:     ctx.user,
		Addr:     ctx.host,
		Port:     uint(ctx.port),
		Auth:     goph.Password(*ctx.pass),
		Timeout:  goph.DefaultTimeout,
		Callback: callback,
	})
	return err
}

func (ctx *Context) Login() error {
	if ctx.client != nil {
		return nil
	}

	if ctx.pass != nil {
		return ctx.loginWithPass()
	}

	return ctx.loginWithPrivateKey()
}

func (ctx *Context) Upload(files ...string) error {
	if ctx.client == nil {
		if err := ctx.Login(); err != nil {
			return err
		}
	}
	if len(files) == 0 {
		return nil
	}

	if err := ctx.Execute(fmt.Sprintf("mkdir -p $TMPDIR/%s", ctx.id)); err != nil {
		return err
	}
	for _, file := range files {
		if err := ctx.client.Upload(file, fmt.Sprintf("$TMPDIR/%s/%s", ctx.id, file)); err != nil {
			return err
		}
	}
	return nil
}

func (ctx *Context) Execute(cmd string) error {
	if ctx.client == nil {
		if err := ctx.Login(); err != nil {
			return err
		}
	}

	// Execute your command.
	sess, err := ctx.client.NewSession()
	if err != nil {
		return err
	}
	defer sess.Close()

	sess.Stdout = ctx.output

	log.Infof("[%s] %s", ctx.rawHost, cmd)

	return sess.Run(cmd)
}