<commit_before>package command

import (
	"os"
	"os/exec"
	"sort"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
)

type argsRule struct {
	cmd           cli.Command
	ui            cli.Ui
	requires      []string
	adminRequires map[string]struct{} // need admin rights: prompt password
	conditions    map[string][]string
}

func validateArgs(cmd cli.Command, ui cli.Ui) *argsRule {
	return &argsRule{
		cmd:           cmd,
		ui:            ui,
		requires:      make([]string, 0),
		adminRequires: make(map[string]struct{}),
		conditions:    make(map[string][]string),
	}
}

func (this *argsRule) require(option ...string) *argsRule {
	this.requires = append(this.requires, option...)
	return this
}

func (this *argsRule) on(whenOption string, requiredOption ...string) *argsRule {
	if _, present := this.conditions[whenOption]; !present {
		this.conditions[whenOption] = make([]string, 0)
	}
	this.conditions[whenOption] = append(this.conditions[whenOption],
		requiredOption...)
	return this
}

func (this *argsRule) requireAdminRights(option ...string) *argsRule {
	for _, opt := range option {
		this.adminRequires[opt] = struct{}{}
	}
	return this
}

func (this *argsRule) invalid(args []string) bool {
	argSet := make(map[string]struct{}, len(args))
	for _, arg := range args {
		argSet[arg] = struct{}{}
	}

	// required
	for _, req := range this.requires {
		if _, present := argSet[req]; !present {
			this.ui.Error(color.Red("%s required", req))
			this.ui.Output(this.cmd.Help())
			return true
		}
	}

	// admin required
	adminAuthRequired := false
	for _, arg := range args {
		if _, present := this.adminRequires[arg]; present {
			adminAuthRequired = true
			break
		}
	}
	if adminAuthRequired {
		pass, err := this.ui.AskSecret("password for admin: ")
		this.ui.Output("")
		if err != nil {
			this.ui.Error(err.Error())
			return true
		}
		if !Authenticator("", pass) {
			this.ui.Error("invalid admin password, bye!")
			return true
		}
	}

	// conditions
	for when, requires := range this.conditions {
		if _, present := argSet[when]; present {
			for _, req := range requires {
				if _, found := argSet[req]; !found {
					this.ui.Error(color.Red("%s required when %s present",
						req, when))
					this.ui.Output(this.cmd.Help())
					return true
				}
			}
		}
	}

	return false
}

func patternMatched(s, pattern string) bool {
	if pattern != "" && !strings.Contains(s, pattern) {
		return false
	}

	return true
}

func refreshScreen() {
	c := exec.Command("clear")
	c.Stdout = os.Stdout
	c.Run()
}

func ensureZoneValid(zone string) {
	ctx.ZoneZkAddrs(zone) // will panic if zone not found
}

func forAllZones(fn func(zkzone *zk.ZkZone)) {
	for _, zone := range ctx.SortedZones() {
		zkAddrs := ctx.ZoneZkAddrs(zone)
		if strings.TrimSpace(zkAddrs) == "" {
			continue
		}

		zkzone := zk.NewZkZone(zk.DefaultConfig(zone, zkAddrs))
		fn(zkzone)
	}
}

func swallow(err error) {
	if err != nil {
		panic(err)
	}
}

type sortedStrMap struct {
	keys []string
	vals []interface{}
}

// TODO map[string]interface{}
func sortStrMap(m map[string]int) sortedStrMap {
	sortedKeys := make([]string, 0, len(m))
	for key, _ := range m {
		sortedKeys = append(sortedKeys, key)
	}
	sort.Strings(sortedKeys)

	r := sortedStrMap{
		keys: sortedKeys,
		vals: make([]interface{}, len(m)),
	}
	for idx, key := range sortedKeys {
		r.vals[idx] = m[key]
	}

	return r
}
<commit_msg>GK_PASS env provided will skip the admin passwd iteractive auth<commit_after>package command

import (
	"os"
	"os/exec"
	"sort"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
)

type argsRule struct {
	cmd           cli.Command
	ui            cli.Ui
	requires      []string
	adminRequires map[string]struct{} // need admin rights: prompt password
	conditions    map[string][]string
}

func validateArgs(cmd cli.Command, ui cli.Ui) *argsRule {
	return &argsRule{
		cmd:           cmd,
		ui:            ui,
		requires:      make([]string, 0),
		adminRequires: make(map[string]struct{}),
		conditions:    make(map[string][]string),
	}
}

func (this *argsRule) require(option ...string) *argsRule {
	this.requires = append(this.requires, option...)
	return this
}

func (this *argsRule) on(whenOption string, requiredOption ...string) *argsRule {
	if _, present := this.conditions[whenOption]; !present {
		this.conditions[whenOption] = make([]string, 0)
	}
	this.conditions[whenOption] = append(this.conditions[whenOption],
		requiredOption...)
	return this
}

func (this *argsRule) requireAdminRights(option ...string) *argsRule {
	for _, opt := range option {
		this.adminRequires[opt] = struct{}{}
	}
	return this
}

func (this *argsRule) invalid(args []string) bool {
	argSet := make(map[string]struct{}, len(args))
	for _, arg := range args {
		argSet[arg] = struct{}{}
	}

	// required
	for _, req := range this.requires {
		if _, present := argSet[req]; !present {
			this.ui.Error(color.Red("%s required", req))
			this.ui.Output(this.cmd.Help())
			return true
		}
	}

	// admin required
	adminAuthRequired := false
	for _, arg := range args {
		if _, present := this.adminRequires[arg]; present {
			adminAuthRequired = true
			break
		}
	}
	if adminAuthRequired {
		if pass := os.Getenv("GK_PASS"); Authenticator("", pass) {
			return false
		}

		pass, err := this.ui.AskSecret("password for admin(or GK_PASS): ")
		this.ui.Output("")
		if err != nil {
			this.ui.Error(err.Error())
			return true
		}
		if !Authenticator("", pass) {
			this.ui.Error("invalid admin password, bye!")
			return true
		}
	}

	// conditions
	for when, requires := range this.conditions {
		if _, present := argSet[when]; present {
			for _, req := range requires {
				if _, found := argSet[req]; !found {
					this.ui.Error(color.Red("%s required when %s present",
						req, when))
					this.ui.Output(this.cmd.Help())
					return true
				}
			}
		}
	}

	return false
}

func patternMatched(s, pattern string) bool {
	if pattern != "" && !strings.Contains(s, pattern) {
		return false
	}

	return true
}

func refreshScreen() {
	c := exec.Command("clear")
	c.Stdout = os.Stdout
	c.Run()
}

func ensureZoneValid(zone string) {
	ctx.ZoneZkAddrs(zone) // will panic if zone not found
}

func forAllZones(fn func(zkzone *zk.ZkZone)) {
	for _, zone := range ctx.SortedZones() {
		zkAddrs := ctx.ZoneZkAddrs(zone)
		if strings.TrimSpace(zkAddrs) == "" {
			continue
		}

		zkzone := zk.NewZkZone(zk.DefaultConfig(zone, zkAddrs))
		fn(zkzone)
	}
}

func swallow(err error) {
	if err != nil {
		panic(err)
	}
}

type sortedStrMap struct {
	keys []string
	vals []interface{}
}

// TODO map[string]interface{}
func sortStrMap(m map[string]int) sortedStrMap {
	sortedKeys := make([]string, 0, len(m))
	for key, _ := range m {
		sortedKeys = append(sortedKeys, key)
	}
	sort.Strings(sortedKeys)

	r := sortedStrMap{
		keys: sortedKeys,
		vals: make([]interface{}, len(m)),
	}
	for idx, key := range sortedKeys {
		r.vals[idx] = m[key]
	}

	return r
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"math"
	"net"
	"sort"
	"strings"
	"time"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/termui"
	"github.com/nsf/termbox-go"
)

