<commit_before>package command

import (
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/mackerelio/mackerel-agent/agent"
	"github.com/mackerelio/mackerel-agent/config"
	"github.com/mackerelio/mackerel-agent/logging"
	"github.com/mackerelio/mackerel-agent/mackerel"
	"github.com/mackerelio/mackerel-agent/spec"
)

var logger = logging.GetLogger("command")

const idFileName = "id"

func idFilePath(root string) string {
	return filepath.Join(root, idFileName)
}

func loadHostID(root string) (string, error) {
	content, err := ioutil.ReadFile(idFilePath(root))
	if err != nil {
		return "", err
	}
	return string(content), nil
}

func saveHostID(root string, id string) error {
	err := os.MkdirAll(root, 0755)
	if err != nil {
		return err
	}

	file, err := os.Create(idFilePath(root))
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.Write([]byte(id))
	if err != nil {
		return err
	}

	return nil
}

// prepareHost collects specs of the host and sends them to Mackerel server.
// A unique host-id is returned by the server if one is not specified.
func prepareHost(root string, api *mackerel.API, roleFullnames []string) (*mackerel.Host, error) {
	// XXX this configuration should be moved to under spec/linux
	os.Setenv("PATH", "/sbin:/usr/sbin:/bin:/usr/bin:"+os.Getenv("PATH"))
	os.Setenv("LANG", "C") // prevent changing outputs of some command, e.g. ifconfig.

	hostname, meta, interfaces, err := collectHostSpecs()
	if err != nil {
		return nil, fmt.Errorf("error while collecting host specs: %s", err.Error())
	}

	var result *mackerel.Host
	if hostID, err := loadHostID(root); err != nil { // create
		logger.Debugf("Registering new host on mackerel...")
		createdHostID, err := api.CreateHost(hostname, meta, interfaces, roleFullnames)
		if err != nil {
			return nil, fmt.Errorf("Failed to register this host: %s", err.Error())
		}

		result, err = api.FindHost(createdHostID)
		if err != nil {
			return nil, fmt.Errorf("Failed to find this host on mackerel: %s", err.Error())
		}
	} else { // update
		result, err = api.FindHost(hostID)
		if err != nil {
			return nil, fmt.Errorf("Failed to find this host on mackerel (You may want to delete file \"%s\" to register this host to an another organization): %s", idFilePath(root), err.Error())
		}
		err := api.UpdateHost(hostID, hostname, meta, interfaces, roleFullnames)
		if err != nil {
			return nil, fmt.Errorf("Failed to update this host: %s", err.Error())
		}
	}

	err = saveHostID(root, result.ID)
	if err != nil {
		return nil, fmt.Errorf("Failed to save host ID: %s", err.Error())
	}

	return result, nil
}

// Interval between each updating host specs.
var specsUpdateInterval = 1 * time.Hour

func delayByHost(host *mackerel.Host) int {
	s := sha1.Sum([]byte(host.ID))
	return int(s[len(s)-1]) % int(config.PostMetricsInterval.Seconds())
}

type context struct {
	ag   *agent.Agent
	conf *config.Config
	host *mackerel.Host
	api  *mackerel.API
}

type postValue struct {
	values   []*mackerel.CreatingMetricsValue
	retryCnt int
}

func newPostValue(values []*mackerel.CreatingMetricsValue) *postValue {
	return &postValue{values, 0}
}

type loopState uint8

const (
	loopStateFirst loopState = iota
	loopStateDefault
	loopStateQueued
	loopStateHadError
	loopStateTerminating
)

func loop(c *context, termCh chan struct{}) int {
	quit := make(chan struct{})
	defer close(quit) // broadcast terminating

	// Periodically update host specs.
	go updateHostSpecsLoop(c, quit)

	postQueue := make(chan *postValue, c.conf.Connection.PostMetricsBufferSize)
	go enqueueLoop(c, postQueue, quit)

	postDelaySeconds := delayByHost(c.host)
	initialDelay := postDelaySeconds / 2
	logger.Debugf("wait %d seconds before initial posting.", initialDelay)
	select {
	case <-termCh:
		return 0
	case <-time.After(time.Duration(initialDelay) * time.Second):
		c.ag.InitPluginGenerators(c.api)
	}

	lState := loopStateFirst
	for {
		select {
		case <-termCh:
			if lState == loopStateTerminating {
				return 1
			}
			lState = loopStateTerminating
			if len(postQueue) <= 0 {
				return 0
			}
		case v := <-postQueue:
			origPostValues := [](*postValue){v}
			if len(postQueue) > 0 {
				// Bulk posting. However at most "two" metrics are to be posted, so postQueue isn't always empty yet.
				logger.Debugf("Merging datapoints with next queued ones")
				nextValues := <-postQueue
				origPostValues = append(origPostValues, nextValues)
			}

			delaySeconds := 0
			switch lState {
			case loopStateFirst: // request immediately to create graph defs of host
				// nop
			case loopStateQueued:
				delaySeconds = c.conf.Connection.PostMetricsDequeueDelaySeconds
			case loopStateHadError:
				// TODO: better interval calculation. exponential backoff or so.
				delaySeconds = c.conf.Connection.PostMetricsRetryDelaySeconds
			case loopStateTerminating:
				// dequeue and post every one second when terminating.
				delaySeconds = 1
			default:
				// Sending data at every 0 second from all hosts causes request flooding.
				// To prevent flooding, this loop sleeps for some seconds
				// which is specific to the ID of the host running agent on.
				// The sleep second is up to 60s (to be exact up to `config.Postmetricsinterval.Seconds()`.
				elapsedSeconds := int(time.Now().Unix() % int64(config.PostMetricsInterval.Seconds()))
				if postDelaySeconds > elapsedSeconds {
					delaySeconds = postDelaySeconds - elapsedSeconds
				}
			}

			// determin next loopState before sleeping
			if lState != loopStateTerminating {
				if len(postQueue) > 0 {
					lState = loopStateQueued
				} else {
					lState = loopStateDefault
				}
			}

			logger.Debugf("Sleep %d seconds before posting.", delaySeconds)
			select {
			case <-time.After(time.Duration(delaySeconds) * time.Second):
				// nop
			case <-termCh:
				if lState == loopStateTerminating {
					return 1
				}
				lState = loopStateTerminating
			}

			postValues := [](*mackerel.CreatingMetricsValue){}
			for _, v := range origPostValues {
				postValues = append(postValues, v.values...)
			}
			err := c.api.PostMetricsValues(postValues)
			if err != nil {
				logger.Errorf("Failed to post metrics value (will retry): %s", err.Error())
				if lState != loopStateTerminating {
					lState = loopStateHadError
				}
				go func() {
					for _, v := range origPostValues {
						v.retryCnt++
						// It is difficult to distinguish the error is server error or data error.
						// So, if retryCnt exceeded the configured limit, postValue is considered invalid and abandoned.
						if v.retryCnt > c.conf.Connection.PostMetricsRetryMax {
							json, err := json.Marshal(v.values)
							if err != nil {
								logger.Errorf("Something wrong with post values. marshaling failed.")
							} else {
								logger.Errorf("Post values may be invalid and abandoned: %s", string(json))
							}
							continue
						}
						postQueue <- v
					}
				}()
				continue
			}
			logger.Debugf("Posting metrics succeeded.")

			if lState == loopStateTerminating && len(postQueue) <= 0 {
				return 0
			}
		}
	}
}

func updateHostSpecsLoop(c *context, quit chan struct{}) {
	for {
		select {
		case <-quit:
			return
		case <-time.After(specsUpdateInterval):
			UpdateHostSpecs(c.conf, c.api, c.host)
		}
	}
}

func enqueueLoop(c *context, postQueue chan *postValue, quit chan struct{}) {
	metricsResult := c.ag.Watch()
	for {
		select {
		case <-quit:
			return
		case result := <-metricsResult:
			created := float64(result.Created.Unix())
			creatingValues := [](*mackerel.CreatingMetricsValue){}
			for name, value := range (map[string]float64)(result.Values) {
				creatingValues = append(
					creatingValues,
					&mackerel.CreatingMetricsValue{
						HostID: c.host.ID,
						Name:   name,
						Time:   created,
						Value:  value,
					},
				)
			}
			logger.Debugf("Enqueuing task to post metrics.")
			postQueue <- newPostValue(creatingValues)
		}
	}
}

// collectHostSpecs collects host specs (correspond to "name", "meta" and "interfaces" fields in API v0)
func collectHostSpecs() (string, map[string]interface{}, []map[string]interface{}, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return "", nil, nil, fmt.Errorf("failed to obtain hostname: %s", err.Error())
	}

	meta := spec.Collect(specGenerators())

	interfacesSpec, err := interfaceGenerator().Generate()
	if err != nil {
		return "", nil, nil, fmt.Errorf("failed to collect interfaces: %s", err.Error())
	}

	interfaces, _ := interfacesSpec.([]map[string]interface{})

	return hostname, meta, interfaces, nil
}

// UpdateHostSpecs updates the host information that is already registered on Mackerel.
func UpdateHostSpecs(conf *config.Config, api *mackerel.API, host *mackerel.Host) {
	logger.Debugf("Updating host specs...")

	hostname, meta, interfaces, err := collectHostSpecs()
	if err != nil {
		logger.Errorf("While collecting host specs: %s", err)
		return
	}

	err = api.UpdateHost(host.ID, hostname, meta, interfaces, conf.Roles)
	if err != nil {
		logger.Errorf("Error while updating host specs: %s", err)
	} else {
		logger.Debugf("Host specs sent.")
	}
}

// Prepare sets up API and registers the host data to the Mackerel server.
// Use returned values to call Run().
func Prepare(conf *config.Config) (*mackerel.API, *mackerel.Host, error) {
	api, err := mackerel.NewAPI(conf.Apibase, conf.Apikey, conf.Verbose)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to prepare an api: %s", err.Error())
	}

	host, err := prepareHost(conf.Root, api, conf.Roles)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to preapre host: %s", err.Error())
	}

	return api, host, nil
}

// Run starts the main metric collecting logic and this function will never return.
func Run(conf *config.Config, api *mackerel.API, host *mackerel.Host, termCh chan struct{}) int {
	logger.Infof("Start: apibase = %s, hostName = %s, hostID = %s", conf.Apibase, host.Name, host.ID)

	ag := &agent.Agent{
		MetricsGenerators: metricsGenerators(conf),
		PluginGenerators:  pluginGenerators(conf),
	}

	c := &context{
		ag:   ag,
		conf: conf,
		host: host,
		api:  api,
	}

	return loop(c, termCh)
}
<commit_msg>Filter invalid float values<commit_after>package command

import (
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"time"

	"github.com/mackerelio/mackerel-agent/agent"
	"github.com/mackerelio/mackerel-agent/config"
	"github.com/mackerelio/mackerel-agent/logging"
	"github.com/mackerelio/mackerel-agent/mackerel"
	"github.com/mackerelio/mackerel-agent/spec"
)

var logger = logging.GetLogger("command")

const idFileName = "id"

func idFilePath(root string) string {
	return filepath.Join(root, idFileName)
}

func loadHostID(root string) (string, error) {
	content, err := ioutil.ReadFile(idFilePath(root))
	if err != nil {
		return "", err
	}
	return string(content), nil
}

func saveHostID(root string, id string) error {
	err := os.MkdirAll(root, 0755)
	if err != nil {
		return err
	}

	file, err := os.Create(idFilePath(root))
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.Write([]byte(id))
	if err != nil {
		return err
	}

	return nil
}

// prepareHost collects specs of the host and sends them to Mackerel server.
// A unique host-id is returned by the server if one is not specified.
func prepareHost(root string, api *mackerel.API, roleFullnames []string) (*mackerel.Host, error) {
	// XXX this configuration should be moved to under spec/linux
	os.Setenv("PATH", "/sbin:/usr/sbin:/bin:/usr/bin:"+os.Getenv("PATH"))
	os.Setenv("LANG", "C") // prevent changing outputs of some command, e.g. ifconfig.

	hostname, meta, interfaces, err := collectHostSpecs()
	if err != nil {
		return nil, fmt.Errorf("error while collecting host specs: %s", err.Error())
	}

	var result *mackerel.Host
	if hostID, err := loadHostID(root); err != nil { // create
		logger.Debugf("Registering new host on mackerel...")
		createdHostID, err := api.CreateHost(hostname, meta, interfaces, roleFullnames)
		if err != nil {
			return nil, fmt.Errorf("Failed to register this host: %s", err.Error())
		}

		result, err = api.FindHost(createdHostID)
		if err != nil {
			return nil, fmt.Errorf("Failed to find this host on mackerel: %s", err.Error())
		}
	} else { // update
		result, err = api.FindHost(hostID)
		if err != nil {
			return nil, fmt.Errorf("Failed to find this host on mackerel (You may want to delete file \"%s\" to register this host to an another organization): %s", idFilePath(root), err.Error())
		}
		err := api.UpdateHost(hostID, hostname, meta, interfaces, roleFullnames)
		if err != nil {
			return nil, fmt.Errorf("Failed to update this host: %s", err.Error())
		}
	}

	err = saveHostID(root, result.ID)
	if err != nil {
		return nil, fmt.Errorf("Failed to save host ID: %s", err.Error())
	}

	return result, nil
}

