<commit_before>// Package cli provides all methods to control command line functions
package cli

import (
	"encoding/json"
	"fmt"
	"github.com/mehrdadrad/mylg/banner"
	"gopkg.in/readline.v1"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"strings"
)

const usage = `Usage:
	The myLG tool developed to troubleshoot networking situations.
	The vi/emacs mode,almost all basic features is supported. press tab to see what options are available.

	connect <provider name>     connects to external looking glass, press tab to see the menu
	node <city/country name>    connects to specific node at current looking glass, press tab to see the available nodes
	local                       back to local
	lg                          change mode to extenal looking glass
	ns                          change mode to name server looking up
	ping                        ping ip address or domain name
	dig                         name server looking up
	whois                       resolve AS number/IP/CIDR to holder (provides by ripe ncc)
	hping                       Ping through HTTP/HTTPS w/ GET/HEAD methods
	scan                        scan tcp ports (you can provide range >scan host minport maxport)
	dump                        prints out a description of the contents of packets on a network interface
	disc                        discover all the devices on a LAN
	peering                     peering information (provides by peeringdb.com)
	web                         web dashboard - opens dashboard at your default browser
	`

// Readline structure
type Readline struct {
	instance  *readline.Instance
	completer *readline.PrefixCompleter
	prompt    string
	next      chan struct{}
}

var (
	// validation command regex
	CMDReg, _ = regexp.Compile(
		`(ping|trace|bgp|lg|ns|dig|dump|disc|whois|peering|scan|hping|connect|node|local|mode|help|web|exit|quit)\s{0,1}(.*)`)
)

// Init set readline imain items
func Init(prompt, version string) *Readline {
	var (
		r         Readline
		err       error
		completer = readline.NewPrefixCompleter(
			readline.PcItem("ping"),
			readline.PcItem("trace"),
			readline.PcItem("bgp"),
			readline.PcItem("hping"),
			readline.PcItem("connect"),
			readline.PcItem("node"),
			readline.PcItem("local"),
			readline.PcItem("lg"),
			readline.PcItem("ns"),
			readline.PcItem("dig"),
			readline.PcItem("whois"),
			readline.PcItem("scan"),
			readline.PcItem("dump"),
			readline.PcItem("disc"),
			readline.PcItem("peering"),
			readline.PcItem("help"),
			readline.PcItem("web"),
			readline.PcItem("exit"),
		)
	)
	r.completer = completer
	r.instance, err = readline.NewEx(&readline.Config{
		Prompt:          prompt + "> ",
		HistoryFile:     "/tmp/myping",
		InterruptPrompt: "^C",
		EOFPrompt:       "exit",
		AutoComplete:    completer,
	})
	if err != nil {
		panic(err)
	}
	banner.Println(version) // print banner
	go checkUpdate(version) // check update version
	r.prompt = prompt       // init local prompt
	return &r
}

// RemoveItemCompleter removes subitem(s) from a specific main item
func (r *Readline) RemoveItemCompleter(pcItem string) {
	child := []readline.PrefixCompleterInterface{}
	for _, p := range r.completer.Children {
		if strings.TrimSpace(string(p.GetName())) != pcItem {
			child = append(child, p)
		}
	}

	r.completer.Children = child

}

// AddCompleter updates subitem(s) from a specific main item
func (r *Readline) AddCompleter(pcItem string, pcSubItems []string) {
	var pc readline.PrefixCompleter
	c := []readline.PrefixCompleterInterface{}
	for _, item := range pcSubItems {
		c = append(c, readline.PcItem(item))
	}
	pc.Name = []rune(pcItem + " ")
	pc.Children = c
	r.completer.Children = append(r.completer.Children, &pc)
}

// UpdateCompleter updates subitem(s) from a specific main item
func (r *Readline) UpdateCompleter(pcItem string, pcSubItems []string) {
	child := []readline.PrefixCompleterInterface{}
	var pc readline.PrefixCompleter
	for _, p := range r.completer.Children {
		if strings.TrimSpace(string(p.GetName())) == pcItem {
			c := []readline.PrefixCompleterInterface{}
			for _, item := range pcSubItems {
				c = append(c, readline.PcItem(item))
			}
			pc.Name = []rune(pcItem + " ")
			pc.Children = c
			child = append(child, &pc)
		} else {
			child = append(child, p)
		}
	}
	if len(pc.Name) < 1 {
		// todo adding new
	}
	r.completer.Children = child
}

// SetPrompt set readline prompt and store it
func (r *Readline) SetPrompt(p string) {
	p = strings.ToLower(p)
	r.prompt = p
	r.instance.SetPrompt(p + "> ")
}

// UpdatePromptN appends readline prompt
func (r *Readline) UpdatePromptN(p string, n int) {
	var parts []string
	p = strings.ToLower(p)
	parts = strings.SplitAfterN(r.prompt, "/", n)
	if n <= len(parts) && n > -1 {
		parts[n-1] = p
		r.prompt = strings.Join(parts, "")
	} else {
		r.prompt += "/" + p
	}
	r.instance.SetPrompt(r.prompt + "> ")
}

// GetPrompt returns the current prompt string
func (r *Readline) GetPrompt() string {
	return r.prompt
}

// Refresh prompt
func (r *Readline) Refresh() {
	r.instance.Refresh()
}

// SetVim set mode to vim
func (r *Readline) SetVim() {
	if !r.instance.IsVimMode() {
		r.instance.SetVimMode(true)
		println("mode changed to vim")
	} else {
		println("mode already is vim")
	}
}

// SetEmacs set mode to emacs
func (r *Readline) SetEmacs() {
	if r.instance.IsVimMode() {
		r.instance.SetVimMode(false)
		println("mode changed to emacs")
	} else {
		println("mode already is emacs")
	}
}

// Next trigers to read next line
func (r *Readline) Next() {
	r.next <- struct{}{}
}

// Run the main loop
func (r *Readline) Run(cmd chan<- string, next chan struct{}) {
	r.next = next
	func() {
		for {
			line, err := r.instance.Readline()
			if err != nil { // io.EOF, readline.ErrInterrupt
				break
			}
			cmd <- line
			if _, ok := <-next; !ok {
				break
			}
		}
	}()
}

// Close the readline instance
func (r *Readline) Close(next chan struct{}) {
	r.instance.Close()
}

// Help print out the main help
func (r *Readline) Help() {
	fmt.Println(usage)
}

// checkUpdate checks if any new version is available
func checkUpdate(version string) {
	type mylg struct {
		Version string
	}
	var appCtl mylg

	if version == "test" {
		return
	}

	resp, err := http.Get("http://mylg.io/appctl/mylg")
	if err != nil {
		println("error: check update has been failed ")
		return
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		println("error: check update has been failed (2)" + err.Error())
		return
	}
	err = json.Unmarshal(body, &appCtl)
	if err != nil {
		println(err.Error())
		return
	}
	if version != appCtl.Version {
		fmt.Printf("New version is available (v%s) at http://mylg.io/download\n", appCtl.Version)
	}
}

//Flag parses the command arguments syntax:
// -flag=x
// -flag x
// help
func Flag(args string) (string, map[string]interface{}) {
	var (
		r   = make(map[string]interface{}, 10)
		err error
	)
	args = strings.TrimSpace(args)
	re := regexp.MustCompile(`(?i)-([a-z]+)={0,1}\s{0,1}([0-9|a-z|\-|'"{}:\/]*)`)
	f := re.FindAllStringSubmatch(args, -1)
	for _, kv := range f {
		if len(kv) > 1 {
			// trim extra characters (' and ") from value
			kv[2] = strings.Trim(kv[2], "'")
			kv[2] = strings.Trim(kv[2], `"`)
			r[kv[1]], err = strconv.Atoi(kv[2])
			if err != nil {
				r[kv[1]] = kv[2]
			}
			args = strings.Replace(args, kv[0], "", -1)
		}
	}
	if m, _ := regexp.MatchString(`(?i)help$`, args); m {
		r["help"] = true
	}
	args = strings.TrimSpace(args)
	return args, r
}

// SetFlag returns command option(s)
func SetFlag(flag map[string]interface{}, option string, v interface{}) interface{} {
	if sValue, ok := flag[option]; ok {
		switch v.(type) {
		case int:
			return sValue.(int)
		case string:
			return sValue.(string)
		case bool:
			return true
		default:
			return sValue.(string)
		}
	} else {
		return v
	}
}
<commit_msg>enhance SetFlag to support int/float as string<commit_after>// Package cli provides all methods to control command line functions
package cli

import (
	"encoding/json"
	"fmt"
	"github.com/mehrdadrad/mylg/banner"
	"gopkg.in/readline.v1"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"strings"
)

const usage = `Usage:
	The myLG tool developed to troubleshoot networking situations.
	The vi/emacs mode,almost all basic features is supported. press tab to see what options are available.

	connect <provider name>     connects to external looking glass, press tab to see the menu
	node <city/country name>    connects to specific node at current looking glass, press tab to see the available nodes
	local                       back to local
	lg                          change mode to extenal looking glass
	ns                          change mode to name server looking up
	ping                        ping ip address or domain name
	dig                         name server looking up
	whois                       resolve AS number/IP/CIDR to holder (provides by ripe ncc)
	hping                       Ping through HTTP/HTTPS w/ GET/HEAD methods
	scan                        scan tcp ports (you can provide range >scan host minport maxport)
	dump                        prints out a description of the contents of packets on a network interface
	disc                        discover all the devices on a LAN
	peering                     peering information (provides by peeringdb.com)
	web                         web dashboard - opens dashboard at your default browser
	`

// Readline structure
type Readline struct {
	instance  *readline.Instance
	completer *readline.PrefixCompleter
	prompt    string
	next      chan struct{}
}

var (
	// validation command regex
	CMDReg, _ = regexp.Compile(
		`(ping|trace|bgp|lg|ns|dig|dump|disc|whois|peering|scan|hping|connect|node|local|mode|help|web|exit|quit)\s{0,1}(.*)`)
)

// Init set readline imain items
func Init(prompt, version string) *Readline {
	var (
		r         Readline
		err       error
		completer = readline.NewPrefixCompleter(
			readline.PcItem("ping"),
			readline.PcItem("trace"),
			readline.PcItem("bgp"),
			readline.PcItem("hping"),
			readline.PcItem("connect"),
			readline.PcItem("node"),
			readline.PcItem("local"),
			readline.PcItem("lg"),
			readline.PcItem("ns"),
			readline.PcItem("dig"),
			readline.PcItem("whois"),
			readline.PcItem("scan"),
			readline.PcItem("dump"),
			readline.PcItem("disc"),
			readline.PcItem("peering"),
			readline.PcItem("help"),
			readline.PcItem("web"),
			readline.PcItem("exit"),
		)
	)
	r.completer = completer
	r.instance, err = readline.NewEx(&readline.Config{
		Prompt:          prompt + "> ",
		HistoryFile:     "/tmp/myping",
		InterruptPrompt: "^C",
		EOFPrompt:       "exit",
		AutoComplete:    completer,
	})
	if err != nil {
		panic(err)
	}
	banner.Println(version) // print banner
	go checkUpdate(version) // check update version
	r.prompt = prompt       // init local prompt
	return &r
}

// RemoveItemCompleter removes subitem(s) from a specific main item
func (r *Readline) RemoveItemCompleter(pcItem string) {
	child := []readline.PrefixCompleterInterface{}
	for _, p := range r.completer.Children {
		if strings.TrimSpace(string(p.GetName())) != pcItem {
			child = append(child, p)
		}
	}

	r.completer.Children = child

}

// AddCompleter updates subitem(s) from a specific main item
func (r *Readline) AddCompleter(pcItem string, pcSubItems []string) {
	var pc readline.PrefixCompleter
	c := []readline.PrefixCompleterInterface{}
	for _, item := range pcSubItems {
		c = append(c, readline.PcItem(item))
	}
	pc.Name = []rune(pcItem + " ")
	pc.Children = c
	r.completer.Children = append(r.completer.Children, &pc)
}