type Zktop struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Zktop) Run(args []string) (exitCode int) {
	var (
		zone            string
		graph           bool
		refreshInterval time.Duration
	)
	cmdFlags := flag.NewFlagSet("zktop", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.DurationVar(&refreshInterval, "i", time.Second*10, "")
	cmdFlags.BoolVar(&graph, "g", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 2
	}

	if graph {
		var zkzones = make([]*zk.ZkZone, 0)
		if zone == "" {
			forSortedZones(func(zkzone *zk.ZkZone) {
				zkzones = append(zkzones, zkzone)
			})
		} else {
			zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
			zkzones = append(zkzones, zkzone)
		}

		this.draw(zkzones)
		return
	}

	for {
		refreshScreen()

		if zone == "" {
			forSortedZones(func(zkzone *zk.ZkZone) {
				this.displayZoneTop(zkzone)
			})
		} else {
			zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
			this.displayZoneTop(zkzone)
		}

		time.Sleep(refreshInterval)
	}

	return
}

func (this *Zktop) displayZoneTop(zkzone *zk.ZkZone) {
	this.Ui.Output(color.Green(zkzone.Name()))
	header := "VER             SERVER           PORT M      OUTST        RECVD         SENT CONNS  ZNODES LAT(MIN/AVG/MAX)"
	this.Ui.Output(header)

	stats := zkzone.RunZkFourLetterCommand("stat")
	sortedHosts := make([]string, 0, len(stats))
	for hp, _ := range stats {
		sortedHosts = append(sortedHosts, hp)
	}
	sort.Strings(sortedHosts)

	for _, hostPort := range sortedHosts {
		host, port, err := net.SplitHostPort(hostPort)
		if err != nil {
			panic(err)
		}

		stat := this.parsedStat(stats[hostPort])
		if stat.mode == "" {
			stat.mode = color.Red("E")
		} else if stat.mode == "L" {
			stat.mode = color.Green("L")
		}
		this.Ui.Output(fmt.Sprintf("%-15s %-15s %5s %1s %10s %12s %12s %5s %7s %s",
			stat.ver,
			host, port,
			stat.mode,
			stat.outstanding,
			stat.received,
			stat.sent,
			stat.connections,
			stat.znodes,
			stat.latency,
		))
	}
}

func (this *Zktop) draw(zkzones []*zk.ZkZone) {
	err := termui.Init()
	if err != nil {
		panic(err)
	}
	defer termui.Close()

	termui.UseTheme("helloworld")

	sinps := (func() []float64 {
		n := 220
		ps := make([]float64, n)
		for i := range ps {
			ps[i] = 1 + math.Sin(float64(i)/5)
		}
		return ps
	})()

	lc0 := termui.NewLineChart()
	lc0.Border.Label = "zk"
	lc0.Data = sinps
	lc0.Width = 50
	lc0.Height = 12
	lc0.X = 0
	lc0.Y = 0
	lc0.AxesColor = termui.ColorWhite
	lc0.LineColor = termui.ColorGreen | termui.AttrBold

	termui.Render(lc0)
	termbox.PollEvent()
}

type zkStat struct {
	ver            string
	latency        string
	connections    string
	outstanding    string
	mode           string
	znodes         string
	received, sent string
}

func (this *Zktop) parsedStat(s string) (stat zkStat) {
	lines := strings.Split(s, "\n")
	for _, l := range lines {
		switch {
		case strings.HasPrefix(l, "Zookeeper version:"):
			p := strings.SplitN(l, ":", 2)
			p = strings.SplitN(p[1], ",", 2)
			stat.ver = strings.TrimSpace(p[0])

		case strings.HasPrefix(l, "Latency"):
			stat.latency = this.extractStatValue(l)

		case strings.HasPrefix(l, "Sent"):
			stat.sent = this.extractStatValue(l)

		case strings.HasPrefix(l, "Received"):
			stat.received = this.extractStatValue(l)

		case strings.HasPrefix(l, "Connections"):
			stat.connections = this.extractStatValue(l)

		case strings.HasPrefix(l, "Mode"):
			stat.mode = strings.ToUpper(this.extractStatValue(l)[:1])

		case strings.HasPrefix(l, "Node count"):
			stat.znodes = this.extractStatValue(l)

		case strings.HasPrefix(l, "Outstanding"):
			stat.outstanding = this.extractStatValue(l)

		}
	}
	return
}

func (this *Zktop) extractStatValue(l string) string {
	p := strings.SplitN(l, ":", 2)
	return strings.TrimSpace(p[1])
}

func (*Zktop) Synopsis() string {
	return "Unix “top” like utility for ZooKeeper"
}

func (this *Zktop) Help() string {
	help := fmt.Sprintf(`
Usage: %s zktop [options]

    Unix “top” like utility for ZooKeeper

Options:

    -z zone

    -g
      Draws zk connections in graph. TODO

    -i interval
      Refresh interval in seconds.
      e,g. 5s

`, this.Cmd)
	return strings.TrimSpace(help)
}
<commit_msg>highlight zk leader ip instead of mode<commit_after>package command

import (
	"flag"
	"fmt"
	"math"
	"net"
	"sort"
	"strings"
	"time"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/termui"
	"github.com/nsf/termbox-go"
)

type Zktop struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Zktop) Run(args []string) (exitCode int) {
	var (
		zone            string
		graph           bool
		refreshInterval time.Duration
	)
	cmdFlags := flag.NewFlagSet("zktop", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.DurationVar(&refreshInterval, "i", time.Second*10, "")
	cmdFlags.BoolVar(&graph, "g", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 2
	}

	if graph {
		var zkzones = make([]*zk.ZkZone, 0)
		if zone == "" {
			forSortedZones(func(zkzone *zk.ZkZone) {
				zkzones = append(zkzones, zkzone)
			})
		} else {
			zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
			zkzones = append(zkzones, zkzone)
		}

		this.draw(zkzones)
		return
	}

	for {
		refreshScreen()

		if zone == "" {
			forSortedZones(func(zkzone *zk.ZkZone) {
				this.displayZoneTop(zkzone)
			})
		} else {
			zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
			this.displayZoneTop(zkzone)
		}

		time.Sleep(refreshInterval)
	}

	return
}

func (this *Zktop) displayZoneTop(zkzone *zk.ZkZone) {
	this.Ui.Output(color.Green(zkzone.Name()))
	header := "VER             SERVER           PORT M      OUTST        RECVD         SENT CONNS  ZNODES LAT(MIN/AVG/MAX)"
	this.Ui.Output(header)

	stats := zkzone.RunZkFourLetterCommand("stat")
	sortedHosts := make([]string, 0, len(stats))
	for hp, _ := range stats {
		sortedHosts = append(sortedHosts, hp)
	}
	sort.Strings(sortedHosts)

	for _, hostPort := range sortedHosts {
		host, port, err := net.SplitHostPort(hostPort)
		if err != nil {
			panic(err)
		}

		stat := this.parsedStat(stats[hostPort])
		if stat.mode == "" {
			stat.mode = color.Red("E")
			host = color.Red(host)
		} else if stat.mode == "L" {
			host = color.Green(host)
		}
		this.Ui.Output(fmt.Sprintf("%-15s %-15s %5s %1s %10s %12s %12s %5s %7s %s",
			stat.ver,
			host, port,
			stat.mode,
			stat.outstanding,
			stat.received,
			stat.sent,
			stat.connections,
			stat.znodes,
			stat.latency,
		))
	}
}

func (this *Zktop) draw(zkzones []*zk.ZkZone) {
	err := termui.Init()
	if err != nil {
		panic(err)
	}
	defer termui.Close()

	termui.UseTheme("helloworld")

	sinps := (func() []float64 {
		n := 220
		ps := make([]float64, n)
		for i := range ps {
			ps[i] = 1 + math.Sin(float64(i)/5)
		}
		return ps
	})()

	lc0 := termui.NewLineChart()
	lc0.Border.Label = "zk"
	lc0.Data = sinps
	lc0.Width = 50
	lc0.Height = 12
	lc0.X = 0
	lc0.Y = 0
	lc0.AxesColor = termui.ColorWhite
	lc0.LineColor = termui.ColorGreen | termui.AttrBold

	termui.Render(lc0)
	termbox.PollEvent()
}

type zkStat struct {
	ver            string
	latency        string
	connections    string
	outstanding    string
	mode           string
	znodes         string
	received, sent string
}

func (this *Zktop) parsedStat(s string) (stat zkStat) {
	lines := strings.Split(s, "\n")
	for _, l := range lines {
		switch {
		case strings.HasPrefix(l, "Zookeeper version:"):
			p := strings.SplitN(l, ":", 2)
			p = strings.SplitN(p[1], ",", 2)
			stat.ver = strings.TrimSpace(p[0])

		case strings.HasPrefix(l, "Latency"):
			stat.latency = this.extractStatValue(l)

		case strings.HasPrefix(l, "Sent"):
			stat.sent = this.extractStatValue(l)

		case strings.HasPrefix(l, "Received"):
			stat.received = this.extractStatValue(l)

		case strings.HasPrefix(l, "Connections"):
			stat.connections = this.extractStatValue(l)

		case strings.HasPrefix(l, "Mode"):
			stat.mode = strings.ToUpper(this.extractStatValue(l)[:1])

		case strings.HasPrefix(l, "Node count"):
			stat.znodes = this.extractStatValue(l)

		case strings.HasPrefix(l, "Outstanding"):
			stat.outstanding = this.extractStatValue(l)

		}
	}
	return
}

func (this *Zktop) extractStatValue(l string) string {
	p := strings.SplitN(l, ":", 2)
	return strings.TrimSpace(p[1])
}

func (*Zktop) Synopsis() string {
	return "Unix “top” like utility for ZooKeeper"
}

func (this *Zktop) Help() string {
	help := fmt.Sprintf(`
Usage: %s zktop [options]

    Unix “top” like utility for ZooKeeper

Options:

    -z zone

    -g
      Draws zk connections in graph. TODO

    -i interval
      Refresh interval in seconds.
      e,g. 5s

`, this.Cmd)
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"regexp"
	"strings"

	"github.com/tkrajina/golongfuncs/internal"
)

func main() {
	ty := make([]string, len(internal.AllTypes))
	for n := range internal.AllTypes {
		ty[n] = string(internal.AllTypes[n])
	}

	var ignoreRegexp, types string

	var params internal.CmdParams
	flag.StringVar(&types, "type", string(internal.Lines), "Type of stats, valid types are: "+strings.Join(ty, ", "))
	flag.Float64Var(&params.Treshold, "treshold", 0, "Min value, functions with value less than this will be ignored")
	flag.IntVar(&params.MinLines, "min-lines", 10, "Functions shorter than this will be ignored")
	flag.IntVar(&params.Top, "top", 25, "Show only top n functions")
	flag.BoolVar(&params.IncludeTests, "include-tests", false, "Include tests")
	flag.BoolVar(&params.IncludeVendor, "include-vendor", false, "Include vendored files")
	flag.StringVar(&ignoreRegexp, "ignore", "", "Regexp for files/directories to ignore")
	flag.Parse()

	paths := flag.Args()
	if len(paths) == 0 {
		paths = append(paths, "./...")
	}

	prepareParams(&params, types, ignoreRegexp)
	stats := internal.Do(params, paths)
	printStats(params, stats)
}

func prepareParams(params *internal.CmdParams, types, ignoreRegexp string) {
	var err error
	params.Types, err = internal.ParseTypes(types)
	if err != nil {
		internal.PrintUsage("Invalid type(s) '%s'", types)
	}
	if len(ignoreRegexp) > 0 {
		r, err := regexp.Compile(ignoreRegexp)
		if err != nil {
			internal.PrintUsage("Invalid ignore regexp '%s'", ignoreRegexp)
		}
		params.Ignore = r
	}
}

func printStats(params internal.CmdParams, stats []internal.FunctionStats) {
	count := 0
	for _, st := range stats {
		val, err := st.Get(params.Types[0])
		if err != nil {
			internal.PrintUsage("Invalid type %s\n", params.Types[0])
		}
		lines, _ := st.Get(internal.Lines)
		if val >= params.Treshold && int(lines) >= params.MinLines {
			loc := st.Location
			if len(loc) >= 38 {
				loc = "..." + loc[len(loc)-35:]
			}
			fmt.Printf("%60s %-40s", st.FuncWithRecv(), loc)
			printSingleStat(params.Types[0], val)
			count += 1
			if len(params.Types) > 1 {
				for i := 1; i < len(params.Types); i++ {
					val, _ := st.Get(params.Types[i])
					printSingleStat(params.Types[i], val)
				}
			}
			fmt.Println()
		}
		if count >= params.Top {
			return
		}
	}
}

func printSingleStat(ty internal.FuncMeasurement, val float64) {
	format := fmt.Sprintf("%%%ds", len(string(ty))+8)
	fmt.Printf(format, fmt.Sprintf("%s=%.1f", ty, val))
}
<commit_msg>Shorten function names<commit_after>package main

import (
	"flag"
	"fmt"
	"regexp"
	"strings"

	"github.com/tkrajina/golongfuncs/internal"
)

func main() {
	ty := make([]string, len(internal.AllTypes))
	for n := range internal.AllTypes {
		ty[n] = string(internal.AllTypes[n])
	}

	var ignoreRegexp, types string

	var params internal.CmdParams
	flag.StringVar(&types, "type", string(internal.Lines), "Type of stats, valid types are: "+strings.Join(ty, ", "))
	flag.Float64Var(&params.Treshold, "treshold", 0, "Min value, functions with value less than this will be ignored")
	flag.IntVar(&params.MinLines, "min-lines", 10, "Functions shorter than this will be ignored")
	flag.IntVar(&params.Top, "top", 25, "Show only top n functions")
	flag.BoolVar(&params.IncludeTests, "include-tests", false, "Include tests")
	flag.BoolVar(&params.IncludeVendor, "include-vendor", false, "Include vendored files")
	flag.StringVar(&ignoreRegexp, "ignore", "", "Regexp for files/directories to ignore")
	flag.Parse()

	paths := flag.Args()
	if len(paths) == 0 {
		paths = append(paths, "./...")
	}

	prepareParams(&params, types, ignoreRegexp)
	stats := internal.Do(params, paths)
	printStats(params, stats)
}

func prepareParams(params *internal.CmdParams, types, ignoreRegexp string) {
	var err error
	params.Types, err = internal.ParseTypes(types)
	if err != nil {
		internal.PrintUsage("Invalid type(s) '%s'", types)
	}
	if len(ignoreRegexp) > 0 {
		r, err := regexp.Compile(ignoreRegexp)
		if err != nil {
			internal.PrintUsage("Invalid ignore regexp '%s'", ignoreRegexp)
		}
		params.Ignore = r
	}
}

func printStats(params internal.CmdParams, stats []internal.FunctionStats) {
	count := 0
	for _, st := range stats {
		val, err := st.Get(params.Types[0])
		if err != nil {
			internal.PrintUsage("Invalid type %s\n", params.Types[0])
		}
		lines, _ := st.Get(internal.Lines)
		if val >= params.Treshold && int(lines) >= params.MinLines {
			fmt.Printf("%40s %-40s", shortenTo(st.FuncWithRecv(), 40), shortenTo(st.Location, 40))
			printSingleStat(params.Types[0], val)
			count += 1
			if len(params.Types) > 1 {
				for i := 1; i < len(params.Types); i++ {
					val, _ := st.Get(params.Types[i])
					printSingleStat(params.Types[i], val)
				}
			}
			fmt.Println()
		}
		if count >= params.Top {
			return
		}
	}
}

func shortenTo(str string, l int) string {
	if len(str) > l {
		return "..." + str[len(str)-l+5:]
	}
	return str
}

func printSingleStat(ty internal.FuncMeasurement, val float64) {
	format := fmt.Sprintf("%%%ds", len(string(ty))+8)
	fmt.Printf(format, fmt.Sprintf("%s=%.1f", ty, val))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/chromium/hstspreload"
	"github.com/chromium/hstspreload/chromiumpreload"
)

func main() {
	args := os.Args[1:]

	if len(args) < 2 {
		fmt.Printf(`hstspreload is a tool for checking conditions to be added to Chromium 's
HSTS preload list. See hstspreload.appspot.com for more details.

Usage:

  hstspreload command argument

The commands are:

  preloadabledomain (+d) Check the TLS configuration and headers of a domain for
                         preload requirements.
  removabledomain   (-d) Check the headers of a domain for removal requirements.
  preloadableheader (+h) Check an HSTS header for preload requirements
  removableheader   (-h) Check an HSTS header for removal requirements
  status                 Check the preload status of a domain

Examples:

  hstspreload +d wikipedia.org
  hstspreload +h "max-age=10886400; includeSubDomains; preload"
  hstspreload -h "max-age=10886400; includeSubDomains"

Return code:

  0    Passed all checks.
  1    Error (failed at least one requirement).
  2    Had warnings, but passed all requirements.
  3    Invalid commandline arguments
  4    Displayed help

`)
		os.Exit(4)
		return
	}

	var header *string
	var issues hstspreload.Issues

	switch args[0] {
	case "+h", "preloadableheader":
		issues = preloadableHeader(args[1])

	case "-h", "removableheader":
		issues = removableHeader(args[1])

	case "+d", "preloadabledomain":
		header, issues = preloadableDomain(args[1])

	case "-d", "removabledomain":
		header, issues = removableDomain(args[1])

	case "status":
		l, err := chromiumpreload.GetLatest()
		if err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err)
		}
		m := chromiumpreload.PreloadEntriesToMap(l)
		fmt.Printf("Status: %v", m[chromiumpreload.Domain(args[1])])
		os.Exit(0)

	default:
		fmt.Printf("Unknown command: %s\n", args[0])
		os.Exit(3)
	}

	// Wrap this in a function to (statically) enforce a return code.
	showResult := func() int {
		if header != nil {
			fmt.Printf("Observed header: %s%s%s\n", bold, *header, reset)
		}

		fmt.Printf("\n")
		switch {
		case len(issues.Errors) > 0:
			return 1

		case len(issues.Warnings) > 0:
			return 2

		default:
			fmt.Printf("%sSatisfies requirements.%s\n\n", green, reset)
			return 0
		}
	}
	exitCode := showResult()

	printList(issues.Errors, "Error", red)
	printList(issues.Warnings, "Warning", yellow)

	os.Exit(exitCode)
}