// Interval between each updating host specs.
var specsUpdateInterval = 1 * time.Hour

func delayByHost(host *mackerel.Host) int {
	s := sha1.Sum([]byte(host.ID))
	return int(s[len(s)-1]) % int(config.PostMetricsInterval.Seconds())
}

type context struct {
	ag   *agent.Agent
	conf *config.Config
	host *mackerel.Host
	api  *mackerel.API
}

type postValue struct {
	values   []*mackerel.CreatingMetricsValue
	retryCnt int
}

func newPostValue(values []*mackerel.CreatingMetricsValue) *postValue {
	return &postValue{values, 0}
}

type loopState uint8

const (
	loopStateFirst loopState = iota
	loopStateDefault
	loopStateQueued
	loopStateHadError
	loopStateTerminating
)

func loop(c *context, termCh chan struct{}) int {
	quit := make(chan struct{})
	defer close(quit) // broadcast terminating

	// Periodically update host specs.
	go updateHostSpecsLoop(c, quit)

	postQueue := make(chan *postValue, c.conf.Connection.PostMetricsBufferSize)
	go enqueueLoop(c, postQueue, quit)

	postDelaySeconds := delayByHost(c.host)
	initialDelay := postDelaySeconds / 2
	logger.Debugf("wait %d seconds before initial posting.", initialDelay)
	select {
	case <-termCh:
		return 0
	case <-time.After(time.Duration(initialDelay) * time.Second):
		c.ag.InitPluginGenerators(c.api)
	}

	lState := loopStateFirst
	for {
		select {
		case <-termCh:
			if lState == loopStateTerminating {
				return 1
			}
			lState = loopStateTerminating
			if len(postQueue) <= 0 {
				return 0
			}
		case v := <-postQueue:
			origPostValues := [](*postValue){v}
			if len(postQueue) > 0 {
				// Bulk posting. However at most "two" metrics are to be posted, so postQueue isn't always empty yet.
				logger.Debugf("Merging datapoints with next queued ones")
				nextValues := <-postQueue
				origPostValues = append(origPostValues, nextValues)
			}

			delaySeconds := 0
			switch lState {
			case loopStateFirst: // request immediately to create graph defs of host
				// nop
			case loopStateQueued:
				delaySeconds = c.conf.Connection.PostMetricsDequeueDelaySeconds
			case loopStateHadError:
				// TODO: better interval calculation. exponential backoff or so.
				delaySeconds = c.conf.Connection.PostMetricsRetryDelaySeconds
			case loopStateTerminating:
				// dequeue and post every one second when terminating.
				delaySeconds = 1
			default:
				// Sending data at every 0 second from all hosts causes request flooding.
				// To prevent flooding, this loop sleeps for some seconds
				// which is specific to the ID of the host running agent on.
				// The sleep second is up to 60s (to be exact up to `config.Postmetricsinterval.Seconds()`.
				elapsedSeconds := int(time.Now().Unix() % int64(config.PostMetricsInterval.Seconds()))
				if postDelaySeconds > elapsedSeconds {
					delaySeconds = postDelaySeconds - elapsedSeconds
				}
			}

			// determin next loopState before sleeping
			if lState != loopStateTerminating {
				if len(postQueue) > 0 {
					lState = loopStateQueued
				} else {
					lState = loopStateDefault
				}
			}

			logger.Debugf("Sleep %d seconds before posting.", delaySeconds)
			select {
			case <-time.After(time.Duration(delaySeconds) * time.Second):
				// nop
			case <-termCh:
				if lState == loopStateTerminating {
					return 1
				}
				lState = loopStateTerminating
			}

			postValues := [](*mackerel.CreatingMetricsValue){}
			for _, v := range origPostValues {
				postValues = append(postValues, v.values...)
			}
			err := c.api.PostMetricsValues(postValues)
			if err != nil {
				logger.Errorf("Failed to post metrics value (will retry): %s", err.Error())
				if lState != loopStateTerminating {
					lState = loopStateHadError
				}
				go func() {
					for _, v := range origPostValues {
						v.retryCnt++
						// It is difficult to distinguish the error is server error or data error.
						// So, if retryCnt exceeded the configured limit, postValue is considered invalid and abandoned.
						if v.retryCnt > c.conf.Connection.PostMetricsRetryMax {
							json, err := json.Marshal(v.values)
							if err != nil {
								logger.Errorf("Something wrong with post values. marshaling failed.")
							} else {
								logger.Errorf("Post values may be invalid and abandoned: %s", string(json))
							}
							continue
						}
						postQueue <- v
					}
				}()
				continue
			}
			logger.Debugf("Posting metrics succeeded.")

			if lState == loopStateTerminating && len(postQueue) <= 0 {
				return 0
			}
		}
	}
}

func updateHostSpecsLoop(c *context, quit chan struct{}) {
	for {
		select {
		case <-quit:
			return
		case <-time.After(specsUpdateInterval):
			UpdateHostSpecs(c.conf, c.api, c.host)
		}
	}
}

func enqueueLoop(c *context, postQueue chan *postValue, quit chan struct{}) {
	metricsResult := c.ag.Watch()
	for {
		select {
		case <-quit:
			return
		case result := <-metricsResult:
			created := float64(result.Created.Unix())
			creatingValues := [](*mackerel.CreatingMetricsValue){}
			for name, value := range (map[string]float64)(result.Values) {
				if math.IsNaN(value) || math.IsInf(value, 0) {
					logger.Warningf("Invalid value: hostID = %s, name = %s, value = %f\n is not sent.", c.host.ID, name, value)
					continue
				}

				creatingValues = append(
					creatingValues,
					&mackerel.CreatingMetricsValue{
						HostID: c.host.ID,
						Name:   name,
						Time:   created,
						Value:  value,
					},
				)
			}
			logger.Debugf("Enqueuing task to post metrics.")
			postQueue <- newPostValue(creatingValues)
		}
	}
}

// collectHostSpecs collects host specs (correspond to "name", "meta" and "interfaces" fields in API v0)
func collectHostSpecs() (string, map[string]interface{}, []map[string]interface{}, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return "", nil, nil, fmt.Errorf("failed to obtain hostname: %s", err.Error())
	}

	meta := spec.Collect(specGenerators())

	interfacesSpec, err := interfaceGenerator().Generate()
	if err != nil {
		return "", nil, nil, fmt.Errorf("failed to collect interfaces: %s", err.Error())
	}

	interfaces, _ := interfacesSpec.([]map[string]interface{})

	return hostname, meta, interfaces, nil
}

// UpdateHostSpecs updates the host information that is already registered on Mackerel.
func UpdateHostSpecs(conf *config.Config, api *mackerel.API, host *mackerel.Host) {
	logger.Debugf("Updating host specs...")

	hostname, meta, interfaces, err := collectHostSpecs()
	if err != nil {
		logger.Errorf("While collecting host specs: %s", err)
		return
	}

	err = api.UpdateHost(host.ID, hostname, meta, interfaces, conf.Roles)
	if err != nil {
		logger.Errorf("Error while updating host specs: %s", err)
	} else {
		logger.Debugf("Host specs sent.")
	}
}

// Prepare sets up API and registers the host data to the Mackerel server.
// Use returned values to call Run().
func Prepare(conf *config.Config) (*mackerel.API, *mackerel.Host, error) {
	api, err := mackerel.NewAPI(conf.Apibase, conf.Apikey, conf.Verbose)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to prepare an api: %s", err.Error())
	}

	host, err := prepareHost(conf.Root, api, conf.Roles)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to preapre host: %s", err.Error())
	}

	return api, host, nil
}

// Run starts the main metric collecting logic and this function will never return.
func Run(conf *config.Config, api *mackerel.API, host *mackerel.Host, termCh chan struct{}) int {
	logger.Infof("Start: apibase = %s, hostName = %s, hostID = %s", conf.Apibase, host.Name, host.ID)

	ag := &agent.Agent{
		MetricsGenerators: metricsGenerators(conf),
		PluginGenerators:  pluginGenerators(conf),
	}

	c := &context{
		ag:   ag,
		conf: conf,
		host: host,
		api:  api,
	}

	return loop(c, termCh)
}
<|endoftext|>


<commit_before>package command

import (
	"bufio"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"os/exec"
	"path/filepath"
	"strings"
)

// Command is our structured data/object for Command
type Command struct {
	alias      string
	command    string
	verbose    bool
	workingdir string
}

const pathDrush = "/usr/local/bin/drush"

// NewDrushCommand creates a new container for []Command objects
func NewDrushCommand() *Command {
	return &Command{}
}

// Set changes all given values for a Drush command object.
func (drush *Command) Set(alias string, command string, verbose bool) {
	drush.alias = alias
	drush.command = command
	drush.verbose = verbose
	drush.workingdir = "."
}

// GetWorkingDir returns the specified working directory used with executed Drush commands.
func (drush *Command) GetWorkingDir() string {
	return drush.workingdir
}

// SetWorkingDir sets the specified working directory used with executed Drush commands.
func (drush *Command) SetWorkingDir(value string) {
	drush.workingdir = value
}

// GetAlias returns the alias used to executed Drush commands.
func (drush *Command) GetAlias() string {
	return drush.alias
}

// SetAlias changes the alias used to executed Drush commands.
func (drush *Command) SetAlias(value string) {
	drush.alias = value
}

// GetCommand returns the command string on executed Drush commands.
func (drush *Command) GetCommand() string {
	return drush.command
}

// SetCommand changes the command string on executed Drush commands.
func (drush *Command) SetCommand(value string) {
	drush.command = value
}

// GetVerbose returns the verbosity setting on executed Drush commands.
func (drush *Command) GetVerbose() bool {
	return drush.verbose
}

// SetVerbose changes the verbosity setting on executed Drush commands.
func (drush *Command) SetVerbose(value bool) {
	drush.verbose = value
}

// LiveOutput returns, and prints the live output of the executing program
// This will wait for completion before proceeding.
func (drush *Command) LiveOutput() error {
	if strings.Contains(drush.alias, "@") == true {
		drush.alias = strings.Replace(drush.alias, "@", "", -1)
	}
	if drush.alias != "" {
		drush.alias = fmt.Sprintf("@%v", drush.alias)
	}
	if drush.verbose == true {
		drush.alias = fmt.Sprintf("%v --verbose", drush.alias)
	}
	args := fmt.Sprintf("%v %v", drush.alias, drush.command)

	comm := new(exec.Cmd)
	comm = exec.Command("sh", "-c", "cd "+drush.workingdir+" && "+pathDrush+" "+args)
	Pipe, _ := comm.StderrPipe()
	scanner := bufio.NewScanner(Pipe)
	go func() {
		for scanner.Scan() {
			if strings.Contains(scanner.Text(), "[error]") || strings.Contains(scanner.Text(), "[warning]") {
				log.Warnf("%s", scanner.Text())
			} else {
				log.Infof("%s", scanner.Text())
			}
		}
	}()
	err := comm.Start()
	err = comm.Wait()
	return err
}

// Output gets the output from a single Command object, does not support []Command items.
func (drush *Command) Output() ([]string, error) {
	comm, err := drush.Run()
	response := filepath.SplitList(string(comm))
	return response, err
}

// CombinedOutput will return the CombinedOutput of a command.
func (drush *Command) CombinedOutput() ([]byte, error) {
	if strings.Contains(drush.alias, "@") == true {
		drush.alias = strings.Replace(drush.alias, "@", "", -1)
	}
	if drush.alias != "" {
		drush.alias = fmt.Sprintf("@%v", drush.alias)
	}
	if drush.verbose == true {
		drush.alias = fmt.Sprintf("%v --verbose", drush.alias)
	}
	args := fmt.Sprintf("%v %v", drush.alias, drush.command)
	if drush.GetWorkingDir() != "." {
		comm, err := exec.Command("sh", "-c", "cd "+drush.workingdir+" && "+pathDrush+" "+args).CombinedOutput()
		return comm, err
	}
	comm, err := exec.Command("sh", "-c", pathDrush+" "+args).CombinedOutput()
	return comm, err
}

// Run runs an individual Command object, does not support []Command items.
func (drush *Command) Run() ([]byte, error) {
	if strings.Contains(drush.alias, "@") == true {
		drush.alias = strings.Replace(drush.alias, "@", "", -1)
	}
	if drush.alias != "" {
		drush.alias = fmt.Sprintf("@%v", drush.alias)
	}
	if drush.verbose == true {
		drush.alias = fmt.Sprintf("%v --verbose", drush.alias)
	}
	args := fmt.Sprintf("%v %v", drush.alias, drush.command)
	if drush.GetWorkingDir() != "." {
		comm, err := exec.Command("sh", "-c", "cd "+drush.workingdir+" && "+pathDrush+" "+args).CombinedOutput()
		return comm, err
	}
	comm, err := exec.Command("sh", "-c", pathDrush+" "+args).CombinedOutput()
	return comm, err
}

