<commit_before>package command

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/config"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/metrics"
)

var (
	stats *peekStats
)

type peekStats struct {
	MsgCountPerSecond metrics.Meter
	MsgBytesPerSecond metrics.Meter
}

func newPeekStats() *peekStats {
	this := &peekStats{
		MsgCountPerSecond: metrics.NewMeter(),
		MsgBytesPerSecond: metrics.NewMeter(),
	}

	metrics.Register("msg.count.per.second", this.MsgCountPerSecond)
	metrics.Register("msg.bytes.per.second", this.MsgBytesPerSecond)
	return this
}

func (this *peekStats) start() {
	metrics.Log(metrics.DefaultRegistry, time.Second*10,
		log.New(os.Stdout, "metrics: ", log.Lmicroseconds))
}

type Peek struct {
	Ui cli.Ui
}

func (this *Peek) Run(args []string) (exitCode int) {
	var (
		cluster     string
		zone        string
		topic       string
		partitionId int
		neat        bool
	)
	cmdFlags := flag.NewFlagSet("peek", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.StringVar(&topic, "t", "", "")
	cmdFlags.IntVar(&partitionId, "p", 0, "")
	cmdFlags.BoolVar(&neat, "n", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).require("-z").invalid(args) {
		return 2
	}

	stats := newPeekStats()
	go stats.start()

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, config.ZonePath(zone)))
	msgChan := make(chan *sarama.ConsumerMessage, 2000) // msg aggerator channel
	if cluster == "" {
		zkzone.WithinClusters(func(name string, path string) {
			zkcluster := zkzone.NewCluster(name)
			this.consumeCluster(zkcluster, topic, partitionId, msgChan)
		})
	} else {
		zkcluster := zkzone.NewCluster(cluster)
		this.consumeCluster(zkcluster, topic, partitionId, msgChan)
	}

	var msg *sarama.ConsumerMessage
	for {
		select {
		case msg = <-msgChan:
			stats.MsgCountPerSecond.Mark(1)
			stats.MsgBytesPerSecond.Mark(int64(len(msg.Value)))

			if !neat {
				this.Ui.Output(fmt.Sprintf("%s %s %s", color.Green(topic),
					gofmt.Comma(msg.Offset), string(msg.Value)))
			}
		}
	}

	return
}

func (this *Peek) consumeCluster(zkcluster *zk.ZkCluster, topic string,
	partitionId int, msgChan chan *sarama.ConsumerMessage) {
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}
	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}
	//defer kfk.Close() // FIXME how to close it

	if topic == "" {
		// peek all topics
		topics, err := kfk.Topics()
		if err != nil {
			this.Ui.Output(err.Error())
			return
		}

		for _, t := range topics {
			go this.consumeTopic(kfk, t, int32(partitionId), msgChan)
		}
	} else {
		go this.consumeTopic(kfk, topic, int32(partitionId), msgChan)
	}
}

func (this *Peek) consumeTopic(kfk sarama.Client, topic string, partitionId int32,
	msgCh chan *sarama.ConsumerMessage) {
	consumer, err := sarama.NewConsumerFromClient(kfk)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()

	if partitionId == -1 {
		// all partitions
		partitions, err := kfk.Partitions(topic)
		if err != nil {
			panic(err)
		}

		for _, p := range partitions {
			go this.consumePartition(kfk, consumer, topic, p, msgCh)
		}

	} else {
		this.consumePartition(kfk, consumer, topic, partitionId, msgCh)
	}

}

func (this *Peek) consumePartition(kfk sarama.Client, consumer sarama.Consumer,
	topic string, partitionId int32, msgCh chan *sarama.ConsumerMessage) {
	p, err := consumer.ConsumePartition(topic, partitionId, sarama.OffsetNewest)
	if err != nil {
		panic(err)
	}
	defer p.Close()

	for {
		select {
		case msg := <-p.Messages():
			msgCh <- msg
		}
	}
}

func (*Peek) Synopsis() string {
	return "Peek kafka cluster messages ongoing"
}

func (*Peek) Help() string {
	help := `
Usage: gafka peek -z zone [options]

	Peek kafka cluster messages ongoing

Options:

  -c cluster

  -t topic

  -p partition id
  	If -1, peek all partitions of a topic

  -n
  	Neat mode, only display statastics instead of message content
`
	return strings.TrimSpace(help)
}
<commit_msg>buf fix: peek not showing msg topic name<commit_after>package command

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/config"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/metrics"
)

var (
	stats *peekStats
)

type peekStats struct {
	MsgCountPerSecond metrics.Meter
	MsgBytesPerSecond metrics.Meter
}

func newPeekStats() *peekStats {
	this := &peekStats{
		MsgCountPerSecond: metrics.NewMeter(),
		MsgBytesPerSecond: metrics.NewMeter(),
	}

	metrics.Register("msg.count.per.second", this.MsgCountPerSecond)
	metrics.Register("msg.bytes.per.second", this.MsgBytesPerSecond)
	return this
}

func (this *peekStats) start() {
	metrics.Log(metrics.DefaultRegistry, time.Second*10,
		log.New(os.Stdout, "metrics: ", log.Lmicroseconds))
}

type Peek struct {
	Ui cli.Ui
}

func (this *Peek) Run(args []string) (exitCode int) {
	var (
		cluster     string
		zone        string
		topic       string
		partitionId int
		neat        bool
	)
	cmdFlags := flag.NewFlagSet("peek", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.StringVar(&topic, "t", "", "")
	cmdFlags.IntVar(&partitionId, "p", 0, "")
	cmdFlags.BoolVar(&neat, "n", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).require("-z").invalid(args) {
		return 2
	}

	stats := newPeekStats()
	go stats.start()

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, config.ZonePath(zone)))
	msgChan := make(chan *sarama.ConsumerMessage, 2000) // msg aggerator channel
	if cluster == "" {
		zkzone.WithinClusters(func(name string, path string) {
			zkcluster := zkzone.NewCluster(name)
			this.consumeCluster(zkcluster, topic, partitionId, msgChan)
		})
	} else {
		zkcluster := zkzone.NewCluster(cluster)
		this.consumeCluster(zkcluster, topic, partitionId, msgChan)
	}

	var msg *sarama.ConsumerMessage
	for {
		select {
		case msg = <-msgChan:
			stats.MsgCountPerSecond.Mark(1)
			stats.MsgBytesPerSecond.Mark(int64(len(msg.Value)))

			if !neat {
				this.Ui.Output(fmt.Sprintf("%s %s  %s", color.Green(msg.Topic),
					gofmt.Comma(msg.Offset), string(msg.Value)))
			}
		}
	}

	return
}

func (this *Peek) consumeCluster(zkcluster *zk.ZkCluster, topic string,
	partitionId int, msgChan chan *sarama.ConsumerMessage) {
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}
	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}
	//defer kfk.Close() // FIXME how to close it

	if topic == "" {
		// peek all topics
		topics, err := kfk.Topics()
		if err != nil {
			this.Ui.Output(err.Error())
			return
		}

		for _, t := range topics {
			go this.consumeTopic(kfk, t, int32(partitionId), msgChan)
		}
	} else {
		go this.consumeTopic(kfk, topic, int32(partitionId), msgChan)
	}
}

func (this *Peek) consumeTopic(kfk sarama.Client, topic string, partitionId int32,
	msgCh chan *sarama.ConsumerMessage) {
	consumer, err := sarama.NewConsumerFromClient(kfk)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()

	if partitionId == -1 {
		// all partitions
		partitions, err := kfk.Partitions(topic)
		if err != nil {
			panic(err)
		}

		for _, p := range partitions {
			go this.consumePartition(kfk, consumer, topic, p, msgCh)
		}

	} else {
		this.consumePartition(kfk, consumer, topic, partitionId, msgCh)
	}

}

func (this *Peek) consumePartition(kfk sarama.Client, consumer sarama.Consumer,
	topic string, partitionId int32, msgCh chan *sarama.ConsumerMessage) {
	p, err := consumer.ConsumePartition(topic, partitionId, sarama.OffsetNewest)
	if err != nil {
		panic(err)
	}
	defer p.Close()

	for {
		select {
		case msg := <-p.Messages():
			msgCh <- msg
		}
	}
}

func (*Peek) Synopsis() string {
	return "Peek kafka cluster messages ongoing"
}

func (*Peek) Help() string {
	help := `
Usage: gafka peek -z zone [options]

	Peek kafka cluster messages ongoing

Options:

  -c cluster

  -t topic

  -p partition id
  	-1 will peek all partitions of a topic

  -n
  	Neat mode, only display statastics instead of message content
`
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	"github.com/Shopify/sarama"
)

func subLoop() {
	c, _ := sarama.NewConsumer([]string{
		"10.209.18.65:11004",
		"10.209.18.65:11004",
	}, sarama.NewConfig())
	s, _ := c.ConsumePartition(options.topic, 0, sarama.OffsetOldest)
	defer s.Close()

	loaded := false
	for {
		select {
		case <-time.After(time.Second * 5):
			loaded = true
			close(ready)

		case msg := <-s.Messages():
			hook := decode(msg.Value)

			lock.Lock()
			events = append(events, hook)
			lock.Unlock()
			if loaded {
				newEvt <- struct{}{}
			}

		case <-quit:
			return
		}
	}

}
<commit_msg>bug fix: close a closed chan<commit_after>package main

import (
	"time"

	"github.com/Shopify/sarama"
)

func subLoop() {
	c, _ := sarama.NewConsumer([]string{
		"10.209.18.65:11004",
		"10.209.18.65:11004",
	}, sarama.NewConfig())
	s, _ := c.ConsumePartition(options.topic, 0, sarama.OffsetOldest)
	defer s.Close()

	loaded := false
	for {
		select {
		case <-time.After(time.Second * 5):
			if !loaded {
				loaded = true
				close(ready)
			}

		case msg := <-s.Messages():
			hook := decode(msg.Value)

			lock.Lock()
			events = append(events, hook)
			lock.Unlock()
			if loaded {
				newEvt <- struct{}{}
			}

		case <-quit:
			return
		}
	}

}
<|endoftext|>


<commit_before>package command

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/pipestream"
	"github.com/ryanuber/columnize"
)

type Brokers struct {
	Ui  cli.Ui
	Cmd string

	staleOnly    bool
	showVersions bool
	ipInNumber   bool
	cluster      string
}