// UpdateCompleter updates subitem(s) from a specific main item
func (r *Readline) UpdateCompleter(pcItem string, pcSubItems []string) {
	child := []readline.PrefixCompleterInterface{}
	var pc readline.PrefixCompleter
	for _, p := range r.completer.Children {
		if strings.TrimSpace(string(p.GetName())) == pcItem {
			c := []readline.PrefixCompleterInterface{}
			for _, item := range pcSubItems {
				c = append(c, readline.PcItem(item))
			}
			pc.Name = []rune(pcItem + " ")
			pc.Children = c
			child = append(child, &pc)
		} else {
			child = append(child, p)
		}
	}
	if len(pc.Name) < 1 {
		// todo adding new
	}
	r.completer.Children = child
}

// SetPrompt set readline prompt and store it
func (r *Readline) SetPrompt(p string) {
	p = strings.ToLower(p)
	r.prompt = p
	r.instance.SetPrompt(p + "> ")
}

// UpdatePromptN appends readline prompt
func (r *Readline) UpdatePromptN(p string, n int) {
	var parts []string
	p = strings.ToLower(p)
	parts = strings.SplitAfterN(r.prompt, "/", n)
	if n <= len(parts) && n > -1 {
		parts[n-1] = p
		r.prompt = strings.Join(parts, "")
	} else {
		r.prompt += "/" + p
	}
	r.instance.SetPrompt(r.prompt + "> ")
}

// GetPrompt returns the current prompt string
func (r *Readline) GetPrompt() string {
	return r.prompt
}

// Refresh prompt
func (r *Readline) Refresh() {
	r.instance.Refresh()
}

// SetVim set mode to vim
func (r *Readline) SetVim() {
	if !r.instance.IsVimMode() {
		r.instance.SetVimMode(true)
		println("mode changed to vim")
	} else {
		println("mode already is vim")
	}
}

// SetEmacs set mode to emacs
func (r *Readline) SetEmacs() {
	if r.instance.IsVimMode() {
		r.instance.SetVimMode(false)
		println("mode changed to emacs")
	} else {
		println("mode already is emacs")
	}
}

// Next trigers to read next line
func (r *Readline) Next() {
	r.next <- struct{}{}
}

// Run the main loop
func (r *Readline) Run(cmd chan<- string, next chan struct{}) {
	r.next = next
	func() {
		for {
			line, err := r.instance.Readline()
			if err != nil { // io.EOF, readline.ErrInterrupt
				break
			}
			cmd <- line
			if _, ok := <-next; !ok {
				break
			}
		}
	}()
}

// Close the readline instance
func (r *Readline) Close(next chan struct{}) {
	r.instance.Close()
}

// Help print out the main help
func (r *Readline) Help() {
	fmt.Println(usage)
}

// checkUpdate checks if any new version is available
func checkUpdate(version string) {
	type mylg struct {
		Version string
	}
	var appCtl mylg

	if version == "test" {
		return
	}

	resp, err := http.Get("http://mylg.io/appctl/mylg")
	if err != nil {
		println("error: check update has been failed ")
		return
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		println("error: check update has been failed (2)" + err.Error())
		return
	}
	err = json.Unmarshal(body, &appCtl)
	if err != nil {
		println(err.Error())
		return
	}
	if version != appCtl.Version {
		fmt.Printf("New version is available (v%s) at http://mylg.io/download\n", appCtl.Version)
	}
}

//Flag parses the command arguments syntax:
// -flag=x
// -flag x
// help
func Flag(args string) (string, map[string]interface{}) {
	var (
		r   = make(map[string]interface{}, 10)
		err error
	)
	args = strings.TrimSpace(args)
	re := regexp.MustCompile(`(?i)-([a-z]+)={0,1}\s{0,1}([0-9|a-z|\-|'"{}:\/]*)`)
	f := re.FindAllStringSubmatch(args, -1)
	for _, kv := range f {
		if len(kv) > 1 {
			// trim extra characters (' and ") from value
			kv[2] = strings.Trim(kv[2], "'")
			kv[2] = strings.Trim(kv[2], `"`)
			r[kv[1]], err = strconv.Atoi(kv[2])
			if err != nil {
				r[kv[1]] = kv[2]
			}
			args = strings.Replace(args, kv[0], "", -1)
		}
	}
	if m, _ := regexp.MatchString(`(?i)help$`, args); m {
		r["help"] = true
	}
	args = strings.TrimSpace(args)
	return args, r
}

// SetFlag returns command option(s)
func SetFlag(flag map[string]interface{}, option string, v interface{}) interface{} {
	if sValue, ok := flag[option]; ok {
		switch v.(type) {
		case int:
			return sValue.(int)
		case string:
			switch sValue.(type) {
			case string:
				return sValue.(string)
			case int:
				str := strconv.FormatInt(int64(sValue.(int)), 10)
				return str
			case float64:
				str := strconv.FormatFloat(sValue.(float64), 'f', -1, 64)
				return str
			}
		case bool:
			return true
		default:
			return sValue.(string)
		}
	}
	return v
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"strings"

	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/jobspec"
	"github.com/hashicorp/nomad/scheduler"
	"github.com/mitchellh/colorstring"
)

type PlanCommand struct {
	Meta
	color *colorstring.Colorize
}

func (c *PlanCommand) Help() string {
	helpText := `
Usage: nomad plan [options] <file>

General Options:

  ` + generalOptionsUsage() + `

Run Options:

  -diff
    Defaults to true, but can be toggled off to omit diff output.

  -no-color
    Disable colored output.
`
	return strings.TrimSpace(helpText)
}

func (c *PlanCommand) Synopsis() string {
	return "Dry-run a job update to determine its effects"
}

func (c *PlanCommand) Run(args []string) int {
	var diff bool

	flags := c.Meta.FlagSet("plan", FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.BoolVar(&diff, "diff", true, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	// Check that we got exactly one job
	args = flags.Args()
	if len(args) != 1 {
		c.Ui.Error(c.Help())
		return 1
	}
	file := args[0]

	// Parse the job file
	job, err := jobspec.ParseFile(file)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error parsing job file %s: %s", file, err))
		return 1
	}

	// Initialize any fields that need to be.
	job.InitFields()

	// Check that the job is valid
	if err := job.Validate(); err != nil {
		c.Ui.Error(fmt.Sprintf("Error validating job: %s", err))
		return 1
	}

	// Convert it to something we can use
	apiJob, err := convertStructJob(job)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error converting job: %s", err))
		return 1
	}

	// Get the HTTP client
	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	// Submit the job
	resp, _, err := client.Jobs().Plan(apiJob, diff, nil)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error during plan: %s", err))
		return 1
	}

	if diff {
		c.Ui.Output(c.Colorize().Color(strings.TrimSpace(formatJobDiff(resp.Diff))))
	}

	return 0
}

func formatJobDiff(job *api.JobDiff) string {
	out := fmt.Sprintf("%s[bold]Job: %q\n", getDiffString(job.Type), job.ID)

	if job.Type == "Edited" {
		for _, field := range job.Fields {
			out += fmt.Sprintf("%s\n", formatFieldDiff(field, ""))
		}

		for _, object := range job.Objects {
			out += fmt.Sprintf("%s\n", formatObjectDiff(object, ""))
		}
	}

	for _, tg := range job.TaskGroups {
		out += fmt.Sprintf("%s\n", formatTaskGroupDiff(tg))
	}

	return out
}

func formatTaskGroupDiff(tg *api.TaskGroupDiff) string {
	out := fmt.Sprintf("%s[bold]Task Group: %q", getDiffString(tg.Type), tg.Name)

	// Append the updates
	if l := len(tg.Updates); l > 0 {
		updates := make([]string, 0, l)
		for updateType, count := range tg.Updates {
			var color string
			switch updateType {
			case scheduler.UpdateTypeIgnore:
			case scheduler.UpdateTypeCreate:
				color = "[green]"
			case scheduler.UpdateTypeDestroy:
				color = "[red]"
			case scheduler.UpdateTypeMigrate:
				color = "[blue]"
			case scheduler.UpdateTypeInplaceUpdate:
				color = "[light_yellow]"
			case scheduler.UpdateTypeDestructiveUpdate:
				color = "[yellow]"
			}
			updates = append(updates, fmt.Sprintf("[reset]%s%d %s", color, count, updateType))
		}
		out += fmt.Sprintf(" (%s[reset])\n", strings.Join(updates, ", "))
	} else {
		out += "[reset]\n"
	}

	if tg.Type == "Edited" {
		for _, field := range tg.Fields {
			out += fmt.Sprintf("%s\n", formatFieldDiff(field, "  "))
		}

		for _, object := range tg.Objects {
			out += fmt.Sprintf("%s\n", formatObjectDiff(object, "  "))
		}
	}

	for _, task := range tg.Tasks {
		out += fmt.Sprintf("%s\n", formatTaskDiff(task))
	}

	return out
}

func formatTaskDiff(task *api.TaskDiff) string {
	out := fmt.Sprintf("  %s[bold]Task: %q", getDiffString(task.Type), task.Name)
	if len(task.Annotations) != 0 {
		out += fmt.Sprintf(" [reset](%s)", strings.Join(task.Annotations, ", "))
	}

	if task.Type != "Edited" {
		return out
	} else {
		out += "\n"
	}

	if task.Type == "Edited" {
		for _, field := range task.Fields {
			out += fmt.Sprintf("%s\n", formatFieldDiff(field, "    "))
		}

		for _, object := range task.Objects {
			out += fmt.Sprintf("%s\n", formatObjectDiff(object, "    "))
		}
	}

	return out
}

func formatFieldDiff(diff *api.FieldDiff, prefix string) string {
	out := prefix
	switch diff.Type {
	case "Added":
		out += fmt.Sprintf("%s%s: %q", getDiffString(diff.Type), diff.Name, diff.New)
	case "Deleted":
		out += fmt.Sprintf("%s%s: %q", getDiffString(diff.Type), diff.Name, diff.Old)
	case "Edited":
		out += fmt.Sprintf("%s%s: %q => %q", getDiffString(diff.Type), diff.Name, diff.Old, diff.New)
	default:
		out += fmt.Sprintf("%s: %q", diff.Name, diff.New)
	}

	if len(diff.Annotations) != 0 {
		out += fmt.Sprintf(" (%s)", strings.Join(diff.Annotations, ", "))
	}

	return out
}

func formatObjectDiff(diff *api.ObjectDiff, prefix string) string {
	out := fmt.Sprintf("%s%s%s {\n", prefix, getDiffString(diff.Type), diff.Name)

	newPrefix := prefix + "  "
	numFields := len(diff.Fields)
	numObjects := len(diff.Objects)
	haveObjects := numObjects != 0
	for i, field := range diff.Fields {
		out += formatFieldDiff(field, newPrefix)
		if i+1 != numFields || haveObjects {
			out += "\n"
		}
	}

	for i, object := range diff.Objects {
		out += formatObjectDiff(object, newPrefix)
		if i+1 != numObjects {
			out += "\n"
		}
	}

	return fmt.Sprintf("%s\n%s}", out, prefix)
}

func getDiffString(diffType string) string {
	switch diffType {
	case "Added":
		return "[green]+[reset] "
	case "Deleted":
		return "[red]-[reset] "
	case "Edited":
		return "[light_yellow]+/-[reset] "
	default:
		return ""
	}
}
<commit_msg>verbose output<commit_after>package command

import (
	"fmt"
	"strings"

	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/jobspec"
	"github.com/hashicorp/nomad/scheduler"
	"github.com/mitchellh/colorstring"
)

type PlanCommand struct {
	Meta
	color *colorstring.Colorize
}

func (c *PlanCommand) Help() string {
	helpText := `
Usage: nomad plan [options] <file>

General Options:

  ` + generalOptionsUsage() + `

Run Options:

  -diff
    Defaults to true, but can be toggled off to omit diff output.

  -no-color
    Disable colored output.

  -verbose
    Increased diff verbosity

`
	return strings.TrimSpace(helpText)
}