// DrushDatabaseSync executes a database synchronisation task from a source to destination with the use of Drush.
func DrushDatabaseSync(srcAlias, destAlias string) {
	/* So our binary and this function combined support two-way traffic...  */
	x := NewDrushCommand()
	srcAlias = strings.Replace(srcAlias, "@", "", -1)
	destAlias = strings.Replace(destAlias, "@", "", -1)
	x.Set("", fmt.Sprintf("sql-sync @%v @%v -y", srcAlias, destAlias), true)
	_, err := x.Output()
	if err == nil {
		log.Infoln("Syncronised databases complete.")
	} else {
		log.Errorln("Could not syncronise databases.")
	}
}

// DrushFilesSync executes a file synchronisation task from a source to destination with the use of Drush.
func DrushFilesSync(srcAlias, destAlias string) {
	x := NewDrushCommand()
	srcAlias = strings.Replace(srcAlias, "@", "", -1)
	destAlias = strings.Replace(destAlias, "@", "", -1)
	x.Set("", fmt.Sprintf("--yes rsync --exclude-other-sites --exclude-conf @%v:%%files @%v:%%files", srcAlias, destAlias), true)
	_, err := x.Output()
	if err == nil {
		log.Infoln("Synced public file system.")
	} else {
		log.Warnln("Public file system has not been synced.")
	}
	x.Set("", fmt.Sprintf("--yes rsync --exclude-other-sites --exclude-conf @%v:%%private @%v:%%private", srcAlias, destAlias), true)
	_, err = x.Output()
	if err == nil {
		log.Infoln("Synced private file system.")
	} else {
		log.Warnln("Private file system has not been synced.")
	}
}

// DrushClearCache performs a cache clear task on an input site alias with the use of Drush.
func DrushClearCache(alias string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set(alias, "cc all", false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not clear caches.", err)
	} else {
		log.Infoln("Caches cleared.")
	}
}

// DrushRebuildRegistry performs a registry rebuild task on an input site alias with the use of Drush.
func DrushRebuildRegistry(alias string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set(alias, "rr", false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not rebuild registry.", err)
	} else {
		log.Infoln("Rebuilt registry.")
	}
}

// DrushUpdateDatabase performs a database update task on an input site alias with the use of Drush.
func DrushUpdateDatabase(alias string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set(alias, "updb -y", false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not update database.", err)
	} else {
		log.Infoln("Updated database where possible.")
	}
}

// DrushDownloadToPath performs a database update task on an input site alias with the use of Drush.
func DrushDownloadToPath(path, project string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set("", "pm-download --yes "+project+" --destination="+path, false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not download module ", project, err)
	} else {
		log.Infoln("Downloaded module", project)
	}
}

// DrushDownloadToAlias performs a database update task on an input site alias with the use of Drush.
func DrushDownloadToAlias(alias, project string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set(alias, "pm-download --yes "+project, false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not download module ", project, err)
	} else {
		log.Infoln("Downloaded module", project)
	}
}

// DrushVariableSet Runs drush vset with a given variable name and value.
func DrushVariableSet(alias, variable_name, variable_value string) {
	srcAlias := strings.Replace(alias, "@", "", -1)
	x := NewDrushCommand()
	x.Set(srcAlias, fmt.Sprintf("vset %v %v", variable_name, variable_value), true)
	drushOut, err := x.Output()
	if err == nil {
		log.Infof("Succesfully set files path %v to %v via Drush\n")
	} else {
		log.Errorf("Could not set files path %v to %v via Drush %v\n", variable_value, drushOut)
	}
}<commit_msg>#22: Improve reported output of vset function.<commit_after>package command

import (
	"bufio"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"os/exec"
	"path/filepath"
	"strings"
)

// Command is our structured data/object for Command
type Command struct {
	alias      string
	command    string
	verbose    bool
	workingdir string
}

const pathDrush = "/usr/local/bin/drush"

// NewDrushCommand creates a new container for []Command objects
func NewDrushCommand() *Command {
	return &Command{}
}

// Set changes all given values for a Drush command object.
func (drush *Command) Set(alias string, command string, verbose bool) {
	drush.alias = alias
	drush.command = command
	drush.verbose = verbose
	drush.workingdir = "."
}

// GetWorkingDir returns the specified working directory used with executed Drush commands.
func (drush *Command) GetWorkingDir() string {
	return drush.workingdir
}

// SetWorkingDir sets the specified working directory used with executed Drush commands.
func (drush *Command) SetWorkingDir(value string) {
	drush.workingdir = value
}

// GetAlias returns the alias used to executed Drush commands.
func (drush *Command) GetAlias() string {
	return drush.alias
}

// SetAlias changes the alias used to executed Drush commands.
func (drush *Command) SetAlias(value string) {
	drush.alias = value
}

// GetCommand returns the command string on executed Drush commands.
func (drush *Command) GetCommand() string {
	return drush.command
}

// SetCommand changes the command string on executed Drush commands.
func (drush *Command) SetCommand(value string) {
	drush.command = value
}

// GetVerbose returns the verbosity setting on executed Drush commands.
func (drush *Command) GetVerbose() bool {
	return drush.verbose
}

// SetVerbose changes the verbosity setting on executed Drush commands.
func (drush *Command) SetVerbose(value bool) {
	drush.verbose = value
}

// LiveOutput returns, and prints the live output of the executing program
// This will wait for completion before proceeding.
func (drush *Command) LiveOutput() error {
	if strings.Contains(drush.alias, "@") == true {
		drush.alias = strings.Replace(drush.alias, "@", "", -1)
	}
	if drush.alias != "" {
		drush.alias = fmt.Sprintf("@%v", drush.alias)
	}
	if drush.verbose == true {
		drush.alias = fmt.Sprintf("%v --verbose", drush.alias)
	}
	args := fmt.Sprintf("%v %v", drush.alias, drush.command)

	comm := new(exec.Cmd)
	comm = exec.Command("sh", "-c", "cd "+drush.workingdir+" && "+pathDrush+" "+args)
	Pipe, _ := comm.StderrPipe()
	scanner := bufio.NewScanner(Pipe)
	go func() {
		for scanner.Scan() {
			if strings.Contains(scanner.Text(), "[error]") || strings.Contains(scanner.Text(), "[warning]") {
				log.Warnf("%s", scanner.Text())
			} else {
				log.Infof("%s", scanner.Text())
			}
		}
	}()
	err := comm.Start()
	err = comm.Wait()
	return err
}

// Output gets the output from a single Command object, does not support []Command items.
func (drush *Command) Output() ([]string, error) {
	comm, err := drush.Run()
	response := filepath.SplitList(string(comm))
	return response, err
}

// CombinedOutput will return the CombinedOutput of a command.
func (drush *Command) CombinedOutput() ([]byte, error) {
	if strings.Contains(drush.alias, "@") == true {
		drush.alias = strings.Replace(drush.alias, "@", "", -1)
	}
	if drush.alias != "" {
		drush.alias = fmt.Sprintf("@%v", drush.alias)
	}
	if drush.verbose == true {
		drush.alias = fmt.Sprintf("%v --verbose", drush.alias)
	}
	args := fmt.Sprintf("%v %v", drush.alias, drush.command)
	if drush.GetWorkingDir() != "." {
		comm, err := exec.Command("sh", "-c", "cd "+drush.workingdir+" && "+pathDrush+" "+args).CombinedOutput()
		return comm, err
	}
	comm, err := exec.Command("sh", "-c", pathDrush+" "+args).CombinedOutput()
	return comm, err
}

// Run runs an individual Command object, does not support []Command items.
func (drush *Command) Run() ([]byte, error) {
	if strings.Contains(drush.alias, "@") == true {
		drush.alias = strings.Replace(drush.alias, "@", "", -1)
	}
	if drush.alias != "" {
		drush.alias = fmt.Sprintf("@%v", drush.alias)
	}
	if drush.verbose == true {
		drush.alias = fmt.Sprintf("%v --verbose", drush.alias)
	}
	args := fmt.Sprintf("%v %v", drush.alias, drush.command)
	if drush.GetWorkingDir() != "." {
		comm, err := exec.Command("sh", "-c", "cd "+drush.workingdir+" && "+pathDrush+" "+args).CombinedOutput()
		return comm, err
	}
	comm, err := exec.Command("sh", "-c", pathDrush+" "+args).CombinedOutput()
	return comm, err
}

// DrushDatabaseSync executes a database synchronisation task from a source to destination with the use of Drush.
func DrushDatabaseSync(srcAlias, destAlias string) {
	/* So our binary and this function combined support two-way traffic...  */
	x := NewDrushCommand()
	srcAlias = strings.Replace(srcAlias, "@", "", -1)
	destAlias = strings.Replace(destAlias, "@", "", -1)
	x.Set("", fmt.Sprintf("sql-sync @%v @%v -y", srcAlias, destAlias), true)
	_, err := x.Output()
	if err == nil {
		log.Infoln("Syncronised databases complete.")
	} else {
		log.Errorln("Could not syncronise databases.")
	}
}

// DrushFilesSync executes a file synchronisation task from a source to destination with the use of Drush.
func DrushFilesSync(srcAlias, destAlias string) {
	x := NewDrushCommand()
	srcAlias = strings.Replace(srcAlias, "@", "", -1)
	destAlias = strings.Replace(destAlias, "@", "", -1)
	x.Set("", fmt.Sprintf("--yes rsync --exclude-other-sites --exclude-conf @%v:%%files @%v:%%files", srcAlias, destAlias), true)
	_, err := x.Output()
	if err == nil {
		log.Infoln("Synced public file system.")
	} else {
		log.Warnln("Public file system has not been synced.")
	}
	x.Set("", fmt.Sprintf("--yes rsync --exclude-other-sites --exclude-conf @%v:%%private @%v:%%private", srcAlias, destAlias), true)
	_, err = x.Output()
	if err == nil {
		log.Infoln("Synced private file system.")
	} else {
		log.Warnln("Private file system has not been synced.")
	}
}

// DrushClearCache performs a cache clear task on an input site alias with the use of Drush.
func DrushClearCache(alias string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set(alias, "cc all", false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not clear caches.", err)
	} else {
		log.Infoln("Caches cleared.")
	}
}

// DrushRebuildRegistry performs a registry rebuild task on an input site alias with the use of Drush.
func DrushRebuildRegistry(alias string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set(alias, "rr", false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not rebuild registry.", err)
	} else {
		log.Infoln("Rebuilt registry.")
	}
}

// DrushUpdateDatabase performs a database update task on an input site alias with the use of Drush.
func DrushUpdateDatabase(alias string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set(alias, "updb -y", false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not update database.", err)
	} else {
		log.Infoln("Updated database where possible.")
	}
}

// DrushDownloadToPath performs a database update task on an input site alias with the use of Drush.
func DrushDownloadToPath(path, project string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set("", "pm-download --yes "+project+" --destination="+path, false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not download module ", project, err)
	} else {
		log.Infoln("Downloaded module", project)
	}
}

// DrushDownloadToAlias performs a database update task on an input site alias with the use of Drush.
func DrushDownloadToAlias(alias, project string) {
	drushCommand := NewDrushCommand()
	drushCommand.Set(alias, "pm-download --yes "+project, false)
	_, err := drushCommand.Output()
	if err != nil {
		log.Warnln("Could not download module ", project, err)
	} else {
		log.Infoln("Downloaded module", project)
	}
}

// DrushVariableSet Runs drush vset with a given variable name and value.
func DrushVariableSet(alias, variable_name, variable_value string) {
	srcAlias := strings.Replace(alias, "@", "", -1)
	x := NewDrushCommand()
	x.Set(srcAlias, fmt.Sprintf("vset %v %v", variable_name, variable_value), true)
	drushOut, err := x.Output()
	if err == nil {
		log.Infof("Succesfully set %v to %v via Drush\n", variable_name, variable_value)
	} else {
		log.Errorf("Could not set %v to %v via Drush: %v\n", variable_name, variable_value, drushOut)
	}
}<|endoftext|>


<commit_before>package command

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	plugin "github.com/hashicorp/go-plugin"
	tfplugin "github.com/hashicorp/terraform/plugin"
	"github.com/hashicorp/terraform/plugin/discovery"
	"github.com/hashicorp/terraform/terraform"
	"github.com/kardianos/osext"
)

// multiVersionProviderResolver is an implementation of
// terraform.ResourceProviderResolver that matches the given version constraints
// against a set of versioned provider plugins to find the newest version of
// each that satisfies the given constraints.
type multiVersionProviderResolver struct {
	Available discovery.PluginMetaSet
}

func choosePlugins(avail discovery.PluginMetaSet, reqd discovery.PluginRequirements) map[string]discovery.PluginMeta {
	candidates := avail.ConstrainVersions(reqd)
	ret := map[string]discovery.PluginMeta{}
	for name, metas := range candidates {
		if len(metas) == 0 {
			continue
		}
		ret[name] = metas.Newest()
	}
	return ret
}