func preloadableHeader(header string) (issues hstspreload.Issues) {
	warnIfNotHeader(header)

	fmt.Printf(
		"Checking header \"%s%s%s\" for preload requirements...\n",
		bold, header, reset)

	return hstspreload.PreloadableHeaderString(header)
}

func removableHeader(header string) (issues hstspreload.Issues) {
	warnIfNotHeader(header)

	fmt.Printf(
		"Checking header \"%s%s%s\" for removal requirements...\n",
		bold, header, reset)

	return hstspreload.RemovableHeaderString(header)
}

func preloadableDomain(domain string) (header *string, issues hstspreload.Issues) {
	mustBeDomain(domain)

	fmt.Printf(
		"Checking domain %s%s%s for preload requirements...\n",
		underline, domain, reset)

	return hstspreload.PreloadableDomain(domain)
}

func removableDomain(domain string) (header *string, issues hstspreload.Issues) {
	mustBeDomain(domain)

	fmt.Printf(
		"Checking domain %s%s%s for removal requirements...\n",
		underline, domain, reset)

	return hstspreload.RemovableDomain(domain)
}

func warnIfNotHeader(str string) {
	if probablyURL(str) {
		fmt.Fprintf(os.Stderr,
			"Warning: please supply an HSTS header string (it appears you supplied a URL).\n")
	}
	if probablyDomain(str) {
		fmt.Fprintf(os.Stderr,
			"Warning: please supply an HSTS header string (it appears you supplied a domain).\n")
	}
}

func mustBeDomain(str string) {
	if probablyHeader(str) {
		fmt.Fprintf(os.Stderr,
			"Invalid argument: please supply a domain (example.com), not a header string.\n")
		os.Exit(3)
	}

	if probablyURL(str) {
		fmt.Fprintf(os.Stderr,
			"Invalid argument: please supply a domain (example.com) rather than a URL (https://example.com/index.html).\n")
		os.Exit(3)
	}
}

func probablyHeader(str string) bool {
	return strings.Contains(str, ";") || strings.Contains(str, " ")
}

func probablyURL(str string) bool {
	return strings.HasPrefix(str, "http") || strings.Contains(str, ":") || strings.Contains(str, "/")
}

func probablyDomain(str string) bool {
	return strings.Contains(str, ".") && !strings.Contains(str, " ")
}

func printList(list []hstspreload.Issue, title string, fs string) {
	if len(list) == 0 {
		return
	}

	titlePluralized := title
	if len(list) != 1 {
		titlePluralized += "s"
	}
	fmt.Printf("%s%s:%s\n", fs, titlePluralized, reset)

	for i, is := range list {
		fmt.Printf(
			"\n%d. %s%s%s [%s]\n%s\n",
			i+1, fs, is.Summary, reset, is.Code, is.Message)
	}

	fmt.Printf("\n")
}
<commit_msg>cmd/hstspreload: Avoid introducing a new variable to pluralize section titles.<commit_after>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/chromium/hstspreload"
	"github.com/chromium/hstspreload/chromiumpreload"
)

func main() {
	args := os.Args[1:]

	if len(args) < 2 {
		fmt.Printf(`hstspreload is a tool for checking conditions to be added to Chromium 's
HSTS preload list. See hstspreload.appspot.com for more details.

Usage:

  hstspreload command argument

The commands are:

  preloadabledomain (+d) Check the TLS configuration and headers of a domain for
                         preload requirements.
  removabledomain   (-d) Check the headers of a domain for removal requirements.
  preloadableheader (+h) Check an HSTS header for preload requirements
  removableheader   (-h) Check an HSTS header for removal requirements
  status                 Check the preload status of a domain

Examples:

  hstspreload +d wikipedia.org
  hstspreload +h "max-age=10886400; includeSubDomains; preload"
  hstspreload -h "max-age=10886400; includeSubDomains"

Return code:

  0    Passed all checks.
  1    Error (failed at least one requirement).
  2    Had warnings, but passed all requirements.
  3    Invalid commandline arguments
  4    Displayed help

`)
		os.Exit(4)
		return
	}

	var header *string
	var issues hstspreload.Issues

	switch args[0] {
	case "+h", "preloadableheader":
		issues = preloadableHeader(args[1])

	case "-h", "removableheader":
		issues = removableHeader(args[1])

	case "+d", "preloadabledomain":
		header, issues = preloadableDomain(args[1])

	case "-d", "removabledomain":
		header, issues = removableDomain(args[1])

	case "status":
		l, err := chromiumpreload.GetLatest()
		if err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err)
		}
		m := chromiumpreload.PreloadEntriesToMap(l)
		fmt.Printf("Status: %v", m[chromiumpreload.Domain(args[1])])
		os.Exit(0)

	default:
		fmt.Printf("Unknown command: %s\n", args[0])
		os.Exit(3)
	}

	// Wrap this in a function to (statically) enforce a return code.
	showResult := func() int {
		if header != nil {
			fmt.Printf("Observed header: %s%s%s\n", bold, *header, reset)
		}

		fmt.Printf("\n")
		switch {
		case len(issues.Errors) > 0:
			return 1

		case len(issues.Warnings) > 0:
			return 2

		default:
			fmt.Printf("%sSatisfies requirements.%s\n\n", green, reset)
			return 0
		}
	}
	exitCode := showResult()

	printList(issues.Errors, "Error", red)
	printList(issues.Warnings, "Warning", yellow)

	os.Exit(exitCode)
}