func (this *Brokers) Run(args []string) (exitCode int) {
	var (
		zone  string
		debug bool
	)
	cmdFlags := flag.NewFlagSet("brokers", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.StringVar(&this.cluster, "c", "", "")
	cmdFlags.BoolVar(&debug, "debug", false, "")
	cmdFlags.BoolVar(&this.ipInNumber, "n", false, "")
	cmdFlags.BoolVar(&this.staleOnly, "stale", false, "")
	cmdFlags.BoolVar(&this.showVersions, "versions", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.showVersions {
		this.doShowVersions()
		return
	}

	if debug {
		sarama.Logger = log.New(os.Stderr, color.Magenta("[sarama]"), log.LstdFlags)
	}

	if zone != "" {
		ensureZoneValid(zone)

		zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
		this.displayZoneBrokers(zkzone)

		return
	}

	// print all brokers on all zones by default
	forSortedZones(func(zkzone *zk.ZkZone) {
		this.displayZoneBrokers(zkzone)
	})

	return
}

func (this *Brokers) maxBrokerId(zkzone *zk.ZkZone, clusterName string) int {
	var maxBrokerId int
	zkzone.ForSortedBrokers(func(cluster string, liveBrokers map[string]*zk.BrokerZnode) {
		if cluster == clusterName {
			for _, b := range liveBrokers {
				id, _ := strconv.Atoi(b.Id)
				if id > maxBrokerId {
					maxBrokerId = id
				}
			}
		}
	})

	return maxBrokerId
}

func (this *Brokers) displayZoneBrokers(zkzone *zk.ZkZone) {
	lines := make([]string, 0)
	header := "Zone|Cluster|Id|Broker|Uptime"
	lines = append(lines, header)

	n := 0
	zkzone.ForSortedBrokers(func(cluster string, liveBrokers map[string]*zk.BrokerZnode) {
		outputs := this.clusterBrokers(zkzone.Name(), cluster, liveBrokers)
		n += len(outputs)
		lines = append(lines, outputs...)
	})
	this.Ui.Info(fmt.Sprintf("%d brokers in zone[%s]", n, zkzone.Name()))
	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (this *Brokers) clusterBrokers(zone, cluster string, brokers map[string]*zk.BrokerZnode) []string {
	if !patternMatched(cluster, this.cluster) {
		return nil
	}

	if brokers == nil || len(brokers) == 0 {
		return []string{fmt.Sprintf("%s|%s|%s|%s|%s",
			zone, cluster, " ", color.Red("empty brokers"), " ")}
	}

	lines := make([]string, 0, len(brokers))
	if this.staleOnly {
		// try each broker's aliveness
		for brokerId, broker := range brokers {
			cf := sarama.NewConfig()
			cf.Net.ReadTimeout = time.Second * 4
			cf.Net.WriteTimeout = time.Second * 4
			kfk, err := sarama.NewClient([]string{broker.Addr()}, cf)
			if err != nil {
				lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s",
					zone, cluster,
					brokerId, broker.Addr(),
					gofmt.PrettySince(broker.Uptime())))
			} else {
				kfk.Close()
			}
		}

		return lines
	}

	// sort by broker id
	sortedBrokerIds := make([]string, 0, len(brokers))
	for brokerId, _ := range brokers {
		sortedBrokerIds = append(sortedBrokerIds, brokerId)
	}
	sort.Strings(sortedBrokerIds)

	for _, brokerId := range sortedBrokerIds {
		b := brokers[brokerId]
		uptime := gofmt.PrettySince(b.Uptime())
		if time.Since(b.Uptime()) < time.Hour*24*7 {
			uptime = color.Green(uptime)
		}
		if this.ipInNumber {
			lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s",
				zone, cluster,
				brokerId, b.Addr(),
				gofmt.PrettySince(b.Uptime())))
		} else {
			lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s",
				zone, cluster,
				brokerId, b.NamedAddr(),
				gofmt.PrettySince(b.Uptime())))
		}

	}
	return lines
}

func (this *Brokers) doShowVersions() {
	kafkaVerExp := regexp.MustCompile(`/kafka_(?P<ver>[-\d.]*)\.jar`)
	processExp := regexp.MustCompile(`kfk_(?P<process>\S*)/config/server.properties`)

	cmd := pipestream.New("/usr/bin/consul", "exec",
		"pgrep", "-lf", "java",
		"|", "grep", "-w", "kafka",
		"|", "grep", "-vw", "grep")
	err := cmd.Open()
	swallow(err)
	defer cmd.Close()

	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)

	var (
		line     string
		lastLine string
	)
	hosts := make(map[string]struct{})
	lines := make([]string, 0)
	header := "Process|Host|Version"
	lines = append(lines, header)
	records := make(map[string]map[string]string) // {process: {host: ver}}
	for scanner.Scan() {
		line = scanner.Text()
		if strings.Contains(line, "finished with exit code") {
			continue
		}
		if strings.Contains(line, "node(s) completed") {
			continue
		}

		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue
		}

		if _, err := strconv.Atoi(fields[1]); err != nil {
			// field 1 should be pid
			// if not pid, it continues with last line
			line = lastLine + strings.Join(fields[1:], " ")

			// redo fields
			fields := strings.Fields(line)
			if len(fields) < 2 {
				continue
			}
		}

		lastLine = line

		// version
		matched := kafkaVerExp.FindStringSubmatch(line)
		if len(matched) < 2 {
			continue
		}
		ver := matched[1]

		// process name
		matched = processExp.FindStringSubmatch(line)
		if len(matched) < 2 {
			continue
		}
		process := matched[1]

		// got a valid process record
		if _, present := records[process]; !present {
			records[process] = make(map[string]string)
		}
		host := fields[0][0 : len(fields[0])-1] // discard the ending ':'
		hosts[host] = struct{}{}
		records[process][host] = ver
	}
	swallow(scanner.Err())

	sortedProceses := make([]string, 0, len(records))
	for proc, _ := range records {
		sortedProceses = append(sortedProceses, proc)
	}
	sort.Strings(sortedProceses)

	procsWithSingleInstance := make([]string, 0)
	for _, proc := range sortedProceses {
		sortedHosts := make([]string, 0, len(records[proc]))
		for host, _ := range records[proc] {
			sortedHosts = append(sortedHosts, host)
		}
		sort.Strings(sortedHosts)

		if len(sortedHosts) < 2 {
			procsWithSingleInstance = append(procsWithSingleInstance, proc)
		}

		for _, host := range sortedHosts {
			lines = append(lines, fmt.Sprintf("%s|%s|%s", proc, host, records[proc][host]))
		}
	}

	this.Ui.Output(columnize.SimpleFormat(lines))
	this.Ui.Output("")
	this.Ui.Output(fmt.Sprintf("TOTAL %d processes running on %d hosts",
		len(lines)-1, len(hosts)))
	if len(procsWithSingleInstance) > 0 {
		this.Ui.Output(fmt.Sprintf("\nProcess with 1 SPOF: "))
		this.Ui.Warn(fmt.Sprintf("%v", procsWithSingleInstance))
	}
}

func (*Brokers) Synopsis() string {
	return "Print online brokers from Zookeeper"
}

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

    Print online brokers from Zookeeper

Options:

    -z zone
      Only print brokers within a zone

    -c cluster name
      Only print brokers of this cluster

    -versions
      Display kafka instances versions by host
      Precondition: you MUST install consul on each broker host

    -debug

    -n
      Show network addresses as numbers

    -stale
      Only print stale brokers: found in zk but not connectable