func (r *multiVersionProviderResolver) ResolveProviders(
	reqd discovery.PluginRequirements,
) (map[string]terraform.ResourceProviderFactory, []error) {
	factories := make(map[string]terraform.ResourceProviderFactory, len(reqd))
	var errs []error

	chosen := choosePlugins(r.Available, reqd)
	for name, req := range reqd {
		if newest, available := chosen[name]; available {
			digest, err := newest.SHA256()
			if err != nil {
				errs = append(errs, fmt.Errorf("provider.%s: failed to load plugin to verify its signature: %s", name, err))
				continue
			}
			if !reqd[name].AcceptsSHA256(digest) {
				errs = append(errs, fmt.Errorf("provider.%s: new or changed plugin executable", name))
				continue
			}

			client := tfplugin.Client(newest)
			factories[name] = providerFactory(client)
		} else {
			msg := fmt.Sprintf("provider.%s: no suitable version installed", name)

			required := req.Versions.String()
			// no version is unconstrained
			if required == "" {
				required = "(any version)"
			}

			foundVersions := []string{}
			for meta := range r.Available.WithName(name) {
				foundVersions = append(foundVersions, fmt.Sprintf("%q", meta.Version))
			}

			found := "none"
			if len(foundVersions) > 0 {
				found = strings.Join(foundVersions, ", ")
			}

			msg += fmt.Sprintf("\n  version requirements: %q\n  versions installed: %s", required, found)

			errs = append(errs, errors.New(msg))
		}
	}

	return factories, errs
}

// store the user-supplied path for plugin discovery
func (m *Meta) storePluginPath(pluginPath []string) error {
	if len(pluginPath) == 0 {
		return nil
	}

	js, err := json.MarshalIndent(pluginPath, "", "  ")
	if err != nil {
		return err
	}

	// if this fails, so will WriteFile
	os.MkdirAll(m.DataDir(), 0755)

	return ioutil.WriteFile(filepath.Join(m.DataDir(), PluginPathFile), js, 0644)
}

// Load the user-defined plugin search path into Meta.pluginPath if the file
// exists.
func (m *Meta) loadPluginPath() ([]string, error) {
	js, err := ioutil.ReadFile(filepath.Join(m.DataDir(), PluginPathFile))
	if os.IsNotExist(err) {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	var pluginPath []string
	if err := json.Unmarshal(js, &pluginPath); err != nil {
		return nil, err
	}

	return pluginPath, nil
}

// the default location for automatically installed plugins
func (m *Meta) pluginDir() string {
	return filepath.Join(m.DataDir(), "plugins", fmt.Sprintf("%s_%s", runtime.GOOS, runtime.GOARCH))
}

// pluginDirs return a list of directories to search for plugins.
//
// Earlier entries in this slice get priority over later when multiple copies
// of the same plugin version are found, but newer versions always override
// older versions where both satisfy the provider version constraints.
func (m *Meta) pluginDirs(includeAutoInstalled bool) []string {
	// user defined paths take precedence
	if len(m.pluginPath) > 0 {
		return m.pluginPath
	}

	// When searching the following directories, earlier entries get precedence
	// if the same plugin version is found twice, but newer versions will
	// always get preference below regardless of where they are coming from.
	// TODO: Add auto-install dir, default vendor dir and optional override
	// vendor dir(s).
	dirs := []string{"."}

	// Look in the same directory as the Terraform executable.
	// If found, this replaces what we found in the config path.
	exePath, err := osext.Executable()
	if err != nil {
		log.Printf("[ERROR] Error discovering exe directory: %s", err)
	} else {
		dirs = append(dirs, filepath.Dir(exePath))
	}

	// add the user vendor directory
	dirs = append(dirs, DefaultPluginVendorDir)

	if includeAutoInstalled {
		dirs = append(dirs, m.pluginDir())
	}
	dirs = append(dirs, m.GlobalPluginDirs...)

	return dirs
}

// providerPluginSet returns the set of valid providers that were discovered in
// the defined search paths.
func (m *Meta) providerPluginSet() discovery.PluginMetaSet {
	plugins := discovery.FindPlugins("provider", m.pluginDirs(true))
	plugins, _ = plugins.ValidateVersions()

	for p := range plugins {
		log.Printf("[DEBUG] found valid plugin: %q", p.Name)
	}

	return plugins
}

// providerPluginAutoInstalledSet returns the set of providers that exist
// within the auto-install directory.
func (m *Meta) providerPluginAutoInstalledSet() discovery.PluginMetaSet {
	plugins := discovery.FindPlugins("provider", []string{m.pluginDir()})
	plugins, _ = plugins.ValidateVersions()

	for p := range plugins {
		log.Printf("[DEBUG] found valid plugin: %q", p.Name)
	}

	return plugins
}

// providerPluginManuallyInstalledSet returns the set of providers that exist
// in all locations *except* the auto-install directory.
func (m *Meta) providerPluginManuallyInstalledSet() discovery.PluginMetaSet {
	plugins := discovery.FindPlugins("provider", m.pluginDirs(false))
	plugins, _ = plugins.ValidateVersions()

	for p := range plugins {
		log.Printf("[DEBUG] found valid plugin: %q", p.Name)
	}

	return plugins
}

func (m *Meta) providerResolver() terraform.ResourceProviderResolver {
	return &multiVersionProviderResolver{
		Available: m.providerPluginSet(),
	}
}

// filter the requirements returning only the providers that we can't resolve
func (m *Meta) missingPlugins(avail discovery.PluginMetaSet, reqd discovery.PluginRequirements) discovery.PluginRequirements {
	missing := make(discovery.PluginRequirements)

	for n, r := range reqd {
		log.Printf("[DEBUG] plugin requirements: %q=%q", n, r.Versions)
	}

	candidates := avail.ConstrainVersions(reqd)

	for name, versionSet := range reqd {
		if metas := candidates[name]; metas.Count() == 0 {
			missing[name] = versionSet
		}
	}

	return missing
}

func (m *Meta) provisionerFactories() map[string]terraform.ResourceProvisionerFactory {
	dirs := m.pluginDirs(true)
	plugins := discovery.FindPlugins("provisioner", dirs)
	plugins, _ = plugins.ValidateVersions()

	// For now our goal is to just find the latest version of each plugin
	// we have on the system. All provisioners should be at version 0.0.0
	// currently, so there should actually only be one instance of each plugin
	// name here, even though the discovery interface forces us to pretend
	// that might not be true.

	factories := make(map[string]terraform.ResourceProvisionerFactory)

	// Wire up the internal provisioners first. These might be overridden
	// by discovered provisioners below.
	for name := range InternalProvisioners {
		client, err := internalPluginClient("provisioner", name)
		if err != nil {
			log.Printf("[WARN] failed to build command line for internal plugin %q: %s", name, err)
			continue
		}
		factories[name] = provisionerFactory(client)
	}

	byName := plugins.ByName()
	for name, metas := range byName {
		// Since we validated versions above and we partitioned the sets
		// by name, we're guaranteed that the metas in our set all have
		// valid versions and that there's at least one meta.
		newest := metas.Newest()
		client := tfplugin.Client(newest)
		factories[name] = provisionerFactory(client)
	}

	return factories
}

func internalPluginClient(kind, name string) (*plugin.Client, error) {
	cmdLine, err := BuildPluginCommandString(kind, name)
	if err != nil {
		return nil, err
	}

	// See the docstring for BuildPluginCommandString for why we need to do
	// this split here.
	cmdArgv := strings.Split(cmdLine, TFSPACE)

	cfg := &plugin.ClientConfig{
		Cmd:             exec.Command(cmdArgv[0], cmdArgv[1:]...),
		HandshakeConfig: tfplugin.Handshake,
		Managed:         true,
		Plugins:         tfplugin.PluginMap,
	}

	return plugin.NewClient(cfg), nil
}

func providerFactory(client *plugin.Client) terraform.ResourceProviderFactory {
	return func() (terraform.ResourceProvider, error) {
		// Request the RPC client so we can get the provider
		// so we can build the actual RPC-implemented provider.
		rpcClient, err := client.Client()
		if err != nil {
			return nil, err
		}

		raw, err := rpcClient.Dispense(tfplugin.ProviderPluginName)
		if err != nil {
			return nil, err
		}

		return raw.(terraform.ResourceProvider), nil
	}
}

func provisionerFactory(client *plugin.Client) terraform.ResourceProvisionerFactory {
	return func() (terraform.ResourceProvisioner, error) {
		// Request the RPC client so we can get the provisioner
		// so we can build the actual RPC-implemented provisioner.
		rpcClient, err := client.Client()
		if err != nil {
			return nil, err
		}

		raw, err := rpcClient.Dispense(tfplugin.ProvisionerPluginName)
		if err != nil {
			return nil, err
		}

		return raw.(terraform.ResourceProvisioner), nil
	}
}
<commit_msg>load Meta.PluginOverrides in dicovery<commit_after>package command

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	plugin "github.com/hashicorp/go-plugin"
	tfplugin "github.com/hashicorp/terraform/plugin"
	"github.com/hashicorp/terraform/plugin/discovery"
	"github.com/hashicorp/terraform/terraform"
	"github.com/kardianos/osext"
)

// multiVersionProviderResolver is an implementation of
// terraform.ResourceProviderResolver that matches the given version constraints
// against a set of versioned provider plugins to find the newest version of
// each that satisfies the given constraints.
type multiVersionProviderResolver struct {
	Available discovery.PluginMetaSet
}

func choosePlugins(avail discovery.PluginMetaSet, reqd discovery.PluginRequirements) map[string]discovery.PluginMeta {
	candidates := avail.ConstrainVersions(reqd)
	ret := map[string]discovery.PluginMeta{}
	for name, metas := range candidates {
		if len(metas) == 0 {
			continue
		}
		ret[name] = metas.Newest()
	}
	return ret
}

func (r *multiVersionProviderResolver) ResolveProviders(
	reqd discovery.PluginRequirements,
) (map[string]terraform.ResourceProviderFactory, []error) {
	factories := make(map[string]terraform.ResourceProviderFactory, len(reqd))
	var errs []error

	chosen := choosePlugins(r.Available, reqd)
	for name, req := range reqd {
		if newest, available := chosen[name]; available {
			digest, err := newest.SHA256()
			if err != nil {
				errs = append(errs, fmt.Errorf("provider.%s: failed to load plugin to verify its signature: %s", name, err))
				continue
			}
			if !reqd[name].AcceptsSHA256(digest) {
				errs = append(errs, fmt.Errorf("provider.%s: new or changed plugin executable", name))
				continue
			}

			client := tfplugin.Client(newest)
			factories[name] = providerFactory(client)
		} else {
			msg := fmt.Sprintf("provider.%s: no suitable version installed", name)

			required := req.Versions.String()
			// no version is unconstrained
			if required == "" {
				required = "(any version)"
			}

			foundVersions := []string{}
			for meta := range r.Available.WithName(name) {
				foundVersions = append(foundVersions, fmt.Sprintf("%q", meta.Version))
			}

			found := "none"
			if len(foundVersions) > 0 {
				found = strings.Join(foundVersions, ", ")
			}

			msg += fmt.Sprintf("\n  version requirements: %q\n  versions installed: %s", required, found)

			errs = append(errs, errors.New(msg))
		}
	}

	return factories, errs
}

// store the user-supplied path for plugin discovery
func (m *Meta) storePluginPath(pluginPath []string) error {
	if len(pluginPath) == 0 {
		return nil
	}

	js, err := json.MarshalIndent(pluginPath, "", "  ")
	if err != nil {
		return err
	}

	// if this fails, so will WriteFile
	os.MkdirAll(m.DataDir(), 0755)

	return ioutil.WriteFile(filepath.Join(m.DataDir(), PluginPathFile), js, 0644)
}