func (c *PlanCommand) Synopsis() string {
	return "Dry-run a job update to determine its effects"
}

func (c *PlanCommand) Run(args []string) int {
	var diff, verbose bool

	flags := c.Meta.FlagSet("plan", FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.BoolVar(&diff, "diff", true, "")
	flags.BoolVar(&verbose, "verbose", false, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	// Check that we got exactly one job
	args = flags.Args()
	if len(args) != 1 {
		c.Ui.Error(c.Help())
		return 1
	}
	file := args[0]

	// Parse the job file
	job, err := jobspec.ParseFile(file)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error parsing job file %s: %s", file, err))
		return 1
	}

	// Initialize any fields that need to be.
	job.InitFields()

	// Check that the job is valid
	if err := job.Validate(); err != nil {
		c.Ui.Error(fmt.Sprintf("Error validating job: %s", err))
		return 1
	}

	// Convert it to something we can use
	apiJob, err := convertStructJob(job)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error converting job: %s", err))
		return 1
	}

	// Get the HTTP client
	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	// Submit the job
	resp, _, err := client.Jobs().Plan(apiJob, diff, nil)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error during plan: %s", err))
		return 1
	}

	if diff {
		c.Ui.Output(c.Colorize().Color(strings.TrimSpace(formatJobDiff(resp.Diff, verbose))))
	}

	return 0
}

func formatJobDiff(job *api.JobDiff, verbose bool) string {
	out := fmt.Sprintf("%s[bold]Job: %q\n", getDiffString(job.Type), job.ID)

	if job.Type == "Edited" || verbose {
		for _, field := range job.Fields {
			out += fmt.Sprintf("%s\n", formatFieldDiff(field, "", verbose))
		}

		for _, object := range job.Objects {
			out += fmt.Sprintf("%s\n", formatObjectDiff(object, "", verbose))
		}
	}

	for _, tg := range job.TaskGroups {
		out += fmt.Sprintf("%s\n", formatTaskGroupDiff(tg, verbose))
	}

	return out
}

func formatTaskGroupDiff(tg *api.TaskGroupDiff, verbose bool) string {
	out := fmt.Sprintf("%s[bold]Task Group: %q", getDiffString(tg.Type), tg.Name)

	// Append the updates
	if l := len(tg.Updates); l > 0 {
		updates := make([]string, 0, l)
		for updateType, count := range tg.Updates {
			var color string
			switch updateType {
			case scheduler.UpdateTypeIgnore:
			case scheduler.UpdateTypeCreate:
				color = "[green]"
			case scheduler.UpdateTypeDestroy:
				color = "[red]"
			case scheduler.UpdateTypeMigrate:
				color = "[blue]"
			case scheduler.UpdateTypeInplaceUpdate:
				color = "[cyan]"
			case scheduler.UpdateTypeDestructiveUpdate:
				color = "[yellow]"
			}
			updates = append(updates, fmt.Sprintf("[reset]%s%d %s", color, count, updateType))
		}
		out += fmt.Sprintf(" (%s[reset])\n", strings.Join(updates, ", "))
	} else {
		out += "[reset]\n"
	}

	if tg.Type == "Edited" || verbose {
		for _, field := range tg.Fields {
			out += fmt.Sprintf("%s\n", formatFieldDiff(field, "  ", verbose))
		}

		for _, object := range tg.Objects {
			out += fmt.Sprintf("%s\n", formatObjectDiff(object, "  ", verbose))
		}
	}

	for _, task := range tg.Tasks {
		out += fmt.Sprintf("%s\n", formatTaskDiff(task, verbose))
	}

	return out
}

func formatTaskDiff(task *api.TaskDiff, verbose bool) string {
	out := fmt.Sprintf("  %s[bold]Task: %q", getDiffString(task.Type), task.Name)
	if len(task.Annotations) != 0 {
		out += fmt.Sprintf(" [reset](%s)", strings.Join(task.Annotations, ", "))
	}

	if task.Type == "None" {
		return out
	} else if (task.Type == "Deleted" || task.Type == "Added") && !verbose {
		return out
	} else {
		out += "\n"
	}

	for _, field := range task.Fields {
		out += fmt.Sprintf("%s\n", formatFieldDiff(field, "    ", verbose))
	}

	for _, object := range task.Objects {
		out += fmt.Sprintf("%s\n", formatObjectDiff(object, "    ", verbose))
	}

	return out
}

func formatFieldDiff(diff *api.FieldDiff, prefix string, verbose bool) string {
	out := prefix
	switch diff.Type {
	case "Added":
		out += fmt.Sprintf("%s%s: %q", getDiffString(diff.Type), diff.Name, diff.New)
	case "Deleted":
		out += fmt.Sprintf("%s%s: %q", getDiffString(diff.Type), diff.Name, diff.Old)
	case "Edited":
		out += fmt.Sprintf("%s%s: %q => %q", getDiffString(diff.Type), diff.Name, diff.Old, diff.New)
	default:
		out += fmt.Sprintf("%s: %q", diff.Name, diff.New)
	}

	if len(diff.Annotations) != 0 {
		out += fmt.Sprintf(" (%s)", strings.Join(diff.Annotations, ", "))
	}

	return out
}

func formatObjectDiff(diff *api.ObjectDiff, prefix string, verbose bool) string {
	out := fmt.Sprintf("%s%s%s {\n", prefix, getDiffString(diff.Type), diff.Name)

	newPrefix := prefix + "  "
	numFields := len(diff.Fields)
	numObjects := len(diff.Objects)
	haveObjects := numObjects != 0
	for i, field := range diff.Fields {
		out += formatFieldDiff(field, newPrefix, verbose)
		if i+1 != numFields || haveObjects {
			out += "\n"
		}
	}

	for i, object := range diff.Objects {
		out += formatObjectDiff(object, newPrefix, verbose)
		if i+1 != numObjects {
			out += "\n"
		}
	}

	return fmt.Sprintf("%s\n%s}", out, prefix)
}

func getDiffString(diffType string) string {
	switch diffType {
	case "Added":
		return "[green]+[reset] "
	case "Deleted":
		return "[red]-[reset] "
	case "Edited":
		return "[light_yellow]+/-[reset] "
	default:
		return ""
	}
}
<|endoftext|>


<commit_before>package commands

import (
	"bufio"
	"errors"
	"fmt"
	"golang.org/x/sys/unix"
	"os"
	fp "path/filepath"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/docker/docker/pkg/fileutils"
	"github.com/docker/docker/pkg/mount"

	"github.com/yuuki1/droot/errwrap"
	"github.com/yuuki1/droot/log"
	"github.com/yuuki1/droot/osutil"
)

var CommandArgRun = "--root ROOT_DIR [--user USER] [--group GROUP] [--bind SRC-PATH[:DEST-PATH]] [--robind SRC-PATH[:DEST-PATH]] [--no-dropcaps] -- COMMAND"
var CommandRun = cli.Command{
	Name:   "run",
	Usage:  "Run an extracted docker image from s3",
	Action: fatalOnError(doRun),
	Flags: []cli.Flag{
		cli.StringFlag{Name: "root, r", Usage: "Root directory path for chrooting"},
		cli.StringFlag{Name: "user, u", Usage: "User (ID or name) to switch before running the program"},
		cli.StringFlag{Name: "group, g", Usage: "Group (ID or name) to switch to"},
		cli.StringSliceFlag{
			Name:  "bind, b",
			Value: &cli.StringSlice{},
			Usage: "Bind mount directory (can be specifies multiple times)",
		},
		cli.StringSliceFlag{
			Name:  "robind",
			Value: &cli.StringSlice{},
			Usage: "Readonly bind mount directory (can be specifies multiple times)",
		},
		cli.BoolFlag{
			Name:  "copy-files, cp",
			Usage: "Copy host files to container such as /etc/group, /etc/passwd, /etc/resolv.conf, /etc/hosts",
		},
		cli.BoolFlag{Name: "no-dropcaps", Usage: "Provide COMMAND's process in chroot with root permission (dangerous)"},
	},
}

var copyFiles = []string{
	"etc/group",
	"etc/passwd",
	"etc/resolv.conf",
	"etc/hosts",
}

var keepCaps = map[uint]bool{
	0:  true, // CAP_CHOWN
	1:  true, // CAP_DAC_OVERRIDE
	2:  true, // CAP_DAC_READ_SEARCH
	3:  true, // CAP_FOWNER
	6:  true, // CAP_SETGID
	7:  true, // CAP_SETUID
	10: true, // CAP_NET_BIND_SERVICE
}

func doRun(c *cli.Context) error {
	command := c.Args()
	if len(command) < 1 {
		cli.ShowCommandHelp(c, "run")
		return errors.New("command required")
	}

	rootDir := c.String("root")
	if rootDir == "" {
		cli.ShowCommandHelp(c, "run")
		return errors.New("--root option required")
	}

	if !osutil.ExistsDir(rootDir) {
		return fmt.Errorf("No such directory %s:", rootDir)
	}

	var err error
	uid, gid := os.Getuid(), os.Getgid()

	if group := c.String("group"); group != "" {
		if gid, err = osutil.LookupGroup(group); err != nil {
			return fmt.Errorf("Failed to lookup group: %s", err)
		}
	}
	if user := c.String("user"); user != "" {
		if uid, err = osutil.LookupUser(user); err != nil {
			return fmt.Errorf("Failed to lookup user: %s", err)
		}
	}

	// copy files
	if c.Bool("copy-files") {
		for _, f := range copyFiles {
			srcFile, destFile := fp.Join("/", f), fp.Join(rootDir, f)
			if err := osutil.Cp(srcFile, destFile); err != nil {
				return fmt.Errorf("Failed to copy %s: %s", f, err)
			}
			if err := os.Lchown(destFile, uid, gid); err != nil {
				return fmt.Errorf("Failed to lchown %s: %s", f, err)
			}
		}
	}

	// mount -t proc none {{rootDir}}/proc
	if err := mount.Mount("none", fp.Join(rootDir, "/proc"), "proc", ""); err != nil {
		return fmt.Errorf("Failed to mount /proc: %s", err)
	}
	// mount --rbind /sys {{rootDir}}/sys
	if err := mount.Mount("/sys", fp.Join(rootDir, "/sys"), "none", "rbind"); err != nil {
		return fmt.Errorf("Failed to mount /sys: %s", err)
	}

	for _, dir := range c.StringSlice("bind") {
		if err := bindMount(dir, rootDir, false); err != nil {
			return fmt.Errorf("Failed to bind mount %s: %s", dir, err)
		}
	}
	for _, dir := range c.StringSlice("robind") {
		if err := bindMount(dir, rootDir, true); err != nil {
			return fmt.Errorf("Failed to robind mount %s: %s", dir, err)
		}
	}

	// create symlinks
	if err := osutil.Symlink("../run/lock", fp.Join(rootDir, "/var/lock")); err != nil {
		return fmt.Errorf("Failed to symlink lock file: %s", err)
	}

	if err := createDevices(rootDir, uid, gid); err != nil {
		return fmt.Errorf("Failed to create devices: %s", err)
	}

	log.Debug("chroot", rootDir, command)

	if err := unix.Chroot(rootDir); err != nil {
		return fmt.Errorf("Failed to chroot: %s", err)
	}
	if err := unix.Chdir("/"); err != nil {
		return fmt.Errorf("Failed to chdir /: %s", err)
	}

	if !c.Bool("no-dropcaps") {
		log.Debug("drop capabilities")
		if err := osutil.DropCapabilities(keepCaps); err != nil {
			return fmt.Errorf("Failed to drop capabilities: %s", err)
		}
	}

	log.Debug("setgid", gid)
	if err := osutil.Setgid(gid); err != nil {
		return fmt.Errorf("Failed to set group %d: %s", gid, err)
	}
	log.Debug("setuid", uid)
	if err := osutil.Setuid(uid); err != nil {
		return fmt.Errorf("Failed to set user %d: %s", uid, err)
	}

	env, err := getEnvironFromEnvFile("/.drootenv")
	if err != nil {
		return fmt.Errorf("Failed to environ from '/.drootenv'")
	}

	return osutil.Execv(command[0], command[0:], env)
}