`, this.Cmd)
	return strings.TrimSpace(help)
}
<commit_msg>output considers stale option<commit_after>package command

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/pipestream"
	"github.com/ryanuber/columnize"
)

type Brokers struct {
	Ui  cli.Ui
	Cmd string

	staleOnly    bool
	showVersions bool
	ipInNumber   bool
	cluster      string
}

func (this *Brokers) Run(args []string) (exitCode int) {
	var (
		zone  string
		debug bool
	)
	cmdFlags := flag.NewFlagSet("brokers", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.StringVar(&this.cluster, "c", "", "")
	cmdFlags.BoolVar(&debug, "debug", false, "")
	cmdFlags.BoolVar(&this.ipInNumber, "n", false, "")
	cmdFlags.BoolVar(&this.staleOnly, "stale", false, "")
	cmdFlags.BoolVar(&this.showVersions, "versions", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.showVersions {
		this.doShowVersions()
		return
	}

	if debug {
		sarama.Logger = log.New(os.Stderr, color.Magenta("[sarama]"), log.LstdFlags)
	}

	if zone != "" {
		ensureZoneValid(zone)

		zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
		this.displayZoneBrokers(zkzone)

		return
	}

	// print all brokers on all zones by default
	forSortedZones(func(zkzone *zk.ZkZone) {
		this.displayZoneBrokers(zkzone)
	})

	return
}

func (this *Brokers) maxBrokerId(zkzone *zk.ZkZone, clusterName string) int {
	var maxBrokerId int
	zkzone.ForSortedBrokers(func(cluster string, liveBrokers map[string]*zk.BrokerZnode) {
		if cluster == clusterName {
			for _, b := range liveBrokers {
				id, _ := strconv.Atoi(b.Id)
				if id > maxBrokerId {
					maxBrokerId = id
				}
			}
		}
	})

	return maxBrokerId
}

func (this *Brokers) displayZoneBrokers(zkzone *zk.ZkZone) {
	lines := make([]string, 0)
	header := "Zone|Cluster|Id|Broker|Uptime"
	lines = append(lines, header)

	n := 0
	zkzone.ForSortedBrokers(func(cluster string, liveBrokers map[string]*zk.BrokerZnode) {
		outputs := this.clusterBrokers(zkzone.Name(), cluster, liveBrokers)
		n += len(outputs)
		lines = append(lines, outputs...)
	})
	if this.staleOnly {
		this.Ui.Info(fmt.Sprintf("%d problematic brokers in zone[%s]", n, zkzone.Name()))
	} else {
		this.Ui.Info(fmt.Sprintf("%d brokers in zone[%s]", n, zkzone.Name()))
	}
	if len(lines) > 1 {
		// lines has header
		this.Ui.Output(columnize.SimpleFormat(lines))
	}
}

func (this *Brokers) clusterBrokers(zone, cluster string, brokers map[string]*zk.BrokerZnode) []string {
	if !patternMatched(cluster, this.cluster) {
		return nil
	}

	if brokers == nil || len(brokers) == 0 {
		return []string{fmt.Sprintf("%s|%s|%s|%s|%s",
			zone, cluster, " ", color.Red("empty brokers"), " ")}
	}

	lines := make([]string, 0, len(brokers))
	if this.staleOnly {
		// try each broker's aliveness
		for brokerId, broker := range brokers {
			cf := sarama.NewConfig()
			cf.Net.ReadTimeout = time.Second * 4
			cf.Net.WriteTimeout = time.Second * 4
			kfk, err := sarama.NewClient([]string{broker.Addr()}, cf)
			if err != nil {
				lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s",
					zone, cluster,
					brokerId, broker.Addr(),
					gofmt.PrettySince(broker.Uptime())))
			} else {
				kfk.Close()
			}
		}

		return lines
	}

	// sort by broker id
	sortedBrokerIds := make([]string, 0, len(brokers))
	for brokerId, _ := range brokers {
		sortedBrokerIds = append(sortedBrokerIds, brokerId)
	}
	sort.Strings(sortedBrokerIds)

	for _, brokerId := range sortedBrokerIds {
		b := brokers[brokerId]
		uptime := gofmt.PrettySince(b.Uptime())
		if time.Since(b.Uptime()) < time.Hour*24*7 {
			uptime = color.Green(uptime)
		}
		if this.ipInNumber {
			lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s",
				zone, cluster,
				brokerId, b.Addr(),
				gofmt.PrettySince(b.Uptime())))
		} else {
			lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s",
				zone, cluster,
				brokerId, b.NamedAddr(),
				gofmt.PrettySince(b.Uptime())))
		}

	}
	return lines
}

func (this *Brokers) doShowVersions() {
	kafkaVerExp := regexp.MustCompile(`/kafka_(?P<ver>[-\d.]*)\.jar`)
	processExp := regexp.MustCompile(`kfk_(?P<process>\S*)/config/server.properties`)

	cmd := pipestream.New("/usr/bin/consul", "exec",
		"pgrep", "-lf", "java",
		"|", "grep", "-w", "kafka",
		"|", "grep", "-vw", "grep")
	err := cmd.Open()
	swallow(err)
	defer cmd.Close()

	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)

	var (
		line     string
		lastLine string
	)
	hosts := make(map[string]struct{})
	lines := make([]string, 0)
	header := "Process|Host|Version"
	lines = append(lines, header)
	records := make(map[string]map[string]string) // {process: {host: ver}}
	for scanner.Scan() {
		line = scanner.Text()
		if strings.Contains(line, "finished with exit code") {
			continue
		}
		if strings.Contains(line, "node(s) completed") {
			continue
		}

		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue
		}

		if _, err := strconv.Atoi(fields[1]); err != nil {
			// field 1 should be pid
			// if not pid, it continues with last line
			line = lastLine + strings.Join(fields[1:], " ")

			// redo fields
			fields := strings.Fields(line)
			if len(fields) < 2 {
				continue
			}
		}

		lastLine = line

		// version
		matched := kafkaVerExp.FindStringSubmatch(line)
		if len(matched) < 2 {
			continue
		}
		ver := matched[1]

		// process name
		matched = processExp.FindStringSubmatch(line)
		if len(matched) < 2 {
			continue
		}
		process := matched[1]

		// got a valid process record
		if _, present := records[process]; !present {
			records[process] = make(map[string]string)
		}
		host := fields[0][0 : len(fields[0])-1] // discard the ending ':'
		hosts[host] = struct{}{}
		records[process][host] = ver
	}
	swallow(scanner.Err())

	sortedProceses := make([]string, 0, len(records))
	for proc, _ := range records {
		sortedProceses = append(sortedProceses, proc)
	}
	sort.Strings(sortedProceses)

	procsWithSingleInstance := make([]string, 0)
	for _, proc := range sortedProceses {
		sortedHosts := make([]string, 0, len(records[proc]))
		for host, _ := range records[proc] {
			sortedHosts = append(sortedHosts, host)
		}
		sort.Strings(sortedHosts)

		if len(sortedHosts) < 2 {
			procsWithSingleInstance = append(procsWithSingleInstance, proc)
		}

		for _, host := range sortedHosts {
			lines = append(lines, fmt.Sprintf("%s|%s|%s", proc, host, records[proc][host]))
		}
	}

	this.Ui.Output(columnize.SimpleFormat(lines))
	this.Ui.Output("")
	this.Ui.Output(fmt.Sprintf("TOTAL %d processes running on %d hosts",
		len(lines)-1, len(hosts)))
	if len(procsWithSingleInstance) > 0 {
		this.Ui.Output(fmt.Sprintf("\nProcess with 1 SPOF: "))
		this.Ui.Warn(fmt.Sprintf("%v", procsWithSingleInstance))
	}
}

func (*Brokers) Synopsis() string {
	return "Print online brokers from Zookeeper"
}

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

    Print online brokers from Zookeeper

Options:

    -z zone
      Only print brokers within a zone

    -c cluster name
      Only print brokers of this cluster

    -versions
      Display kafka instances versions by host
      Precondition: you MUST install consul on each broker host

    -debug

    -n
      Show network addresses as numbers

    -stale
      Only print stale brokers: found in zk but not connectable

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


<commit_before>package command

import (
	"encoding/csv"
	"flag"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gocli"
	"github.com/ryanuber/columnize"
)

type Haproxy struct {
	Ui  cli.Ui
	Cmd string

	zone string

	cols    []string
	colsMap map[string]struct{}
}

func (this *Haproxy) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("haproxy", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.DefaultZone(), "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	// initialize the output columns
	this.cols = []string{
		"# pxname", "svname",
		"bin", "bout",
		"scur", "req_rate", "rate",
		"hrsp_1xx", "hrsp_2xx", "hrsp_3xx", "hrsp_4xx", "hrsp_5xx",
	}
	this.colsMap = make(map[string]struct{})
	for _, c := range this.cols {
		this.colsMap[c] = struct{}{}
	}

	zone := ctx.Zone(this.zone)
	for _, uri := range zone.HaProxyStatsUri {
		this.fetchStats(uri)
	}

	return
}

func (*Haproxy) Synopsis() string {
	return "Query ehaproxy cluster for load stats"
}

func (this *Haproxy) fetchStats(statsUri string) {
	client := http.Client{Timeout: time.Second * 30}
	resp, err := client.Get(statsUri)
	swallow(err)
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		swallow(fmt.Errorf("fetch[%s] stats got status: %d", resp.StatusCode))
	}

	reader := csv.NewReader(resp.Body)
	records, err := reader.ReadAll()
	swallow(err)

	this.Ui.Info(statsUri)
	cols := make(map[int]string) // col:name

	lines := []string{strings.Join(this.cols, "|")}
	for i, row := range records {
		if i == 0 {
			// header
			for j, col := range row {
				cols[j] = col
			}
			continue
		}

		if (row[0] != "pub" && row[0] != "sub") || row[1] == "BACKEND" || row[1] == "FRONTEND" {
			continue
		}

		var vals []string
		for j, col := range row {
			if _, present := this.colsMap[cols[j]]; !present {
				continue
			}

			vals = append(vals, col)
		}

		lines = append(lines, strings.Join(vals, "|"))
	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

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

    %s

Options:

    -z zone

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<commit_msg>display host only<commit_after>package command

import (
	"encoding/csv"
	"flag"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gocli"
	"github.com/ryanuber/columnize"
)

type Haproxy struct {
	Ui  cli.Ui
	Cmd string

	zone string

	cols    []string
	colsMap map[string]struct{}
}

func (this *Haproxy) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("haproxy", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.DefaultZone(), "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	// initialize the output columns
	this.cols = []string{
		"# pxname", "svname",
		"bin", "bout",
		"scur", "req_rate", "rate",
		"hrsp_1xx", "hrsp_2xx", "hrsp_3xx", "hrsp_4xx", "hrsp_5xx",
	}
	this.colsMap = make(map[string]struct{})
	for _, c := range this.cols {
		this.colsMap[c] = struct{}{}
	}

	zone := ctx.Zone(this.zone)
	for _, uri := range zone.HaProxyStatsUri {
		this.fetchStats(uri)
	}

	return
}

func (*Haproxy) Synopsis() string {
	return "Query ehaproxy cluster for load stats"
}

func (this *Haproxy) fetchStats(statsUri string) {
	client := http.Client{Timeout: time.Second * 30}
	resp, err := client.Get(statsUri)
	swallow(err)
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		swallow(fmt.Errorf("fetch[%s] stats got status: %d", resp.StatusCode))
	}

	reader := csv.NewReader(resp.Body)
	records, err := reader.ReadAll()
	swallow(err)

	u, err := url.Parse(statsUri)
	swallow(err)
	this.Ui.Info(u.Host)
	cols := make(map[int]string) // col:name

	lines := []string{strings.Join(this.cols, "|")}
	for i, row := range records {
		if i == 0 {
			// header
			for j, col := range row {
				cols[j] = col
			}
			continue
		}

		if (row[0] != "pub" && row[0] != "sub") || row[1] == "BACKEND" || row[1] == "FRONTEND" {
			continue
		}

		var vals []string
		for j, col := range row {
			if _, present := this.colsMap[cols[j]]; !present {
				continue
			}

			vals = append(vals, col)
		}

		lines = append(lines, strings.Join(vals, "|"))
	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

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

    %s

Options:

    -z zone

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


<commit_before>package main

import (
	"context"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"math/big"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/Cloud-Foundations/keymaster/lib/certgen"
	"github.com/Cloud-Foundations/keymaster/lib/instrumentedwriter"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/aws/arn"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/organizations"
	"github.com/aws/aws-sdk-go-v2/service/sts"
)

const (
	awsAccountListInterval = time.Minute * 5
)

type assumeRoleCredentialsProvider struct {
	credentials aws.Credentials
	roleArn     *string
	stsClient   *sts.Client
}

type getCallerIdentityResult struct {
	Arn string
}

type getCallerIdentityResponse struct {
	GetCallerIdentityResult getCallerIdentityResult
}

type parsedArnType struct {
	parsedArn arn.ARN
	role      string
}

func awsListAccounts(ctx context.Context, orgClient *organizations.Client) (
	map[string]struct{}, error) {
	output, err := orgClient.ListAccounts(ctx,
		&organizations.ListAccountsInput{})
	if err != nil {
		return nil, err
	}
	list := make(map[string]struct{}, len(output.Accounts))
	for _, account := range output.Accounts {
		list[*account.Id] = struct{}{}
	}
	return list, nil
}

func getCallerIdentity(header http.Header,
	validator func(presignedUrl string) (*url.URL, error)) (
	*parsedArnType, error) {
	claimedArn := header.Get("claimed-arn")
	presignedMethod := header.Get("presigned-method")
	presignedUrl := header.Get("presigned-url")
	if claimedArn == "" || presignedUrl == "" || presignedMethod == "" {
		return nil, fmt.Errorf("missing presigned request data")
	}
	validatedUrl, err := validator(presignedUrl)
	if err != nil {
		return nil, err
	}
	presignedUrl = validatedUrl.String()
	validateReq, err := http.NewRequest(presignedMethod, presignedUrl, nil)
	if err != nil {
		return nil, err
	}
	validateResp, err := http.DefaultClient.Do(validateReq)
	if err != nil {
		return nil, err
	}
	defer validateResp.Body.Close()
	if validateResp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("verification request failed")
	}
	body, err := ioutil.ReadAll(validateResp.Body)
	if err != nil {
		return nil, err
	}
	var callerIdentity getCallerIdentityResponse
	if err := xml.Unmarshal(body, &callerIdentity); err != nil {
		return nil, err
	}
	parsedArn, err := arn.Parse(callerIdentity.GetCallerIdentityResult.Arn)
	if err != nil {
		return nil, err
	}
	// Normalise to the actual role ARN, rather than an ARN showing how the
	// credentials were obtained. This mirrors the way AWS policy documents are
	// written.
	parsedArn.Region = ""
	parsedArn.Service = "iam"
	splitResource := strings.Split(parsedArn.Resource, "/")
	if len(splitResource) < 2 || splitResource[0] != "assumed-role" {
		return nil, fmt.Errorf("invalid resource: %s", parsedArn.Resource)
	}
	parsedArn.Resource = "role/" + splitResource[1]
	if parsedArn.String() != claimedArn {
		return nil, fmt.Errorf("validated ARN: %s != claimed ARN: %s",
			parsedArn.String(), claimedArn)
	}
	return &parsedArnType{
		parsedArn: parsedArn,
		role:      splitResource[1],
	}, nil
}

// validateStsPresignedUrl will validate if the URL is a valid AWS URL.
// It returns the parsed, validated URL so that the caller can rebuild the URL
// (to hopefully silence code security scanners which are dumb).
func validateStsPresignedUrl(presignedUrl string) (*url.URL, error) {
	parsedPresignedUrl, err := url.Parse(presignedUrl)
	if err != nil {
		return nil, err
	}
	if parsedPresignedUrl.Scheme != "https" {
		return nil, fmt.Errorf("invalid scheme: %s", parsedPresignedUrl.Scheme)
	}
	if parsedPresignedUrl.Path != "/" {
		return nil, fmt.Errorf("invalid path: %s", parsedPresignedUrl.Path)
	}
	if !strings.HasPrefix(parsedPresignedUrl.RawQuery,
		"Action=GetCallerIdentity&") {
		return nil,
			fmt.Errorf("invalid action: %s", parsedPresignedUrl.RawQuery)
	}
	splitHost := strings.Split(parsedPresignedUrl.Host, ".")
	if len(splitHost) != 4 ||
		splitHost[0] != "sts" ||
		splitHost[2] != "amazonaws" ||
		splitHost[3] != "com" {
		return nil, fmt.Errorf("malformed presigned URL host")
	}
	return parsedPresignedUrl, nil
}

func (p *assumeRoleCredentialsProvider) Retrieve(ctx context.Context) (
	aws.Credentials, error) {
	if time.Until(p.credentials.Expires) > time.Minute {
		return p.credentials, nil
	}
	output, err := p.stsClient.AssumeRole(ctx, &sts.AssumeRoleInput{
		RoleArn:         p.roleArn,
		RoleSessionName: aws.String("keymaster"),
	})
	if err != nil {
		return aws.Credentials{}, err
	}
	p.credentials = aws.Credentials{
		AccessKeyID:     *output.Credentials.AccessKeyId,
		CanExpire:       true,
		Expires:         *output.Credentials.Expiration,
		SecretAccessKey: *output.Credentials.SecretAccessKey,
		SessionToken:    *output.Credentials.SessionToken,
	}
	return p.credentials, nil
}

func (state *RuntimeState) checkAwsRolesEnabled() bool {
	if len(state.Config.AwsCerts.AllowedAccounts) > 0 {
		return true
	}
	if state.Config.AwsCerts.ListAccountsRole != "" {
		return true
	}
	return false
}

func (state *RuntimeState) configureAwsRoles() error {
	if len(state.Config.AwsCerts.AllowedAccounts) > 0 {
		state.Config.AwsCerts.allowedAccounts =
			make(map[string]struct{})
		for _, id := range state.Config.AwsCerts.AllowedAccounts {
			if _, err := strconv.ParseUint(id, 10, 64); err != nil {
				return fmt.Errorf("accountID: %s is not a number", id)
			}
			state.Config.AwsCerts.allowedAccounts[id] = struct{}{}
		}
	}
	if state.Config.AwsCerts.ListAccountsRole != "" {
		ctx := context.TODO()
		awsConfig, err := awsconfig.LoadDefaultConfig(ctx,
			awsconfig.WithEC2IMDSRegion())
		if err != nil {
			return err
		}
		credsProvider := &assumeRoleCredentialsProvider{
			roleArn:   aws.String(state.Config.AwsCerts.ListAccountsRole),
			stsClient: sts.NewFromConfig(awsConfig),
		}
		awsConfig, err = awsconfig.LoadDefaultConfig(ctx,
			awsconfig.WithEC2IMDSRegion(),
			awsconfig.WithCredentialsProvider(credsProvider))
		if err != nil {
			return err
		}
		orgClient := organizations.NewFromConfig(awsConfig)
		state.Config.AwsCerts.organisationAccounts, err =
			awsListAccounts(ctx, orgClient)
		if err != nil {
			return err
		}
		go state.refreshAwsAccounts(ctx, orgClient)
	}
	return nil
}

func (state *RuntimeState) checkAwsAccountAllowed(accountId string) bool {
	if _, ok := state.Config.AwsCerts.allowedAccounts[accountId]; ok {
		return true
	}
	if _, ok := state.Config.AwsCerts.organisationAccounts[accountId]; ok {
		return true
	}
	if _, ok := state.Config.AwsCerts.allowedAccounts["*"]; ok {
		return true
	}
	return false
}

func (state *RuntimeState) refreshAwsAccounts(ctx context.Context,
	orgClient *organizations.Client) {
	for {
		time.Sleep(awsAccountListInterval)
		if list, err := awsListAccounts(ctx, orgClient); err != nil {
			state.logger.Println(err)
		} else {
			state.Config.AwsCerts.organisationAccounts = list
		}
	}
}

func (state *RuntimeState) requestAwsRoleCertificateHandler(
	w http.ResponseWriter, r *http.Request) {
	state.logger.Debugln(1, "Entered requestAwsRoleCertificateHandler()")
	if state.sendFailureToClientIfLocked(w, r) {
		return
	}
	if r.Method != "POST" {
		state.writeFailureResponse(w, r, http.StatusMethodNotAllowed, "")
		return
	}
	// First extract and validate AWS credentials claim.
	callerArn, err := getCallerIdentity(r.Header, validateStsPresignedUrl)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusUnauthorized,
			"verification request failed")
		return
	}
	if !state.checkAwsAccountAllowed(callerArn.parsedArn.AccountID) {
		state.logger.Printf("AWS account: %s not allowed\n",
			callerArn.parsedArn.AccountID)
		state.writeFailureResponse(w, r, http.StatusUnauthorized,
			"AWS account not allowed")
	}
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusInternalServerError,
			"error reading body")
		return
	}
	// Now extract the public key PEM data.
	block, _ := pem.Decode(body)
	if block == nil {
		state.logger.Println("unable to decode PEM block")
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"invalid PEM block")
		return
	}
	if block.Type != "PUBLIC KEY" {
		state.logger.Printf("unsupport PEM type: %s\n", block.Type)
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"unsupported PEM type")
		return
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusBadRequest, "invalid DER")
		return
	}
	strong, err := certgen.ValidatePublicKeyStrength(pub)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"cannot check key strength")
		return
	}
	if !strong {
		state.writeFailureResponse(w, r, http.StatusBadRequest, "key too weak")
		return
	}
	certDER, commonName, err := state.generateRoleCert(pub, callerArn)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusInternalServerError,
			"cannot generate certificate")
		return
	}
	w.(*instrumentedwriter.LoggingWriter).SetUsername(commonName)
	pem.Encode(w, &pem.Block{Bytes: certDER, Type: "CERTIFICATE"})
}

// Returns certificate DER and CommonName.
func (state *RuntimeState) generateRoleCert(publicKey interface{},
	callerArn *parsedArnType) ([]byte, string, error) {
	commonName := fmt.Sprintf("aws:iam:%s:%s",
		callerArn.parsedArn.AccountID, callerArn.role)
	subject := pkix.Name{
		CommonName:   commonName,
		Organization: []string{"keymaster"},
	}
	arnUrl, err := url.Parse(callerArn.parsedArn.String())
	if err != nil {
		return nil, "", err
	}
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, "", err
	}
	now := time.Now()
	template := x509.Certificate{
		SerialNumber:          serialNumber,
		Subject:               subject,
		NotBefore:             now,
		NotAfter:              now.Add(time.Hour * 24),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageKeyAgreement,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
		IsCA:                  false,
		URIs:                  []*url.URL{arnUrl},
	}
	caCert, err := x509.ParseCertificate(state.caCertDer)
	if err != nil {
		return nil, "", err
	}
	certDER, err := x509.CreateCertificate(rand.Reader, &template, caCert,
		publicKey, state.Signer)
	if err != nil {
		return nil, "", err
	}
	state.logger.Debugf(1, "generateRoleCert: ARN=`%s`, expires=%s",
		callerArn.parsedArn.String(), template.NotAfter)
	metricLogCertDuration("x509", "granted",
		float64(time.Until(template.NotAfter).Seconds()))
	return certDER, commonName, nil
}
<commit_msg>Code review comments.<commit_after>package main

import (
	"context"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"math/big"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/Cloud-Foundations/keymaster/lib/certgen"
	"github.com/Cloud-Foundations/keymaster/lib/instrumentedwriter"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/aws/arn"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/organizations"
	"github.com/aws/aws-sdk-go-v2/service/sts"
)

const (
	awsAccountListInterval = time.Minute * 5
)

type assumeRoleCredentialsProvider struct {
	credentials aws.Credentials
	roleArn     *string
	stsClient   *sts.Client
}

type getCallerIdentityResult struct {
	Arn string
}

type getCallerIdentityResponse struct {
	GetCallerIdentityResult getCallerIdentityResult
}

type parsedArnType struct {
	parsedArn arn.ARN
	role      string
}

func awsListAccounts(ctx context.Context, orgClient *organizations.Client) (
	map[string]struct{}, error) {
	output, err := orgClient.ListAccounts(ctx,
		&organizations.ListAccountsInput{})
	if err != nil {
		return nil, err
	}
	list := make(map[string]struct{}, len(output.Accounts))
	for _, account := range output.Accounts {
		list[*account.Id] = struct{}{}
	}
	return list, nil
}

func getCallerIdentity(header http.Header,
	validator func(presignedUrl string) (*url.URL, error)) (
	*parsedArnType, error) {
	claimedArn := header.Get("claimed-arn")
	presignedMethod := header.Get("presigned-method")
	presignedUrl := header.Get("presigned-url")
	if claimedArn == "" || presignedUrl == "" || presignedMethod == "" {
		return nil, fmt.Errorf("missing presigned request data")
	}
	validatedUrl, err := validator(presignedUrl)
	if err != nil {
		return nil, err
	}
	presignedUrl = validatedUrl.String()
	validateReq, err := http.NewRequest(presignedMethod, presignedUrl, nil)
	if err != nil {
		return nil, err
	}
	validateResp, err := http.DefaultClient.Do(validateReq)
	if err != nil {
		return nil, err
	}
	defer validateResp.Body.Close()
	if validateResp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("verification request failed")
	}
	body, err := ioutil.ReadAll(validateResp.Body)
	if err != nil {
		return nil, err
	}
	var callerIdentity getCallerIdentityResponse
	if err := xml.Unmarshal(body, &callerIdentity); err != nil {
		return nil, err
	}
	parsedArn, err := arn.Parse(callerIdentity.GetCallerIdentityResult.Arn)
	if err != nil {
		return nil, err
	}
	// Normalise to the actual role ARN, rather than an ARN showing how the
	// credentials were obtained. This mirrors the way AWS policy documents are
	// written.
	parsedArn.Region = ""
	parsedArn.Service = "iam"
	splitResource := strings.Split(parsedArn.Resource, "/")
	if len(splitResource) < 2 || splitResource[0] != "assumed-role" {
		return nil, fmt.Errorf("invalid resource: %s", parsedArn.Resource)
	}
	parsedArn.Resource = "role/" + splitResource[1]
	if parsedArn.String() != claimedArn {
		return nil, fmt.Errorf("validated ARN: %s != claimed ARN: %s",
			parsedArn.String(), claimedArn)
	}
	return &parsedArnType{
		parsedArn: parsedArn,
		role:      splitResource[1],
	}, nil
}

// validateStsPresignedUrl will validate if the URL is a valid AWS URL.
// It returns the parsed, validated URL so that the caller can rebuild the URL
// (to hopefully silence code security scanners which are dumb).
func validateStsPresignedUrl(presignedUrl string) (*url.URL, error) {
	parsedPresignedUrl, err := url.Parse(presignedUrl)
	if err != nil {
		return nil, err
	}
	if parsedPresignedUrl.Scheme != "https" {
		return nil, fmt.Errorf("invalid scheme: %s", parsedPresignedUrl.Scheme)
	}
	if parsedPresignedUrl.Path != "/" {
		return nil, fmt.Errorf("invalid path: %s", parsedPresignedUrl.Path)
	}
	if !strings.HasPrefix(parsedPresignedUrl.RawQuery,
		"Action=GetCallerIdentity&") {
		return nil,
			fmt.Errorf("invalid action: %s", parsedPresignedUrl.RawQuery)
	}
	splitHost := strings.Split(parsedPresignedUrl.Host, ".")
	if len(splitHost) != 4 ||
		splitHost[0] != "sts" ||
		splitHost[2] != "amazonaws" ||
		splitHost[3] != "com" {
		return nil, fmt.Errorf("malformed presigned URL host")
	}
	return parsedPresignedUrl, nil
}

func (p *assumeRoleCredentialsProvider) Retrieve(ctx context.Context) (
	aws.Credentials, error) {
	if time.Until(p.credentials.Expires) > time.Minute {
		return p.credentials, nil
	}
	output, err := p.stsClient.AssumeRole(ctx, &sts.AssumeRoleInput{
		RoleArn:         p.roleArn,
		RoleSessionName: aws.String("keymaster"),
	})
	if err != nil {
		return aws.Credentials{}, err
	}
	p.credentials = aws.Credentials{
		AccessKeyID:     *output.Credentials.AccessKeyId,
		CanExpire:       true,
		Expires:         *output.Credentials.Expiration,
		SecretAccessKey: *output.Credentials.SecretAccessKey,
		SessionToken:    *output.Credentials.SessionToken,
	}
	return p.credentials, nil
}

func (state *RuntimeState) checkAwsRolesEnabled() bool {
	if len(state.Config.AwsCerts.AllowedAccounts) > 0 {
		return true
	}
	if state.Config.AwsCerts.ListAccountsRole != "" {
		return true
	}
	return false
}

func (state *RuntimeState) configureAwsRoles() error {
	if len(state.Config.AwsCerts.AllowedAccounts) > 0 {
		state.Config.AwsCerts.allowedAccounts =
			make(map[string]struct{})
		for _, id := range state.Config.AwsCerts.AllowedAccounts {
			if _, err := strconv.ParseUint(id, 10, 64); err != nil {
				return fmt.Errorf("accountID: %s is not a number", id)
			}
			state.Config.AwsCerts.allowedAccounts[id] = struct{}{}
		}
	}
	if state.Config.AwsCerts.ListAccountsRole != "" {
		ctx := context.TODO()
		awsConfig, err := awsconfig.LoadDefaultConfig(ctx,
			awsconfig.WithEC2IMDSRegion())
		if err != nil {
			return err
		}
		credsProvider := &assumeRoleCredentialsProvider{
			roleArn:   aws.String(state.Config.AwsCerts.ListAccountsRole),
			stsClient: sts.NewFromConfig(awsConfig),
		}
		awsConfig, err = awsconfig.LoadDefaultConfig(ctx,
			awsconfig.WithEC2IMDSRegion(),
			awsconfig.WithCredentialsProvider(credsProvider))
		if err != nil {
			return err
		}
		orgClient := organizations.NewFromConfig(awsConfig)
		state.Config.AwsCerts.organisationAccounts, err =
			awsListAccounts(ctx, orgClient)
		if err != nil {
			return err
		}
		go state.refreshAwsAccounts(ctx, orgClient)
	}
	return nil
}

func (state *RuntimeState) checkAwsAccountAllowed(accountId string) bool {
	if _, ok := state.Config.AwsCerts.allowedAccounts[accountId]; ok {
		return true
	}
	if _, ok := state.Config.AwsCerts.organisationAccounts[accountId]; ok {
		return true
	}
	if _, ok := state.Config.AwsCerts.allowedAccounts["*"]; ok {
		return true
	}
	return false
}

func (state *RuntimeState) refreshAwsAccounts(ctx context.Context,
	orgClient *organizations.Client) {
	for {
		time.Sleep(awsAccountListInterval)
		if list, err := awsListAccounts(ctx, orgClient); err != nil {
			state.logger.Println(err)
		} else {
			state.Config.AwsCerts.organisationAccounts = list
		}
	}
}

func (state *RuntimeState) requestAwsRoleCertificateHandler(
	w http.ResponseWriter, r *http.Request) {
	state.logger.Debugln(1, "Entered requestAwsRoleCertificateHandler()")
	if state.sendFailureToClientIfLocked(w, r) {
		return
	}
	if r.Method != "POST" {
		state.writeFailureResponse(w, r, http.StatusMethodNotAllowed, "")
		return
	}
	// First extract and validate AWS credentials claim.
	callerArn, err := getCallerIdentity(r.Header, validateStsPresignedUrl)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusUnauthorized,
			"verification request failed")
		return
	}
	if !state.checkAwsAccountAllowed(callerArn.parsedArn.AccountID) {
		state.logger.Printf("AWS account: %s not allowed\n",
			callerArn.parsedArn.AccountID)
		state.writeFailureResponse(w, r, http.StatusUnauthorized,
			"AWS account not allowed")
	}
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusInternalServerError,
			"error reading body")
		return
	}
	// Now extract the public key PEM data.
	block, _ := pem.Decode(body)
	if block == nil {
		state.logger.Println("unable to decode PEM block")
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"invalid PEM block")
		return
	}
	if block.Type != "PUBLIC KEY" {
		state.logger.Printf("unsupport PEM type: %s\n", block.Type)
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"unsupported PEM type")
		return
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusBadRequest, "invalid DER")
		return
	}
	strong, err := certgen.ValidatePublicKeyStrength(pub)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"cannot check key strength")
		return
	}
	if !strong {
		state.writeFailureResponse(w, r, http.StatusBadRequest, "key too weak")
		return
	}
	certDER, commonName, err := state.generateRoleCert(pub, callerArn)
	if err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusInternalServerError,
			"cannot generate certificate")
		return
	}
	w.(*instrumentedwriter.LoggingWriter).SetUsername(commonName)
	pem.Encode(w, &pem.Block{Bytes: certDER, Type: "CERTIFICATE"})
}

// Returns certificate DER and CommonName.
func (state *RuntimeState) generateRoleCert(publicKey interface{},
	callerArn *parsedArnType) ([]byte, string, error) {
	commonName := fmt.Sprintf("aws:iam:%s:%s",
		callerArn.parsedArn.AccountID, callerArn.role)
	subject := pkix.Name{
		CommonName:   commonName,
		Organization: []string{"keymaster"},
	}
	arnUrl, err := url.Parse(callerArn.parsedArn.String())
	if err != nil {
		return nil, "", err
	}
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, "", err
	}
	now := time.Now()
	template := x509.Certificate{
		SerialNumber:          serialNumber,
		Subject:               subject,
		NotBefore:             now,
		NotAfter:              now.Add(time.Hour * 24),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageKeyAgreement,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
		IsCA:                  false,
		URIs:                  []*url.URL{arnUrl},
	}
	caCert, err := x509.ParseCertificate(state.caCertDer)
	if err != nil {
		return nil, "", err
	}
	certDER, err := x509.CreateCertificate(rand.Reader, &template, caCert,
		publicKey, state.Signer)
	if err != nil {
		return nil, "", err
	}
	state.logger.Printf("Generated x509 Certificate for ARN=`%s`, expires=%s",
		callerArn.parsedArn.String(), template.NotAfter)
	metricLogCertDuration("x509", "granted",
		float64(time.Until(template.NotAfter).Seconds()))
	go func(username string, certType string) {
		metricsMutex.Lock()
		defer metricsMutex.Unlock()
		certGenCounter.WithLabelValues(username, certType).Inc()
	}(commonName, "x509")
	return certDER, commonName, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"testing"
	"time"
)

// ResetFlags restores the command line flags to their default
// values.  This does NOT reset the result of Parsed().
func ResetFlags() {
	flag.VisitAll(func(f *flag.Flag) {
		f.Value.Set(f.DefValue)
	})
}

func Test_queueFor(t *testing.T) {
	// Set the flag values?
	flag.CommandLine.Parse([]string{"-queue", "base-", "-dry_run"})
	defer ResetFlags()

	tme, _ := time.Parse("2006/01/02", "2017/09/01")
	queue := queueFor(tme)
	if queue != "base-2" {
		t.Error("bad queue: ", queue)
	}
}

type countingHTTP struct {
	count   int
	lastURL string
}

func (h *countingHTTP) Get(url string) (resp *http.Response, err error) {
	h.count++
	h.lastURL = url
	resp = &http.Response{}
	resp.Body = ioutil.NopCloser(bytes.NewReader([]byte{}))
	resp.Status = "200 OK"
	resp.StatusCode = 200
	return
}

// NOTE: This test depends on the contents of the test-archive-mlab-sandbox
// bucket.  If it fails, check whether that bucket has been modified.
func Test_month(t *testing.T) {
	flag.CommandLine.Parse([]string{
		"-queue", "base-", "-dry_run", "-bucket=test-archive-mlab-sandbox",
		"-month=2017/10"})
	defer ResetFlags()

	fake := countingHTTP{}
	setup(&fake)
	run()
	if fake.count != 100 {
		t.Error("Should have made 100 http requests.")
	}
}

// NOTE: This test depends on the contents of the test-archive-mlab-sandbox
// bucket.  If it fails, check whether that bucket has been modified.
func Example_month() {
	log.SetFlags(0)
	flag.CommandLine.Parse([]string{
		"-queue", "base-", "-dry_run", "-bucket=test-archive-mlab-sandbox",
		"-month=2017/10"})
	defer ResetFlags()

	log.SetOutput(os.Stdout) // Redirect to stdout so test framework sees it.
	setup(nil)
	run()
	log.SetOutput(os.Stderr) // restore
	// Unordered output:
	// Adding  ndt/2017/10/28/  to  base-3
	// Adding  ndt/2017/10/29/  to  base-4
	// Added  10  tasks to  base-3
	// Added  90  tasks to  base-4
}

// NOTE: This test depends on the contents of the test-archive-mlab-sandbox
// bucket.  If it fails, check whether that bucket has been modified.
func Example_day() {
	log.SetFlags(0)
	flag.CommandLine.Parse([]string{
		"-queue", "base-", "-dry_run", "-bucket=test-archive-mlab-sandbox",
		"-day=2017/09/24"})
	defer ResetFlags()

	log.SetOutput(os.Stdout) // Redirect to stdout so test framework sees it.
	setup(nil)
	run()
	log.SetOutput(os.Stderr) // restore
	// Output:
	// Adding  ndt/2017/09/24/  to  base-1
	// Added  76  tasks to  base-1
}
<commit_msg>fix outdated function name<commit_after>package main

import (
	"bytes"
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"testing"
	"time"
)

// ResetFlags restores the command line flags to their default
// values.  This does NOT reset the result of Parsed().
func ResetFlags() {
	flag.VisitAll(func(f *flag.Flag) {
		f.Value.Set(f.DefValue)
	})
}

func Test_queueFor(t *testing.T) {
	// Set the flag values?
	flag.CommandLine.Parse([]string{"-queue", "base-", "-dry_run"})
	defer ResetFlags()

	tme, _ := time.Parse("2006/01/02", "2017/09/01")
	queue := queueForDate(tme)
	if queue != "base-2" {
		t.Error("bad queue: ", queue)
	}
}

type countingHTTP struct {
	count   int
	lastURL string
}

func (h *countingHTTP) Get(url string) (resp *http.Response, err error) {
	h.count++
	h.lastURL = url
	resp = &http.Response{}
	resp.Body = ioutil.NopCloser(bytes.NewReader([]byte{}))
	resp.Status = "200 OK"
	resp.StatusCode = 200
	return
}

// NOTE: This test depends on the contents of the test-archive-mlab-sandbox
// bucket.  If it fails, check whether that bucket has been modified.
func Test_month(t *testing.T) {
	flag.CommandLine.Parse([]string{
		"-queue", "base-", "-dry_run", "-bucket=test-archive-mlab-sandbox",
		"-month=2017/10"})
	defer ResetFlags()

	fake := countingHTTP{}
	setup(&fake)
	run()
	if fake.count != 100 {
		t.Error("Should have made 100 http requests.")
	}
}

// NOTE: This test depends on the contents of the test-archive-mlab-sandbox
// bucket.  If it fails, check whether that bucket has been modified.
func Example_month() {
	log.SetFlags(0)
	flag.CommandLine.Parse([]string{
		"-queue", "base-", "-dry_run", "-bucket=test-archive-mlab-sandbox",
		"-month=2017/10"})
	defer ResetFlags()

	log.SetOutput(os.Stdout) // Redirect to stdout so test framework sees it.
	setup(nil)
	run()
	log.SetOutput(os.Stderr) // restore
	// Unordered output:
	// Adding  ndt/2017/10/28/  to  base-3
	// Adding  ndt/2017/10/29/  to  base-4
	// Added  10  tasks to  base-3
	// Added  90  tasks to  base-4
}

// NOTE: This test depends on the contents of the test-archive-mlab-sandbox
// bucket.  If it fails, check whether that bucket has been modified.
func Example_day() {
	log.SetFlags(0)
	flag.CommandLine.Parse([]string{
		"-queue", "base-", "-dry_run", "-bucket=test-archive-mlab-sandbox",
		"-day=2017/09/24"})
	defer ResetFlags()

	log.SetOutput(os.Stdout) // Redirect to stdout so test framework sees it.
	setup(nil)
	run()
	log.SetOutput(os.Stderr) // restore
	// Output:
	// Adding  ndt/2017/09/24/  to  base-1
	// Added  76  tasks to  base-1
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/nelhage/taktician/ai"
	"github.com/nelhage/taktician/playtak"
	"github.com/nelhage/taktician/tak"
)

const (
	minThink = 2 * time.Second
	maxThink = time.Minute

	defaultLevel = 2
)

var commandRE = regexp.MustCompile(`^([^ :]+):\s*([^ ]+)\s*(.*)$`)

type Friendly struct {
	client *playtak.Client
	ai     *ai.MinimaxAI
	g      *Game

	level    int
	levelSet time.Time
}

func (f *Friendly) NewGame(g *Game) {
	if time.Now().Sub(f.levelSet) > 1*time.Hour {
		f.level = defaultLevel
	}
	f.g = g
	f.ai = ai.NewMinimax(f.Config())
	f.client.SendCommand("Shout",
		fmt.Sprintf("[FriendlyBot@level %d] Good luck %s!",
			f.level, g.opponent,
		))
}

func (f *Friendly) GameOver() {
	f.g = nil
}

func (f *Friendly) GetMove(p *tak.Position, mine, theirs time.Duration) tak.Move {
	deadline := time.After(minThink)
	m := f.ai.GetMove(p, maxThink)
	<-deadline
	return m
}

func (f *Friendly) HandleChat(who string, msg string) {
	log.Printf("chat from=%q msg=%q", who, msg)
	if f.g != nil && who != f.g.opponent {
		return
	}
	gs := commandRE.FindStringSubmatch(msg)
	if gs == nil {
		return
	}
	if !strings.EqualFold(gs[1], *user) {
		return
	}
	switch gs[2] {
	case "level":
		if gs[3] == "max" {
			f.level = 100
			f.levelSet = time.Now()
			f.client.SendCommand("Shout", "OK! I'll play as best as I can!")
			break
		}
		l, e := strconv.ParseUint(strings.Trim(gs[3], " "), 10, 64)
		if e != nil {
			log.Printf("bad level: %v", e)
			return
		}
		if int(l) < 1 || int(l) > len(levels)+1 {
			f.client.SendCommand("Shout", fmt.Sprintf("I only know about levels up to %d", len(levels)+1))
		}
		f.level = int(l)
		f.levelSet = time.Now()
		f.client.SendCommand("Shout", fmt.Sprintf("OK! I'll play at level %d from now on.", l))
		if f.g != nil {
			f.ai = ai.NewMinimax(f.Config())
		}
	}
}

func (f *Friendly) Config() ai.MinimaxConfig {
	cfg := ai.MinimaxConfig{
		Size:  f.g.size,
		Debug: *debug,

		NoSort:  !*sort,
		NoTable: !*table,
	}
	cfg.Depth, cfg.Evaluate = f.levelSettings(f.g.size, f.level)

	return cfg
}

var (
	easyWeights = ai.Weights{
		TopFlat: 100,
		Tempo:   50,
	}
	medWeights = ai.Weights{
		TopFlat:  200,
		Standing: 200,
		Capstone: 200,
		Flat:     100,
		Tempo:    100,
		Groups:   [8]int{0, 0, 0, 100, 200},
	}
)

var levels = []struct {
	depth   int
	weights ai.Weights
}{
	{3, easyWeights},
	{4, easyWeights},
	{2, ai.DefaultWeights},
	{3, medWeights},
	{4, medWeights},
	{3, ai.DefaultWeights},
	{5, easyWeights},
	{5, medWeights},
	{4, ai.DefaultWeights},
	{5, ai.DefaultWeights},
}

func (f *Friendly) levelSettings(size int, level int) (int, ai.EvaluationFunc) {
	if level == 0 {
		level = 3
	}
	if level > len(levels)+1 {
		return 7, ai.DefaultEvaluate
	}
	s := levels[level-1]
	return s.depth, ai.MakeEvaluator(&s.weights)
}
<commit_msg>tweak levels slightly<commit_after>package main

import (
	"fmt"
	"log"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/nelhage/taktician/ai"
	"github.com/nelhage/taktician/playtak"
	"github.com/nelhage/taktician/tak"
)

const (
	minThink = 2 * time.Second
	maxThink = time.Minute

	defaultLevel = 2
)

var commandRE = regexp.MustCompile(`^([^ :]+):\s*([^ ]+)\s*(.*)$`)

type Friendly struct {
	client *playtak.Client
	ai     *ai.MinimaxAI
	g      *Game

	level    int
	levelSet time.Time
}

func (f *Friendly) NewGame(g *Game) {
	if time.Now().Sub(f.levelSet) > 1*time.Hour {
		f.level = defaultLevel
	}
	f.g = g
	f.ai = ai.NewMinimax(f.Config())
	f.client.SendCommand("Shout",
		fmt.Sprintf("[FriendlyBot@level %d] Good luck %s!",
			f.level, g.opponent,
		))
}

func (f *Friendly) GameOver() {
	f.g = nil
}

func (f *Friendly) GetMove(p *tak.Position, mine, theirs time.Duration) tak.Move {
	deadline := time.After(minThink)
	m := f.ai.GetMove(p, maxThink)
	<-deadline
	return m
}

func (f *Friendly) HandleChat(who string, msg string) {
	log.Printf("chat from=%q msg=%q", who, msg)
	if f.g != nil && who != f.g.opponent {
		return
	}
	gs := commandRE.FindStringSubmatch(msg)
	if gs == nil {
		return
	}
	if !strings.EqualFold(gs[1], *user) {
		return
	}
	switch gs[2] {
	case "level":
		if gs[3] == "max" {
			f.level = 100
			f.levelSet = time.Now()
			f.client.SendCommand("Shout", "OK! I'll play as best as I can!")
			break
		}
		l, e := strconv.ParseUint(strings.Trim(gs[3], " "), 10, 64)
		if e != nil {
			log.Printf("bad level: %v", e)
			return
		}
		if int(l) < 1 || int(l) > len(levels)+1 {
			f.client.SendCommand("Shout", fmt.Sprintf("I only know about levels up to %d", len(levels)+1))
		}
		f.level = int(l)
		f.levelSet = time.Now()
		f.client.SendCommand("Shout", fmt.Sprintf("OK! I'll play at level %d from now on.", l))
		if f.g != nil {
			f.ai = ai.NewMinimax(f.Config())
		}
	}
}

func (f *Friendly) Config() ai.MinimaxConfig {
	cfg := ai.MinimaxConfig{
		Size:  f.g.size,
		Debug: *debug,

		NoSort:  !*sort,
		NoTable: !*table,
	}
	cfg.Depth, cfg.Evaluate = f.levelSettings(f.g.size, f.level)

	return cfg
}

var (
	easyWeights = ai.Weights{
		TopFlat: 100,
		Tempo:   50,
	}
	medWeights = ai.Weights{
		TopFlat:  200,
		Standing: 200,
		Capstone: 200,
		Flat:     100,
		Tempo:    100,
		Groups:   [8]int{0, 0, 0, 100, 200},
	}
)

var levels = []struct {
	depth   int
	weights ai.Weights
}{
	{2, easyWeights},
	{2, medWeights},
	{2, ai.DefaultWeights},
	{3, medWeights},
	{4, medWeights},
	{3, ai.DefaultWeights},
	{5, easyWeights},
	{5, medWeights},
	{4, ai.DefaultWeights},
	{5, ai.DefaultWeights},
}

func (f *Friendly) levelSettings(size int, level int) (int, ai.EvaluationFunc) {
	if level == 0 {
		level = 3
	}
	if level > len(levels)+1 {
		return 7, ai.DefaultEvaluate
	}
	s := levels[level-1]
	return s.depth, ai.MakeEvaluator(&s.weights)
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"github.com/notegio/openrelay/funds"
	"github.com/notegio/openrelay/types"
	"github.com/notegio/openrelay/config"
	"log"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
)

func hexToBytes(address string) [20]byte {
	slice, err := hex.DecodeString(address)
	if err != nil {
		return [20]byte{}
	}
	output := [20]byte{}
	copy(output[:], slice[:])
	return output
}

func main() {
	rpcURL := os.Args[1]
	orderFile := os.Args[2]
	feeToken, err := config.NewRpcFeeToken(rpcURL)
	if err != nil {
		log.Fatalf("Error creating RpcOrderValidator: '%v'", err.Error())
	}
	tokenProxy, err := config.NewRpcTokenProxy(rpcURL)
	if err != nil {
		log.Fatalf("Error creating RpcOrderValidator: '%v'", err.Error())
	}
	fundChecker, err := funds.NewRpcOrderValidator(rpcURL, feeToken, tokenProxy)
	if err != nil {
		log.Fatalf(err.Error())
	}
	newOrder := types.Order{}
	if orderData, err := ioutil.ReadFile(orderFile); err == nil {
		if err := json.Unmarshal(orderData, &newOrder); err != nil {
			log.Fatalf(err.Error())
		}
	}
	fmt.Println(fundChecker.ValidateOrder(&newOrder))
}
<commit_msg>Update order validator<commit_after>package main

import (
	"os"
	"github.com/notegio/openrelay/funds"
	"github.com/notegio/openrelay/types"
	"github.com/notegio/openrelay/config"
	"log"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
)

func hexToBytes(address string) [20]byte {
	slice, err := hex.DecodeString(address)
	if err != nil {
		return [20]byte{}
	}
	output := [20]byte{}
	copy(output[:], slice[:])
	return output
}

func main() {
	rpcURL := os.Args[1]
	orderFile := os.Args[2]
	feeToken, err := config.NewRpcFeeToken(rpcURL)
	if err != nil {
		log.Fatalf("Error creating RpcOrderValidator: '%v'", err.Error())
	}
	tokenProxy, err := config.NewRpcTokenProxy(rpcURL)
	if err != nil {
		log.Fatalf("Error creating RpcOrderValidator: '%v'", err.Error())
	}
	fundChecker, err := funds.NewRpcOrderValidator(rpcURL, feeToken, tokenProxy)
	if err != nil {
		log.Fatalf(err.Error())
	}
	filledLookup, err := funds.NewRpcFilledLookup(rpcURL)
	if err != nil {
		log.Fatalf(err.Error())
	}
	newOrder := types.Order{}
	if orderData, err := ioutil.ReadFile(orderFile); err == nil {
		if err := json.Unmarshal(orderData, &newOrder); err != nil {
			log.Fatalf(err.Error())
		}
	}
	newOrder.TakerTokenAmountFilled, err = filledLookup.GetAmountFilled(&newOrder)
	if err != nil {
		log.Fatalf(err.Error())
	}
	newOrder.TakerTokenAmountCancelled, err = filledLookup.GetAmountCancelled(&newOrder)
	if err != nil {
		log.Fatalf(err.Error())
	}
	fmt.Println(fundChecker.ValidateOrder(&newOrder))
}
<|endoftext|>


<commit_before>package kcp_test

import (
	"crypto/rand"
	"io"
	"net"
	"testing"
	"time"

	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/testing/assert"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/internet/authenticators/srtp"
	. "v2ray.com/core/transport/internet/kcp"
)

type NoOpWriteCloser struct{}

func (this *NoOpWriteCloser) Write(b []byte) (int, error) {
	return len(b), nil
}

func (this *NoOpWriteCloser) Close() error {
	return nil
}

func TestConnectionReadTimeout(t *testing.T) {
	assert := assert.On(t)

	conn := NewConnection(1, &NoOpWriteCloser{}, nil, nil, NewSimpleAuthenticator(), &Config{})
	conn.SetReadDeadline(time.Now().Add(time.Second))

	b := make([]byte, 1024)
	nBytes, err := conn.Read(b)
	assert.Int(nBytes).Equals(0)
	assert.Error(err).IsNotNil()

	conn.Terminate()
}

func TestConnectionReadWrite(t *testing.T) {
	assert := assert.On(t)

	upReader, upWriter := io.Pipe()
	downReader, downWriter := io.Pipe()

	auth := internet.NewAuthenticatorChain(srtp.SRTPFactory{}.Create(nil), NewSimpleAuthenticator())

	connClient := NewConnection(1, upWriter, &net.UDPAddr{IP: v2net.LocalHostIP.IP(), Port: 1}, &net.UDPAddr{IP: v2net.LocalHostIP.IP(), Port: 2}, auth, &Config{})
	connClient.FetchInputFrom(downReader)

	connServer := NewConnection(1, downWriter, &net.UDPAddr{IP: v2net.LocalHostIP.IP(), Port: 2}, &net.UDPAddr{IP: v2net.LocalHostIP.IP(), Port: 1}, auth, &Config{})
	connServer.FetchInputFrom(upReader)

	totalWritten := 1024 * 1024
	clientSend := make([]byte, totalWritten)
	rand.Read(clientSend)
	go func() {
		nBytes, err := connClient.Write(clientSend)
		assert.Int(nBytes).Equals(totalWritten)
		assert.Error(err).IsNil()
	}()

	serverReceived := make([]byte, totalWritten)
	totalRead := 0
	for totalRead < totalWritten {
		nBytes, err := connServer.Read(serverReceived[totalRead:])
		assert.Error(err).IsNil()
		totalRead += nBytes
	}
	assert.Bytes(serverReceived).Equals(clientSend)

	connClient.Close()
	connServer.Close()

	for connClient.State() != StateTerminated || connServer.State() != StateTerminated {
		time.Sleep(time.Second)
	}
}
<commit_msg>remove buggy test case<commit_after>package kcp_test

import (
	"testing"
	"time"

	"v2ray.com/core/testing/assert"
	. "v2ray.com/core/transport/internet/kcp"
)

type NoOpWriteCloser struct{}

func (this *NoOpWriteCloser) Write(b []byte) (int, error) {
	return len(b), nil
}

func (this *NoOpWriteCloser) Close() error {
	return nil
}

func TestConnectionReadTimeout(t *testing.T) {
	assert := assert.On(t)

	conn := NewConnection(1, &NoOpWriteCloser{}, nil, nil, NewSimpleAuthenticator(), &Config{})
	conn.SetReadDeadline(time.Now().Add(time.Second))

	b := make([]byte, 1024)
	nBytes, err := conn.Read(b)
	assert.Int(nBytes).Equals(0)
	assert.Error(err).IsNotNil()

	conn.Terminate()
}
<|endoftext|>


<commit_before>// Copyright 2013 Artur Grabowski. All rights reserved.
// Use of this source code is governed by a ISC-style
// license that can be found in the LICENSE file.
package index

import (
	"bconf"
)

type Index struct {
	br     *blob_reader
	Docs   map[uint32][]byte
	Attrs  map[string][]IbDoc
	Meta   bconf.Bconf
	header string
}

func Open(name string) (*Index, error) {
	var in Index
	var err error

	in.br, err = open_blob_reader(name)
	if err != nil {
		return nil, err
	}

	in.Docs = make(map[uint32][]byte)
	for _, d := range in.br.get_documents() {
		in.Docs[d.Doc.Id] = in.br.get_document_data(&d)
	}

	in.Attrs = make(map[string][]IbDoc)
	for _, a := range in.br.get_invattrs() {
		in.Attrs[in.br.get_attr_name(&a)] = in.br.get_attr_docs(&a)
	}

	in.Meta.LoadJson(in.br.get_meta())

	in.Header() // Pre-cache the header to avoid race conditions.

	return &in, nil
}

func (in Index) Header() string {
	if in.header == "" {
		in.Meta.GetNode("attr", "order").ForeachSorted(func(k, v string) {
			if in.header != "" {
				in.header += "\t"
			}
			in.header += v
		})
	}
	return in.header
}

func (in Index) Close() {
	in.br.close()
}
<commit_msg>Use the bconf from github.<commit_after>// Copyright 2013 Artur Grabowski. All rights reserved.
// Use of this source code is governed by a ISC-style
// license that can be found in the LICENSE file.
package index

import (
	"github.com/art4711/bconf"
)

type Index struct {
	br     *blob_reader
	Docs   map[uint32][]byte
	Attrs  map[string][]IbDoc
	Meta   bconf.Bconf
	header string
}

func Open(name string) (*Index, error) {
	var in Index
	var err error

	in.br, err = open_blob_reader(name)
	if err != nil {
		return nil, err
	}

	in.Docs = make(map[uint32][]byte)
	for _, d := range in.br.get_documents() {
		in.Docs[d.Doc.Id] = in.br.get_document_data(&d)
	}

	in.Attrs = make(map[string][]IbDoc)
	for _, a := range in.br.get_invattrs() {
		in.Attrs[in.br.get_attr_name(&a)] = in.br.get_attr_docs(&a)
	}

	in.Meta.LoadJson(in.br.get_meta())

	in.Header() // Pre-cache the header to avoid race conditions.

	return &in, nil
}

func (in Index) Header() string {
	if in.header == "" {
		in.Meta.GetNode("attr", "order").ForeachSorted(func(k, v string) {
			if in.header != "" {
				in.header += "\t"
			}
			in.header += v
		})
	}
	return in.header
}

func (in Index) Close() {
	in.br.close()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/jpfairbanks/cse6140/project/hashes"
	"math/rand"
)

//CMSketch: A streaming data structure for tracking the elements of
//a long vector in a space efficient manner.
//Provides update and query operations
type CMSketch struct {
	Depth   int64
	Width   int64
	Hash    []hashes.Hash
	Counter Matrix
}

//NewCMSketch: Allocates a new sketch
//Does not create the hashes. We leave that up to you so that
//the consumer can pick hashes according to different schema.
func NewCMSketch(Depth, Width int64) CMSketch {
	Hash := make([]hashes.Hash, Depth)
	Counter := NewMatrix(Depth, Width)
	cms := CMSketch{Depth, Width, Hash, Counter}
	return cms
}

//RandomHashes: Make uniformly random hashes from a user provided
//stream of random numbers.
func RandomHashes(r *rand.Rand, Depth int64) []hashes.Hash {
	hslice := make([]hashes.Hash, Depth)
	for i, _ := range hslice {
		hslice[i] = hashes.Rand(r)
	}
	return hslice
}

//UpdateSerial: Insert a single item into the sketch with a count.
//If counts can be negative, then you mist estimate differently
func (cms *CMSketch) UpdateSerial(position int64, count int64) {
	for i, h := range cms.Hash {
		cms.Counter.Add(int64(i), h.Apply(position)%cms.Width, count)
	}
}

//AddSignal: Updates the counter and signals completion on a channel
func (cms *CMSketch) AddSignal(row int64, position int64,
	count int64, ch chan int64) {
	h := cms.Hash[row]
	cms.Counter.Add(row, h.Apply(position)%cms.Width, count)
	ch <- 1
}

//UpdateDepthParallel: Insert a single item into the sketch with a count.
//If counts can be negative, then you mist estimate differently
func (cms *CMSketch) UpdateDepthParallel(position int64, count int64) {
	ch := make(chan int64, cms.Depth)
	var i int64
	for i = 0; i < cms.Depth; i++ {
		go cms.AddSignal(i, position, count, ch)
	}
	for i = 0; i < cms.Depth; i++ {
		<-ch
	}
}

//PointQuery: query the value at position
func (cms *CMSketch) PointQuery(position int64) int64 {
	temp := make([]int64, cms.Depth)
	var h hashes.Hash
	var i int64
	for i = 0; i < cms.Depth; i++ {
		h = cms.Hash[i]
		temp[i] = cms.Counter.Read(i, h.Apply(position)%cms.Width)
	}
	mini := temp[0]
	for i = 0; i < cms.Depth; i++ {
		if mini > temp[i] {
			mini = temp[i]
		}
	}
	return mini
}

func main() {
	fmt.Println("starting main")
	src := rand.NewSource(0)
	r := rand.New(src)
	var Depth, Width int64
	Depth = 160
	Width = 480
	hslice := RandomHashes(r, Depth)
	cms := NewCMSketch(Depth, Width)
	cms.Hash = hslice
	fmt.Printf("%s\n", cms.Counter.String())
	fmt.Printf("Inserting\n")
	cms.UpdateSerial(1, 1)
	fmt.Printf("%s\n", cms.Counter.String())
	fmt.Printf("Inserting\n")
	var j, z int64
	var s, v float64
	var imax uint64
	s = 1.2
	v = 1.0
	imax = 2 << 10
	zipfer := rand.NewZipf(r, s, v, imax)
	set := make(map[int64]int64)
	for j = 0; j < Depth*Width*10; j++ {
		z = int64(zipfer.Uint64())
		set[z] += 1
		fmt.Println(z)
		cms.UpdateSerial(z, 1)
	}
	fmt.Printf("%s\n", cms.Counter.String())
	var qj int64
	for j, cj := range set {
		qj = cms.PointQuery(j)
		fmt.Printf("results:%d %d %d %f\n", j, qj, cj, float64(qj)/float64(cj))
	}
}
<commit_msg>add logging and command line flags<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/jpfairbanks/cse6140/project/hashes"
	"log"
	"math/rand"
)

//CMSketch: A streaming data structure for tracking the elements of
//a long vector in a space efficient manner.
//Provides update and query operations
type CMSketch struct {
	Depth   int64
	Width   int64
	Hash    []hashes.Hash
	Counter Matrix
}

//NewCMSketch: Allocates a new sketch
//Does not create the hashes. We leave that up to you so that
//the consumer can pick hashes according to different schema.
func NewCMSketch(Depth, Width int64) CMSketch {
	Hash := make([]hashes.Hash, Depth)
	Counter := NewMatrix(Depth, Width)
	cms := CMSketch{Depth, Width, Hash, Counter}
	return cms
}

//RandomHashes: Make uniformly random hashes from a user provided
//stream of random numbers.
func RandomHashes(r *rand.Rand, Depth int64) []hashes.Hash {
	hslice := make([]hashes.Hash, Depth)
	for i, _ := range hslice {
		hslice[i] = hashes.Rand(r)
	}
	return hslice
}

//UpdateSerial: Insert a single item into the sketch with a count.
//If counts can be negative, then you mist estimate differently
func (cms *CMSketch) UpdateSerial(position int64, count int64) {
	for i, h := range cms.Hash {
		cms.Counter.Add(int64(i), h.Apply(position)%cms.Width, count)
	}
}

//AddSignal: Updates the counter and signals completion on a channel
func (cms *CMSketch) AddSignal(row int64, position int64,
	count int64, ch chan int64) {
	h := cms.Hash[row]
	cms.Counter.Add(row, h.Apply(position)%cms.Width, count)
	ch <- 1
}

//UpdateDepthParallel: Insert a single item into the sketch with a count.
//If counts can be negative, then you mist estimate differently
func (cms *CMSketch) UpdateDepthParallel(position int64, count int64) {
	ch := make(chan int64, cms.Depth)
	var i int64
	for i = 0; i < cms.Depth; i++ {
		go cms.AddSignal(i, position, count, ch)
	}
	for i = 0; i < cms.Depth; i++ {
		<-ch
	}
}

//PointQuery: query the value at position
func (cms *CMSketch) PointQuery(position int64) int64 {
	temp := make([]int64, cms.Depth)
	var h hashes.Hash
	var i int64
	for i = 0; i < cms.Depth; i++ {
		h = cms.Hash[i]
		temp[i] = cms.Counter.Read(i, h.Apply(position)%cms.Width)
	}
	mini := temp[0]
	for i = 0; i < cms.Depth; i++ {
		if mini > temp[i] {
			mini = temp[i]
		}
	}
	return mini
}

var depthPtr, widthPtr, efactorPtr *int64

func init() {
	depthPtr = flag.Int64("depth", 50, "sets the number of rows (and hash functions) in the CMSketch")
	widthPtr = flag.Int64("width", 80, "sets the number of columns in the CMSketch")
	efactorPtr = flag.Int64("efactor", 10, "number of elements = depth*width*efactor")
	flag.Parse()
}

func main() {
	//log.New(os.Stderr, "", log.LstdFlags)
	log.Printf("starting main\n")
	src := rand.NewSource(0)
	r := rand.New(src)
	var Depth, Width, efactor, numElements int64
	Depth = *depthPtr
	Width = *widthPtr
	efactor = *efactorPtr
	numElements = Depth * Width * efactor
	log.Printf("params:Depth:%d\n", Depth)
	log.Printf("params:Width:%d\n", Width)
	log.Printf("params:efactor:%d\n", efactor)
	log.Printf("params:numElements:%d\n", numElements)
	hslice := RandomHashes(r, Depth)
	cms := NewCMSketch(Depth, Width)
	cms.Hash = hslice
	fmt.Printf("%s\n", cms.Counter.String())
	log.Printf("Inserting\n")
	cms.UpdateSerial(1, 1)
	fmt.Printf("%s\n", cms.Counter.String())
	log.Printf("Inserting\n")
	//Make the zipf distribution of random input
	var j, z int64
	var s, v float64
	var imax uint64
	s = 1.2
	v = 1.0
	imax = 2 << 10
	zipfer := rand.NewZipf(r, s, v, imax)

	//Use set to store the exact answers
	set := make(map[int64]int64)
	for j = 0; j < numElements; j++ {
		z = int64(zipfer.Uint64())
		set[z] += 1
		fmt.Println(z)
		cms.UpdateSerial(z, 1)
	}
	fmt.Printf("%s\n", cms.Counter.String())
	var qj int64 //approximate answers
	var totalLoss float64
	loss := func(cj, qj float64) float64 {
		return (cj - qj) * (cj - qj)
	}
	for j, cj := range set {
		qj = cms.PointQuery(j)
		fmt.Printf("results:%d %d %d %f\n", j, qj, cj, float64(qj)/float64(cj))
		totalLoss += loss(float64(cj), float64(qj))
	}
	fmt.Printf("Total Loss: %f/%d\n", totalLoss, numElements)
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package psbt

// The Creator has only one function:
// takes a list of inputs and outputs and converts them to
// the simplest Psbt, i.e. one with no data appended to inputs or outputs,
// but only the raw unsigned transaction in the global section.

import (
	"github.com/btcsuite/btcd/wire"
)

// Creator holds a reference to a created Psbt struct.
type Creator struct {
	Cpsbt *Psbt
}

// CreatePsbt , on provision of an input and output 'skeleton' for
// the transaction, returns a Creator struct.
// Note that we require OutPoints and not TxIn structs, as we will
// only populate the txid:n information, *not* any scriptSig/witness
// information. The values of nLockTime, nSequence (per input) and
// transaction version (must be 1 of 2) must be specified here. Note
// that the default nSequence value is wire.MaxTxInSequenceNum.
func (c *Creator) CreatePsbt(inputs []*wire.OutPoint,
	outputs []*wire.TxOut, Version int32, nLockTime uint32,
	nSequences []uint32) error {
	// Create the new struct; the input and output lists will be empty,
	// the unsignedTx object must be constructed and serialized,
	// and that serialization should be entered as the only entry for
	// the globalKVPairs list.

	// Check the version is a valid Bitcoin tx version; the nLockTime
	// can be any valid uint32. There must be one sequence number per
	// input.
	if !(Version == 1 || Version == 2) || len(nSequences) != len(inputs) {
		return ErrInvalidPsbtFormat
	}
	unsignedTx := wire.NewMsgTx(Version)
	unsignedTx.LockTime = nLockTime

	for i, in := range inputs {
		unsignedTx.AddTxIn(&wire.TxIn{
			PreviousOutPoint: *in,
			Sequence:         nSequences[i],
		})
	}
	for _, out := range outputs {
		unsignedTx.AddTxOut(out)
	}

	// The input and output lists are empty, but there is a list of those
	// two lists, and each one must be of length matching the unsigned
	// transaction; the unknown list can be nil.
	pInputs := make([]PInput, len(unsignedTx.TxIn))
	pOutputs := make([]POutput, len(unsignedTx.TxOut))
	c.Cpsbt = &Psbt{
		UnsignedTx: unsignedTx,
		Inputs:     pInputs,
		Outputs:    pOutputs,
		Unknowns:   nil,
	}

	// This new Psbt is "raw" and contains no key-value fields,
	// so sanity checking with c.Cpsbt.SanityCheck() is not required.
	return nil
}
<commit_msg>psbt: remove Creator struct, create New function as entry point into package<commit_after>// Copyright (c) 2018 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package psbt

import (
	"github.com/btcsuite/btcd/wire"
)

// MinTxVersion is the lowest transaction version that we'll permit.
const MinTxVersion = 1

// New on provision of an input and output 'skeleton' for the transaction, a
// new partially populated PBST packet. The populated packet will include the
// unsigned transaction, and the set of known inputs and outputs contained
// within the unsigned transaction.  The values of nLockTime, nSequence (per
// input) and transaction version (must be 1 of 2) must be specified here. Note
// that the default nSequence value is wire.MaxTxInSequenceNum.  Referencing
// the PSBT BIP, this function serves the roles of teh Creator.
func New(inputs []*wire.OutPoint,
	outputs []*wire.TxOut, version int32, nLockTime uint32,
	nSequences []uint32) (*Packet, error) {

	// Create the new struct; the input and output lists will be empty, the
	// unsignedTx object must be constructed and serialized, and that
	// serialization should be entered as the only entry for the
	// globalKVPairs list.
	//
	// Ensure that the version of the transaction is greater then our
	// minimum allowed transaction version. There must be one sequence
	// number per input.
	if version < MinTxVersion || len(nSequences) != len(inputs) {
		return nil, ErrInvalidPsbtFormat
	}

	unsignedTx := wire.NewMsgTx(version)
	unsignedTx.LockTime = nLockTime
	for i, in := range inputs {
		unsignedTx.AddTxIn(&wire.TxIn{
			PreviousOutPoint: *in,
			Sequence:         nSequences[i],
		})
	}
	for _, out := range outputs {
		unsignedTx.AddTxOut(out)
	}

	// The input and output lists are empty, but there is a list of those
	// two lists, and each one must be of length matching the unsigned
	// transaction; the unknown list can be nil.
	pInputs := make([]PInput, len(unsignedTx.TxIn))
	pOutputs := make([]POutput, len(unsignedTx.TxOut))

	// This new Psbt is "raw" and contains no key-value fields, so sanity
	// checking with c.Cpsbt.SanityCheck() is not required.
	return &Packet{
		UnsignedTx: unsignedTx,
		Inputs:     pInputs,
		Outputs:    pOutputs,
		Unknowns:   nil,
	}, nil
}
<|endoftext|>