// Load the user-defined plugin search path into Meta.pluginPath if the file
// exists.
func (m *Meta) loadPluginPath() ([]string, error) {
	js, err := ioutil.ReadFile(filepath.Join(m.DataDir(), PluginPathFile))
	if os.IsNotExist(err) {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	var pluginPath []string
	if err := json.Unmarshal(js, &pluginPath); err != nil {
		return nil, err
	}

	return pluginPath, nil
}

// the default location for automatically installed plugins
func (m *Meta) pluginDir() string {
	return filepath.Join(m.DataDir(), "plugins", fmt.Sprintf("%s_%s", runtime.GOOS, runtime.GOARCH))
}

// pluginDirs return a list of directories to search for plugins.
//
// Earlier entries in this slice get priority over later when multiple copies
// of the same plugin version are found, but newer versions always override
// older versions where both satisfy the provider version constraints.
func (m *Meta) pluginDirs(includeAutoInstalled bool) []string {
	// user defined paths take precedence
	if len(m.pluginPath) > 0 {
		return m.pluginPath
	}

	// When searching the following directories, earlier entries get precedence
	// if the same plugin version is found twice, but newer versions will
	// always get preference below regardless of where they are coming from.
	// TODO: Add auto-install dir, default vendor dir and optional override
	// vendor dir(s).
	dirs := []string{"."}

	// Look in the same directory as the Terraform executable.
	// If found, this replaces what we found in the config path.
	exePath, err := osext.Executable()
	if err != nil {
		log.Printf("[ERROR] Error discovering exe directory: %s", err)
	} else {
		dirs = append(dirs, filepath.Dir(exePath))
	}

	// add the user vendor directory
	dirs = append(dirs, DefaultPluginVendorDir)

	if includeAutoInstalled {
		dirs = append(dirs, m.pluginDir())
	}
	dirs = append(dirs, m.GlobalPluginDirs...)

	return dirs
}

// providerPluginSet returns the set of valid providers that were discovered in
// the defined search paths.
func (m *Meta) providerPluginSet() discovery.PluginMetaSet {
	plugins := discovery.FindPlugins("provider", m.pluginDirs(true))

	// Add providers defined in the legacy .terraformrc,
	if m.PluginOverrides != nil {
		plugins = plugins.OverridePaths(m.PluginOverrides.Providers)
	}

	plugins, _ = plugins.ValidateVersions()

	for p := range plugins {
		log.Printf("[DEBUG] found valid plugin: %q", p.Name)
	}

	return plugins
}

// providerPluginAutoInstalledSet returns the set of providers that exist
// within the auto-install directory.
func (m *Meta) providerPluginAutoInstalledSet() discovery.PluginMetaSet {
	plugins := discovery.FindPlugins("provider", []string{m.pluginDir()})
	plugins, _ = plugins.ValidateVersions()

	for p := range plugins {
		log.Printf("[DEBUG] found valid plugin: %q", p.Name)
	}

	return plugins
}

// providerPluginManuallyInstalledSet returns the set of providers that exist
// in all locations *except* the auto-install directory.
func (m *Meta) providerPluginManuallyInstalledSet() discovery.PluginMetaSet {
	plugins := discovery.FindPlugins("provider", m.pluginDirs(false))

	// Add providers defined in the legacy .terraformrc,
	if m.PluginOverrides != nil {
		plugins = plugins.OverridePaths(m.PluginOverrides.Providers)
	}

	plugins, _ = plugins.ValidateVersions()

	for p := range plugins {
		log.Printf("[DEBUG] found valid plugin: %q", p.Name)
	}

	return plugins
}

func (m *Meta) providerResolver() terraform.ResourceProviderResolver {
	return &multiVersionProviderResolver{
		Available: m.providerPluginSet(),
	}
}

// filter the requirements returning only the providers that we can't resolve
func (m *Meta) missingPlugins(avail discovery.PluginMetaSet, reqd discovery.PluginRequirements) discovery.PluginRequirements {
	missing := make(discovery.PluginRequirements)

	for n, r := range reqd {
		log.Printf("[DEBUG] plugin requirements: %q=%q", n, r.Versions)
	}

	candidates := avail.ConstrainVersions(reqd)

	for name, versionSet := range reqd {
		if metas := candidates[name]; metas.Count() == 0 {
			missing[name] = versionSet
		}
	}

	return missing
}

func (m *Meta) provisionerFactories() map[string]terraform.ResourceProvisionerFactory {
	dirs := m.pluginDirs(true)
	plugins := discovery.FindPlugins("provisioner", dirs)
	plugins, _ = plugins.ValidateVersions()

	// For now our goal is to just find the latest version of each plugin
	// we have on the system. All provisioners should be at version 0.0.0
	// currently, so there should actually only be one instance of each plugin
	// name here, even though the discovery interface forces us to pretend
	// that might not be true.

	factories := make(map[string]terraform.ResourceProvisionerFactory)

	// Wire up the internal provisioners first. These might be overridden
	// by discovered provisioners below.
	for name := range InternalProvisioners {
		client, err := internalPluginClient("provisioner", name)
		if err != nil {
			log.Printf("[WARN] failed to build command line for internal plugin %q: %s", name, err)
			continue
		}
		factories[name] = provisionerFactory(client)
	}

	byName := plugins.ByName()
	for name, metas := range byName {
		// Since we validated versions above and we partitioned the sets
		// by name, we're guaranteed that the metas in our set all have
		// valid versions and that there's at least one meta.
		newest := metas.Newest()
		client := tfplugin.Client(newest)
		factories[name] = provisionerFactory(client)
	}

	return factories
}

func internalPluginClient(kind, name string) (*plugin.Client, error) {
	cmdLine, err := BuildPluginCommandString(kind, name)
	if err != nil {
		return nil, err
	}

	// See the docstring for BuildPluginCommandString for why we need to do
	// this split here.
	cmdArgv := strings.Split(cmdLine, TFSPACE)

	cfg := &plugin.ClientConfig{
		Cmd:             exec.Command(cmdArgv[0], cmdArgv[1:]...),
		HandshakeConfig: tfplugin.Handshake,
		Managed:         true,
		Plugins:         tfplugin.PluginMap,
	}

	return plugin.NewClient(cfg), nil
}

func providerFactory(client *plugin.Client) terraform.ResourceProviderFactory {
	return func() (terraform.ResourceProvider, error) {
		// Request the RPC client so we can get the provider
		// so we can build the actual RPC-implemented provider.
		rpcClient, err := client.Client()
		if err != nil {
			return nil, err
		}

		raw, err := rpcClient.Dispense(tfplugin.ProviderPluginName)
		if err != nil {
			return nil, err
		}

		return raw.(terraform.ResourceProvider), nil
	}
}

func provisionerFactory(client *plugin.Client) terraform.ResourceProvisionerFactory {
	return func() (terraform.ResourceProvisioner, error) {
		// Request the RPC client so we can get the provisioner
		// so we can build the actual RPC-implemented provisioner.
		rpcClient, err := client.Client()
		if err != nil {
			return nil, err
		}

		raw, err := rpcClient.Dispense(tfplugin.ProvisionerPluginName)
		if err != nil {
			return nil, err
		}

		return raw.(terraform.ResourceProvisioner), nil
	}
}
<|endoftext|>


<commit_before>package commands

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/concourse/atc"
	"github.com/concourse/fly/commands/internal/flaghelpers"
	"github.com/concourse/fly/pty"
	"github.com/concourse/fly/rc"
	"github.com/concourse/go-concourse/concourse"
	"github.com/mgutz/ansi"
	"github.com/tedsuo/rata"
	"github.com/vito/go-interact/interact"
)

type HijackCommand struct {
	Job      flaghelpers.JobFlag      `short:"j" long:"job"   value-name:"PIPELINE/JOB"   description:"Name of a job to hijack"`
	Check    flaghelpers.ResourceFlag `short:"c" long:"check" value-name:"PIPELINE/CHECK" description:"Name of a resource's checking container to hijack"`
	Build    string                   `short:"b" long:"build"                             description:"Build number within the job, or global build ID"`
	StepName string                   `short:"s" long:"step"                              description:"Name of step to hijack (e.g. build, unit, resource name)"`
}

func remoteCommand(argv []string) (string, []string) {
	var path string
	var args []string

	switch len(argv) {
	case 0:
		path = "bash"
	case 1:
		path = argv[0]
	default:
		path = argv[0]
		args = argv[1:]
	}

	return path, args
}

type containerLocator interface {
	locate(containerFingerprint) (map[string]string, error)
}

type stepContainerLocator struct {
	client concourse.Client
}

func (locator stepContainerLocator) locate(fingerprint containerFingerprint) (map[string]string, error) {
	reqValues := map[string]string{}

	if fingerprint.jobName != "" {
		reqValues["pipeline_name"] = fingerprint.pipelineName
		reqValues["job_name"] = fingerprint.jobName
		if fingerprint.buildNameOrID != "" {
			reqValues["build_name"] = fingerprint.buildNameOrID
		}
	} else if fingerprint.buildNameOrID != "" {
		reqValues["build-id"] = fingerprint.buildNameOrID
	} else {
		build, err := GetBuild(locator.client, "", "", "")
		if err != nil {
			return reqValues, err
		}
		reqValues["build-id"] = strconv.Itoa(build.ID)
	}
	if fingerprint.stepName != "" {
		reqValues["step_name"] = fingerprint.stepName
	}

	return reqValues, nil
}

type checkContainerLocator struct{}

func (locator checkContainerLocator) locate(fingerprint containerFingerprint) (map[string]string, error) {
	reqValues := map[string]string{}

	reqValues["type"] = "check"
	if fingerprint.checkName != "" {
		reqValues["resource_name"] = fingerprint.checkName
	}
	if fingerprint.pipelineName != "" {
		reqValues["pipeline_name"] = fingerprint.pipelineName
	}

	return reqValues, nil
}

type containerFingerprint struct {
	pipelineName  string
	jobName       string
	buildNameOrID string

	stepName string

	checkName string
}

func locateContainer(client concourse.Client, fingerprint containerFingerprint) (map[string]string, error) {
	var locator containerLocator

	if fingerprint.checkName == "" {
		locator = stepContainerLocator{
			client: client,
		}
	} else {
		locator = checkContainerLocator{}
	}

	return locator.locate(fingerprint)
}

func constructRequest(reqGenerator *rata.RequestGenerator, spec atc.HijackProcessSpec, id string, token *rc.TargetToken) *http.Request {
	payload, err := json.Marshal(spec)
	if err != nil {
		log.Fatalln("failed to marshal process spec:", err)
	}

	hijackReq, err := reqGenerator.CreateRequest(
		atc.HijackContainer,
		rata.Params{"id": id},
		bytes.NewBuffer(payload),
	)
	if err != nil {
		log.Fatalln("failed to create hijack request:", err)
	}

	if token != nil {
		hijackReq.Header.Add("Authorization", token.Type+" "+token.Value)
	}

	return hijackReq
}

func getContainerIDs(c *HijackCommand) []atc.Container {
	var pipelineName string
	if c.Job.PipelineName != "" {
		pipelineName = c.Job.PipelineName
	} else {
		pipelineName = c.Check.PipelineName
	}

	buildNameOrID := c.Build
	stepName := c.StepName
	jobName := c.Job.JobName
	check := c.Check.ResourceName

	fingerprint := containerFingerprint{
		pipelineName:  pipelineName,
		jobName:       jobName,
		buildNameOrID: buildNameOrID,
		stepName:      stepName,
		checkName:     check,
	}

	connection, err := rc.TargetConnection(Fly.Target)
	if err != nil {
		log.Fatalln("failed to create client:", err)
	}
	client := concourse.NewClient(connection)

	reqValues, err := locateContainer(client, fingerprint)
	if err != nil {
		log.Fatalln(err)
	}

	containers, err := client.ListContainers(reqValues)
	if err != nil {
		log.Fatalln("failed to get containers:", err)
	}
	return containers
}

func (command *HijackCommand) Execute(args []string) error {
	target, err := rc.SelectTarget(Fly.Target)
	if err != nil {
		log.Fatalln(err)
		return nil
	}

	containers := getContainerIDs(command)

	var chosenContainer atc.Container
	if len(containers) == 0 {
		fmt.Fprintln(os.Stderr, "no containers matched your search parameters! they may have expired if your build hasn't recently finished")
		os.Exit(1)
	} else if len(containers) > 1 {
		var choices []interact.Choice
		for _, container := range containers {
			var infos []string

			infos = append(infos, fmt.Sprintf("worker: %s", container.WorkerName))
			if container.PipelineName != "" {
				infos = append(infos, fmt.Sprintf("pipeline: %s", container.PipelineName))
			}
			if container.JobName != "" {
				infos = append(infos, fmt.Sprintf("job: %s", container.JobName))
			}
			if container.BuildName != "" {
				infos = append(infos, fmt.Sprintf("build #: %s", container.BuildName))
			}
			if container.BuildID != 0 {
				infos = append(infos, fmt.Sprintf("build id: %d", container.BuildID))
			}

			if container.StepType != "" {
				infos = append(infos, fmt.Sprintf("type: %s", container.StepType))
				infos = append(infos, fmt.Sprintf("name: %s", container.StepName))
			} else {
				infos = append(infos, "type: check")
				infos = append(infos, fmt.Sprintf("name: %s", container.ResourceName))
			}

			choices = append(choices, interact.Choice{
				Display: strings.Join(infos, ", "),
				Value:   container,
			})
		}

		err = interact.NewInteraction("choose a container", choices...).Resolve(&chosenContainer)
		if err == io.EOF {
			os.Exit(0)
		}

		if err != nil {
			return err
		}
	} else {
		chosenContainer = containers[0]
	}

	path, args := remoteCommand(args)
	privileged := true

	reqGenerator := rata.NewRequestGenerator(target.API, atc.Routes)
	tlsConfig := &tls.Config{InsecureSkipVerify: target.Insecure}

	var ttySpec *atc.HijackTTYSpec
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		ttySpec = &atc.HijackTTYSpec{
			WindowSize: atc.HijackWindowSize{
				Columns: cols,
				Rows:    rows,
			},
		}
	}

	envVariables := append(chosenContainer.EnvironmentVariables, "TERM="+os.Getenv("TERM"))

	spec := atc.HijackProcessSpec{
		Path: path,
		Args: args,
		Env:  envVariables,
		User: "root",
		Dir:  chosenContainer.WorkingDirectory,

		Privileged: privileged,
		TTY:        ttySpec,
	}

	hijackReq := constructRequest(reqGenerator, spec, chosenContainer.ID, target.Token)
	hijackResult := performHijack(hijackReq, tlsConfig)
	os.Exit(hijackResult)

	return nil
}