func bindMount(bindDir string, rootDir string, readonly bool) error {
	var srcDir, destDir string

	d := strings.SplitN(bindDir, ":", 2)
	if len(d) < 2 {
		srcDir = d[0]
	} else {
		srcDir, destDir = d[0], d[1]
	}
	if destDir == "" {
		destDir = srcDir
	}

	ok, err := osutil.IsDirEmpty(srcDir)
	if err != nil {
		return err
	}
	if ok {
		if _, err := os.Create(fp.Join(srcDir, ".droot.keep")); err != nil {
			return errwrap.Wrapf(err, "Failed to create .droot.keep: {{err}}")
		}
	}

	containerDir := fp.Join(rootDir, destDir)

	if err := fileutils.CreateIfNotExists(containerDir, true); err != nil { // mkdir -p
		return errwrap.Wrapff(err, "Failed to create directory: %s: {{err}}", containerDir)
	}

	ok, err = osutil.IsDirEmpty(containerDir)
	if err != nil {
		return err
	}
	if ok {
		log.Debug("bind mount", bindDir, "to", containerDir)
		if err := mount.Mount(srcDir, containerDir, "none", "bind,rw"); err != nil {
			return errwrap.Wrapff(err, "Failed to bind mount %s: {{err}}", containerDir)
		}

		if readonly {
			log.Debug("robind mount", bindDir, "to", containerDir)
			if err := mount.Mount(srcDir, containerDir, "none", "remount,ro,bind"); err != nil {
				return errwrap.Wrapff(err, "Failed to robind mount %s: {{err}}", containerDir)
			}
		}
	}

	return nil
}

func createDevices(rootDir string, uid, gid int) error {
	nullDir := fp.Join(rootDir, os.DevNull)
	if err := osutil.Mknod(nullDir, unix.S_IFCHR|uint32(os.FileMode(0666)), 1*256+3); err != nil {
		return err
	}

	if err := os.Lchown(nullDir, uid, gid); err != nil {
		return errwrap.Wrapff(err, "Failed to lchown %s: {{err}}", nullDir)
	}

	zeroDir := fp.Join(rootDir, "/dev/zero")
	if err := osutil.Mknod(zeroDir, unix.S_IFCHR|uint32(os.FileMode(0666)), 1*256+3); err != nil {
		return err
	}

	if err := os.Lchown(zeroDir, uid, gid); err != nil {
		return errwrap.Wrapff(err, "Failed to lchown %s:", zeroDir)
	}

	for _, f := range []string{"/dev/random", "/dev/urandom"} {
		randomDir := fp.Join(rootDir, f)
		if err := osutil.Mknod(randomDir, unix.S_IFCHR|uint32(os.FileMode(0666)), 1*256+9); err != nil {
			return err
		}

		if err := os.Lchown(randomDir, uid, gid); err != nil {
			return errwrap.Wrapff(err, "Failed to lchown %s: {{err}}", randomDir)
		}
	}

	return nil
}

func getEnvironFromEnvFile(filename string) ([]string, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var env []string
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		l := strings.Trim(scanner.Text(), " \n\t")
		if len(l) == 0 {
			continue
		}
		if len(strings.Split(l, "=")) != 2 { // line should be `key=value`
			continue
		}
		env = append(env, l)
	}

	return env, nil
}
<commit_msg>Fix bug that run command failed when .drootenv doesn't exist<commit_after>package commands

import (
	"bufio"
	"errors"
	"fmt"
	"golang.org/x/sys/unix"
	"os"
	fp "path/filepath"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/docker/docker/pkg/fileutils"
	"github.com/docker/docker/pkg/mount"

	"github.com/yuuki1/droot/errwrap"
	"github.com/yuuki1/droot/log"
	"github.com/yuuki1/droot/osutil"
)

var CommandArgRun = "--root ROOT_DIR [--user USER] [--group GROUP] [--bind SRC-PATH[:DEST-PATH]] [--robind SRC-PATH[:DEST-PATH]] [--no-dropcaps] -- COMMAND"
var CommandRun = cli.Command{
	Name:   "run",
	Usage:  "Run an extracted docker image from s3",
	Action: fatalOnError(doRun),
	Flags: []cli.Flag{
		cli.StringFlag{Name: "root, r", Usage: "Root directory path for chrooting"},
		cli.StringFlag{Name: "user, u", Usage: "User (ID or name) to switch before running the program"},
		cli.StringFlag{Name: "group, g", Usage: "Group (ID or name) to switch to"},
		cli.StringSliceFlag{
			Name:  "bind, b",
			Value: &cli.StringSlice{},
			Usage: "Bind mount directory (can be specifies multiple times)",
		},
		cli.StringSliceFlag{
			Name:  "robind",
			Value: &cli.StringSlice{},
			Usage: "Readonly bind mount directory (can be specifies multiple times)",
		},
		cli.BoolFlag{
			Name:  "copy-files, cp",
			Usage: "Copy host files to container such as /etc/group, /etc/passwd, /etc/resolv.conf, /etc/hosts",
		},
		cli.BoolFlag{Name: "no-dropcaps", Usage: "Provide COMMAND's process in chroot with root permission (dangerous)"},
	},
}

var copyFiles = []string{
	"etc/group",
	"etc/passwd",
	"etc/resolv.conf",
	"etc/hosts",
}

var keepCaps = map[uint]bool{
	0:  true, // CAP_CHOWN
	1:  true, // CAP_DAC_OVERRIDE
	2:  true, // CAP_DAC_READ_SEARCH
	3:  true, // CAP_FOWNER
	6:  true, // CAP_SETGID
	7:  true, // CAP_SETUID
	10: true, // CAP_NET_BIND_SERVICE
}

func doRun(c *cli.Context) error {
	command := c.Args()
	if len(command) < 1 {
		cli.ShowCommandHelp(c, "run")
		return errors.New("command required")
	}

	rootDir := c.String("root")
	if rootDir == "" {
		cli.ShowCommandHelp(c, "run")
		return errors.New("--root option required")
	}

	if !osutil.ExistsDir(rootDir) {
		return fmt.Errorf("No such directory %s:", rootDir)
	}

	var err error
	uid, gid := os.Getuid(), os.Getgid()

	if group := c.String("group"); group != "" {
		if gid, err = osutil.LookupGroup(group); err != nil {
			return fmt.Errorf("Failed to lookup group: %s", err)
		}
	}
	if user := c.String("user"); user != "" {
		if uid, err = osutil.LookupUser(user); err != nil {
			return fmt.Errorf("Failed to lookup user: %s", err)
		}
	}

	// copy files
	if c.Bool("copy-files") {
		for _, f := range copyFiles {
			srcFile, destFile := fp.Join("/", f), fp.Join(rootDir, f)
			if err := osutil.Cp(srcFile, destFile); err != nil {
				return fmt.Errorf("Failed to copy %s: %s", f, err)
			}
			if err := os.Lchown(destFile, uid, gid); err != nil {
				return fmt.Errorf("Failed to lchown %s: %s", f, err)
			}
		}
	}

	// mount -t proc none {{rootDir}}/proc
	if err := mount.Mount("none", fp.Join(rootDir, "/proc"), "proc", ""); err != nil {
		return fmt.Errorf("Failed to mount /proc: %s", err)
	}
	// mount --rbind /sys {{rootDir}}/sys
	if err := mount.Mount("/sys", fp.Join(rootDir, "/sys"), "none", "rbind"); err != nil {
		return fmt.Errorf("Failed to mount /sys: %s", err)
	}

	for _, dir := range c.StringSlice("bind") {
		if err := bindMount(dir, rootDir, false); err != nil {
			return fmt.Errorf("Failed to bind mount %s: %s", dir, err)
		}
	}
	for _, dir := range c.StringSlice("robind") {
		if err := bindMount(dir, rootDir, true); err != nil {
			return fmt.Errorf("Failed to robind mount %s: %s", dir, err)
		}
	}

	// create symlinks
	if err := osutil.Symlink("../run/lock", fp.Join(rootDir, "/var/lock")); err != nil {
		return fmt.Errorf("Failed to symlink lock file: %s", err)
	}

	if err := createDevices(rootDir, uid, gid); err != nil {
		return fmt.Errorf("Failed to create devices: %s", err)
	}

	log.Debug("chroot", rootDir, command)

	if err := unix.Chroot(rootDir); err != nil {
		return fmt.Errorf("Failed to chroot: %s", err)
	}
	if err := unix.Chdir("/"); err != nil {
		return fmt.Errorf("Failed to chdir /: %s", err)
	}

	if !c.Bool("no-dropcaps") {
		log.Debug("drop capabilities")
		if err := osutil.DropCapabilities(keepCaps); err != nil {
			return fmt.Errorf("Failed to drop capabilities: %s", err)
		}
	}

	log.Debug("setgid", gid)
	if err := osutil.Setgid(gid); err != nil {
		return fmt.Errorf("Failed to set group %d: %s", gid, err)
	}
	log.Debug("setuid", uid)
	if err := osutil.Setuid(uid); err != nil {
		return fmt.Errorf("Failed to set user %d: %s", uid, err)
	}

	var env []string
	if osutil.ExistsFile("/.drootenv") {
		env, err = getEnvironFromEnvFile("/.drootenv")
		if err != nil {
			return fmt.Errorf("Failed to read environ from '/.drootenv'")
		}
	}
	return osutil.Execv(command[0], command[0:], env)
}

func bindMount(bindDir string, rootDir string, readonly bool) error {
	var srcDir, destDir string

	d := strings.SplitN(bindDir, ":", 2)
	if len(d) < 2 {
		srcDir = d[0]
	} else {
		srcDir, destDir = d[0], d[1]
	}
	if destDir == "" {
		destDir = srcDir
	}

	ok, err := osutil.IsDirEmpty(srcDir)
	if err != nil {
		return err
	}
	if ok {
		if _, err := os.Create(fp.Join(srcDir, ".droot.keep")); err != nil {
			return errwrap.Wrapf(err, "Failed to create .droot.keep: {{err}}")
		}
	}

	containerDir := fp.Join(rootDir, destDir)

	if err := fileutils.CreateIfNotExists(containerDir, true); err != nil { // mkdir -p
		return errwrap.Wrapff(err, "Failed to create directory: %s: {{err}}", containerDir)
	}

	ok, err = osutil.IsDirEmpty(containerDir)
	if err != nil {
		return err
	}
	if ok {
		log.Debug("bind mount", bindDir, "to", containerDir)
		if err := mount.Mount(srcDir, containerDir, "none", "bind,rw"); err != nil {
			return errwrap.Wrapff(err, "Failed to bind mount %s: {{err}}", containerDir)
		}

		if readonly {
			log.Debug("robind mount", bindDir, "to", containerDir)
			if err := mount.Mount(srcDir, containerDir, "none", "remount,ro,bind"); err != nil {
				return errwrap.Wrapff(err, "Failed to robind mount %s: {{err}}", containerDir)
			}
		}
	}

	return nil
}

func createDevices(rootDir string, uid, gid int) error {
	nullDir := fp.Join(rootDir, os.DevNull)
	if err := osutil.Mknod(nullDir, unix.S_IFCHR|uint32(os.FileMode(0666)), 1*256+3); err != nil {
		return err
	}

	if err := os.Lchown(nullDir, uid, gid); err != nil {
		return errwrap.Wrapff(err, "Failed to lchown %s: {{err}}", nullDir)
	}

	zeroDir := fp.Join(rootDir, "/dev/zero")
	if err := osutil.Mknod(zeroDir, unix.S_IFCHR|uint32(os.FileMode(0666)), 1*256+3); err != nil {
		return err
	}

	if err := os.Lchown(zeroDir, uid, gid); err != nil {
		return errwrap.Wrapff(err, "Failed to lchown %s:", zeroDir)
	}

	for _, f := range []string{"/dev/random", "/dev/urandom"} {
		randomDir := fp.Join(rootDir, f)
		if err := osutil.Mknod(randomDir, unix.S_IFCHR|uint32(os.FileMode(0666)), 1*256+9); err != nil {
			return err
		}

		if err := os.Lchown(randomDir, uid, gid); err != nil {
			return errwrap.Wrapff(err, "Failed to lchown %s: {{err}}", randomDir)
		}
	}

	return nil
}