func preloadableHeader(header string) (issues hstspreload.Issues) {
	warnIfNotHeader(header)

	fmt.Printf(
		"Checking header \"%s%s%s\" for preload requirements...\n",
		bold, header, reset)

	return hstspreload.PreloadableHeaderString(header)
}

func removableHeader(header string) (issues hstspreload.Issues) {
	warnIfNotHeader(header)

	fmt.Printf(
		"Checking header \"%s%s%s\" for removal requirements...\n",
		bold, header, reset)

	return hstspreload.RemovableHeaderString(header)
}

func preloadableDomain(domain string) (header *string, issues hstspreload.Issues) {
	mustBeDomain(domain)

	fmt.Printf(
		"Checking domain %s%s%s for preload requirements...\n",
		underline, domain, reset)

	return hstspreload.PreloadableDomain(domain)
}

func removableDomain(domain string) (header *string, issues hstspreload.Issues) {
	mustBeDomain(domain)

	fmt.Printf(
		"Checking domain %s%s%s for removal requirements...\n",
		underline, domain, reset)

	return hstspreload.RemovableDomain(domain)
}

func warnIfNotHeader(str string) {
	if probablyURL(str) {
		fmt.Fprintf(os.Stderr,
			"Warning: please supply an HSTS header string (it appears you supplied a URL).\n")
	}
	if probablyDomain(str) {
		fmt.Fprintf(os.Stderr,
			"Warning: please supply an HSTS header string (it appears you supplied a domain).\n")
	}
}

func mustBeDomain(str string) {
	if probablyHeader(str) {
		fmt.Fprintf(os.Stderr,
			"Invalid argument: please supply a domain (example.com), not a header string.\n")
		os.Exit(3)
	}

	if probablyURL(str) {
		fmt.Fprintf(os.Stderr,
			"Invalid argument: please supply a domain (example.com) rather than a URL (https://example.com/index.html).\n")
		os.Exit(3)
	}
}

func probablyHeader(str string) bool {
	return strings.Contains(str, ";") || strings.Contains(str, " ")
}

func probablyURL(str string) bool {
	return strings.HasPrefix(str, "http") || strings.Contains(str, ":") || strings.Contains(str, "/")
}

func probablyDomain(str string) bool {
	return strings.Contains(str, ".") && !strings.Contains(str, " ")
}

func printList(list []hstspreload.Issue, title string, fs string) {
	if len(list) == 0 {
		return
	}

	if len(list) != 1 {
		title += "s"
	}
	fmt.Printf("%s%s:%s\n", fs, title, reset)

	for i, is := range list {
		fmt.Printf(
			"\n%d. %s%s%s [%s]\n%s\n",
			i+1, fs, is.Summary, reset, is.Code, is.Message)
	}

	fmt.Printf("\n")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/juju"
	"strings"
)

// GetEnvironmentCommand is able to output either the entire environment or
// the requested value in a format of the user's choosing.
type GetEnvironmentCommand struct {
	EnvCommandBase
	key string
	out cmd.Output
}

const getEnvHelpDoc = `
If no extra args passed on the command line, all configuration keys and values
for the environment are output using the selected formatter.

A single environment value can be output by adding the environment key name to
the end of the command line.

e.g. $ juju get-environment default-series
     precise
`

func (c *GetEnvironmentCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "get-environment",
		Args:    "[<environment key>]",
		Purpose: "view environment values",
		Doc:     strings.TrimSpace(getEnvHelpDoc),
		Aliases: []string{"get-env"},
	}
}

func (c *GetEnvironmentCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	c.out.AddFlags(f, "smart", cmd.DefaultFormatters)
}

func (c *GetEnvironmentCommand) Init(args []string) (err error) {
	c.key, err = cmd.ZeroOrOneArgs(args)
	return
}

func (c *GetEnvironmentCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()

	// Get the existing environment config from the state.
	config, err := conn.State.EnvironConfig()
	if err != nil {
		return err
	}
	attrs := config.AllAttrs()

	// If no key specified, write out the whole lot.
	if c.key == "" {
		return c.out.Write(ctx, attrs)
	}

	value, found := attrs[c.key]
	if found {
		return c.out.Write(ctx, value)
	}

	return fmt.Errorf("Key %q not found in %q environment.", c.key, config.Name())
}

type attributes map[string]interface{}

// SetEnvironment
type SetEnvironmentCommand struct {
	EnvCommandBase
	values attributes
}

const setEnvHelpDoc = `
Updates the environment of a running Juju instance.  Multiple key/value pairs
can be passed on as command line arguments.
`

func (c *SetEnvironmentCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "set-environment",
		Args:    "key=[value] ...",
		Purpose: "replace environment values",
		Doc:     strings.TrimSpace(setEnvHelpDoc),
		Aliases: []string{"set-env"},
	}
}

// SetFlags handled entirely by EnvCommandBase

func (c *SetEnvironmentCommand) Init(args []string) (err error) {
	if len(args) == 0 {
		return fmt.Errorf("No key, value pairs specified")
	}
	// TODO(thumper) look to have a common library of functions for dealing
	// with key=value pairs.
	c.values = make(attributes)
	for i, arg := range args {
		bits := strings.SplitN(arg, "=", 2)
		if len(bits) < 2 {
			return fmt.Errorf(`Missing "=" in arg %d: %q`, i+1, arg)
		}
		key := bits[0]
		if _, exists := c.values[key]; exists {
			return fmt.Errorf(`Key %q specified more than once`, key)
		}
		c.values[key] = bits[1]
	}
	return nil
}

func (c *SetEnvironmentCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()

	// Here is the magic around setting the attributes:

	// Get the existing environment config from the state.
	oldConfig, err := conn.State.EnvironConfig()
	if err != nil {
		return err
	}
	// Apply the attributes specified for the command to the state config.
	newConfig, err := oldConfig.Apply(c.values)
	if err != nil {
		return err
	}
	// Now validate this new config against the existing config via the provider.
	provider := conn.Environ.Provider()
	newProviderConfig, err := provider.Validate(newConfig, oldConfig)
	if err != nil {
		return err
	}
	// Now try to apply the new validated config.
	return conn.State.SetEnvironConfig(newProviderConfig)
}
<commit_msg>Add a TODO.<commit_after>package main

import (
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/juju"
	"strings"
)

// GetEnvironmentCommand is able to output either the entire environment or
// the requested value in a format of the user's choosing.
type GetEnvironmentCommand struct {
	EnvCommandBase
	key string
	out cmd.Output
}

const getEnvHelpDoc = `
If no extra args passed on the command line, all configuration keys and values
for the environment are output using the selected formatter.

A single environment value can be output by adding the environment key name to
the end of the command line.

e.g. $ juju get-environment default-series
     precise
`

func (c *GetEnvironmentCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "get-environment",
		Args:    "[<environment key>]",
		Purpose: "view environment values",
		Doc:     strings.TrimSpace(getEnvHelpDoc),
		Aliases: []string{"get-env"},
	}
}

func (c *GetEnvironmentCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	c.out.AddFlags(f, "smart", cmd.DefaultFormatters)
}

func (c *GetEnvironmentCommand) Init(args []string) (err error) {
	c.key, err = cmd.ZeroOrOneArgs(args)
	return
}

func (c *GetEnvironmentCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()

	// Get the existing environment config from the state.
	config, err := conn.State.EnvironConfig()
	if err != nil {
		return err
	}
	attrs := config.AllAttrs()

	// If no key specified, write out the whole lot.
	if c.key == "" {
		return c.out.Write(ctx, attrs)
	}

	value, found := attrs[c.key]
	if found {
		return c.out.Write(ctx, value)
	}

	return fmt.Errorf("Key %q not found in %q environment.", c.key, config.Name())
}

type attributes map[string]interface{}

// SetEnvironment
type SetEnvironmentCommand struct {
	EnvCommandBase
	values attributes
}

const setEnvHelpDoc = `
Updates the environment of a running Juju instance.  Multiple key/value pairs
can be passed on as command line arguments.
`

func (c *SetEnvironmentCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "set-environment",
		Args:    "key=[value] ...",
		Purpose: "replace environment values",
		Doc:     strings.TrimSpace(setEnvHelpDoc),
		Aliases: []string{"set-env"},
	}
}

// SetFlags handled entirely by EnvCommandBase

func (c *SetEnvironmentCommand) Init(args []string) (err error) {
	if len(args) == 0 {
		return fmt.Errorf("No key, value pairs specified")
	}
	// TODO(thumper) look to have a common library of functions for dealing
	// with key=value pairs.
	c.values = make(attributes)
	for i, arg := range args {
		bits := strings.SplitN(arg, "=", 2)
		if len(bits) < 2 {
			return fmt.Errorf(`Missing "=" in arg %d: %q`, i+1, arg)
		}
		key := bits[0]
		if _, exists := c.values[key]; exists {
			return fmt.Errorf(`Key %q specified more than once`, key)
		}
		c.values[key] = bits[1]
	}
	return nil
}

func (c *SetEnvironmentCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()

	// Here is the magic around setting the attributes:
	// TODO(thumper): get this magic under test somewhere, and update other call-sites to use it.
	// Get the existing environment config from the state.
	oldConfig, err := conn.State.EnvironConfig()
	if err != nil {
		return err
	}
	// Apply the attributes specified for the command to the state config.
	newConfig, err := oldConfig.Apply(c.values)
	if err != nil {
		return err
	}
	// Now validate this new config against the existing config via the provider.
	provider := conn.Environ.Provider()
	newProviderConfig, err := provider.Validate(newConfig, oldConfig)
	if err != nil {
		return err
	}
	// Now try to apply the new validated config.
	return conn.State.SetEnvironConfig(newProviderConfig)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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 main

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

	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kops/cmd/kops/util"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/registry"
	"k8s.io/kops/util/pkg/tables"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
)