func performHijack(hijackReq *http.Request, tlsConfig *tls.Config) int {
	conn, err := dialEndpoint(hijackReq.URL, tlsConfig)
	if err != nil {
		log.Fatalln("failed to dial hijack endpoint:", err)
	}

	clientConn := httputil.NewClientConn(conn, nil)

	resp, err := clientConn.Do(hijackReq)
	if err != nil {
		log.Fatalln("failed to hijack:", err)
	}

	if resp.StatusCode != http.StatusOK {
		handleBadResponse("hijacking", resp)
	}

	return hijack(clientConn.Hijack())
}

func hijack(conn net.Conn, br *bufio.Reader) int {
	var in io.Reader

	term, err := pty.OpenRawTerm()
	if err == nil {
		defer term.Restore()

		in = term
	} else {
		in = os.Stdin
	}

	encoder := json.NewEncoder(conn)
	decoder := json.NewDecoder(br)

	resized := pty.ResizeNotifier()

	go func() {
		for {
			<-resized
			// TODO json race
			sendSize(encoder)
		}
	}()

	go io.Copy(&stdinWriter{encoder}, in)

	var exitStatus int
	for {
		var output atc.HijackOutput
		err := decoder.Decode(&output)
		if err != nil {
			break
		}

		if output.ExitStatus != nil {
			exitStatus = *output.ExitStatus
		} else if len(output.Error) > 0 {
			fmt.Fprintf(os.Stderr, "%s\n", ansi.Color(output.Error, "red+b"))
			exitStatus = 255
		} else if len(output.Stdout) > 0 {
			os.Stdout.Write(output.Stdout)
		} else if len(output.Stderr) > 0 {
			os.Stderr.Write(output.Stderr)
		}
	}

	return exitStatus
}

func sendSize(enc *json.Encoder) {
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		enc.Encode(atc.HijackInput{
			TTYSpec: &atc.HijackTTYSpec{
				WindowSize: atc.HijackWindowSize{
					Columns: cols,
					Rows:    rows,
				},
			},
		})
	}
}

type stdinWriter struct {
	enc *json.Encoder
}

func (w *stdinWriter) Write(d []byte) (int, error) {
	err := w.enc.Encode(atc.HijackInput{
		Stdin: d,
	})
	if err != nil {
		return 0, err
	}

	return len(d), nil
}

var canonicalPortMap = map[string]string{
	"http":  "80",
	"https": "443",
}

func dialEndpoint(url *url.URL, tlsConfig *tls.Config) (net.Conn, error) {
	addr := canonicalAddr(url)

	if url.Scheme == "https" {
		return tls.Dial("tcp", addr, tlsConfig)
	}

	return net.Dial("tcp", addr)
}

func canonicalAddr(url *url.URL) string {
	host, port, err := net.SplitHostPort(url.Host)
	if err != nil {
		if strings.Contains(err.Error(), "missing port in address") {
			host = url.Host
			port = canonicalPortMap[url.Scheme]
		} else {
			log.Fatalln("invalid host:", err)
		}
	}

	return net.JoinHostPort(host, port)
}
<commit_msg>shorten hijack containers list entries<commit_after>package commands

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/concourse/atc"
	"github.com/concourse/fly/commands/internal/flaghelpers"
	"github.com/concourse/fly/pty"
	"github.com/concourse/fly/rc"
	"github.com/concourse/go-concourse/concourse"
	"github.com/mgutz/ansi"
	"github.com/tedsuo/rata"
	"github.com/vito/go-interact/interact"
)

type HijackCommand struct {
	Job      flaghelpers.JobFlag      `short:"j" long:"job"   value-name:"PIPELINE/JOB"   description:"Name of a job to hijack"`
	Check    flaghelpers.ResourceFlag `short:"c" long:"check" value-name:"PIPELINE/CHECK" description:"Name of a resource's checking container to hijack"`
	Build    string                   `short:"b" long:"build"                             description:"Build number within the job, or global build ID"`
	StepName string                   `short:"s" long:"step"                              description:"Name of step to hijack (e.g. build, unit, resource name)"`
}

func remoteCommand(argv []string) (string, []string) {
	var path string
	var args []string

	switch len(argv) {
	case 0:
		path = "bash"
	case 1:
		path = argv[0]
	default:
		path = argv[0]
		args = argv[1:]
	}

	return path, args
}

type containerLocator interface {
	locate(containerFingerprint) (map[string]string, error)
}

type stepContainerLocator struct {
	client concourse.Client
}

func (locator stepContainerLocator) locate(fingerprint containerFingerprint) (map[string]string, error) {
	reqValues := map[string]string{}

	if fingerprint.jobName != "" {
		reqValues["pipeline_name"] = fingerprint.pipelineName
		reqValues["job_name"] = fingerprint.jobName
		if fingerprint.buildNameOrID != "" {
			reqValues["build_name"] = fingerprint.buildNameOrID
		}
	} else if fingerprint.buildNameOrID != "" {
		reqValues["build-id"] = fingerprint.buildNameOrID
	} else {
		build, err := GetBuild(locator.client, "", "", "")
		if err != nil {
			return reqValues, err
		}
		reqValues["build-id"] = strconv.Itoa(build.ID)
	}
	if fingerprint.stepName != "" {
		reqValues["step_name"] = fingerprint.stepName
	}

	return reqValues, nil
}

type checkContainerLocator struct{}

func (locator checkContainerLocator) locate(fingerprint containerFingerprint) (map[string]string, error) {
	reqValues := map[string]string{}

	reqValues["type"] = "check"
	if fingerprint.checkName != "" {
		reqValues["resource_name"] = fingerprint.checkName
	}
	if fingerprint.pipelineName != "" {
		reqValues["pipeline_name"] = fingerprint.pipelineName
	}

	return reqValues, nil
}

type containerFingerprint struct {
	pipelineName  string
	jobName       string
	buildNameOrID string

	stepName string

	checkName string
}

func locateContainer(client concourse.Client, fingerprint containerFingerprint) (map[string]string, error) {
	var locator containerLocator

	if fingerprint.checkName == "" {
		locator = stepContainerLocator{
			client: client,
		}
	} else {
		locator = checkContainerLocator{}
	}

	return locator.locate(fingerprint)
}

func constructRequest(reqGenerator *rata.RequestGenerator, spec atc.HijackProcessSpec, id string, token *rc.TargetToken) *http.Request {
	payload, err := json.Marshal(spec)
	if err != nil {
		log.Fatalln("failed to marshal process spec:", err)
	}

	hijackReq, err := reqGenerator.CreateRequest(
		atc.HijackContainer,
		rata.Params{"id": id},
		bytes.NewBuffer(payload),
	)
	if err != nil {
		log.Fatalln("failed to create hijack request:", err)
	}

	if token != nil {
		hijackReq.Header.Add("Authorization", token.Type+" "+token.Value)
	}

	return hijackReq
}

func getContainerIDs(c *HijackCommand) []atc.Container {
	var pipelineName string
	if c.Job.PipelineName != "" {
		pipelineName = c.Job.PipelineName
	} else {
		pipelineName = c.Check.PipelineName
	}

	buildNameOrID := c.Build
	stepName := c.StepName
	jobName := c.Job.JobName
	check := c.Check.ResourceName

	fingerprint := containerFingerprint{
		pipelineName:  pipelineName,
		jobName:       jobName,
		buildNameOrID: buildNameOrID,
		stepName:      stepName,
		checkName:     check,
	}

	connection, err := rc.TargetConnection(Fly.Target)
	if err != nil {
		log.Fatalln("failed to create client:", err)
	}
	client := concourse.NewClient(connection)

	reqValues, err := locateContainer(client, fingerprint)
	if err != nil {
		log.Fatalln(err)
	}

	containers, err := client.ListContainers(reqValues)
	if err != nil {
		log.Fatalln("failed to get containers:", err)
	}
	return containers
}

func (command *HijackCommand) Execute(args []string) error {
	target, err := rc.SelectTarget(Fly.Target)
	if err != nil {
		log.Fatalln(err)
		return nil
	}

	containers := getContainerIDs(command)

	var chosenContainer atc.Container
	if len(containers) == 0 {
		fmt.Fprintln(os.Stderr, "no containers matched your search parameters! they may have expired if your build hasn't recently finished")
		os.Exit(1)
	} else if len(containers) > 1 {
		var choices []interact.Choice
		for _, container := range containers {
			var infos []string

			if container.JobName != "" {
				infos = append(infos, fmt.Sprintf("build: #%s", container.BuildName))
			} else {
				infos = append(infos, fmt.Sprintf("build id: %d", container.BuildID))
			}

			if container.StepType != "" {
				infos = append(infos, fmt.Sprintf("step: %s", container.StepName))
				infos = append(infos, fmt.Sprintf("type: %s", container.StepType))
			} else {
				infos = append(infos, fmt.Sprintf("resource: %s", container.ResourceName))
				infos = append(infos, "type: check")
			}

			choices = append(choices, interact.Choice{
				Display: strings.Join(infos, ", "),
				Value:   container,
			})
		}

		err = interact.NewInteraction("choose a container", choices...).Resolve(&chosenContainer)
		if err == io.EOF {
			os.Exit(0)
		}

		if err != nil {
			return err
		}
	} else {
		chosenContainer = containers[0]
	}

	path, args := remoteCommand(args)
	privileged := true

	reqGenerator := rata.NewRequestGenerator(target.API, atc.Routes)
	tlsConfig := &tls.Config{InsecureSkipVerify: target.Insecure}

	var ttySpec *atc.HijackTTYSpec
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		ttySpec = &atc.HijackTTYSpec{
			WindowSize: atc.HijackWindowSize{
				Columns: cols,
				Rows:    rows,
			},
		}
	}

	envVariables := append(chosenContainer.EnvironmentVariables, "TERM="+os.Getenv("TERM"))

	spec := atc.HijackProcessSpec{
		Path: path,
		Args: args,
		Env:  envVariables,
		User: "root",
		Dir:  chosenContainer.WorkingDirectory,

		Privileged: privileged,
		TTY:        ttySpec,
	}

	hijackReq := constructRequest(reqGenerator, spec, chosenContainer.ID, target.Token)
	hijackResult := performHijack(hijackReq, tlsConfig)
	os.Exit(hijackResult)

	return nil
}

func performHijack(hijackReq *http.Request, tlsConfig *tls.Config) int {
	conn, err := dialEndpoint(hijackReq.URL, tlsConfig)
	if err != nil {
		log.Fatalln("failed to dial hijack endpoint:", err)
	}

	clientConn := httputil.NewClientConn(conn, nil)

	resp, err := clientConn.Do(hijackReq)
	if err != nil {
		log.Fatalln("failed to hijack:", err)
	}

	if resp.StatusCode != http.StatusOK {
		handleBadResponse("hijacking", resp)
	}

	return hijack(clientConn.Hijack())
}

func hijack(conn net.Conn, br *bufio.Reader) int {
	var in io.Reader

	term, err := pty.OpenRawTerm()
	if err == nil {
		defer term.Restore()

		in = term
	} else {
		in = os.Stdin
	}

	encoder := json.NewEncoder(conn)
	decoder := json.NewDecoder(br)

	resized := pty.ResizeNotifier()

	go func() {
		for {
			<-resized
			// TODO json race
			sendSize(encoder)
		}
	}()

	go io.Copy(&stdinWriter{encoder}, in)

	var exitStatus int
	for {
		var output atc.HijackOutput
		err := decoder.Decode(&output)
		if err != nil {
			break
		}

		if output.ExitStatus != nil {
			exitStatus = *output.ExitStatus
		} else if len(output.Error) > 0 {
			fmt.Fprintf(os.Stderr, "%s\n", ansi.Color(output.Error, "red+b"))
			exitStatus = 255
		} else if len(output.Stdout) > 0 {
			os.Stdout.Write(output.Stdout)
		} else if len(output.Stderr) > 0 {
			os.Stderr.Write(output.Stderr)
		}
	}

	return exitStatus
}

func sendSize(enc *json.Encoder) {
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		enc.Encode(atc.HijackInput{
			TTYSpec: &atc.HijackTTYSpec{
				WindowSize: atc.HijackWindowSize{
					Columns: cols,
					Rows:    rows,
				},
			},
		})
	}
}

type stdinWriter struct {
	enc *json.Encoder
}

func (w *stdinWriter) Write(d []byte) (int, error) {
	err := w.enc.Encode(atc.HijackInput{
		Stdin: d,
	})
	if err != nil {
		return 0, err
	}

	return len(d), nil
}

var canonicalPortMap = map[string]string{
	"http":  "80",
	"https": "443",
}

func dialEndpoint(url *url.URL, tlsConfig *tls.Config) (net.Conn, error) {
	addr := canonicalAddr(url)

	if url.Scheme == "https" {
		return tls.Dial("tcp", addr, tlsConfig)
	}

	return net.Dial("tcp", addr)
}