func getEnvironFromEnvFile(filename string) ([]string, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var env []string
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		l := strings.Trim(scanner.Text(), " \n\t")
		if len(l) == 0 {
			continue
		}
		if len(strings.Split(l, "=")) != 2 { // line should be `key=value`
			continue
		}
		env = append(env, l)
	}

	return env, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 WALLIX

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package commands

import (
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/spf13/cobra"
	"github.com/wallix/awless/aws"
	"github.com/wallix/awless/aws/config"
	"github.com/wallix/awless/cloud"
	"github.com/wallix/awless/cloud/properties"
	"github.com/wallix/awless/config"
	"github.com/wallix/awless/console"
	"github.com/wallix/awless/graph"
	"github.com/wallix/awless/logger"
	"github.com/wallix/awless/ssh"
)

var keyPathFlag string
var printSSHConfigFlag bool
var printSSHCLIFlag bool
var privateIPFlag bool
var disableStrictHostKeyCheckingFlag bool

func init() {
	RootCmd.AddCommand(sshCmd)
	sshCmd.Flags().StringVarP(&keyPathFlag, "identity", "i", "", "Set path or name toward the identity (key file) to use to connect through SSH")
	sshCmd.Flags().BoolVar(&printSSHConfigFlag, "print-config", false, "Print SSH configuration for ~/.ssh/config file.")
	sshCmd.Flags().BoolVar(&printSSHCLIFlag, "print-cli", false, "Print the CLI one-liner to connect with SSH. (/usr/bin/ssh user@ip -i ...)")
	sshCmd.Flags().BoolVar(&privateIPFlag, "private", false, "Use private ip to connect to host")
	sshCmd.Flags().BoolVar(&disableStrictHostKeyCheckingFlag, "disable-strict-host-keychecking", false, "Disable the remote host key check from ~/.ssh/known_hosts or ~/.awless/known_hosts file")
}

var sshCmd = &cobra.Command{
	Use:   "ssh [USER@]INSTANCE",
	Short: "Launch a SSH (Secure Shell) session to an instance given an id or alias",
	Example: `  awless ssh i-8d43b21b                       # using the instance id
  awless ssh ec2-user@redis-prod              # using the instance name and specify a user
  awless ssh @redis-prod -i keyname # using a key stored in ~/.ssh/keyname.pem
  awless ssh @redis-prod -i ./path/toward/key # with a keyfile`,
	PersistentPreRun:  applyHooks(initLoggerHook, initAwlessEnvHook, initCloudServicesHook),
	PersistentPostRun: applyHooks(verifyNewVersionHook),

	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) != 1 {
			return fmt.Errorf("instance required")
		}

		connectionCtx, err := initInstanceConnectionContext(args[0], keyPathFlag)
		exitOn(err)

		client, err := ssh.InitClient(connectionCtx.keypath,
			[]string{config.KeysDir, filepath.Join(os.Getenv("HOME"), ".ssh")},
			disableStrictHostKeyCheckingFlag)
		if err != nil && strings.Contains(err.Error(), "cannot find SSH key") && keyPathFlag == "" {
			logger.Info("you may want to specify a key filepath with `-i /path/to/key.pem`")
		}
		exitOn(err)

		client.SetLogger(logger.DefaultLogger)
		client.InteractiveTerminalFunc = console.InteractiveTerminal
		client.IP = connectionCtx.ip

		if connectionCtx.user != "" {
			client, err = client.DialWithUsers(connectionCtx.user)
		} else {
			client, err = client.DialWithUsers(awsconfig.DefaultAMIUsers...)
		}

		if isConnectionRefusedErr(err) {
			logger.Warning("cannot connect to this instance, maybe the system is still booting?")
			return nil
		}

		exitOn(err)

		exitOn(connectionCtx.checkInstanceAccessible())

		if printSSHConfigFlag {
			fmt.Println(client.SSHConfigString(connectionCtx.instanceName))
			return nil
		}

		if printSSHCLIFlag {
			fmt.Println(client.ConnectString())
			return nil
		}

		exitOn(client.Connect())
		return nil
	},
}

func instanceCredentialsFromGraph(g *graph.Graph, inst *graph.Resource, keyFlag string) (keypath string, ip string, err error) {
	if ip, err = getIP(inst); err != nil {
		return
	}

	if keyFlag != "" {
		keypath = keyFlag
	} else {
		keypair, ok := inst.Properties[properties.KeyPair]
		if !ok {
			err = fmt.Errorf("no access key set for instance %s", inst.Id())
			return
		}
		keypath = fmt.Sprint(keypair)
	}

	return
}

func isConnectionRefusedErr(err error) bool {
	return err != nil && strings.Contains(err.Error(), "connection refused")
}

func getIP(inst *graph.Resource) (string, error) {
	var ipKeyType string

	if privateIPFlag {
		ipKeyType = properties.PrivateIP
	} else {
		ipKeyType = properties.PublicIP
	}

	ip, ok := inst.Properties[ipKeyType]

	if !ok {
		return "", fmt.Errorf("no IP address for instance %s", inst.Id())
	}

	return fmt.Sprint(ip), nil
}

type instanceConnectionContext struct {
	user           string
	ip             string
	keypath        string
	instanceName   string
	instance       *graph.Resource
	resourcesGraph *graph.Graph
	myip           net.IP
}

func initInstanceConnectionContext(userhost, keypath string) (*instanceConnectionContext, error) {
	ctx := &instanceConnectionContext{}

	if strings.Contains(userhost, "@") {
		ctx.user = strings.Split(userhost, "@")[0]
		ctx.instanceName = strings.Split(userhost, "@")[1]
	} else {
		ctx.instanceName = userhost
	}

	ctx.fetchConnectionInfo()

	instanceResolvers := []graph.Resolver{&graph.ByProperty{Key: "Name", Value: ctx.instanceName}, &graph.ByType{Typ: cloud.Instance}}
	resources, err := ctx.resourcesGraph.ResolveResources(&graph.And{Resolvers: instanceResolvers})
	exitOn(err)
	switch len(resources) {
	case 0:
		// No instance with that name, use the id
		ctx.instance, err = findResource(ctx.resourcesGraph, ctx.instanceName, cloud.Instance)
		exitOn(err)
	case 1:
		ctx.instance = resources[0]
	default:
		idStatus := graph.Resources(resources).Map(func(r *graph.Resource) string {
			return fmt.Sprintf("%s (%s)", r.Id(), r.Properties[properties.State])
		})
		logger.Infof("Found %d resources with name '%s': %s", len(resources), ctx.instanceName, strings.Join(idStatus, ", "))

		var running []*graph.Resource
		running, err = ctx.resourcesGraph.ResolveResources(&graph.And{Resolvers: append(instanceResolvers, &graph.ByProperty{Key: properties.State, Value: "running"})})
		exitOn(err)

		switch len(running) {
		case 0:
			logger.Warning("None of them is running, cannot connect through SSH")
			return ctx, errors.New("non running instances")
		case 1:
			logger.Infof("Found only one instance running: %s. Will connect to this instance.", running[0].Id())
			ctx.instance = running[0]
		default:
			logger.Warning("Connect through the running ones using their id:")
			for _, res := range running {
				var up string
				if uptime, ok := res.Properties[properties.Launched].(time.Time); ok {
					up = fmt.Sprintf("\t\t(uptime: %s)", console.HumanizeTime(uptime))
				}
				logger.Warningf("\t`awless ssh %s`%s", res.Id(), up)
			}
			return ctx, errors.New("use instances ids")
		}
	}

	keypath, IP, err := instanceCredentialsFromGraph(ctx.resourcesGraph, ctx.instance, keypath)
	if err != nil {
		return nil, err
	}

	ctx.keypath = keypath
	ctx.ip = IP

	return ctx, nil
}

func (ctx *instanceConnectionContext) fetchConnectionInfo() {
	var resourcesGraph, sgroupsGraph *graph.Graph
	var myip net.IP
	var wg sync.WaitGroup
	var errc = make(chan error)

	wg.Add(1)
	go func() {
		var err error
		defer wg.Done()
		resourcesGraph, err = aws.InfraService.FetchByType(cloud.Instance)
		if err != nil {
			errc <- err
		}
	}()

	wg.Add(1)
	go func() {
		var err error
		defer wg.Done()
		sgroupsGraph, err = aws.InfraService.FetchByType(cloud.SecurityGroup)
		if err != nil {
			errc <- err
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		myip = getMyIP()
	}()
	go func() {
		wg.Wait()
		close(errc)
	}()
	for err := range errc {
		if err != nil {
			exitOn(err)
		}
	}
	resourcesGraph.AddGraph(sgroupsGraph)

	ctx.resourcesGraph = resourcesGraph
	ctx.myip = myip
	return
}

func (ctx *instanceConnectionContext) checkInstanceAccessible() (err error) {
	state, ok := ctx.instance.Properties[properties.State]
	if st := fmt.Sprint(state); ok && st != "running" {
		logger.Warningf("this instance is '%s' (cannot ssh to a non running state)", st)
		if st == "stopped" {
			logger.Warningf("you can start it with `awless -f start instance id=%s`", ctx.instance.Id())
		}
		return errors.New("instance not accessible")
	}

	sgroups, ok := ctx.instance.Properties[properties.SecurityGroups].([]string)
	if ok {
		var sshPortOpen, myIPAllowed bool
		for _, id := range sgroups {
			var sgroup *graph.Resource
			sgroup, err = findResource(ctx.resourcesGraph, id, cloud.SecurityGroup)
			if err != nil {
				break
			}

			rules, ok := sgroup.Properties[properties.InboundRules].([]*graph.FirewallRule)
			if ok {
				for _, r := range rules {
					if r.PortRange.Contains(22) {
						sshPortOpen = true
					}
					if ctx.myip != nil && r.Contains(ctx.myip.String()) {
						myIPAllowed = true
					}
				}
			}
		}

		if !sshPortOpen {
			logger.Warning("port 22 is not open on this instance")
			return errors.New("instance not accessible")
		}

		if !myIPAllowed && ctx.myip != nil {
			logger.Warningf("your ip %s is not authorized for this instance. You might want to update the securitygroup with:", ctx.myip)
			var group = "mygroup"
			if len(sgroups) == 1 {
				group = sgroups[0]
			}
			logger.Warningf("`awless update securitygroup id=%s inbound=authorize protocol=tcp cidr=%s/32 portrange=22`", group, ctx.myip)
			return errors.New("instance not accessible")
		}
	}

	return nil
}

func findResource(g *graph.Graph, id, typ string) (*graph.Resource, error) {
	if found, err := g.FindResource(id); found == nil || err != nil {
		return nil, fmt.Errorf("instance '%s' not found", id)
	}

	return g.GetResource(typ, id)
}
<commit_msg>Half fix #100: inst. access diagnostic should not... ...  bail out command<commit_after>/*
Copyright 2017 WALLIX

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package commands

import (
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/spf13/cobra"
	"github.com/wallix/awless/aws"
	"github.com/wallix/awless/aws/config"
	"github.com/wallix/awless/cloud"
	"github.com/wallix/awless/cloud/properties"
	"github.com/wallix/awless/config"
	"github.com/wallix/awless/console"
	"github.com/wallix/awless/graph"
	"github.com/wallix/awless/logger"
	"github.com/wallix/awless/ssh"
)

var keyPathFlag string
var printSSHConfigFlag bool
var printSSHCLIFlag bool
var privateIPFlag bool
var disableStrictHostKeyCheckingFlag bool

func init() {
	RootCmd.AddCommand(sshCmd)
	sshCmd.Flags().StringVarP(&keyPathFlag, "identity", "i", "", "Set path or name toward the identity (key file) to use to connect through SSH")
	sshCmd.Flags().BoolVar(&printSSHConfigFlag, "print-config", false, "Print SSH configuration for ~/.ssh/config file.")
	sshCmd.Flags().BoolVar(&printSSHCLIFlag, "print-cli", false, "Print the CLI one-liner to connect with SSH. (/usr/bin/ssh user@ip -i ...)")
	sshCmd.Flags().BoolVar(&privateIPFlag, "private", false, "Use private ip to connect to host")
	sshCmd.Flags().BoolVar(&disableStrictHostKeyCheckingFlag, "disable-strict-host-keychecking", false, "Disable the remote host key check from ~/.ssh/known_hosts or ~/.awless/known_hosts file")
}

var sshCmd = &cobra.Command{
	Use:   "ssh [USER@]INSTANCE",
	Short: "Launch a SSH (Secure Shell) session to an instance given an id or alias",
	Example: `  awless ssh i-8d43b21b                       # using the instance id
  awless ssh ec2-user@redis-prod              # using the instance name and specify a user
  awless ssh @redis-prod -i keyname # using a key stored in ~/.ssh/keyname.pem
  awless ssh @redis-prod -i ./path/toward/key # with a keyfile`,
	PersistentPreRun:  applyHooks(initLoggerHook, initAwlessEnvHook, initCloudServicesHook),
	PersistentPostRun: applyHooks(verifyNewVersionHook),

	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) != 1 {
			return fmt.Errorf("instance required")
		}

		connectionCtx, err := initInstanceConnectionContext(args[0], keyPathFlag)
		exitOn(err)

		client, err := ssh.InitClient(connectionCtx.keypath,
			[]string{config.KeysDir, filepath.Join(os.Getenv("HOME"), ".ssh")},
			disableStrictHostKeyCheckingFlag)
		if err != nil && strings.Contains(err.Error(), "cannot find SSH key") && keyPathFlag == "" {
			logger.Info("you may want to specify a key filepath with `-i /path/to/key.pem`")
		}
		exitOn(err)

		client.SetLogger(logger.DefaultLogger)
		client.InteractiveTerminalFunc = console.InteractiveTerminal
		client.IP = connectionCtx.ip

		if connectionCtx.user != "" {
			client, err = client.DialWithUsers(connectionCtx.user)
		} else {
			client, err = client.DialWithUsers(awsconfig.DefaultAMIUsers...)
		}

		if isConnectionRefusedErr(err) {
			logger.Warning("cannot connect to this instance, maybe the system is still booting?")
			return nil
		}

		if err != nil {
			if e := connectionCtx.checkInstanceAccessible(); e != nil {
				logger.Error(e.Error())
			}
			exitOn(err)
		}

		if printSSHConfigFlag {
			fmt.Println(client.SSHConfigString(connectionCtx.instanceName))
			return nil
		}

		if printSSHCLIFlag {
			fmt.Println(client.ConnectString())
			return nil
		}

		exitOn(client.Connect())
		return nil
	},
}