var (
	get_cluster_long = templates.LongDesc(i18n.T(`
	Display one or many cluster resources.`))

	get_cluster_example = templates.Examples(i18n.T(`
	# Get all clusters in a state store
	kops get clusters

	# Get a cluster
	kops get cluster k8s-cluster.example.com

	# Get a cluster YAML desired configuration
	kops get cluster k8s-cluster.example.com -o yaml

	# Save a cluster desired configuration to YAML file
	kops get cluster k8s-cluster.example.com -o yaml > cluster-desired-config.yaml
	`))

	get_cluster_short = i18n.T(`Get one or many clusters.`)

	// Warning for --full.  Since we are not using the template from kubectl
	// we have to have zero white space before the comment characters otherwise
	// output to stdout is going to be off.
	get_cluster_full_warning = i18n.T(`
//
//   WARNING: Do not use a '--full' cluster specification to define a Kubernetes installation.
//   You may experience unexpected behavior and other bugs.  Use only the required elements
//   and any modifications that you require.
//
//   Use the following command to retrieve only the required elements:
//   $ kop get cluster -o yaml
//

`)
)

type GetClusterOptions struct {
	*GetOptions

	// FullSpec determines if we should output the completed (fully populated) spec
	FullSpec bool

	// ClusterNames is a list of cluster names to show; if not specified all clusters will be shown
	ClusterNames []string
}

func NewCmdGetCluster(f *util.Factory, out io.Writer, getOptions *GetOptions) *cobra.Command {
	options := GetClusterOptions{
		GetOptions: getOptions,
	}

	cmd := &cobra.Command{
		Use:     "clusters",
		Aliases: []string{"cluster"},
		Short:   get_cluster_short,
		Long:    get_cluster_long,
		Example: get_cluster_example,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 0 {
				options.ClusterNames = append(options.ClusterNames, args...)
			}

			if rootCommand.clusterName != "" {
				if len(args) != 0 {
					exitWithError(fmt.Errorf("cannot mix --name for cluster with positional arguments"))
				}

				options.ClusterNames = append(options.ClusterNames, rootCommand.clusterName)
			}

			err := RunGetClusters(&rootCommand, os.Stdout, &options)
			if err != nil {
				exitWithError(err)
			}
		},
	}

	cmd.Flags().BoolVar(&options.FullSpec, "full", options.FullSpec, "Show fully populated configuration")

	return cmd
}

func RunGetClusters(context Factory, out io.Writer, options *GetClusterOptions) error {
	client, err := context.Clientset()
	if err != nil {
		return err
	}

	clusterList, err := client.ListClusters(metav1.ListOptions{})
	if err != nil {
		return err
	}

	clusters, err := buildClusters(options.ClusterNames, clusterList)
	if err != nil {
		return err
	}

	if len(clusters) == 0 {
		return fmt.Errorf("no clusters found")
	}

	if options.FullSpec {
		var err error
		clusters, err = fullClusterSpecs(clusters)
		if err != nil {
			return err
		}

		fmt.Fprint(out, get_cluster_full_warning)
	}

	var obj []runtime.Object
	if options.output != OutputTable {
		for _, c := range clusters {
			obj = append(obj, c)
		}
	}

	switch options.output {
	case OutputTable:
		return clusterOutputTable(clusters, out)
	case OutputYaml:
		return fullOutputYAML(out, obj...)
	case OutputJSON:
		return fullOutputJSON(out, obj...)
	default:
		return fmt.Errorf("Unknown output format: %q", options.output)
	}
}

func buildClusters(args []string, clusterList *api.ClusterList) ([]*api.Cluster, error) {
	var clusters []*api.Cluster
	if len(args) != 0 {
		m := make(map[string]*api.Cluster)
		for i := range clusterList.Items {
			c := &clusterList.Items[i]
			m[c.ObjectMeta.Name] = c
		}
		for _, clusterName := range args {
			c := m[clusterName]
			if c == nil {
				return nil, fmt.Errorf("cluster not found %q", clusterName)
			}

			clusters = append(clusters, c)
		}
	} else {
		for i := range clusterList.Items {
			c := &clusterList.Items[i]
			clusters = append(clusters, c)
		}
	}

	return clusters, nil
}

func clusterOutputTable(clusters []*api.Cluster, out io.Writer) error {
	t := &tables.Table{}
	t.AddColumn("NAME", func(c *api.Cluster) string {
		return c.ObjectMeta.Name
	})
	t.AddColumn("CLOUD", func(c *api.Cluster) string {
		return c.Spec.CloudProvider
	})
	t.AddColumn("ZONES", func(c *api.Cluster) string {
		zones := sets.NewString()
		for _, s := range c.Spec.Subnets {
			if s.Zone != "" {
				zones.Insert(s.Zone)
			}
		}
		return strings.Join(zones.List(), ",")
	})

	return t.Render(clusters, out, "NAME", "CLOUD", "ZONES")
}

// fullOutputJson outputs the marshalled JSON of a list of clusters and instance groups.  It will handle
// nils for clusters and instanceGroups slices.
func fullOutputJSON(out io.Writer, args ...runtime.Object) error {
	argsLen := len(args)

	if argsLen > 1 {
		if _, err := fmt.Fprint(out, "["); err != nil {
			return err
		}
	}

	for i, arg := range args {
		if i != 0 {
			if _, err := fmt.Fprint(out, ","); err != nil {
				return err
			}
		}
		if err := marshalToWriter(arg, marshalJSON, out); err != nil {
			return err
		}
	}

	if argsLen > 1 {
		if _, err := fmt.Fprint(out, "]"); err != nil {
			return err
		}
	}

	return nil
}

// fullOutputJson outputs the marshalled JSON of a list of clusters and instance groups.  It will handle
// nils for clusters and instanceGroups slices.
func fullOutputYAML(out io.Writer, args ...runtime.Object) error {
	for i, obj := range args {
		if i != 0 {
			if err := writeYAMLSep(out); err != nil {
				return fmt.Errorf("error writing to stdout: %v", err)
			}
		}
		if err := marshalToWriter(obj, marshalYaml, out); err != nil {
			return err
		}
	}
	return nil
}

func fullClusterSpecs(clusters []*api.Cluster) ([]*api.Cluster, error) {
	var fullSpecs []*api.Cluster
	for _, cluster := range clusters {
		configBase, err := registry.ConfigBase(cluster)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		fullSpec := &api.Cluster{}
		err = registry.ReadConfigDeprecated(configBase.Join(registry.PathClusterCompleted), fullSpec)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		fullSpecs = append(fullSpecs, fullSpec)
	}
	return fullSpecs, nil
}
<commit_msg>Fix typo in kops get cluster --full hint<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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 main

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

	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kops/cmd/kops/util"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/registry"
	"k8s.io/kops/util/pkg/tables"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
)

var (
	get_cluster_long = templates.LongDesc(i18n.T(`
	Display one or many cluster resources.`))

	get_cluster_example = templates.Examples(i18n.T(`
	# Get all clusters in a state store
	kops get clusters

	# Get a cluster
	kops get cluster k8s-cluster.example.com

	# Get a cluster YAML desired configuration
	kops get cluster k8s-cluster.example.com -o yaml

	# Save a cluster desired configuration to YAML file
	kops get cluster k8s-cluster.example.com -o yaml > cluster-desired-config.yaml
	`))

	get_cluster_short = i18n.T(`Get one or many clusters.`)

	// Warning for --full.  Since we are not using the template from kubectl
	// we have to have zero white space before the comment characters otherwise
	// output to stdout is going to be off.
	get_cluster_full_warning = i18n.T(`
//
//   WARNING: Do not use a '--full' cluster specification to define a Kubernetes installation.
//   You may experience unexpected behavior and other bugs.  Use only the required elements
//   and any modifications that you require.
//
//   Use the following command to retrieve only the required elements:
//   $ kops get cluster -o yaml
//

`)
)

type GetClusterOptions struct {
	*GetOptions

	// FullSpec determines if we should output the completed (fully populated) spec
	FullSpec bool

	// ClusterNames is a list of cluster names to show; if not specified all clusters will be shown
	ClusterNames []string
}

func NewCmdGetCluster(f *util.Factory, out io.Writer, getOptions *GetOptions) *cobra.Command {
	options := GetClusterOptions{
		GetOptions: getOptions,
	}

	cmd := &cobra.Command{
		Use:     "clusters",
		Aliases: []string{"cluster"},
		Short:   get_cluster_short,
		Long:    get_cluster_long,
		Example: get_cluster_example,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 0 {
				options.ClusterNames = append(options.ClusterNames, args...)
			}

			if rootCommand.clusterName != "" {
				if len(args) != 0 {
					exitWithError(fmt.Errorf("cannot mix --name for cluster with positional arguments"))
				}

				options.ClusterNames = append(options.ClusterNames, rootCommand.clusterName)
			}

			err := RunGetClusters(&rootCommand, os.Stdout, &options)
			if err != nil {
				exitWithError(err)
			}
		},
	}

	cmd.Flags().BoolVar(&options.FullSpec, "full", options.FullSpec, "Show fully populated configuration")

	return cmd
}

func RunGetClusters(context Factory, out io.Writer, options *GetClusterOptions) error {
	client, err := context.Clientset()
	if err != nil {
		return err
	}

	clusterList, err := client.ListClusters(metav1.ListOptions{})
	if err != nil {
		return err
	}

	clusters, err := buildClusters(options.ClusterNames, clusterList)
	if err != nil {
		return err
	}

	if len(clusters) == 0 {
		return fmt.Errorf("no clusters found")
	}

	if options.FullSpec {
		var err error
		clusters, err = fullClusterSpecs(clusters)
		if err != nil {
			return err
		}

		fmt.Fprint(out, get_cluster_full_warning)
	}

	var obj []runtime.Object
	if options.output != OutputTable {
		for _, c := range clusters {
			obj = append(obj, c)
		}
	}

	switch options.output {
	case OutputTable:
		return clusterOutputTable(clusters, out)
	case OutputYaml:
		return fullOutputYAML(out, obj...)
	case OutputJSON:
		return fullOutputJSON(out, obj...)
	default:
		return fmt.Errorf("Unknown output format: %q", options.output)
	}
}

func buildClusters(args []string, clusterList *api.ClusterList) ([]*api.Cluster, error) {
	var clusters []*api.Cluster
	if len(args) != 0 {
		m := make(map[string]*api.Cluster)
		for i := range clusterList.Items {
			c := &clusterList.Items[i]
			m[c.ObjectMeta.Name] = c
		}
		for _, clusterName := range args {
			c := m[clusterName]
			if c == nil {
				return nil, fmt.Errorf("cluster not found %q", clusterName)
			}

			clusters = append(clusters, c)
		}
	} else {
		for i := range clusterList.Items {
			c := &clusterList.Items[i]
			clusters = append(clusters, c)
		}
	}

	return clusters, nil
}