func canonicalAddr(url *url.URL) string {
	host, port, err := net.SplitHostPort(url.Host)
	if err != nil {
		if strings.Contains(err.Error(), "missing port in address") {
			host = url.Host
			port = canonicalPortMap[url.Scheme]
		} else {
			log.Fatalln("invalid host:", err)
		}
	}

	return net.JoinHostPort(host, port)
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"net"
	"os"
	"os/signal"

	"github.com/mcuadros/passage/server"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"gopkg.in/fsnotify.v1"
	"gopkg.in/inconshreveable/log15.v2"
)

type ServerCommand struct {
	LogLevel   string
	LogFile    string
	ConfigFile string
	Config     *server.Config
	Server     *server.Server
	RPCAddr    string
	RPCServer  *server.RPCServer

	done chan bool
}

func NewServerCommand() *ServerCommand {
	return &ServerCommand{
		Server: server.NewServer(),
		done:   make(chan bool),
	}
}

func (c *ServerCommand) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "server",
		Short: "runs a passage server, with multiple SSH tunnels",
		RunE:  c.Execute,
	}

	cmd.Flags().StringVar(&c.ConfigFile, "config", "", "config file (default is $HOME/.passage.yaml)")
	cmd.Flags().StringVar(&c.LogFile, "log-file", "", "log file")
	cmd.Flags().StringVar(&c.LogLevel, "log-level", "info", "max log level enabled")
	cmd.Flags().StringVar(&c.RPCAddr, "rpc-addr", rpcAddrDefault, "passage rpc server address, is an unix socket.")
	return cmd
}

func (c *ServerCommand) Execute(cmd *cobra.Command, args []string) error {
	c.handleSignals()
	if err := c.setupLogging(); err != nil {
		return err
	}

	if err := c.setupServer(); err != nil {
		return err
	}

	if err := c.setupRPCServer(); err != nil {
		return err
	}

	<-c.done
	log15.Info("server stopped successfully")
	return nil
}

func (c *ServerCommand) setupServer() error {
	if err := c.readConfig(); err != nil {
		return err
	}

	if err := c.loadConfig(); err != nil {
		return err
	}

	log15.Info("configuration file loaded", "file", viper.ConfigFileUsed())
	viper.OnConfigChange(func(e fsnotify.Event) {
		log15.Info("configuration file re-loaded", "file", e.Name)
		if err := c.loadConfig(); err != nil {
			log15.Error("unable to read/load config", "error", err.Error())
			return
		}
	})

	return nil
}

func (c *ServerCommand) setupRPCServer() error {
	log15.Debug("rpc server started", "addr", c.RPCAddr)
	a, err := net.ResolveUnixAddr("unix", c.RPCAddr)
	if err != nil {
		return err
	}

	c.RPCServer = server.NewRPCServer(c.Server)
	if err := c.RPCServer.Listen(a); err != nil {
		return err
	}

	return nil
}

func (c *ServerCommand) readConfig() error {
	if c.ConfigFile != "" {
		viper.SetConfigFile(c.ConfigFile)
	}

	viper.WatchConfig()
	viper.SetConfigName(".passage")
	viper.AddConfigPath("$HOME")

	return viper.ReadInConfig()
}

func (c *ServerCommand) loadConfig() error {
	if err := viper.Unmarshal(&c.Config); err != nil {
		return err
	}

	if err := c.Config.Validate(); err != nil {
		return err
	}

	if err := c.Server.Load(c.Config); err != nil {
		return err
	}

	return nil
}

func (c *ServerCommand) setupLogging() error {
	lvl, err := log15.LvlFromString(c.LogLevel)
	if err != nil {
		return fmt.Errorf("unknown log level name %q", c.LogLevel)
	}

	handler := log15.StdoutHandler
	format := log15.LogfmtFormat()

	if c.LogFile != "" {
		handler = log15.MultiHandler(handler, log15.Must.FileHandler(c.LogFile, format))
	}

	handler = log15.LvlFilterHandler(lvl, handler)

	if lvl == log15.LvlDebug {
		handler = log15.CallerFileHandler(log15.LvlFilterHandler(lvl, handler))
	}

	log15.Root().SetHandler(handler)
	return nil
}

func (c *ServerCommand) handleSignals() {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)
	go func() {
		<-signals
		if err := c.stop(); err != nil {
			log15.Error("error stopping services", "err", err.Error())
		}
	}()
}

func (c *ServerCommand) stop() error {
	if err := c.Server.Close(); err != nil {
		return err
	}

	if err := c.RPCServer.Close(); err != nil {
		return err
	}

	c.done <- true
	return nil
}
<commit_msg>update fsnotify import<commit_after>package commands

import (
	"fmt"
	"net"
	"os"
	"os/signal"

	"github.com/mcuadros/passage/server"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/fsnotify/fsnotify"
	"gopkg.in/inconshreveable/log15.v2"
)

type ServerCommand struct {
	LogLevel   string
	LogFile    string
	ConfigFile string
	Config     *server.Config
	Server     *server.Server
	RPCAddr    string
	RPCServer  *server.RPCServer

	done chan bool
}

func NewServerCommand() *ServerCommand {
	return &ServerCommand{
		Server: server.NewServer(),
		done:   make(chan bool),
	}
}

func (c *ServerCommand) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "server",
		Short: "runs a passage server, with multiple SSH tunnels",
		RunE:  c.Execute,
	}

	cmd.Flags().StringVar(&c.ConfigFile, "config", "", "config file (default is $HOME/.passage.yaml)")
	cmd.Flags().StringVar(&c.LogFile, "log-file", "", "log file")
	cmd.Flags().StringVar(&c.LogLevel, "log-level", "info", "max log level enabled")
	cmd.Flags().StringVar(&c.RPCAddr, "rpc-addr", rpcAddrDefault, "passage rpc server address, is an unix socket.")
	return cmd
}

func (c *ServerCommand) Execute(cmd *cobra.Command, args []string) error {
	c.handleSignals()
	if err := c.setupLogging(); err != nil {
		return err
	}

	if err := c.setupServer(); err != nil {
		return err
	}

	if err := c.setupRPCServer(); err != nil {
		return err
	}

	<-c.done
	log15.Info("server stopped successfully")
	return nil
}

func (c *ServerCommand) setupServer() error {
	if err := c.readConfig(); err != nil {
		return err
	}

	if err := c.loadConfig(); err != nil {
		return err
	}

	log15.Info("configuration file loaded", "file", viper.ConfigFileUsed())
	viper.OnConfigChange(func(e fsnotify.Event) {
		log15.Info("configuration file re-loaded", "file", e.Name)
		if err := c.loadConfig(); err != nil {
			log15.Error("unable to read/load config", "error", err.Error())
			return
		}
	})

	return nil
}

func (c *ServerCommand) setupRPCServer() error {
	log15.Debug("rpc server started", "addr", c.RPCAddr)
	a, err := net.ResolveUnixAddr("unix", c.RPCAddr)
	if err != nil {
		return err
	}

	c.RPCServer = server.NewRPCServer(c.Server)
	if err := c.RPCServer.Listen(a); err != nil {
		return err
	}

	return nil
}

func (c *ServerCommand) readConfig() error {
	if c.ConfigFile != "" {
		viper.SetConfigFile(c.ConfigFile)
	}

	viper.WatchConfig()
	viper.SetConfigName(".passage")
	viper.AddConfigPath("$HOME")

	return viper.ReadInConfig()
}

func (c *ServerCommand) loadConfig() error {
	if err := viper.Unmarshal(&c.Config); err != nil {
		return err
	}

	if err := c.Config.Validate(); err != nil {
		return err
	}

	if err := c.Server.Load(c.Config); err != nil {
		return err
	}

	return nil
}

func (c *ServerCommand) setupLogging() error {
	lvl, err := log15.LvlFromString(c.LogLevel)
	if err != nil {
		return fmt.Errorf("unknown log level name %q", c.LogLevel)
	}

	handler := log15.StdoutHandler
	format := log15.LogfmtFormat()

	if c.LogFile != "" {
		handler = log15.MultiHandler(handler, log15.Must.FileHandler(c.LogFile, format))
	}

	handler = log15.LvlFilterHandler(lvl, handler)

	if lvl == log15.LvlDebug {
		handler = log15.CallerFileHandler(log15.LvlFilterHandler(lvl, handler))
	}

	log15.Root().SetHandler(handler)
	return nil
}

func (c *ServerCommand) handleSignals() {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)
	go func() {
		<-signals
		if err := c.stop(); err != nil {
			log15.Error("error stopping services", "err", err.Error())
		}
	}()
}

func (c *ServerCommand) stop() error {
	if err := c.Server.Close(); err != nil {
		return err
	}

	if err := c.RPCServer.Close(); err != nil {
		return err
	}

	c.done <- true
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Pagoda Box Inc
//
// This Source Code Form is subject to the terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not distributed with this file, You can obtain one
// at http://mozilla.org/MPL/2.0/.
//

package commands

//
import (
	"fmt"
	"github.com/kardianos/osext"
	"github.com/nanobox-io/nanobox-golang-stylish"
	"github.com/nanobox-io/nanobox/config"
	// "github.com/nanobox-io/nanobox/util/file"
	"github.com/nanobox-io/nanobox/util/print"
	"github.com/spf13/cobra"
	"os"
	"os/exec"
	"runtime"
	"time"
)

//
var updateCmd = &cobra.Command{
	Use:   "update",
	Short: "Updates the CLI to the newest available version",
	Long:  ``,

	Run: update,
}

// update
func update(ccmd *cobra.Command, args []string) {
	Update()
	fmt.Printf("Nanobox is now up to date (running v%s)", config.VERSION)
}

// Update
func Update() {

	// if there is no nanobox.md5 this is likely the first time nanobox is being
	// run and it should update
	if fi, _ := os.Stat(config.Root + "/nanobox.md5"); fi == nil {
		fmt.Printf(stylish.Bullet("Ensuring nanobox is up to date..."))
		runUpdate()
		return
	}

	//
	match, err := Util.MD5sMatch(config.Root+"/nanobox.md5", "https://s3.amazonaws.com/tools.nanobox.io/cli/nanobox.md5")
	if err != nil {
		Config.Fatal("[commands/update] util.MD5sMatch() failed", err.Error())
	}

	// an error here just means the file doesn't exist (which should never happen
	// since it gets created in the config init at startup)
	fi, _ := os.Stat(config.UpdateFile)

	// if the last update was longer ago than our wait time (14 days), and the md5s
	// dont match, then update
	if time.Since(fi.ModTime()).Hours() >= 336.0 && !match {

		//
		switch print.Prompt("Nanobox is out of date, would you like to update it now (y/N)? ") {

		// don't update by default
		default:
			fmt.Println("You can manually update at any time with 'nanobox update'.")
			return

		// if yes continue to update
		case "Yes", "yes", "Y", "y":
			runUpdate()
		}
	}
}

// runUpdate
func runUpdate() {

	fmt.Printf(stylish.Bullet("Updating nanobox..."))

	// get the path of the current executing CLI
	path, err := osext.Executable()
	if err != nil {
		Config.Fatal("[commands/update] osext.ExecutableFolder() failed", err.Error())
	}

	// download the CLI
	cli, err := os.Create(path)
	if err != nil {
		Config.Fatal("[commands/update] os.Create() failed", err.Error())
	}
	defer cli.Close()

	//
	File.Progress(fmt.Sprintf("https://s3.amazonaws.com/tools.nanobox.io/cli/%v/%v/nanobox", runtime.GOOS, runtime.GOARCH), cli)

	//
	// download the CLI md5
	md5, err := os.Create(config.Root + "/nanobox.md5")
	if err != nil {
		Config.Fatal("[commands/update] os.Create() failed", err.Error())
	}
	defer md5.Close()

	//
	File.Download("https://s3.amazonaws.com/tools.nanobox.io/cli/nanobox.md5", md5)

	// if the new CLI fails to execute, just print a generic message and return
	out, err := exec.Command(path, "-v").Output()
	if err != nil {
		fmt.Printf(stylish.SubBullet("[√] Update successful"))
		return
	}

	fmt.Printf(stylish.SubBullet("[√] Now running %s", string(out)))

	// update the modification time of the .update file
	if err := os.Chtimes(config.UpdateFile, time.Now(), time.Now()); err != nil {
		Config.Fatal("[commands.update] os.Chtimes() failed", err.Error())
	}
}
<commit_msg>adding new line at end of cli update<commit_after>// Copyright (c) 2015 Pagoda Box Inc
//
// This Source Code Form is subject to the terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not distributed with this file, You can obtain one
// at http://mozilla.org/MPL/2.0/.
//

package commands

//
import (
	"fmt"
	"github.com/kardianos/osext"
	"github.com/nanobox-io/nanobox-golang-stylish"
	"github.com/nanobox-io/nanobox/config"
	// "github.com/nanobox-io/nanobox/util/file"
	"github.com/nanobox-io/nanobox/util/print"
	"github.com/spf13/cobra"
	"os"
	"os/exec"
	"runtime"
	"time"
)

//
var updateCmd = &cobra.Command{
	Use:   "update",
	Short: "Updates the CLI to the newest available version",
	Long:  ``,

	Run: update,
}

// update
func update(ccmd *cobra.Command, args []string) {
	Update()
	fmt.Printf("Nanobox is now up to date (running v%s)\n", config.VERSION)
}

// Update
func Update() {

	// if there is no nanobox.md5 this is likely the first time nanobox is being
	// run and it should update
	if fi, _ := os.Stat(config.Root + "/nanobox.md5"); fi == nil {
		fmt.Printf(stylish.Bullet("Ensuring nanobox is up to date..."))
		runUpdate()
		return
	}

	//
	match, err := Util.MD5sMatch(config.Root+"/nanobox.md5", "https://s3.amazonaws.com/tools.nanobox.io/cli/nanobox.md5")
	if err != nil {
		Config.Fatal("[commands/update] util.MD5sMatch() failed", err.Error())
	}

	// an error here just means the file doesn't exist (which should never happen
	// since it gets created in the config init at startup)
	fi, _ := os.Stat(config.UpdateFile)

	// if the last update was longer ago than our wait time (14 days), and the md5s
	// dont match, then update
	if time.Since(fi.ModTime()).Hours() >= 336.0 && !match {

		//
		switch print.Prompt("Nanobox is out of date, would you like to update it now (y/N)? ") {

		// don't update by default
		default:
			fmt.Println("You can manually update at any time with 'nanobox update'.")
			return

		// if yes continue to update
		case "Yes", "yes", "Y", "y":
			runUpdate()
		}
	}
}

// runUpdate
func runUpdate() {

	fmt.Printf(stylish.Bullet("Updating nanobox..."))

	// get the path of the current executing CLI
	path, err := osext.Executable()
	if err != nil {
		Config.Fatal("[commands/update] osext.ExecutableFolder() failed", err.Error())
	}

	// download the CLI
	cli, err := os.Create(path)
	if err != nil {
		Config.Fatal("[commands/update] os.Create() failed", err.Error())
	}
	defer cli.Close()

	//
	File.Progress(fmt.Sprintf("https://s3.amazonaws.com/tools.nanobox.io/cli/%v/%v/nanobox", runtime.GOOS, runtime.GOARCH), cli)

	//
	// download the CLI md5
	md5, err := os.Create(config.Root + "/nanobox.md5")
	if err != nil {
		Config.Fatal("[commands/update] os.Create() failed", err.Error())
	}
	defer md5.Close()

	//
	File.Download("https://s3.amazonaws.com/tools.nanobox.io/cli/nanobox.md5", md5)

	// if the new CLI fails to execute, just print a generic message and return
	out, err := exec.Command(path, "-v").Output()
	if err != nil {
		fmt.Printf(stylish.SubBullet("[√] Update successful"))
		return
	}

	fmt.Printf(stylish.SubBullet("[√] Now running %s", string(out)))

	// update the modification time of the .update file
	if err := os.Chtimes(config.UpdateFile, time.Now(), time.Now()); err != nil {
		Config.Fatal("[commands.update] os.Chtimes() failed", err.Error())
	}
}
<|endoftext|>


<commit_before>//
package commands

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	"github.com/kardianos/osext"
	"github.com/nanobox-io/nanobox-golang-stylish"
	"github.com/spf13/cobra"

	"github.com/nanobox-io/nanobox/config"
	fileutil "github.com/nanobox-io/nanobox/util/file"
	printutil "github.com/nanobox-io/nanobox/util/print"
)