func instanceCredentialsFromGraph(g *graph.Graph, inst *graph.Resource, keyFlag string) (keypath string, ip string, err error) {
	if ip, err = getIP(inst); err != nil {
		return
	}

	if keyFlag != "" {
		keypath = keyFlag
	} else {
		keypair, ok := inst.Properties[properties.KeyPair]
		if !ok {
			err = fmt.Errorf("no access key set for instance %s", inst.Id())
			return
		}
		keypath = fmt.Sprint(keypair)
	}

	return
}

func isConnectionRefusedErr(err error) bool {
	return err != nil && strings.Contains(err.Error(), "connection refused")
}

func getIP(inst *graph.Resource) (string, error) {
	var ipKeyType string

	if privateIPFlag {
		ipKeyType = properties.PrivateIP
	} else {
		ipKeyType = properties.PublicIP
	}

	ip, ok := inst.Properties[ipKeyType]

	if !ok {
		return "", fmt.Errorf("no IP address for instance %s", inst.Id())
	}

	return fmt.Sprint(ip), nil
}

type instanceConnectionContext struct {
	user           string
	ip             string
	keypath        string
	instanceName   string
	instance       *graph.Resource
	resourcesGraph *graph.Graph
	myip           net.IP
}

func initInstanceConnectionContext(userhost, keypath string) (*instanceConnectionContext, error) {
	ctx := &instanceConnectionContext{}

	if strings.Contains(userhost, "@") {
		ctx.user = strings.Split(userhost, "@")[0]
		ctx.instanceName = strings.Split(userhost, "@")[1]
	} else {
		ctx.instanceName = userhost
	}

	ctx.fetchConnectionInfo()

	instanceResolvers := []graph.Resolver{&graph.ByProperty{Key: "Name", Value: ctx.instanceName}, &graph.ByType{Typ: cloud.Instance}}
	resources, err := ctx.resourcesGraph.ResolveResources(&graph.And{Resolvers: instanceResolvers})
	exitOn(err)
	switch len(resources) {
	case 0:
		// No instance with that name, use the id
		ctx.instance, err = findResource(ctx.resourcesGraph, ctx.instanceName, cloud.Instance)
		exitOn(err)
	case 1:
		ctx.instance = resources[0]
	default:
		idStatus := graph.Resources(resources).Map(func(r *graph.Resource) string {
			return fmt.Sprintf("%s (%s)", r.Id(), r.Properties[properties.State])
		})
		logger.Infof("Found %d resources with name '%s': %s", len(resources), ctx.instanceName, strings.Join(idStatus, ", "))

		var running []*graph.Resource
		running, err = ctx.resourcesGraph.ResolveResources(&graph.And{Resolvers: append(instanceResolvers, &graph.ByProperty{Key: properties.State, Value: "running"})})
		exitOn(err)

		switch len(running) {
		case 0:
			logger.Warning("None of them is running, cannot connect through SSH")
			return ctx, errors.New("non running instances")
		case 1:
			logger.Infof("Found only one instance running: %s. Will connect to this instance.", running[0].Id())
			ctx.instance = running[0]
		default:
			logger.Warning("Connect through the running ones using their id:")
			for _, res := range running {
				var up string
				if uptime, ok := res.Properties[properties.Launched].(time.Time); ok {
					up = fmt.Sprintf("\t\t(uptime: %s)", console.HumanizeTime(uptime))
				}
				logger.Warningf("\t`awless ssh %s`%s", res.Id(), up)
			}
			return ctx, errors.New("use instances ids")
		}
	}

	keypath, IP, err := instanceCredentialsFromGraph(ctx.resourcesGraph, ctx.instance, keypath)
	if err != nil {
		return nil, err
	}

	ctx.keypath = keypath
	ctx.ip = IP

	return ctx, nil
}

func (ctx *instanceConnectionContext) fetchConnectionInfo() {
	var resourcesGraph, sgroupsGraph *graph.Graph
	var myip net.IP
	var wg sync.WaitGroup
	var errc = make(chan error)

	wg.Add(1)
	go func() {
		var err error
		defer wg.Done()
		resourcesGraph, err = aws.InfraService.FetchByType(cloud.Instance)
		if err != nil {
			errc <- err
		}
	}()

	wg.Add(1)
	go func() {
		var err error
		defer wg.Done()
		sgroupsGraph, err = aws.InfraService.FetchByType(cloud.SecurityGroup)
		if err != nil {
			errc <- err
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		myip = getMyIP()
	}()
	go func() {
		wg.Wait()
		close(errc)
	}()
	for err := range errc {
		if err != nil {
			exitOn(err)
		}
	}
	resourcesGraph.AddGraph(sgroupsGraph)

	ctx.resourcesGraph = resourcesGraph
	ctx.myip = myip
	return
}

func (ctx *instanceConnectionContext) checkInstanceAccessible() (err error) {
	state, ok := ctx.instance.Properties[properties.State]
	if st := fmt.Sprint(state); ok && st != "running" {
		logger.Warningf("this instance is '%s' (cannot ssh to a non running state)", st)
		if st == "stopped" {
			logger.Warningf("you can start it with `awless -f start instance id=%s`", ctx.instance.Id())
		}
		return errors.New("instance not accessible")
	}

	sgroups, ok := ctx.instance.Properties[properties.SecurityGroups].([]string)
	if ok {
		var sshPortOpen, myIPAllowed bool
		for _, id := range sgroups {
			var sgroup *graph.Resource
			sgroup, err = findResource(ctx.resourcesGraph, id, cloud.SecurityGroup)
			if err != nil {
				break
			}

			rules, ok := sgroup.Properties[properties.InboundRules].([]*graph.FirewallRule)
			if ok {
				for _, r := range rules {
					if r.PortRange.Contains(22) {
						sshPortOpen = true
					}
					if ctx.myip != nil && r.Contains(ctx.myip.String()) {
						myIPAllowed = true
					}
				}
			}
		}

		if !sshPortOpen {
			logger.Warning("port 22 is not open on this instance")
			return errors.New("instance not accessible")
		}

		if !myIPAllowed && ctx.myip != nil {
			logger.Warningf("your ip %s is not authorized for this instance. You might want to update the securitygroup with:", ctx.myip)
			var group = "mygroup"
			if len(sgroups) == 1 {
				group = sgroups[0]
			}
			logger.Warningf("`awless update securitygroup id=%s inbound=authorize protocol=tcp cidr=%s/32 portrange=22`", group, ctx.myip)
			return errors.New("instance not accessible")
		}
	}

	return nil
}

func findResource(g *graph.Graph, id, typ string) (*graph.Resource, error) {
	if found, err := g.FindResource(id); found == nil || err != nil {
		return nil, fmt.Errorf("instance '%s' not found", id)
	}

	return g.GetResource(typ, id)
}
<|endoftext|>


<commit_before>package cmdline

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/disorganizer/brig"
	"github.com/disorganizer/brig/daemon"
	"github.com/disorganizer/brig/repo"
	"github.com/disorganizer/brig/repo/config"
	"github.com/disorganizer/brig/util/colors"
	colorlog "github.com/disorganizer/brig/util/log"
	yamlConfig "github.com/olebedev/config"
	"github.com/tsuibin/goxmpp2/xmpp"
	"github.com/tucnak/climax"
)

func init() {
	log.SetOutput(os.Stderr)

	// Only log the warning severity or above.
	log.SetLevel(log.DebugLevel)

	// Log pretty text
	log.SetFormatter(&colorlog.ColorfulLogFormatter{})
}

///////////////////////
// Utility functions //
///////////////////////

func formatGroup(category string) string {
	return strings.ToUpper(category) + " COMMANDS:"
}

// guessRepoFolder tries to find the repository path
// by using a number of sources.
func guessRepoFolder() string {
	wd := os.Getenv("BRIG_PATH")
	if wd != "" {
		return wd
	}

	wd, err := os.Getwd()
	if err != nil {
		log.Error(err)
	}
	return wd
}

func readPassword() (string, error) {
	repoFolder := guessRepoFolder()
	pwd, err := repo.PromptPasswordMaxTries(4, func(pwd string) bool {
		err := repo.CheckPassword(repoFolder, pwd)
		return err == nil
	})

	return pwd, err
}

///////////////////////
// Handler functions //
///////////////////////

func handleVersion(ctx climax.Context) int {
	fmt.Println(brig.VersionString())
	return 0
}

func handleOpen(ctx climax.Context) int {
	repoFolder := guessRepoFolder()
	pwd, err := readPassword()

	if err != nil {
		log.Errorf("Open failed: %v", err)
		return 1
	}

	if _, err := daemon.Reach(pwd, repoFolder, 6666); err != nil {
		log.Errorf("Unable to start daemon: %v", err)
		return 3
	}
	return 0
}

func handleClose(ctx climax.Context) int {
	// This is currently the same as `brig daemon -q`
	return handleDaemonQuit()
}