func clusterOutputTable(clusters []*api.Cluster, out io.Writer) error {
	t := &tables.Table{}
	t.AddColumn("NAME", func(c *api.Cluster) string {
		return c.ObjectMeta.Name
	})
	t.AddColumn("CLOUD", func(c *api.Cluster) string {
		return c.Spec.CloudProvider
	})
	t.AddColumn("ZONES", func(c *api.Cluster) string {
		zones := sets.NewString()
		for _, s := range c.Spec.Subnets {
			if s.Zone != "" {
				zones.Insert(s.Zone)
			}
		}
		return strings.Join(zones.List(), ",")
	})

	return t.Render(clusters, out, "NAME", "CLOUD", "ZONES")
}

// fullOutputJson outputs the marshalled JSON of a list of clusters and instance groups.  It will handle
// nils for clusters and instanceGroups slices.
func fullOutputJSON(out io.Writer, args ...runtime.Object) error {
	argsLen := len(args)

	if argsLen > 1 {
		if _, err := fmt.Fprint(out, "["); err != nil {
			return err
		}
	}

	for i, arg := range args {
		if i != 0 {
			if _, err := fmt.Fprint(out, ","); err != nil {
				return err
			}
		}
		if err := marshalToWriter(arg, marshalJSON, out); err != nil {
			return err
		}
	}

	if argsLen > 1 {
		if _, err := fmt.Fprint(out, "]"); err != nil {
			return err
		}
	}

	return nil
}

// fullOutputJson outputs the marshalled JSON of a list of clusters and instance groups.  It will handle
// nils for clusters and instanceGroups slices.
func fullOutputYAML(out io.Writer, args ...runtime.Object) error {
	for i, obj := range args {
		if i != 0 {
			if err := writeYAMLSep(out); err != nil {
				return fmt.Errorf("error writing to stdout: %v", err)
			}
		}
		if err := marshalToWriter(obj, marshalYaml, out); err != nil {
			return err
		}
	}
	return nil
}

func fullClusterSpecs(clusters []*api.Cluster) ([]*api.Cluster, error) {
	var fullSpecs []*api.Cluster
	for _, cluster := range clusters {
		configBase, err := registry.ConfigBase(cluster)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		fullSpec := &api.Cluster{}
		err = registry.ReadConfigDeprecated(configBase.Join(registry.PathClusterCompleted), fullSpec)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		fullSpecs = append(fullSpecs, fullSpec)
	}
	return fullSpecs, nil
}
<|endoftext|>


<commit_before>package main

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

	"github.com/omniscale/magnacarto/builder"
	mmlparse "github.com/omniscale/magnacarto/mml"
	"gopkg.in/fsnotify.v1"
)

type Update struct {
	Err        error
	Time       time.Time
	UpdatedMML bool
}

type buildStyleFunc func(mm builder.MapMaker, mml string, mss []string) (string, error)

func notifier(buildStyle buildStyleFunc, mm builder.MapMaker, mml string, mss []string, done <-chan struct{}) (updatec chan Update, errc chan error) {
	updatec = make(chan Update, 1)
	errc = make(chan error, 1)

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		errc <- err
		return
	}

	if err := watcher.Add(mml); err != nil {
		updatec <- Update{Err: err}
		return
	}
	for _, mss := range mss {
		if err := watcher.Add(mss); err != nil {
			updatec <- Update{Err: err}
			return
		}
	}

	if len(mss) == 0 {
		// add mss files from mml to watcher, keep mss empty so we know that
		// we style all mss files
		if err := watchMSSFromMML(watcher, mml); err != nil {
			updatec <- Update{Err: err}
			return
		}
	}

	go func() {
		// dummy event to send initial change message to client
		watcher.Events <- fsnotify.Event{}
	}()
	go func() {
		for {
			select {
			case evt := <-watcher.Events:
				style, err := buildStyle(mm, mml, mss)
				if evt.Name == mml && len(mss) == 0 {
					// update mms files to watch if mml changed and mss files were not set
					if err := watchMSSFromMML(watcher, mml); err != nil {
						updatec <- Update{Err: err}
					}
				}
				// atomic save of some editors will trigger remove event,
				// which will remove the file from the watcher. add back again
				if evt.Name != "" {
					watcher.Add(evt.Name)
				}
				if err != nil {
					updatec <- Update{Err: err}
				} else {
					fi, _ := os.Stat(style)
					updatec <- Update{Time: fi.ModTime(), UpdatedMML: evt.Name == mml}
				}
			case err := <-watcher.Errors:
				errc <- err
			case <-done:
				watcher.Close()
				return
			}
		}
	}()
	return updatec, errc
}
func mssFilesFromMML(mmlFile string) ([]string, error) {
	r, err := os.Open(mmlFile)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	mml, err := mmlparse.Parse(r)
	if err != nil {
		return nil, err
	}
	mssFiles := []string{}
	for _, s := range mml.Stylesheets {
		mssFiles = append(mssFiles, filepath.Join(filepath.Dir(mmlFile), s))
	}
	return mssFiles, nil
}

func watchMSSFromMML(watcher *fsnotify.Watcher, mmlFile string) error {
	mssFiles, err := mssFilesFromMML(mmlFile)
	if err != nil {
		return err
	}
	for _, mssFile := range mssFiles {
		if err := watcher.Add(mssFile); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>magnaserv: make building more reliable when editing with vim<commit_after>package main

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

	"github.com/omniscale/magnacarto/builder"
	mmlparse "github.com/omniscale/magnacarto/mml"
	"gopkg.in/fsnotify.v1"
)

type Update struct {
	Err        error
	Time       time.Time
	UpdatedMML bool
}

type buildStyleFunc func(mm builder.MapMaker, mml string, mss []string) (string, error)

func notifier(buildStyle buildStyleFunc, mm builder.MapMaker, mml string, mss []string, done <-chan struct{}) (updatec chan Update, errc chan error) {
	updatec = make(chan Update, 1)
	errc = make(chan error, 1)

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		errc <- err
		return
	}

	if err := watcher.Add(mml); err != nil {
		updatec <- Update{Err: err}
		return
	}
	for _, mss := range mss {
		if err := watcher.Add(mss); err != nil {
			updatec <- Update{Err: err}
			return
		}
	}

	if len(mss) == 0 {
		// add mss files from mml to watcher, keep mss empty so we know that
		// we style all mss files
		if err := watchMSSFromMML(watcher, mml); err != nil {
			updatec <- Update{Err: err}
			return
		}
	}

	go func() {
		// dummy event to send initial change message to client
		watcher.Events <- fsnotify.Event{}
	}()
	go func() {
		for {
			select {
			case evt := <-watcher.Events:
				// The style is generated on demand for each map request with
				// the selected mss files. We still buile the style here with
				// all mss files to be able to pass any syntax errors back to
				// the client.
				if evt.Op == fsnotify.Rename {
					// Editors like Vim make two renames and the file is not
					// available inbetween. Delay, to avoid file not found
					// errors.
					time.Sleep(100 * time.Millisecond)
				}
				style, err := buildStyle(mm, mml, mss)
				if evt.Name == mml && len(mss) == 0 {
					// update mms files to watch if mml changed and mss files were not set
					if err := watchMSSFromMML(watcher, mml); err != nil {
						updatec <- Update{Err: err}
					}
				}
				// atomic save of some editors will trigger remove event,
				// which will remove the file from the watcher. add back again
				if evt.Name != "" {
					watcher.Add(evt.Name)
				}
				if err != nil {
					updatec <- Update{Err: err}
				} else {
					fi, _ := os.Stat(style)
					updatec <- Update{Time: fi.ModTime(), UpdatedMML: evt.Name == mml}
				}
			case err := <-watcher.Errors:
				errc <- err
			case <-done:
				watcher.Close()
				return
			}
		}
	}()
	return updatec, errc
}
func mssFilesFromMML(mmlFile string) ([]string, error) {
	r, err := os.Open(mmlFile)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	mml, err := mmlparse.Parse(r)
	if err != nil {
		return nil, err
	}
	mssFiles := []string{}
	for _, s := range mml.Stylesheets {
		mssFiles = append(mssFiles, filepath.Join(filepath.Dir(mmlFile), s))
	}
	return mssFiles, nil
}

func watchMSSFromMML(watcher *fsnotify.Watcher, mmlFile string) error {
	mssFiles, err := mssFilesFromMML(mmlFile)
	if err != nil {
		return err
	}
	for _, mssFile := range mssFiles {
		if err := watcher.Add(mssFile); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"
	"syscall"

	"github.com/gentlemanautomaton/signaler"
)

func usage(errmsg string) {
	fmt.Fprintf(os.Stderr,
		"%s\n\n"+
			"usage: %s <command>\n"+
			"       where <command> is one of\n"+
			"       run, list, guardian.\n",
		errmsg, os.Args[0])
	os.Exit(2)
}

func main() {
	interactive, err := isInteractive()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to determine interactive session status: %v\n", err)
		os.Exit(1)
	}

	var (
		app                       = App()
		listCmd, listConf         = ListCommand(app)
		installCmd, installConf   = InstallCommand(app)
		uninstallCmd              = UninstallCommand(app)
		enforceCmd, enforceConf   = EnforceCommand(app)
		guardianCmd, guardianConf = GuardianCommand(app)
		uiCmd                     = UICommand(app)
		runCmd, runConf           = RunCommand(app)
	)

	command, err := app.Parse(os.Args[1:])

	// Non-interactive means we're running as LocalSystem. This typically means
	// we're being invoked as a service, but it could also mean we're being
	// run via "psexec -s -i". We check our own "-i" flag to override
	// invocation via the service framework.
	if !interactive && !enforceConf.Interactive {
		enforceService(*enforceConf, err)
		return
	}

	if err != nil {
		// Special GUI-based error handling for run
		if len(os.Args) > 1 && strings.EqualFold(os.Args[1], "run") {
			runError(err)
		}
		prepareConsole(false)
		app.Fatalf("%s, try --help", err)
	}

	// Shutdown when we receive a termination signal
	shutdown := signaler.New().Capture(os.Interrupt, syscall.SIGTERM)

	// Ensure that we cleanup even if we panic
	defer shutdown.Trigger()

	switch command {
	case uiCmd.FullCommand():
		ui(shutdown.Context())
	case runCmd.FullCommand():
		run(shutdown.Context(), *runConf)
	case listCmd.FullCommand():
		list(shutdown.Context(), *listConf)
	case installCmd.FullCommand():
		install(shutdown.Context(), os.Args[0], *installConf)
	case uninstallCmd.FullCommand():
		uninstall(shutdown.Context())
	case enforceCmd.FullCommand():
		enforceInteractive(shutdown.Context(), *enforceConf)
	case guardianCmd.FullCommand():
		// Run the server
		err := daemon(shutdown, *guardianConf)
		if err != nil {
			os.Exit(2)
		}
	}
}
<commit_msg>cmd: Improved execution as a service detection<commit_after>package main