var updateCmd = &cobra.Command{
	Use:   "update",
	Short: "Updates the CLI to the newest available version",
	Long:  ``,

	Run: update,
}

// update
func update(ccmd *cobra.Command, args []string) {

	update, err := updateAvailable()
	if err != nil {
		fmt.Println("Unable to determing if updates are available (see log for details).")
		Config.Error("[commands/update] updateAvailable() failed", err.Error())
		return
	}

	// if the md5s don't match or it's been forced, update
	switch {
	case update, config.Force:
		if err := runUpdate(); err != nil {
			if _, ok := err.(*os.LinkError); ok {
				fmt.Println(`Nanobox was unable to update, try again with admin privilege (ex. "sudo nanobox update")`)
			} else {
				Config.Fatal("[commands/update] runUpdate() failed", err.Error())
			}
		}
	default:
		fmt.Printf(stylish.SubBullet("[√] Nanobox is up-to-date"))
	}
}

// Update
func Update() {

	update, err := updateAvailable()
	if err != nil {
		fmt.Println("Unable to determing if updates are available (see log for details).")
		Config.Error("[commands/update] updateAvailable() failed", err.Error())
		return
	}

	// stat the update file to get ModTime(); an error here means the file doesn't
	// exist. This is highly unlikely as the file is created if it doesn't exist
	// each time the CLI is run.
	fi, _ := os.Stat(config.UpdateFile)

	// if the md5s don't match and it's 'time' for an update (14 days), OR a force
	// update is issued, update
	if update && time.Since(fi.ModTime()).Hours() >= 336.0 {

		//
		switch printutil.Prompt("Nanobox is out of date, would you like to update it now (y/N)? ") {

		// don't update by default
		default:
			fmt.Println("You can manually update at any time with 'nanobox update'.")

			// if they don't update, assume then that they'll either do it manually or just
			// wait 14 more days
			if err := touchUpdate(); err != nil {
				fmt.Println("Failed to touch update")
				Config.Error("[commands/update] updateAvailable() failed", err.Error())
			}

			return

		// if yes continue to update
		case "Yes", "yes", "Y", "y":
			if err := runUpdate(); err != nil {
				if _, ok := err.(*os.LinkError); ok {
					fmt.Println(`Nanobox was unable to update, try again with admin privilege (ex. "sudo nanobox update")`)
				} else {
					Config.Fatal("[commands/update] runUpdate() failed", err.Error())
				}
			}
		}
	}
}

// updateAvailable
func updateAvailable() (bool, error) {

	// get the path of the current executing CLI
	exe, err := osext.Executable()
	if err != nil {
		return false, err
	}

	// check the current cli md5 against the remote md5; os.Args[0] is used as the
	// final interpolation to determine standard/dev versions
	md5 := fmt.Sprintf("https://s3.amazonaws.com/tools.nanobox.io/cli/%v/%v/%v.md5", runtime.GOOS, runtime.GOARCH, filepath.Base(os.Args[0]))

	match, err := Util.MD5sMatch(exe, md5)
	if err != nil {
		return false, err
	}

	return !match, nil
}

// runUpdate
func runUpdate() error {

	fmt.Printf(stylish.Bullet("Updating nanobox..."))

	// get the path of the current executing CLI
	exe, err := osext.Executable()
	if err != nil {
		return err
	}

	//

	prog := filepath.Base(os.Args[0])
	tmpDir := config.Root + "/tmp"
	tmpFile := tmpDir + "/" + prog

	// create a tmp dir to download the new cli to; don't care about the error here
	// because if the tmp dir already exists we'll just use it
	os.Mkdir(tmpDir, 0755)

	// create a tmp cli in tmp dir
	cli, err := os.Create(tmpFile)
	if err != nil {
		return err
	}
	defer cli.Close()

	// download the new cli
	dl := fmt.Sprintf("https://s3.amazonaws.com/tools.nanobox.io/cli/%v/%v/%v", runtime.GOOS, runtime.GOARCH, prog)
	fileutil.Progress(dl, cli)

	// make new CLI executable
	if err := os.Chmod(tmpFile, 0755); err != nil {
		return err
	}

	// move new executable over current CLI's location
	if err = os.Rename(tmpFile, exe); err != nil {
		return err
	}

	// if the new CLI fails to execute, just print a generic message and return
	out, err := exec.Command(exe, "-v").Output()
	if err != nil {
		fmt.Printf(stylish.SubBullet("[√] Update successful"))
		return nil
	}

	fmt.Printf(stylish.SubBullet("[√] Now running %s", string(out)))

	// update the .update file
	return touchUpdate()
}

// touchUpdate updates the mod time on the ~/.nanobox/.update file
func touchUpdate() error {
	return os.Chtimes(config.UpdateFile, time.Now(), time.Now())
}
<commit_msg>more message updates<commit_after>//
package commands

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	"github.com/kardianos/osext"
	"github.com/nanobox-io/nanobox-golang-stylish"
	"github.com/spf13/cobra"

	"github.com/nanobox-io/nanobox/config"
	fileutil "github.com/nanobox-io/nanobox/util/file"
	printutil "github.com/nanobox-io/nanobox/util/print"
)

var updateCmd = &cobra.Command{
	Use:   "update",
	Short: "Updates the CLI to the newest available version",
	Long:  ``,

	Run: update,
}

// update
func update(ccmd *cobra.Command, args []string) {

	update, err := updateAvailable()
	if err != nil {
		Config.Error("Unable to determing if updates are available", err.Error())
		return
	}

	// if the md5s don't match or it's been forced, update
	switch {
	case update, config.Force:
		if err := runUpdate(); err != nil {
			if _, ok := err.(*os.LinkError); ok {
				fmt.Println(`Nanobox was unable to update, try again with admin privilege (ex. "sudo nanobox update")`)
			} else {
				Config.Fatal("[commands/update] runUpdate() failed", err.Error())
			}
		}
	default:
		fmt.Printf(stylish.SubBullet("[√] Nanobox is up-to-date"))
	}
}

// Update
func Update() {

	update, err := updateAvailable()
	if err != nil {
		Config.Error("Unable to determing if updates are available.", err.Error())
		return
	}

	// stat the update file to get ModTime(); an error here means the file doesn't
	// exist. This is highly unlikely as the file is created if it doesn't exist
	// each time the CLI is run.
	fi, _ := os.Stat(config.UpdateFile)

	// if the md5s don't match and it's 'time' for an update (14 days), OR a force
	// update is issued, update
	if update && time.Since(fi.ModTime()).Hours() >= 336.0 {

		//
		switch printutil.Prompt("Nanobox is out of date, would you like to update it now (y/N)? ") {

		// don't update by default
		default:
			fmt.Println("You can manually update at any time with 'nanobox update'.")

			// if they don't update, assume then that they'll either do it manually or just
			// wait 14 more days
			if err := touchUpdate(); err != nil {
				Config.Error("Failed to touch update", err.Error())
			}

			return

		// if yes continue to update
		case "Yes", "yes", "Y", "y":
			if err := runUpdate(); err != nil {
				if _, ok := err.(*os.LinkError); ok {
					fmt.Println(`Nanobox was unable to update, try again with admin privilege (ex. "sudo nanobox update")`)
				} else {
					Config.Fatal("[commands/update] runUpdate() failed", err.Error())
				}
			}
		}
	}
}

// updateAvailable
func updateAvailable() (bool, error) {

	// get the path of the current executing CLI
	exe, err := osext.Executable()
	if err != nil {
		return false, err
	}

	// check the current cli md5 against the remote md5; os.Args[0] is used as the
	// final interpolation to determine standard/dev versions
	md5 := fmt.Sprintf("https://s3.amazonaws.com/tools.nanobox.io/cli/%v/%v/%v.md5", runtime.GOOS, runtime.GOARCH, filepath.Base(os.Args[0]))

	match, err := Util.MD5sMatch(exe, md5)
	if err != nil {
		return false, err
	}

	return !match, nil
}

// runUpdate
func runUpdate() error {

	fmt.Printf(stylish.Bullet("Updating nanobox..."))

	// get the path of the current executing CLI
	exe, err := osext.Executable()
	if err != nil {
		return err
	}

	//

	prog := filepath.Base(os.Args[0])
	tmpDir := config.Root + "/tmp"
	tmpFile := tmpDir + "/" + prog

	// create a tmp dir to download the new cli to; don't care about the error here
	// because if the tmp dir already exists we'll just use it
	os.Mkdir(tmpDir, 0755)

	// create a tmp cli in tmp dir
	cli, err := os.Create(tmpFile)
	if err != nil {
		return err
	}
	defer cli.Close()

	// download the new cli
	dl := fmt.Sprintf("https://s3.amazonaws.com/tools.nanobox.io/cli/%v/%v/%v", runtime.GOOS, runtime.GOARCH, prog)
	fileutil.Progress(dl, cli)

	// make new CLI executable
	if err := os.Chmod(tmpFile, 0755); err != nil {
		return err
	}

	// move new executable over current CLI's location
	if err = os.Rename(tmpFile, exe); err != nil {
		return err
	}

	// if the new CLI fails to execute, just print a generic message and return
	out, err := exec.Command(exe, "-v").Output()
	if err != nil {
		fmt.Printf(stylish.SubBullet("[√] Update successful"))
		return nil
	}

	fmt.Printf(stylish.SubBullet("[√] Now running %s", string(out)))

	// update the .update file
	return touchUpdate()
}

// touchUpdate updates the mod time on the ~/.nanobox/.update file
func touchUpdate() error {
	return os.Chtimes(config.UpdateFile, time.Now(), time.Now())
}
<|endoftext|>