func handleDaemonPing() int {
	client, err := daemon.Dial(6666)
	if err != nil {
		log.Warning("Unable to dial to daemon: ", err)
		return 1
	}
	defer client.Close()

	for i := 0; i < 100; i++ {
		before := time.Now()
		symbol := colors.Colorize("✔", colors.Green)
		if !client.Ping() {
			symbol = colors.Colorize("✘", colors.Red)
		}

		delay := time.Since(before)

		fmt.Printf("#%02d %s ➔ %s: %s (%v)\n",
			i+1,
			client.LocalAddr().String(),
			client.RemoteAddr().String(),
			symbol, delay)
		time.Sleep(1 * time.Second)
	}

	return 0
}

func handleDaemonQuit() int {
	client, err := daemon.Dial(6666)
	if err != nil {
		log.Warning("Unable to dial to daemon: ", err)
		return 1
	}
	defer client.Close()

	client.Exorcise()
	return 0
}

func handleDaemon(ctx climax.Context) int {
	if ctx.Is("ping") {
		return handleDaemonPing()
	} else if ctx.Is("quit") {
		return handleDaemonQuit()
	}

	pwd, ok := ctx.Get("password")
	if !ok {
		var err error
		pwd, err = readPassword()
		if err != nil {
			log.Errorf("Could not read password: %v", pwd)
			return 1
		}
	}

	repoFolder := guessRepoFolder()
	err := repo.CheckPassword(repoFolder, pwd)
	if err != nil {
		log.Error("Wrong password.")
		return 2
	}

	baal, err := daemon.Summon(pwd, repoFolder, 6666)
	if err != nil {
		log.Warning("Unable to start daemon: ", err)
		return 3
	}

	baal.Serve()
	return 0
}

func handleConfig(ctx climax.Context) int {
	folder := guessRepoFolder()
	cfgPath := filepath.Join(folder, ".brig", "config")

	cfg, err := config.LoadConfig(cfgPath)
	if err != nil {
		log.Errorf("Could not load config: %v", err)
		return 2
	}

	switch len(ctx.Args) {
	case 0:
		yaml, err := yamlConfig.RenderYaml(cfg)
		if err != nil {
			log.Errorf("Unable to render config: %v", err)
			return 3
		}
		fmt.Println(yaml)
	case 1:
		key := ctx.Args[0]
		value, err := cfg.String(key)
		if err != nil {
			log.Errorf("Could not retrieve %s: %v", key, err)
			return 4
		}
		fmt.Println(value)
	case 2:
		key := ctx.Args[0]
		value := ctx.Args[1]
		if err := cfg.Set(key, value); err != nil {
			log.Errorf("Could not set %s: %v", key, err)
			return 5
		}

		if _, err := config.SaveConfig(cfgPath, cfg); err != nil {
			log.Errorf("Could not save config: %v", err)
			return 6
		}
	}

	return 0
}

func handleInit(ctx climax.Context) int {
	if len(ctx.Args) < 1 {
		log.Error("Need your Jabber ID.")
		return 1
	}

	jid := xmpp.JID(ctx.Args[0])
	if jid.Domain() == "" {
		log.Error("Your JabberID needs a domain.")
		return 2
	}

	// Extract the folder from the resource name by default:
	folder := jid.Resource()
	if folder == "" {
		log.Error("Need a resource in your JID.")
		return 3
	}

	if envFolder := os.Getenv("BRIG_PATH"); envFolder != "" {
		folder = envFolder
	}

	if ctx.Is("folder") {
		folder, _ = ctx.Get("folder")
	}

	pwd, err := repo.PromptNewPassword(40.0)
	if err != nil {
		log.Error(err)
		return 4
	}

	repo, err := repo.NewRepository(string(jid), string(pwd), folder)
	if err != nil {
		log.Error(err)
		return 5
	}

	if err := repo.Close(); err != nil {
		log.Errorf("close: %v", err)
		return 6
	}

	if _, err := daemon.Reach(string(pwd), folder, 6666); err != nil {
		log.Errorf("Unable to start daemon: %v", err)
		return 7
	}

	return 0
}

////////////////////////////
// Commandline definition //
////////////////////////////

// RunCmdline starts a brig commandline tool.
func RunCmdline() int {
	demo := climax.New("brig")
	demo.Brief = "brig is a decentralized file syncer based on IPFS and XMPP."
	demo.Version = "unstable"

	repoGroup := demo.AddGroup(formatGroup("repository"))
	xmppGroup := demo.AddGroup(formatGroup("xmpp helper"))
	wdirGroup := demo.AddGroup(formatGroup("working"))
	advnGroup := demo.AddGroup(formatGroup("advanced"))
	miscGroup := demo.AddGroup(formatGroup("misc"))

	commands := []climax.Command{
		climax.Command{
			Name:  "init",
			Brief: "Initialize an empty repository and open it",
			Group: repoGroup,
			Usage: `<JID> [<PATH>]`,
			Help:  `Create an empty repository, open it and associate it with the JID`,
			Flags: []climax.Flag{
				{
					Name:     "--folder",
					Short:    "o",
					Usage:    `--depth="N"`,
					Help:     `Only clone up to this depth of pinned files`,
					Variable: true,
				},
			},
			Examples: []climax.Example{
				{
					Usecase:     `alice@jabber.de/laptop`,
					Description: `Create a folder laptop/ with hidden directories`,
				},
			},
			Handle: func(ctx climax.Context) int {
				return handleInit(ctx)
			},
		},
		climax.Command{
			Name:  "clone",
			Brief: "Clone an repository from somebody else",
			Group: repoGroup,
			Usage: `<OTHER_JID> <YOUR_JID> [<PATH>]`,
			Help:  `...`,
			Flags: []climax.Flag{
				{
					Name:     "--depth",
					Short:    "d",
					Usage:    `--depth="N"`,
					Help:     `Only clone up to this depth of pinned files`,
					Variable: true,
				},
			},
			Examples: []climax.Example{
				{
					Usecase:     `alice@jabber.de/laptop bob@jabber.de/desktop`,
					Description: `Clone Alice' contents`,
				},
			},
			Handle: func(ctx climax.Context) int {
				// TODO: Utils to convert string to int.
				// TODO: Utils to get default value.
				depth, ok := ctx.Get("--depth")
				if !ok {
					depth = "-1"
				}

				fmt.Println(depth)
				return 0
			},
		},
		climax.Command{
			Name:   "open",
			Group:  repoGroup,
			Brief:  "Open an encrypted port. Asks for passphrase.",
			Handle: handleOpen,
		},
		climax.Command{
			Name:   "close",
			Group:  repoGroup,
			Brief:  "Encrypt all metadata in the port and go offline.",
			Handle: handleClose,
		},
		climax.Command{
			Name:  "sync",
			Group: repoGroup,
			Brief: "Sync with all or selected trusted peers.",
		},
		climax.Command{
			Name:  "push",
			Group: repoGroup,
			Brief: "Push your content to all or selected trusted peers.",
		},
		climax.Command{
			Name:  "pull",
			Group: repoGroup,
			Brief: "Pull content from all or selected trusted peers.",
		},
		climax.Command{
			Name:  "watch",
			Group: repoGroup,
			Brief: "Enable or disable watch mode.",
		},
		climax.Command{
			Name:  "discover",
			Group: xmppGroup,
			Brief: "Try to find other brig users near you.",
		},
		climax.Command{
			Name:  "friends",
			Group: xmppGroup,
			Brief: "List your trusted peers.",
		},
		climax.Command{
			Name:  "beg",
			Group: xmppGroup,
			Brief: "Request authorisation from a buddy.",
		},
		climax.Command{
			Name:  "ban",
			Group: xmppGroup,
			Brief: "Discontinue friendship with a peer.",
		},
		climax.Command{
			Name:  "prio",
			Group: xmppGroup,
			Brief: "Change priority of a peer.",
		},
		climax.Command{
			Name:  "status",
			Group: wdirGroup,
			Brief: "Give an overview of brig's current state.",
		},
		climax.Command{
			Name:  "add",
			Group: wdirGroup,
			Brief: "Make file to be managed by brig.",
		},
		climax.Command{
			Name:  "find",
			Group: wdirGroup,
			Brief: "Find filenames in the fleet.",
		},
		climax.Command{
			Name:  "rm",
			Group: wdirGroup,
			Brief: "Remove file from brig's control.",
		},
		climax.Command{
			Name:  "log",
			Group: wdirGroup,
			Brief: "Visualize changelog tree.",
		},
		climax.Command{
			Name:  "checkout",
			Group: wdirGroup,
			Brief: "Attempt to checkout previous version of a file.",
		},
		climax.Command{
			Name:  "fsck",
			Group: advnGroup,
			Brief: "Verify, and possibly fix, broken files.",
		},
		climax.Command{
			Name:  "daemon",
			Group: advnGroup,
			Brief: "Manually run the daemon process.",
			Flags: []climax.Flag{
				{
					Name:  "ping",
					Short: "p",
					Usage: `--ping`,
					Help:  `Ping the dameon to check if it's running.`,
				},
				{
					Name:  "quit",
					Short: "q",
					Usage: `--quit`,
					Help:  `Kill a running daemon.`,
				},
				{
					Name:     "password",
					Short:    "x",
					Usage:    `--password PWD`,
					Help:     `Supply password.`,
					Variable: true,
				},
			},
			Handle: handleDaemon,
		},
		climax.Command{
			Name:  "passwd",
			Group: advnGroup,
			Brief: "Set your XMPP and access password.",
		},
		climax.Command{
			Name:  "yubi",
			Group: advnGroup,
			Brief: "Manage YubiKeys.",
		},
		climax.Command{
			Name:   "config",
			Group:  miscGroup,
			Brief:  "Access, list and modify configuration values.",
			Handle: handleConfig,
		},
		climax.Command{
			Name:  "update",
			Group: miscGroup,
			Brief: "Try to securely update brig.",
		},
		climax.Command{
			Name:  "help",
			Group: miscGroup,
			Brief: "Print some help",
			Usage: "Did you really need help on help?",
		},
		climax.Command{
			Name:   "version",
			Group:  miscGroup,
			Brief:  "Print current version.",
			Usage:  "Print current version.",
			Handle: handleVersion,
		},
	}

	for _, command := range commands {
		demo.AddCommand(command)
	}

	// Help topics:
	demo.AddTopic(climax.Topic{
		Name:  "quickstart",
		Brief: "A very short introduction to brig",
		Text:  "TODO: write.",
	})
	demo.AddTopic(climax.Topic{
		Name:  "tutorial",
		Brief: "A slightly longer introduction.",
		Text:  "TODO: write.",
	})
	demo.AddTopic(climax.Topic{
		Name:  "terms",
		Brief: "Cheat sheet for often used terms.",
		Text:  "TODO: write.",
	})

	return demo.Run()
}
<commit_msg>cmd: make usable from any path inside the repo too.<commit_after>package cmdline

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/disorganizer/brig"
	"github.com/disorganizer/brig/daemon"
	"github.com/disorganizer/brig/repo"
	"github.com/disorganizer/brig/repo/config"
	"github.com/disorganizer/brig/util/colors"
	colorlog "github.com/disorganizer/brig/util/log"
	yamlConfig "github.com/olebedev/config"
	"github.com/tsuibin/goxmpp2/xmpp"
	"github.com/tucnak/climax"
)

func init() {
	log.SetOutput(os.Stderr)

	// Only log the warning severity or above.
	log.SetLevel(log.DebugLevel)

	// Log pretty text
	log.SetFormatter(&colorlog.ColorfulLogFormatter{})
}

///////////////////////
// Utility functions //
///////////////////////

func formatGroup(category string) string {
	return strings.ToUpper(category) + " COMMANDS:"
}

// guessRepoFolder tries to find the repository path
// by using a number of sources.
func guessRepoFolder() string {
	wd := os.Getenv("BRIG_PATH")
	if wd != "" {
		return wd
	}

	wd, err := os.Getwd()
	if err != nil {
		log.Error(err)
	}

	actualPath := repo.FindRepo(wd)
	if actualPath == "" {
		log.Errorf("Unable to find repo in path or any parents:\n%q", wd)
	}

	return actualPath
}