import (
	"fmt"
	"os"
	"strings"
	"syscall"

	"github.com/gentlemanautomaton/signaler"
)

func usage(errmsg string) {
	fmt.Fprintf(os.Stderr,
		"%s\n\n"+
			"usage: %s <command>\n"+
			"       where <command> is one of\n"+
			"       run, list, guardian.\n",
		errmsg, os.Args[0])
	os.Exit(2)
}

func main() {
	interactive, err := isInteractive()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to determine interactive session status: %v\n", err)
		os.Exit(1)
	}

	var (
		app                       = App()
		listCmd, listConf         = ListCommand(app)
		installCmd, installConf   = InstallCommand(app)
		uninstallCmd              = UninstallCommand(app)
		enforceCmd, enforceConf   = EnforceCommand(app)
		guardianCmd, guardianConf = GuardianCommand(app)
		uiCmd                     = UICommand(app)
		runCmd, runConf           = RunCommand(app)
	)

	command, err := app.Parse(os.Args[1:])

	// Non-interactive means we're running as LocalSystem. This typically means
	// we're being invoked as a service, but it could also mean we're being
	// run via "psexec -s -i". We check our own "-i" flag to override
	// invocation via the service framework.
	if !interactive && !enforceConf.Interactive && command == enforceCmd.FullCommand() {
		enforceService(*enforceConf, err)
		return
	}

	if err != nil {
		// Special GUI-based error handling for run
		if len(os.Args) > 1 && strings.EqualFold(os.Args[1], "run") {
			runError(err)
		}
		prepareConsole(false)
		app.Fatalf("%s, try --help", err)
	}

	// Shutdown when we receive a termination signal
	shutdown := signaler.New().Capture(os.Interrupt, syscall.SIGTERM)

	// Ensure that we cleanup even if we panic
	defer shutdown.Trigger()

	switch command {
	case uiCmd.FullCommand():
		ui(shutdown.Context())
	case runCmd.FullCommand():
		run(shutdown.Context(), *runConf)
	case listCmd.FullCommand():
		list(shutdown.Context(), *listConf)
	case installCmd.FullCommand():
		install(shutdown.Context(), os.Args[0], *installConf)
	case uninstallCmd.FullCommand():
		uninstall(shutdown.Context())
	case enforceCmd.FullCommand():
		enforceInteractive(shutdown.Context(), *enforceConf)
	case guardianCmd.FullCommand():
		// Run the server
		err := daemon(shutdown, *guardianConf)
		if err != nil {
			os.Exit(2)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"os"

	"github.com/restic/restic/internal/restic"
	"github.com/spf13/cobra"
)

var cmdStats = &cobra.Command{
	Use:   "stats",
	Short: "Scan the repository and show basic statistics",
	Long: `
The "stats" command walks all snapshots in a repository and accumulates
statistics about the data stored therein. It reports on the number of
unique files and their sizes.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runStats(globalOptions, args)
	},
}

func init() {
	cmdRoot.AddCommand(cmdStats)
}

func runStats(gopts GlobalOptions, args []string) error {
	ctx, cancel := context.WithCancel(gopts.ctx)
	defer cancel()

	repo, err := OpenRepository(gopts)
	if err != nil {
		return err
	}

	if err = repo.LoadIndex(ctx); err != nil {
		return err
	}

	if !gopts.NoLock {
		lock, err := lockRepo(repo)
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	}

	// create a container for the stats, and other state
	// needed while walking the trees
	stats := &statsContainer{uniqueFiles: make(map[fileID]struct{}), idSet: make(restic.IDSet)}

	// iterate every snapshot in the repo
	err = repo.List(ctx, restic.SnapshotFile, func(snapshotID restic.ID, size int64) error {
		snapshot, err := restic.LoadSnapshot(ctx, repo, snapshotID)
		if err != nil {
			return fmt.Errorf("Error loading snapshot %s: %v", snapshotID.Str(), err)
		}
		if snapshot.Tree == nil {
			return fmt.Errorf("snapshot %s has nil tree", snapshot.ID().Str())
		}

		err = walkTree(ctx, repo, *snapshot.Tree, stats)
		if err != nil {
			return fmt.Errorf("walking tree %s: %v", *snapshot.Tree, err)
		}

		return nil
	})

	if gopts.JSON {
		err = json.NewEncoder(os.Stdout).Encode(stats)
		if err != nil {
			return fmt.Errorf("encoding output: %v", err)
		}
		return nil
	}

	Printf("   Cumulative Original Size:   %-5s\n", formatBytes(stats.TotalOriginalSize))
	Printf("  Total Original File Count:   %d\n", stats.TotalCount)
	return nil
}

func walkTree(ctx context.Context, repo restic.Repository, treeID restic.ID, stats *statsContainer) error {
	if stats.idSet.Has(treeID) {
		return nil
	}
	stats.idSet.Insert(treeID)

	tree, err := repo.LoadTree(ctx, treeID)
	if err != nil {
		return fmt.Errorf("loading tree: %v", err)
	}

	for _, node := range tree.Nodes {
		// only count this file if we haven't visited it before
		fid := makeFileID(node)
		if _, ok := stats.uniqueFiles[fid]; !ok {
			// mark the file as visited
			stats.uniqueFiles[fid] = struct{}{}

			// update our stats to account for this node
			stats.TotalOriginalSize += node.Size
			stats.TotalCount++
		}

		// visit subtrees (i.e. directory contents)
		if node.Subtree != nil {
			err = walkTree(ctx, repo, *node.Subtree, stats)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func makeFileID(node *restic.Node) fileID {
	var bb []byte
	for _, c := range node.Content {
		bb = append(bb, []byte(c[:])...)
	}
	return sha256.Sum256(bb)
}

// statsContainer holds information during a walk of a repository
// to collect information about it, as well as state needed
// for a successful and efficient walk.
type statsContainer struct {
	TotalCount        uint64 `json:"total_count"`
	TotalOriginalSize uint64 `json:"total_original_size"`
	idSet             restic.IDSet
	uniqueFiles       map[fileID]struct{}
}

type fileID [32]byte
<commit_msg>Implement four counting modes<commit_after>package main

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"github.com/restic/restic/internal/restic"
	"github.com/spf13/cobra"
)

var cmdStats = &cobra.Command{
	Use:   "stats",
	Short: "Scan the repository and show basic statistics",
	Long: `
The "stats" command walks one or all snapshots in a repository and
accumulates statistics about the data stored therein. It reports on
the number of unique files and their sizes, according to one of
the counting modes as given by a flag.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runStats(globalOptions, args)
	},
}

var countModeFlag []string

func init() {
	cmdRoot.AddCommand(cmdStats)

	f := cmdStats.Flags()
	f.BoolVar(&countModeRestoreSize, "count-restore-size", false, "count the size of files that would be restored (default)")
	f.BoolVar(&countModeUniqueFilesByContent, "count-files-by-contents", false, "count files as unique by their contents")
	f.BoolVar(&countModeBlobsPerFile, "count-blobs-per-file", false, "count sizes of blobs by filename")
	f.BoolVar(&countModeRawData, "count-raw-data", false, "count unique blob sizes irrespective of files referencing them")
	f.StringVar(&snapshotByHost, "host", "", "filter latest snapshot by this hostname")
}

func runStats(gopts GlobalOptions, args []string) error {
	err := verifyStatsInput(gopts, args)
	if err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(gopts.ctx)
	defer cancel()

	repo, err := OpenRepository(gopts)
	if err != nil {
		return err
	}

	if err = repo.LoadIndex(ctx); err != nil {
		return err
	}

	if !gopts.NoLock {
		lock, err := lockRepo(repo)
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	}

	// create a container for the stats (and other needed state)
	stats := &statsContainer{
		uniqueFiles: make(map[fileID]struct{}),
		idSet:       make(restic.IDSet),
		fileBlobs:   make(map[string]restic.IDSet),
		blobs:       restic.NewBlobSet(),
		blobsSeen:   restic.NewBlobSet(),
	}

	if snapshotIDString != "" {
		// scan just a single snapshot

		var sID restic.ID
		if snapshotIDString == "latest" {
			sID, err = restic.FindLatestSnapshot(ctx, repo, []string{}, []restic.TagList{}, snapshotByHost)
			if err != nil {
				Exitf(1, "latest snapshot for criteria not found: %v", err)
			}
		} else {
			sID, err = restic.FindSnapshot(repo, snapshotIDString)
			if err != nil {
				return err
			}
		}

		snapshot, err := restic.LoadSnapshot(ctx, repo, sID)
		if err != nil {
			return err
		}

		err = statsWalkSnapshot(ctx, snapshot, repo, stats)
	} else {
		// iterate every snapshot in the repo
		err = repo.List(ctx, restic.SnapshotFile, func(snapshotID restic.ID, size int64) error {
			snapshot, err := restic.LoadSnapshot(ctx, repo, snapshotID)
			if err != nil {
				return fmt.Errorf("Error loading snapshot %s: %v", snapshotID.Str(), err)
			}
			return statsWalkSnapshot(ctx, snapshot, repo, stats)
		})
	}
	if err != nil {
		return err
	}

	if countModeRawData {
		// the blob handles have been collected, but not yet counted
		for blobHandle := range stats.blobs {
			blobSize, found := repo.LookupBlobSize(blobHandle.ID, blobHandle.Type)
			if !found {
				return fmt.Errorf("blob %v not found", blobHandle)
			}
			stats.TotalSize += uint64(blobSize)
			stats.TotalBlobCount++
		}
	}

	if gopts.JSON {
		err = json.NewEncoder(os.Stdout).Encode(stats)
		if err != nil {
			return fmt.Errorf("encoding output: %v", err)
		}
		return nil
	}

	if stats.TotalBlobCount > 0 {
		Printf("  Total Blob Count:   %d\n", stats.TotalBlobCount)
	}
	if stats.TotalFileCount > 0 {
		Printf("  Total File Count:   %d\n", stats.TotalFileCount)
	}
	Printf("        Total Size:   %-5s\n", formatBytes(stats.TotalSize))

	return nil
}

func statsWalkSnapshot(ctx context.Context, snapshot *restic.Snapshot, repo restic.Repository, stats *statsContainer) error {
	if snapshot.Tree == nil {
		return fmt.Errorf("snapshot %s has nil tree", snapshot.ID().Str())
	}

	if countModeRawData {
		// count just the sizes of unique blobs; we don't need to walk the tree
		// ourselves in this case, since a nifty function does it for us
		return restic.FindUsedBlobs(ctx, repo, *snapshot.Tree, stats.blobs, stats.blobsSeen)
	}

	err := statsWalkTree(ctx, repo, *snapshot.Tree, stats, string(filepath.Separator))
	if err != nil {
		return fmt.Errorf("walking tree %s: %v", *snapshot.Tree, err)
	}
	return nil
}

func statsWalkTree(ctx context.Context, repo restic.Repository, treeID restic.ID, stats *statsContainer, fpath string) error {
	// don't visit a tree we've already walked
	if stats.idSet.Has(treeID) {
		return nil
	}
	stats.idSet.Insert(treeID)

	tree, err := repo.LoadTree(ctx, treeID)
	if err != nil {
		return fmt.Errorf("loading tree: %v", err)
	}

	for _, node := range tree.Nodes {
		if countModeUniqueFilesByContent || countModeBlobsPerFile {
			// only count this file if we haven't visited it before
			fid := makeFileIDByContents(node)
			if _, ok := stats.uniqueFiles[fid]; !ok {
				// mark the file as visited
				stats.uniqueFiles[fid] = struct{}{}

				if countModeUniqueFilesByContent {
					// simply count the size of each unique file (unique by contents only)
					stats.TotalSize += node.Size
					stats.TotalFileCount++
				}
				if countModeBlobsPerFile {
					// count the size of each unique blob reference, which is
					// by unique file (unique by contents and file path)
					for _, blobID := range node.Content {
						// ensure we have this file (by path) in our map; in this
						// mode, a file is unique by both contents and path
						if _, ok := stats.fileBlobs[fpath]; !ok {
							stats.fileBlobs[fpath] = restic.NewIDSet()
							stats.TotalFileCount++
						}
						if _, ok := stats.fileBlobs[fpath][blobID]; !ok {
							// TODO: Is the blob type always 'data' in this case?
							blobSize, found := repo.LookupBlobSize(blobID, restic.DataBlob)
							if !found {
								return fmt.Errorf("blob %s not found for tree %s", blobID, treeID)
							}

							// count the blob's size, then add this blob by this
							// file (path) so we don't double-count it
							stats.TotalSize += uint64(blobSize)
							stats.fileBlobs[fpath].Insert(blobID)

							// this mode also counts total unique blob _references_ per file
							stats.TotalBlobCount++
						}
					}
				}
			}
		}

		if countModeRestoreSize {
			// as this is a file in the snapshot, we can simply count its
			// size without worrying about uniqueness, since duplicate files
			// will still be restored
			stats.TotalSize += node.Size
			stats.TotalFileCount++
		}

		// visit subtrees (i.e. directory contents)
		if node.Subtree != nil {
			err = statsWalkTree(ctx, repo, *node.Subtree, stats, filepath.Join(fpath, node.Name))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// makeFileIDByContents returns a hash of the blob IDs of the
// node's Content in sequence.
func makeFileIDByContents(node *restic.Node) fileID {
	var bb []byte
	for _, c := range node.Content {
		bb = append(bb, []byte(c[:])...)
	}
	return sha256.Sum256(bb)
}

func verifyStatsInput(gopts GlobalOptions, args []string) error {
	// ensure only one counting mode was specified, for clarity
	var countModes int
	if countModeRestoreSize {
		countModes++
	}
	if countModeUniqueFilesByContent {
		countModes++
	}
	if countModeBlobsPerFile {
		countModes++
	}
	if countModeRawData {
		countModes++
	}
	if countModes > 1 {
		return fmt.Errorf("only one counting mode may be used")
	}
	// set a default count mode if none were specified
	if countModes == 0 {
		countModeRestoreSize = true
	}
	// ensure one or none snapshots were specified
	if len(args) > 1 {
		return fmt.Errorf("only one snapshot may be specified")
	}
	// set the snapshot to scan, if one was specified
	if len(args) == 1 {
		snapshotIDString = args[0]
	}
	return nil
}

// statsContainer holds information during a walk of a repository
// to collect information about it, as well as state needed
// for a successful and efficient walk.
type statsContainer struct {
	TotalSize      uint64 `json:"total_size"`
	TotalFileCount uint64 `json:"total_file_count"`
	TotalBlobCount uint64 `json:"total_blob_count,omitempty"`

	// idSet marks visited trees, to avoid repeated walks
	idSet restic.IDSet

	// uniqueFiles marks visited files according to their
	// contents (hashed sequence of content blob IDs)
	uniqueFiles map[fileID]struct{}

	// fileBlobs maps a file name (path) to the set of
	// blobs that have been seen as a part of the file
	fileBlobs map[string]restic.IDSet

	// blobs and blobsSeen are used to count indiviudal
	// unique blobs, independent of references to files
	blobs, blobsSeen restic.BlobSet
}

// fileID is a 256-bit hash that distinguishes unique files.
type fileID [32]byte

var (
	countModeRestoreSize          bool
	countModeUniqueFilesByContent bool
	countModeBlobsPerFile         bool
	countModeRawData              bool

	// the snapshot to scan, as given by the user
	snapshotIDString string

	// snapshotByHost is the host to filter latest
	// snapshot by, if given by user
	snapshotByHost string
)
<|endoftext|>


<commit_before>package main

import (
	"github.com/Sirupsen/logrus"
	"github.com/cerana/cerana/pkg/logrusx"
	"github.com/spf13/pflag"
)

func main() {
	logrus.SetFormatter(&logrusx.MistifyFormatter{})

	config := newConfig(nil, nil)
	pflag.Parse()

	dieOnError(config.loadConfig())
	dieOnError(config.setupLogging())

	sp, err := newStatsPusher(config)
	dieOnError(err)

	dieOnError(sp.run())
	sp.stopOnSignal()
}

func dieOnError(err error) {
	if err != nil {
		logrus.Fatal("encountered an error during startup")
	}
}
<commit_msg>Use logrus.JSONFormatter instead of MistifyFormatter in statspusher<commit_after>package main

import (
	"github.com/Sirupsen/logrus"
	"github.com/cerana/cerana/pkg/logrusx"
	"github.com/spf13/pflag"
)

func main() {
	logrus.SetFormatter(&logrusx.JSONFormatter{})

	config := newConfig(nil, nil)
	pflag.Parse()

	dieOnError(config.loadConfig())
	dieOnError(config.setupLogging())

	sp, err := newStatsPusher(config)
	dieOnError(err)

	dieOnError(sp.run())
	sp.stopOnSignal()
}

func dieOnError(err error) {
	if err != nil {
		logrus.Fatal("encountered an error during startup")
	}
}
<|endoftext|>


<commit_before>package remote

import (
	"fmt"

	"github.com/hashicorp/terraform/state"
)

// Client is the interface that must be implemented for a remote state
// driver. It supports dumb put/get/delete, and the higher level structs
// handle persisting the state properly here.
type Client interface {
	Get() (*Payload, error)
	Put([]byte) error
	Delete() error
}

// ClientLocker is an optional interface that allows a remote state
// backend to enable state lock/unlock.
type ClientLocker interface {
	Client

	Lock(*state.LockInfo) (string, error)
	Unlock(string) error
}

// Payload is the return value from the remote state storage.
type Payload struct {
	MD5  []byte
	Data []byte
}

// Factory is the factory function to create a remote client.
type Factory func(map[string]string) (Client, error)

// NewClient returns a new Client with the given type and configuration.
// The client is looked up in the BuiltinClients variable.
func NewClient(t string, conf map[string]string) (Client, error) {
	f, ok := BuiltinClients[t]
	if !ok {
		return nil, fmt.Errorf("unknown remote client type: %s", t)
	}

	return f(conf)
}

// BuiltinClients is the list of built-in clients that can be used with
// NewClient.
var BuiltinClients = map[string]Factory{
	"artifactory": artifactoryFactory,
	"atlas":       atlasFactory,
	"azure":       azureFactory,
	"etcd":        etcdFactory,
	"gcs":         gcsFactory,
	"http":        httpFactory,
	"local":       fileFactory,
	"s3":          s3Factory,
	"swift":       swiftFactory,
	"manta":       mantaFactory,
}
<commit_msg>state/remote: ClientLocker is just a Client that is a state.Locker<commit_after>package remote

import (
	"fmt"

	"github.com/hashicorp/terraform/state"
)

// Client is the interface that must be implemented for a remote state
// driver. It supports dumb put/get/delete, and the higher level structs
// handle persisting the state properly here.
type Client interface {
	Get() (*Payload, error)
	Put([]byte) error
	Delete() error
}

// ClientLocker is an optional interface that allows a remote state
// backend to enable state lock/unlock.
type ClientLocker interface {
	Client
	state.Locker
}

// Payload is the return value from the remote state storage.
type Payload struct {
	MD5  []byte
	Data []byte
}

// Factory is the factory function to create a remote client.
type Factory func(map[string]string) (Client, error)

// NewClient returns a new Client with the given type and configuration.
// The client is looked up in the BuiltinClients variable.
func NewClient(t string, conf map[string]string) (Client, error) {
	f, ok := BuiltinClients[t]
	if !ok {
		return nil, fmt.Errorf("unknown remote client type: %s", t)
	}

	return f(conf)
}

// BuiltinClients is the list of built-in clients that can be used with
// NewClient.
var BuiltinClients = map[string]Factory{
	"artifactory": artifactoryFactory,
	"atlas":       atlasFactory,
	"azure":       azureFactory,
	"etcd":        etcdFactory,
	"gcs":         gcsFactory,
	"http":        httpFactory,
	"local":       fileFactory,
	"s3":          s3Factory,
	"swift":       swiftFactory,
	"manta":       mantaFactory,
}
<|endoftext|>