func readPassword() (string, error) {
	repoFolder := guessRepoFolder()
	pwd, err := repo.PromptPasswordMaxTries(4, func(pwd string) bool {
		err := repo.CheckPassword(repoFolder, pwd)
		return err == nil
	})

	return pwd, err
}

///////////////////////
// Handler functions //
///////////////////////

func handleVersion(ctx climax.Context) int {
	fmt.Println(brig.VersionString())
	return 0
}

func handleOpen(ctx climax.Context) int {
	repoFolder := guessRepoFolder()
	pwd, err := readPassword()

	if err != nil {
		log.Errorf("Open failed: %v", err)
		return 1
	}

	if _, err := daemon.Reach(pwd, repoFolder, 6666); err != nil {
		log.Errorf("Unable to start daemon: %v", err)
		return 3
	}
	return 0
}

func handleClose(ctx climax.Context) int {
	// This is currently the same as `brig daemon -q`
	return handleDaemonQuit()
}

func handleDaemonPing() int {
	client, err := daemon.Dial(6666)
	if err != nil {
		log.Warning("Unable to dial to daemon: ", err)
		return 1
	}
	defer client.Close()

	for i := 0; i < 100; i++ {
		before := time.Now()
		symbol := colors.Colorize("✔", colors.Green)
		if !client.Ping() {
			symbol = colors.Colorize("✘", colors.Red)
		}

		delay := time.Since(before)

		fmt.Printf("#%02d %s ➔ %s: %s (%v)\n",
			i+1,
			client.LocalAddr().String(),
			client.RemoteAddr().String(),
			symbol, delay)
		time.Sleep(1 * time.Second)
	}

	return 0
}

func handleDaemonQuit() int {
	client, err := daemon.Dial(6666)
	if err != nil {
		log.Warning("Unable to dial to daemon: ", err)
		return 1
	}
	defer client.Close()

	client.Exorcise()
	return 0
}

func handleDaemon(ctx climax.Context) int {
	if ctx.Is("ping") {
		return handleDaemonPing()
	} else if ctx.Is("quit") {
		return handleDaemonQuit()
	}

	pwd, ok := ctx.Get("password")
	if !ok {
		var err error
		pwd, err = readPassword()
		if err != nil {
			log.Errorf("Could not read password: %v", pwd)
			return 1
		}
	}

	repoFolder := guessRepoFolder()
	err := repo.CheckPassword(repoFolder, pwd)
	if err != nil {
		log.Error("Wrong password.")
		return 2
	}

	baal, err := daemon.Summon(pwd, repoFolder, 6666)
	if err != nil {
		log.Warning("Unable to start daemon: ", err)
		return 3
	}

	baal.Serve()
	return 0
}

func handleConfig(ctx climax.Context) int {
	folder := guessRepoFolder()
	cfgPath := filepath.Join(folder, ".brig", "config")

	cfg, err := config.LoadConfig(cfgPath)
	if err != nil {
		log.Errorf("Could not load config: %v", err)
		return 2
	}

	switch len(ctx.Args) {
	case 0:
		yaml, err := yamlConfig.RenderYaml(cfg)
		if err != nil {
			log.Errorf("Unable to render config: %v", err)
			return 3
		}
		fmt.Println(yaml)
	case 1:
		key := ctx.Args[0]
		value, err := cfg.String(key)
		if err != nil {
			log.Errorf("Could not retrieve %s: %v", key, err)
			return 4
		}
		fmt.Println(value)
	case 2:
		key := ctx.Args[0]
		value := ctx.Args[1]
		if err := cfg.Set(key, value); err != nil {
			log.Errorf("Could not set %s: %v", key, err)
			return 5
		}

		if _, err := config.SaveConfig(cfgPath, cfg); err != nil {
			log.Errorf("Could not save config: %v", err)
			return 6
		}
	}

	return 0
}

func handleInit(ctx climax.Context) int {
	if len(ctx.Args) < 1 {
		log.Error("Need your Jabber ID.")
		return 1
	}

	jid := xmpp.JID(ctx.Args[0])
	if jid.Domain() == "" {
		log.Error("Your JabberID needs a domain.")
		return 2
	}

	// Extract the folder from the resource name by default:
	folder := jid.Resource()
	if folder == "" {
		log.Error("Need a resource in your JID.")
		return 3
	}

	if envFolder := os.Getenv("BRIG_PATH"); envFolder != "" {
		folder = envFolder
	}

	if ctx.Is("folder") {
		folder, _ = ctx.Get("folder")
	}

	pwd, err := repo.PromptNewPassword(40.0)
	if err != nil {
		log.Error(err)
		return 4
	}

	repo, err := repo.NewRepository(string(jid), string(pwd), folder)
	if err != nil {
		log.Error(err)
		return 5
	}

	if err := repo.Close(); err != nil {
		log.Errorf("close: %v", err)
		return 6
	}

	if _, err := daemon.Reach(string(pwd), folder, 6666); err != nil {
		log.Errorf("Unable to start daemon: %v", err)
		return 7
	}

	return 0
}

////////////////////////////
// Commandline definition //
////////////////////////////

// RunCmdline starts a brig commandline tool.
func RunCmdline() int {
	demo := climax.New("brig")
	demo.Brief = "brig is a decentralized file syncer based on IPFS and XMPP."
	demo.Version = "unstable"

	repoGroup := demo.AddGroup(formatGroup("repository"))
	xmppGroup := demo.AddGroup(formatGroup("xmpp helper"))
	wdirGroup := demo.AddGroup(formatGroup("working"))
	advnGroup := demo.AddGroup(formatGroup("advanced"))
	miscGroup := demo.AddGroup(formatGroup("misc"))

	commands := []climax.Command{
		climax.Command{
			Name:  "init",
			Brief: "Initialize an empty repository and open it",
			Group: repoGroup,
			Usage: `<JID> [<PATH>]`,
			Help:  `Create an empty repository, open it and associate it with the JID`,
			Flags: []climax.Flag{
				{
					Name:     "--folder",
					Short:    "o",
					Usage:    `--depth="N"`,
					Help:     `Only clone up to this depth of pinned files`,
					Variable: true,
				},
			},
			Examples: []climax.Example{
				{
					Usecase:     `alice@jabber.de/laptop`,
					Description: `Create a folder laptop/ with hidden directories`,
				},
			},
			Handle: func(ctx climax.Context) int {
				return handleInit(ctx)
			},
		},
		climax.Command{
			Name:  "clone",
			Brief: "Clone an repository from somebody else",
			Group: repoGroup,
			Usage: `<OTHER_JID> <YOUR_JID> [<PATH>]`,
			Help:  `...`,
			Flags: []climax.Flag{
				{
					Name:     "--depth",
					Short:    "d",
					Usage:    `--depth="N"`,
					Help:     `Only clone up to this depth of pinned files`,
					Variable: true,
				},
			},
			Examples: []climax.Example{
				{
					Usecase:     `alice@jabber.de/laptop bob@jabber.de/desktop`,
					Description: `Clone Alice' contents`,
				},
			},
			Handle: func(ctx climax.Context) int {
				// TODO: Utils to convert string to int.
				// TODO: Utils to get default value.
				depth, ok := ctx.Get("--depth")
				if !ok {
					depth = "-1"
				}

				fmt.Println(depth)
				return 0
			},
		},
		climax.Command{
			Name:   "open",
			Group:  repoGroup,
			Brief:  "Open an encrypted port. Asks for passphrase.",
			Handle: handleOpen,
		},
		climax.Command{
			Name:   "close",
			Group:  repoGroup,
			Brief:  "Encrypt all metadata in the port and go offline.",
			Handle: handleClose,
		},
		climax.Command{
			Name:  "sync",
			Group: repoGroup,
			Brief: "Sync with all or selected trusted peers.",
		},
		climax.Command{
			Name:  "push",
			Group: repoGroup,
			Brief: "Push your content to all or selected trusted peers.",
		},
		climax.Command{
			Name:  "pull",
			Group: repoGroup,
			Brief: "Pull content from all or selected trusted peers.",
		},
		climax.Command{
			Name:  "watch",
			Group: repoGroup,
			Brief: "Enable or disable watch mode.",
		},
		climax.Command{
			Name:  "discover",
			Group: xmppGroup,
			Brief: "Try to find other brig users near you.",
		},
		climax.Command{
			Name:  "friends",
			Group: xmppGroup,
			Brief: "List your trusted peers.",
		},
		climax.Command{
			Name:  "beg",
			Group: xmppGroup,
			Brief: "Request authorisation from a buddy.",
		},
		climax.Command{
			Name:  "ban",
			Group: xmppGroup,
			Brief: "Discontinue friendship with a peer.",
		},
		climax.Command{
			Name:  "prio",
			Group: xmppGroup,
			Brief: "Change priority of a peer.",
		},
		climax.Command{
			Name:  "status",
			Group: wdirGroup,
			Brief: "Give an overview of brig's current state.",
		},
		climax.Command{
			Name:  "add",
			Group: wdirGroup,
			Brief: "Make file to be managed by brig.",
		},
		climax.Command{
			Name:  "find",
			Group: wdirGroup,
			Brief: "Find filenames in the fleet.",
		},
		climax.Command{
			Name:  "rm",
			Group: wdirGroup,
			Brief: "Remove file from brig's control.",
		},
		climax.Command{
			Name:  "log",
			Group: wdirGroup,
			Brief: "Visualize changelog tree.",
		},
		climax.Command{
			Name:  "checkout",
			Group: wdirGroup,
			Brief: "Attempt to checkout previous version of a file.",
		},
		climax.Command{
			Name:  "fsck",
			Group: advnGroup,
			Brief: "Verify, and possibly fix, broken files.",
		},
		climax.Command{
			Name:  "daemon",
			Group: advnGroup,
			Brief: "Manually run the daemon process.",
			Flags: []climax.Flag{
				{
					Name:  "ping",
					Short: "p",
					Usage: `--ping`,
					Help:  `Ping the dameon to check if it's running.`,
				},
				{
					Name:  "quit",
					Short: "q",
					Usage: `--quit`,
					Help:  `Kill a running daemon.`,
				},
				{
					Name:     "password",
					Short:    "x",
					Usage:    `--password PWD`,
					Help:     `Supply password.`,
					Variable: true,
				},
			},
			Handle: handleDaemon,
		},
		climax.Command{
			Name:  "passwd",
			Group: advnGroup,
			Brief: "Set your XMPP and access password.",
		},
		climax.Command{
			Name:  "yubi",
			Group: advnGroup,
			Brief: "Manage YubiKeys.",
		},
		climax.Command{
			Name:   "config",
			Group:  miscGroup,
			Brief:  "Access, list and modify configuration values.",
			Handle: handleConfig,
		},
		climax.Command{
			Name:  "update",
			Group: miscGroup,
			Brief: "Try to securely update brig.",
		},
		climax.Command{
			Name:  "help",
			Group: miscGroup,
			Brief: "Print some help",
			Usage: "Did you really need help on help?",
		},
		climax.Command{
			Name:   "version",
			Group:  miscGroup,
			Brief:  "Print current version.",
			Usage:  "Print current version.",
			Handle: handleVersion,
		},
	}

	for _, command := range commands {
		demo.AddCommand(command)
	}

	// Help topics:
	demo.AddTopic(climax.Topic{
		Name:  "quickstart",
		Brief: "A very short introduction to brig",
		Text:  "TODO: write.",
	})
	demo.AddTopic(climax.Topic{
		Name:  "tutorial",
		Brief: "A slightly longer introduction.",
		Text:  "TODO: write.",
	})
	demo.AddTopic(climax.Topic{
		Name:  "terms",
		Brief: "Cheat sheet for often used terms.",
		Text:  "TODO: write.",
	})

	return demo.Run()
}
<|endoftext|>