<commit_before>package collector

import (
	"github.com/StackExchange/wmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
)

func init() {
	Factories["thermalzone"] = NewThermalZoneCollector
}

// A ThermalZoneCollector is a Prometheus collector for WMI Win32_PerfRawData_Counters_ThermalZoneInformation metrics
type ThermalZoneCollector struct {
	PercentPassiveLimit *prometheus.Desc
	Temperature         *prometheus.Desc
	ThrottleReasons     *prometheus.Desc
}

// NewThermalZoneCollector ...
func NewThermalZoneCollector() (Collector, error) {
	const subsystem = "thermalzone"
	return &ThermalZoneCollector{
		Temperature: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "temperature_celsius"),
			"(Temperature)",
			[]string{
				"name",
			},
			nil,
		),
		PercentPassiveLimit: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "percent_passive_limit"),
			"(PercentPassiveLimit)",
			[]string{
				"name",
			},
			nil,
		),
		ThrottleReasons: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "throttle_reasons"),
			"(ThrottleReasons)",
			[]string{
				"name",
			},
			nil,
		),
	}, nil
}

// Collect sends the metric values for each metric
// to the provided prometheus Metric channel.
func (c *ThermalZoneCollector) Collect(ch chan<- prometheus.Metric) error {
	if desc, err := c.collect(ch); err != nil {
		log.Error("failed collecting thermalzone metrics:", desc, err)
		return err
	}
	return nil
}

// Win32_PerfRawData_Counters_ThermalZoneInformation docs:
// https://wutils.com/wmi/root/cimv2/win32_perfrawdata_counters_thermalzoneinformation/
type Win32_PerfRawData_Counters_ThermalZoneInformation struct {
	Name string

	HighPrecisionTemperature uint32
	PercentPassiveLimit      uint32
	ThrottleReasons          uint32
}

func (c *ThermalZoneCollector) collect(ch chan<- prometheus.Metric) (*prometheus.Desc, error) {
	var dst []Win32_PerfRawData_Counters_ThermalZoneInformation
	q := queryAll(&dst)
	if err := wmi.Query(q, &dst); err != nil {
		return nil, err
	}

	for _, info := range dst {
		//Divide by 10 and subtract 273.15 to convert decikelvin to celsius
		ch <- prometheus.MustNewConstMetric(
			c.Temperature,
			prometheus.GaugeValue,
			(float64(info.HighPrecisionTemperature)/10.0)-273.15,
			info.Name,
		)

		ch <- prometheus.MustNewConstMetric(
			c.PercentPassiveLimit,
			prometheus.GaugeValue,
			float64(info.PercentPassiveLimit),
			info.Name,
		)

		ch <- prometheus.MustNewConstMetric(
			c.ThrottleReasons,
			prometheus.GaugeValue,
			float64(info.ThrottleReasons),
			info.Name,
		)
	}

	return nil, nil
}
<commit_msg>Adapt thermalzone to new signature<commit_after>package collector

import (
	"github.com/StackExchange/wmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
)

func init() {
	Factories["thermalzone"] = NewThermalZoneCollector
}

// A thermalZoneCollector is a Prometheus collector for WMI Win32_PerfRawData_Counters_ThermalZoneInformation metrics
type thermalZoneCollector struct {
	PercentPassiveLimit *prometheus.Desc
	Temperature         *prometheus.Desc
	ThrottleReasons     *prometheus.Desc
}

// NewThermalZoneCollector ...
func NewThermalZoneCollector() (Collector, error) {
	const subsystem = "thermalzone"
	return &thermalZoneCollector{
		Temperature: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "temperature_celsius"),
			"(Temperature)",
			[]string{
				"name",
			},
			nil,
		),
		PercentPassiveLimit: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "percent_passive_limit"),
			"(PercentPassiveLimit)",
			[]string{
				"name",
			},
			nil,
		),
		ThrottleReasons: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "throttle_reasons"),
			"(ThrottleReasons)",
			[]string{
				"name",
			},
			nil,
		),
	}, nil
}

// Collect sends the metric values for each metric
// to the provided prometheus Metric channel.
func (c *thermalZoneCollector) Collect(ctx *ScrapeContext, ch chan<- prometheus.Metric) error {
	if desc, err := c.collect(ch); err != nil {
		log.Error("failed collecting thermalzone metrics:", desc, err)
		return err
	}
	return nil
}

// Win32_PerfRawData_Counters_ThermalZoneInformation docs:
// https://wutils.com/wmi/root/cimv2/win32_perfrawdata_counters_thermalzoneinformation/
type Win32_PerfRawData_Counters_ThermalZoneInformation struct {
	Name string

	HighPrecisionTemperature uint32
	PercentPassiveLimit      uint32
	ThrottleReasons          uint32
}

func (c *thermalZoneCollector) collect(ch chan<- prometheus.Metric) (*prometheus.Desc, error) {
	var dst []Win32_PerfRawData_Counters_ThermalZoneInformation
	q := queryAll(&dst)
	if err := wmi.Query(q, &dst); err != nil {
		return nil, err
	}

	for _, info := range dst {
		//Divide by 10 and subtract 273.15 to convert decikelvin to celsius
		ch <- prometheus.MustNewConstMetric(
			c.Temperature,
			prometheus.GaugeValue,
			(float64(info.HighPrecisionTemperature)/10.0)-273.15,
			info.Name,
		)

		ch <- prometheus.MustNewConstMetric(
			c.PercentPassiveLimit,
			prometheus.GaugeValue,
			float64(info.PercentPassiveLimit),
			info.Name,
		)

		ch <- prometheus.MustNewConstMetric(
			c.ThrottleReasons,
			prometheus.GaugeValue,
			float64(info.ThrottleReasons),
			info.Name,
		)
	}

	return nil, nil
}
<|endoftext|>


<commit_before>package lightstep

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/lightstep/lightstep-tracer-go/collectorpb"
	"golang.org/x/net/http2"
)

var (
	acceptHeader      = http.CanonicalHeaderKey("Accept")
	contentTypeHeader = http.CanonicalHeaderKey("Content-Type")
)

const (
	collectorHttpMethod = "POST"
	collectorHttpPath   = "/api/v2/reports"
	protoContentType    = "application/octet-stream"
)

// grpcCollectorClient specifies how to send reports back to a LightStep
// collector via grpc.
type httpCollectorClient struct {
	// auth and runtime information
	reporterID  uint64
	accessToken string // accessToken is the access token used for explicit trace collection requests.
	attributes  map[string]string

	reportTimeout time.Duration

	// Remote service that will receive reports.
	url    *url.URL
	client *http.Client

	// converters
	converter *protoConverter
}

type HttpRequest struct {
	http.Request
}

type transportCloser struct {
	transport http2.Transport
}

func (closer *transportCloser) Close() error {
	closer.transport.CloseIdleConnections()

	return nil
}

func newHttpCollectorClient(
	opts Options,
	reporterID uint64,
	attributes map[string]string,
) (*httpCollectorClient, error) {
	url, err := url.Parse(opts.Collector.URL())
	if err != nil {
		fmt.Println("collector config does not produce valid url", err)
		return nil, err
	}
	url.Path = collectorHttpPath

	return &httpCollectorClient{
		reporterID:    reporterID,
		accessToken:   opts.AccessToken,
		attributes:    attributes,
		reportTimeout: opts.ReportTimeout,
		url:           url,
		converter:     newProtoConverter(opts),
	}, nil
}

func (client *httpCollectorClient) ConnectClient() (Connection, error) {
	// The golang http2 client implementation doesn't support plaintext http2 (a.k.a h2c) out of the box.
	// According to https://github.com/golang/go/issues/14141, they don't have plans to.
	// For now, we are falling back to http1 for plaintext.
	// In the future, we might want to add out own h2c implementation (see https://github.com/hkwi/h2c).
	var transport http.RoundTripper
	if client.url.Scheme == "https" {
		transport = &http2.Transport{}
	} else {
		transport = &http.Transport{}
	}

	client.client = &http.Client{
		Transport: transport,
		Timeout:   client.reportTimeout,
	}

	return &transportCloser{}, nil
}

func (client *httpCollectorClient) ShouldReconnect() bool {
	// http2 will handle connection reuse under the hood
	return false
}

func (client *httpCollectorClient) Report(context context.Context, req reportRequest) (collectorResponse, error) {
	if req.httpRequest == nil {
		return nil, fmt.Errorf("httpRequest cannot be null")
	}

	httpResponse, err := client.client.Do(req.httpRequest)
	if err != nil {
		return nil, err
	}
	defer httpResponse.Body.Close()

	response, err := client.toResponse(httpResponse)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (client *httpCollectorClient) Translate(ctx context.Context, buffer *reportBuffer) (reportRequest, error) {

	httpRequest, err := client.toRequest(ctx, buffer)
	if err != nil {
		return reportRequest{}, err
	}
	return reportRequest{
		httpRequest: httpRequest,
	}, nil
}

func (client *httpCollectorClient) toRequest(
	context context.Context,
	buffer *reportBuffer,
) (*http.Request, error) {
	protoRequest := client.converter.toReportRequest(
		client.reporterID,
		client.attributes,
		client.accessToken,
		buffer,
	)

	buf, err := proto.Marshal(protoRequest)
	if err != nil {
		return nil, err
	}

	requestBody := bytes.NewReader(buf)

	request, err := http.NewRequest(collectorHttpMethod, client.url.String(), requestBody)
	if err != nil {
		return nil, err
	}
	request = request.WithContext(context)
	request.Header.Set(contentTypeHeader, protoContentType)
	request.Header.Set(acceptHeader, protoContentType)

	return request, nil
}

func (client *httpCollectorClient) toResponse(response *http.Response) (collectorResponse, error) {
	if response.StatusCode != http.StatusOK {
		return nil, errors.New(fmt.Sprintf("status code (%d) is not ok", response.StatusCode))
	}

	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	protoResponse := &collectorpb.ReportResponse{}
	if err := proto.Unmarshal(body, protoResponse); err != nil {
		return nil, err
	}

	return protoResponse, nil
}
<commit_msg>remove x/net/http2 dependency<commit_after>package lightstep

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/lightstep/lightstep-tracer-go/collectorpb"
)

var (
	acceptHeader      = http.CanonicalHeaderKey("Accept")
	contentTypeHeader = http.CanonicalHeaderKey("Content-Type")
)

const (
	collectorHttpMethod = "POST"
	collectorHttpPath   = "/api/v2/reports"
	protoContentType    = "application/octet-stream"
)

// grpcCollectorClient specifies how to send reports back to a LightStep
// collector via grpc.
type httpCollectorClient struct {
	// auth and runtime information
	reporterID  uint64
	accessToken string // accessToken is the access token used for explicit trace collection requests.
	attributes  map[string]string

	reportTimeout   time.Duration
	reportingPeriod time.Duration

	// Remote service that will receive reports.
	url    *url.URL
	client *http.Client

	// converters
	converter *protoConverter
}

type HttpRequest struct {
	http.Request
}

type transportCloser struct {
	*http.Transport
}

func (closer transportCloser) Close() error {
	closer.CloseIdleConnections()
	return nil
}

func newHttpCollectorClient(
	opts Options,
	reporterID uint64,
	attributes map[string]string,
) (*httpCollectorClient, error) {
	url, err := url.Parse(opts.Collector.URL())
	if err != nil {
		fmt.Println("collector config does not produce valid url", err)
		return nil, err
	}
	url.Path = collectorHttpPath

	return &httpCollectorClient{
		reporterID:      reporterID,
		accessToken:     opts.AccessToken,
		attributes:      attributes,
		reportTimeout:   opts.ReportTimeout,
		reportingPeriod: opts.ReportingPeriod,
		url:             url,
		converter:       newProtoConverter(opts),
	}, nil
}

func (client *httpCollectorClient) ConnectClient() (Connection, error) {
	// Use a transport independent from http.DefaultTransport to provide sane
	// defaults that make sense in the context of the lightstep client.
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   client.reportTimeout / 2,
			DualStack: true,
		}).DialContext,
		DisableCompression:     true,
		IdleConnTimeout:        2 * client.reportingPeriod,
		TLSHandshakeTimeout:    client.reportTimeout / 2,
		ResponseHeaderTimeout:  client.reportTimeout,
		ExpectContinueTimeout:  client.reportTimeout,
		MaxResponseHeaderBytes: 64 * 1024, // 64 KB, just a safeguard
	}

	client.client = &http.Client{
		Transport: transport,
		Timeout:   client.reportTimeout,
	}

	return transportCloser{transport}, nil
}

func (client *httpCollectorClient) ShouldReconnect() bool {
	// http.Transport will handle connection reuse under the hood
	return false
}

func (client *httpCollectorClient) Report(context context.Context, req reportRequest) (collectorResponse, error) {
	if req.httpRequest == nil {
		return nil, fmt.Errorf("httpRequest cannot be null")
	}

	httpResponse, err := client.client.Do(req.httpRequest.WithContext(context))
	if err != nil {
		return nil, err
	}
	defer httpResponse.Body.Close()

	response, err := client.toResponse(httpResponse)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (client *httpCollectorClient) Translate(ctx context.Context, buffer *reportBuffer) (reportRequest, error) {

	httpRequest, err := client.toRequest(ctx, buffer)
	if err != nil {
		return reportRequest{}, err
	}
	return reportRequest{
		httpRequest: httpRequest,
	}, nil
}

func (client *httpCollectorClient) toRequest(
	context context.Context,
	buffer *reportBuffer,
) (*http.Request, error) {
	protoRequest := client.converter.toReportRequest(
		client.reporterID,
		client.attributes,
		client.accessToken,
		buffer,
	)

	buf, err := proto.Marshal(protoRequest)
	if err != nil {
		return nil, err
	}

	requestBody := bytes.NewReader(buf)

	request, err := http.NewRequest(collectorHttpMethod, client.url.String(), requestBody)
	if err != nil {
		return nil, err
	}
	request = request.WithContext(context)
	request.Header.Set(contentTypeHeader, protoContentType)
	request.Header.Set(acceptHeader, protoContentType)

	return request, nil
}

func (client *httpCollectorClient) toResponse(response *http.Response) (collectorResponse, error) {
	if response.StatusCode != http.StatusOK {
		return nil, errors.New(fmt.Sprintf("status code (%d) is not ok", response.StatusCode))
	}

	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	protoResponse := &collectorpb.ReportResponse{}
	if err := proto.Unmarshal(body, protoResponse); err != nil {
		return nil, err
	}

	return protoResponse, nil
}
<|endoftext|>


<commit_before>package webircgateway

import (
	"errors"
	"strconv"
	"strings"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/kiwiirc/webircgateway/pkg/irc"
	"github.com/kiwiirc/webircgateway/pkg/recaptcha"
	"golang.org/x/net/html/charset"
)

/*
 * ProcessLineFromUpstream
 * Processes and makes any changes to a line of data sent from an upstream
 */
func (c *Client) ProcessLineFromUpstream(data string) string {
	client := c

	m, parseErr := irc.ParseLine(data)
	if parseErr != nil {
		return data
	}

	pLen := len(m.Params)

	if pLen > 0 && m.Command == "NICK" && m.Prefix.Nick == c.IrcState.Nick {
		client.IrcState.Nick = m.Params[0]
	}
	if pLen > 0 && m.Command == "001" {
		client.IrcState.Nick = m.Params[0]
		client.State = ClientStateConnected
	}
	if pLen > 0 && m.Command == "005" {
		// If EXTJWT is supported by the IRC server, disable it here
		foundExtJwt := false
		for _, param := range m.Params {
			if strings.HasPrefix(param, "EXTJWT") {
				c.Log(1, "Upstream already supports EXTJWT, disabling feature")
				foundExtJwt = true
			}
		}

		if foundExtJwt {
			c.Features.ExtJwt = false
		}
	}
	if pLen > 0 && m.Command == "JOIN" && m.Prefix.Nick == c.IrcState.Nick {
		channel := irc.NewStateChannel(m.GetParam(0, ""))
		c.IrcState.SetChannel(channel)
	}
	if pLen > 0 && m.Command == "PART" && m.Prefix.Nick == c.IrcState.Nick {
		c.IrcState.RemoveChannel(m.GetParam(0, ""))
	}
	if pLen > 0 && m.Command == "QUIT" && m.Prefix.Nick == c.IrcState.Nick {
		c.IrcState.ClearChannels()
	}
	// :server.com 900 m m!m@irc-3jg.1ab.j4ep8h.IP prawnsalad :You are now logged in as prawnsalad
	if pLen > 0 && m.Command == "900" {
		c.IrcState.Account = m.GetParam(2, "")
	}
	// :server.com 901 itsonlybinary itsonlybinary!itsonlybina@user/itsonlybinary :You are now logged out
	if m.Command == "901" {
		c.IrcState.Account = ""
	}
	// :prawnsalad!prawn@kiwiirc/prawnsalad MODE #kiwiirc-dev +oo notprawn kiwi-n75
	if pLen > 0 && m.Command == "MODE" {
		if strings.HasPrefix(m.GetParam(0, ""), "#") {
			channelName := m.GetParam(0, "")
			modes := m.GetParam(1, "")

			channel := c.IrcState.GetChannel(channelName)
			if channel != nil {
				channel = irc.NewStateChannel(channelName)
				c.IrcState.SetChannel(channel)
			}

			adding := false
			paramIdx := 1
			for i := 0; i < len(modes); i++ {
				mode := string(modes[i])

				if mode == "+" {
					adding = true
				} else if mode == "-" {
					adding = false
				} else {
					paramIdx++
					param := m.GetParam(paramIdx, "")
					if strings.ToLower(param) == strings.ToLower(c.IrcState.Nick) {
						if adding {
							channel.Modes[mode] = ""
						} else {
							delete(channel.Modes, mode)
						}
					}
				}
			}
		}
	}

	// If upstream reports that it supports message-tags natively, disable the wrapping of this feature for
	// this client
	if pLen >= 3 &&
		strings.ToUpper(m.Command) == "CAP" &&
		strings.ToUpper(m.Params[1]) == "LS" {
		// The CAPs could be param 2 or 3 depending on if were using multiple lines to list them all.
		caps := ""
		if pLen >= 4 && m.Params[2] == "*" {
			caps = strings.ToUpper(m.Params[3])
		} else {
			caps = strings.ToUpper(m.Params[2])
		}

		if strings.Contains(caps, "DRAFT/MESSAGE-TAGS-0.2") {
			c.Log(1, "Upstream already supports Messagetags, disabling feature")
			c.Features.Messagetags = false
		}
	}

	if m != nil && client.Features.Messagetags && c.Gateway.messageTags.CanMessageContainClientTags(m) {
		// If we have any message tags stored for this message from a previous PRIVMSG sent
		// by a client, add them back in
		mTags, mTagsExists := c.Gateway.messageTags.GetTagsFromMessage(client, m.Prefix.Nick, m)
		if mTagsExists {
			for k, v := range mTags.Tags {
				m.Tags[k] = v
			}

			data = m.ToLine()
		}
	}

	return data
}

/*
 * ProcessLineFromClient
 * Processes and makes any changes to a line of data sent from a client
 */
func (c *Client) ProcessLineFromClient(line string) (string, error) {
	message, err := irc.ParseLine(line)
	// Just pass any random data upstream
	if err != nil {
		return line, nil
	}

	maybeConnectUpstream := func() {
		if !c.UpstreamStarted && c.IrcState.Username != "" && c.Verified {
			go c.connectUpstream()
		}
	}

	if !c.Verified && strings.ToUpper(message.Command) == "CAPTCHA" {
		verified := false
		if len(message.Params) >= 1 {
			captcha := recaptcha.R{
				Secret: c.Gateway.Config.ReCaptchaSecret,
			}

			verified = captcha.VerifyResponse(message.Params[0])
		}

		if !verified {
			c.SendIrcError("Invalid captcha")
			c.StartShutdown("unverifed")
		} else {
			c.Verified = true
			maybeConnectUpstream()
		}

		return "", nil
	}

	// USER <username> <hostname> <servername> <realname>
	if strings.ToUpper(message.Command) == "USER" && !c.UpstreamStarted {
		if len(message.Params) < 4 {
			return line, errors.New("Invalid USER line")
		}

		if c.Gateway.Config.ClientUsername != "" {
			message.Params[0] = makeClientReplacements(c.Gateway.Config.ClientUsername, c)
		}
		if c.Gateway.Config.ClientRealname != "" {
			message.Params[3] = makeClientReplacements(c.Gateway.Config.ClientRealname, c)
		}

		line = message.ToLine()

		c.IrcState.Username = message.Params[0]
		c.IrcState.RealName = message.Params[3]

		maybeConnectUpstream()
	}

	if strings.ToUpper(message.Command) == "ENCODING" {
		if len(message.Params) > 0 {
			encoding, _ := charset.Lookup(message.Params[0])
			if encoding == nil {
				c.Log(1, "Requested unknown encoding, %s", message.Params[0])
			} else {
				c.Encoding = message.Params[0]
				c.Log(1, "Set encoding to %s", message.Params[0])
			}
		}

		// Don't send the ENCODING command upstream
		return "", nil
	}

	if strings.ToUpper(message.Command) == "HOST" && !c.UpstreamStarted {
		// HOST irc.network.net:6667
		// HOST irc.network.net:+6667

		if !c.Gateway.Config.Gateway {
			return "", nil
		}

		if len(message.Params) == 0 {
			return "", nil
		}

		addr := message.Params[0]
		if addr == "" {
			c.SendIrcError("Missing host")
			c.StartShutdown("missing_host")
			return "", nil
		}

		// Parse host:+port into the c.dest* vars
		portSep := strings.Index(addr, ":")
		if portSep == -1 {
			c.DestHost = addr
			c.DestPort = 6667
			c.DestTLS = false
		} else {
			c.DestHost = addr[0:portSep]
			portParam := addr[portSep+1:]
			if len(portParam) > 0 && portParam[0:1] == "+" {
				c.DestTLS = true
				c.DestPort, err = strconv.Atoi(portParam[1:])
				if err != nil {
					c.DestPort = 6697
				}
			} else {
				c.DestPort, err = strconv.Atoi(portParam[0:])
				if err != nil {
					c.DestPort = 6667
				}
			}
		}

		// Don't send the HOST command upstream
		return "", nil
	}

	// If the client supports CAP, assume the client also supports parsing MessageTags
	// When upstream replies with its CAP listing, we check if message-tags is supported by the IRCd already and if so,
	// we disable this feature flag again to use the IRCds native support.
	if strings.ToUpper(message.Command) == "CAP" && len(message.Params) >= 0 && strings.ToUpper(message.Params[0]) == "LS" {
		c.Log(1, "Enabling client Messagetags feature")
		c.Features.Messagetags = true
	}

	// Check for any client message tags so that we can store them for replaying to other clients
	if c.Features.Messagetags && c.Gateway.messageTags.CanMessageContainClientTags(message) {
		c.Gateway.messageTags.AddTagsFromMessage(c, c.IrcState.Nick, message)
		// Prevent any client tags heading upstream
		for k := range message.Tags {
			if len(k) > 0 && k[0] == '+' {
				delete(message.Tags, k)
			}
		}

		line = message.ToLine()
	}

	if c.Features.Messagetags && message.Command == "TAGMSG" {
		if len(message.Params) == 0 {
			return "", nil
		}

		// We can't be 100% sure what this users correct mask is, so just send the nick
		message.Prefix.Nick = c.IrcState.Nick
		message.Prefix.Hostname = ""
		message.Prefix.Username = ""

		thisHost := strings.ToLower(c.UpstreamConfig.Hostname)
		target := message.Params[0]
		for val := range c.Gateway.Clients.Iter() {
			curClient := val.Val.(*Client)
			sameHost := strings.ToLower(curClient.UpstreamConfig.Hostname) == thisHost
			if !sameHost {
				continue
			}

			// Only send the message on to either the target nick, or the clients in a set channel
			curNick := strings.ToLower(curClient.IrcState.Nick)
			if target != curNick && !curClient.IrcState.HasChannel(target) {
				continue
			}

			curClient.SendClientSignal("data", message.ToLine())
		}

		return "", nil
	}

	if c.Features.ExtJwt && strings.ToUpper(message.Command) == "EXTJWT" {
		tokenFor := message.GetParam(0, "")

		tokenM := irc.Message{}
		tokenM.Command = "EXTJWT"
		tokenData := jwt.MapClaims{
			"exp":         time.Now().UTC().Add(1 * time.Minute).Unix(),
			"iss":         c.UpstreamConfig.Hostname,
			"nick":        c.IrcState.Nick,
			"account":     c.IrcState.Account,
			"net_modes":   []string{},
			"channel":     "",
			"joined":      false,
			"time_joined": 0,
			"modes":       []string{},
		}

		// Use the NetworkCommonAddress if a plugin as assigned one.
		// This allows plugins to associate different upstream hosts to the same network
		if c.UpstreamConfig.NetworkCommonAddress != "" {
			tokenData["iss"] = c.UpstreamConfig.NetworkCommonAddress
		}

		if tokenFor == "" {
			tokenM.Params = append(tokenM.Params, "*")
		} else {
			tokenM.Params = append(tokenM.Params, tokenFor)

			tokenForChan := c.IrcState.GetChannel(tokenFor)
			if tokenForChan != nil {
				tokenData["time_joined"] = tokenForChan.Joined.Unix()
				tokenData["channel"] = tokenForChan.Name
				tokenData["joined"] = true

				modes := []string{}
				for mode := range tokenForChan.Modes {
					modes = append(modes, mode)
				}
				tokenData["modes"] = modes
			} else {
				tokenData["channel"] = tokenFor
			}
		}

		token := jwt.NewWithClaims(jwt.SigningMethodHS256, tokenData)
		tokenSigned, tokenSignedErr := token.SignedString([]byte(c.Gateway.Config.Secret))
		if tokenSignedErr != nil {
			c.Log(3, "Error creating JWT token. %s", tokenSignedErr.Error())
			println(tokenSignedErr.Error())
		}

		tokenM.Params = append(tokenM.Params, tokenSigned)

		c.SendClientSignal("data", tokenM.ToLine())

		return "", nil
	}

	return line, nil
}
<commit_msg>Fix crash when client issues a CAP command with no parameters<commit_after>package webircgateway

import (
	"errors"
	"strconv"
	"strings"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/kiwiirc/webircgateway/pkg/irc"
	"github.com/kiwiirc/webircgateway/pkg/recaptcha"
	"golang.org/x/net/html/charset"
)

/*
 * ProcessLineFromUpstream
 * Processes and makes any changes to a line of data sent from an upstream
 */
func (c *Client) ProcessLineFromUpstream(data string) string {
	client := c

	m, parseErr := irc.ParseLine(data)
	if parseErr != nil {
		return data
	}

	pLen := len(m.Params)

	if pLen > 0 && m.Command == "NICK" && m.Prefix.Nick == c.IrcState.Nick {
		client.IrcState.Nick = m.Params[0]
	}
	if pLen > 0 && m.Command == "001" {
		client.IrcState.Nick = m.Params[0]
		client.State = ClientStateConnected
	}
	if pLen > 0 && m.Command == "005" {
		// If EXTJWT is supported by the IRC server, disable it here
		foundExtJwt := false
		for _, param := range m.Params {
			if strings.HasPrefix(param, "EXTJWT") {
				c.Log(1, "Upstream already supports EXTJWT, disabling feature")
				foundExtJwt = true
			}
		}

		if foundExtJwt {
			c.Features.ExtJwt = false
		}
	}
	if pLen > 0 && m.Command == "JOIN" && m.Prefix.Nick == c.IrcState.Nick {
		channel := irc.NewStateChannel(m.GetParam(0, ""))
		c.IrcState.SetChannel(channel)
	}
	if pLen > 0 && m.Command == "PART" && m.Prefix.Nick == c.IrcState.Nick {
		c.IrcState.RemoveChannel(m.GetParam(0, ""))
	}
	if pLen > 0 && m.Command == "QUIT" && m.Prefix.Nick == c.IrcState.Nick {
		c.IrcState.ClearChannels()
	}
	// :server.com 900 m m!m@irc-3jg.1ab.j4ep8h.IP prawnsalad :You are now logged in as prawnsalad
	if pLen > 0 && m.Command == "900" {
		c.IrcState.Account = m.GetParam(2, "")
	}
	// :server.com 901 itsonlybinary itsonlybinary!itsonlybina@user/itsonlybinary :You are now logged out
	if m.Command == "901" {
		c.IrcState.Account = ""
	}
	// :prawnsalad!prawn@kiwiirc/prawnsalad MODE #kiwiirc-dev +oo notprawn kiwi-n75
	if pLen > 0 && m.Command == "MODE" {
		if strings.HasPrefix(m.GetParam(0, ""), "#") {
			channelName := m.GetParam(0, "")
			modes := m.GetParam(1, "")

			channel := c.IrcState.GetChannel(channelName)
			if channel != nil {
				channel = irc.NewStateChannel(channelName)
				c.IrcState.SetChannel(channel)
			}

			adding := false
			paramIdx := 1
			for i := 0; i < len(modes); i++ {
				mode := string(modes[i])

				if mode == "+" {
					adding = true
				} else if mode == "-" {
					adding = false
				} else {
					paramIdx++
					param := m.GetParam(paramIdx, "")
					if strings.ToLower(param) == strings.ToLower(c.IrcState.Nick) {
						if adding {
							channel.Modes[mode] = ""
						} else {
							delete(channel.Modes, mode)
						}
					}
				}
			}
		}
	}

	// If upstream reports that it supports message-tags natively, disable the wrapping of this feature for
	// this client
	if pLen >= 3 &&
		strings.ToUpper(m.Command) == "CAP" &&
		strings.ToUpper(m.Params[1]) == "LS" {
		// The CAPs could be param 2 or 3 depending on if were using multiple lines to list them all.
		caps := ""
		if pLen >= 4 && m.Params[2] == "*" {
			caps = strings.ToUpper(m.Params[3])
		} else {
			caps = strings.ToUpper(m.Params[2])
		}

		if strings.Contains(caps, "DRAFT/MESSAGE-TAGS-0.2") {
			c.Log(1, "Upstream already supports Messagetags, disabling feature")
			c.Features.Messagetags = false
		}
	}

	if m != nil && client.Features.Messagetags && c.Gateway.messageTags.CanMessageContainClientTags(m) {
		// If we have any message tags stored for this message from a previous PRIVMSG sent
		// by a client, add them back in
		mTags, mTagsExists := c.Gateway.messageTags.GetTagsFromMessage(client, m.Prefix.Nick, m)
		if mTagsExists {
			for k, v := range mTags.Tags {
				m.Tags[k] = v
			}

			data = m.ToLine()
		}
	}

	return data
}

/*
 * ProcessLineFromClient
 * Processes and makes any changes to a line of data sent from a client
 */
func (c *Client) ProcessLineFromClient(line string) (string, error) {
	message, err := irc.ParseLine(line)
	// Just pass any random data upstream
	if err != nil {
		return line, nil
	}

	maybeConnectUpstream := func() {
		if !c.UpstreamStarted && c.IrcState.Username != "" && c.Verified {
			go c.connectUpstream()
		}
	}

	if !c.Verified && strings.ToUpper(message.Command) == "CAPTCHA" {
		verified := false
		if len(message.Params) >= 1 {
			captcha := recaptcha.R{
				Secret: c.Gateway.Config.ReCaptchaSecret,
			}

			verified = captcha.VerifyResponse(message.Params[0])
		}

		if !verified {
			c.SendIrcError("Invalid captcha")
			c.StartShutdown("unverifed")
		} else {
			c.Verified = true
			maybeConnectUpstream()
		}

		return "", nil
	}

	// USER <username> <hostname> <servername> <realname>
	if strings.ToUpper(message.Command) == "USER" && !c.UpstreamStarted {
		if len(message.Params) < 4 {
			return line, errors.New("Invalid USER line")
		}

		if c.Gateway.Config.ClientUsername != "" {
			message.Params[0] = makeClientReplacements(c.Gateway.Config.ClientUsername, c)
		}
		if c.Gateway.Config.ClientRealname != "" {
			message.Params[3] = makeClientReplacements(c.Gateway.Config.ClientRealname, c)
		}

		line = message.ToLine()

		c.IrcState.Username = message.Params[0]
		c.IrcState.RealName = message.Params[3]

		maybeConnectUpstream()
	}

	if strings.ToUpper(message.Command) == "ENCODING" {
		if len(message.Params) > 0 {
			encoding, _ := charset.Lookup(message.Params[0])
			if encoding == nil {
				c.Log(1, "Requested unknown encoding, %s", message.Params[0])
			} else {
				c.Encoding = message.Params[0]
				c.Log(1, "Set encoding to %s", message.Params[0])
			}
		}

		// Don't send the ENCODING command upstream
		return "", nil
	}

	if strings.ToUpper(message.Command) == "HOST" && !c.UpstreamStarted {
		// HOST irc.network.net:6667
		// HOST irc.network.net:+6667

		if !c.Gateway.Config.Gateway {
			return "", nil
		}

		if len(message.Params) == 0 {
			return "", nil
		}

		addr := message.Params[0]
		if addr == "" {
			c.SendIrcError("Missing host")
			c.StartShutdown("missing_host")
			return "", nil
		}

		// Parse host:+port into the c.dest* vars
		portSep := strings.Index(addr, ":")
		if portSep == -1 {
			c.DestHost = addr
			c.DestPort = 6667
			c.DestTLS = false
		} else {
			c.DestHost = addr[0:portSep]
			portParam := addr[portSep+1:]
			if len(portParam) > 0 && portParam[0:1] == "+" {
				c.DestTLS = true
				c.DestPort, err = strconv.Atoi(portParam[1:])
				if err != nil {
					c.DestPort = 6697
				}
			} else {
				c.DestPort, err = strconv.Atoi(portParam[0:])
				if err != nil {
					c.DestPort = 6667
				}
			}
		}

		// Don't send the HOST command upstream
		return "", nil
	}

	// If the client supports CAP, assume the client also supports parsing MessageTags
	// When upstream replies with its CAP listing, we check if message-tags is supported by the IRCd already and if so,
	// we disable this feature flag again to use the IRCds native support.
	if strings.ToUpper(message.Command) == "CAP" && len(message.Params) > 0 && strings.ToUpper(message.Params[0]) == "LS" {
		c.Log(1, "Enabling client Messagetags feature")
		c.Features.Messagetags = true
	}

	// Check for any client message tags so that we can store them for replaying to other clients
	if c.Features.Messagetags && c.Gateway.messageTags.CanMessageContainClientTags(message) {
		c.Gateway.messageTags.AddTagsFromMessage(c, c.IrcState.Nick, message)
		// Prevent any client tags heading upstream
		for k := range message.Tags {
			if len(k) > 0 && k[0] == '+' {
				delete(message.Tags, k)
			}
		}

		line = message.ToLine()
	}

	if c.Features.Messagetags && message.Command == "TAGMSG" {
		if len(message.Params) == 0 {
			return "", nil
		}

		// We can't be 100% sure what this users correct mask is, so just send the nick
		message.Prefix.Nick = c.IrcState.Nick
		message.Prefix.Hostname = ""
		message.Prefix.Username = ""

		thisHost := strings.ToLower(c.UpstreamConfig.Hostname)
		target := message.Params[0]
		for val := range c.Gateway.Clients.Iter() {
			curClient := val.Val.(*Client)
			sameHost := strings.ToLower(curClient.UpstreamConfig.Hostname) == thisHost
			if !sameHost {
				continue
			}

			// Only send the message on to either the target nick, or the clients in a set channel
			curNick := strings.ToLower(curClient.IrcState.Nick)
			if target != curNick && !curClient.IrcState.HasChannel(target) {
				continue
			}

			curClient.SendClientSignal("data", message.ToLine())
		}

		return "", nil
	}

	if c.Features.ExtJwt && strings.ToUpper(message.Command) == "EXTJWT" {
		tokenFor := message.GetParam(0, "")

		tokenM := irc.Message{}
		tokenM.Command = "EXTJWT"
		tokenData := jwt.MapClaims{
			"exp":         time.Now().UTC().Add(1 * time.Minute).Unix(),
			"iss":         c.UpstreamConfig.Hostname,
			"nick":        c.IrcState.Nick,
			"account":     c.IrcState.Account,
			"net_modes":   []string{},
			"channel":     "",
			"joined":      false,
			"time_joined": 0,
			"modes":       []string{},
		}

		// Use the NetworkCommonAddress if a plugin as assigned one.
		// This allows plugins to associate different upstream hosts to the same network
		if c.UpstreamConfig.NetworkCommonAddress != "" {
			tokenData["iss"] = c.UpstreamConfig.NetworkCommonAddress
		}

		if tokenFor == "" {
			tokenM.Params = append(tokenM.Params, "*")
		} else {
			tokenM.Params = append(tokenM.Params, tokenFor)

			tokenForChan := c.IrcState.GetChannel(tokenFor)
			if tokenForChan != nil {
				tokenData["time_joined"] = tokenForChan.Joined.Unix()
				tokenData["channel"] = tokenForChan.Name
				tokenData["joined"] = true

				modes := []string{}
				for mode := range tokenForChan.Modes {
					modes = append(modes, mode)
				}
				tokenData["modes"] = modes
			} else {
				tokenData["channel"] = tokenFor
			}
		}

		token := jwt.NewWithClaims(jwt.SigningMethodHS256, tokenData)
		tokenSigned, tokenSignedErr := token.SignedString([]byte(c.Gateway.Config.Secret))
		if tokenSignedErr != nil {
			c.Log(3, "Error creating JWT token. %s", tokenSignedErr.Error())
			println(tokenSignedErr.Error())
		}

		tokenM.Params = append(tokenM.Params, tokenSigned)

		c.SendClientSignal("data", tokenM.ToLine())

		return "", nil
	}

	return line, nil
}
<|endoftext|>


<commit_before>package collision

type OnHit func(id1, id2 int)

// CallOnHits will send a signal to the passed in channel
// when it has completed all collision functions in the hitmap.
func CallOnHits(s *Space, m map[int]OnHit, doneCh chan bool) {
	progCh := make(chan bool)
	hits := Hits(s)
	for _, s2 := range hits {
		go func(s, s2 *Space, m map[int]OnHit, progCh chan bool) {
			if fn, ok := m[s2.Label]; ok {
				fn(int(s.CID), int(s2.CID))
				progCh <- true
				return
			} else {
				progCh <- false
			}

		}(s, s2, m, progCh)
	}
	// This waits to send our signal that we've
	// finished until we've counted signals for
	// each collision entity
	hitFlag := false
	for range hits {
		hitFlag = hitFlag || <-progCh
	}
	doneCh <- hitFlag
}
<commit_msg>More bug fixes and tweaks. Goons do not just sit there during caution now. Players can no longer walk through walls if they spray oil on them. Players can no longer spray oil into the void.<commit_after>package collision

type OnHit func(id1, id2 int)

// CallOnHits will send a signal to the passed in channel
// when it has completed all collision functions in the hitmap.
func CallOnHits(s *Space, m map[int]OnHit, doneCh chan bool) {
	progCh := make(chan bool)
	hits := Hits(s)
	for _, s2 := range hits {
		go func(s, s2 *Space, m map[int]OnHit, progCh chan bool) {
			if fn, ok := m[s2.Label]; ok {
				fn(int(s.CID), int(s2.CID))
				progCh <- true
			} else {
				progCh <- false
			}
		}(s, s2, m, progCh)
	}
	// This waits to send our signal that we've
	// finished until we've counted signals for
	// each collision entity
	hitFlag := false
	for range hits {
		v := <-progCh
		hitFlag = hitFlag || v
	}
	doneCh <- hitFlag
}
<|endoftext|>


<commit_before>package command

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

// DrushVariableSet Runs drush vset with a given variable name and value.
func DrushVariableSet(alias, variableName, variableValue string) {
	srcAlias := strings.Replace(alias, "@", "", -1)
	x := NewDrushCommand()
	x.Set(srcAlias, fmt.Sprintf("vset %v %v", variableName, variableValue), false)
	drushOut, err := x.Output()
	if err == nil {
		log.Infof("Successfully set %v to %v via Drush", variableName, variableValue)
	} else {
		log.Errorf("Could not set %v to %v via Drush: %v", variableName, variableValue, drushOut)
	}
}

// DrushVariableGet Runs drush vget with an exact variable name from the alias.
func DrushVariableGet(alias, variableName string) string {
	srcAlias := strings.Replace(alias, "@", "", -1)
	x := NewDrushCommand()
	x.Set(srcAlias, fmt.Sprintf("vget --exact %v", variableName), false)
	drushOut, err := x.Output()
	drushOutString := fmt.Sprintf("%s", drushOut)
	if strings.Contains(drushOutString, "No matching variable found") {
		log.Warnf("Variable %v was not found", variableName)
	} else if err == nil {
		log.Infof("Successfully retreived %v via Drush", variableName)
		drushOutString = strings.Replace(drushOutString, "[", "", -1)
		drushOutString = strings.Replace(drushOutString, "]", "", -1)
		drushOutString = strings.Replace(drushOutString, "\n", "", -1)
		return drushOutString
	} else {
		log.Errorf("Could not retreived %v via Drush: %v", variableName, drushOut)
	}
	return ""
}
<commit_msg>Add a rawoutput function to tap directly into os.stdout<commit_after>package command

import (
	"bufio"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	log "github.com/Sirupsen/logrus"
)

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

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

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

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

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

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

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

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

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

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

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

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

func (drush *Command) RawOutput() error {
	if strings.Contains(drush.alias, "@") == true {
		drush.alias = strings.Replace(drush.alias, "@", "", -1)
	}
	if drush.alias != "" {
		drush.alias = fmt.Sprintf("@%v", drush.alias)
	}
	if drush.verbose == true {
		drush.command = fmt.Sprintf("%v --verbose", drush.command)
	}
	args := fmt.Sprintf("%v %v", drush.alias, drush.command)

	comm := new(exec.Cmd)
	comm = exec.Command("sh", "-c", pathDrush+" "+" "+args)
	comm.Dir = drush.GetWorkingDir()
	comm.Stderr = os.Stderr
	comm.Stdout = os.Stdout
	err := comm.Start()
	comm.Wait()
	return err
}

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

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

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

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

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

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

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

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

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

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

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

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

// DrushVariableSet Runs drush vset with a given variable name and value.
func DrushVariableSet(alias, variableName, variableValue string) {
	srcAlias := strings.Replace(alias, "@", "", -1)
	x := NewDrushCommand()
	x.Set(srcAlias, fmt.Sprintf("vset %v %v", variableName, variableValue), false)
	drushOut, err := x.Output()
	if err == nil {
		log.Infof("Successfully set %v to %v via Drush", variableName, variableValue)
	} else {
		log.Errorf("Could not set %v to %v via Drush: %v", variableName, variableValue, drushOut)
	}
}

// DrushVariableGet Runs drush vget with an exact variable name from the alias.
func DrushVariableGet(alias, variableName string) string {
	srcAlias := strings.Replace(alias, "@", "", -1)
	x := NewDrushCommand()
	x.Set(srcAlias, fmt.Sprintf("vget --exact %v", variableName), false)
	drushOut, err := x.Output()
	drushOutString := fmt.Sprintf("%s", drushOut)
	if strings.Contains(drushOutString, "No matching variable found") {
		log.Warnf("Variable %v was not found", variableName)
	} else if err == nil {
		log.Infof("Successfully retreived %v via Drush", variableName)
		drushOutString = strings.Replace(drushOutString, "[", "", -1)
		drushOutString = strings.Replace(drushOutString, "]", "", -1)
		drushOutString = strings.Replace(drushOutString, "\n", "", -1)
		return drushOutString
	} else {
		log.Errorf("Could not retreived %v via Drush: %v", variableName, drushOut)
	}
	return ""
}
<|endoftext|>


<commit_before>package command

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

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

var logger = logging.GetLogger("command")

const idFileName = "id"

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

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

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

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

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

	return nil
}

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

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

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

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

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

	return result, nil
}

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

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

type queueState int

const (
	queueStateFirst queueState = iota
	queueStateDefault
	queueStateQueued
	queueStateRetrying
)

func loop(ag *agent.Agent, conf *config.Config, api *mackerel.API, host *mackerel.Host) {
	metricsResult := ag.Watch()

	postQueue := make(chan []*mackerel.CreatingMetricsValue, conf.Connection.Post_Metrics_Buffer_Size)
	go func() {
		postDelaySeconds := delayByHost(host)
		qState := queueStateFirst
		for values := range postQueue {
			orgPostValues := [][](*mackerel.CreatingMetricsValue){values}
			if len(postQueue) > 0 {
				// Bulk posting. However at most "two" metrics are to be posted, so postQueue isn't always empty yet.
				logger.Debugf("Merging datapoints with next queued ones")
				nextValues := <-postQueue
				orgPostValues = append(orgPostValues, nextValues)
			}
			postValues := [](*mackerel.CreatingMetricsValue){}
			for _, v := range orgPostValues {
				postValues = append(postValues, v...)
			}

			delaySeconds := 0
			switch qState {
			case queueStateFirst: // request immediately to create graph defs of host
				// nop
			case queueStateQueued:
				delaySeconds = conf.Connection.Post_Metrics_Dequeue_Delay_Seconds
			case queueStateRetrying:
				delaySeconds = conf.Connection.Post_Metrics_Retry_Delay_Seconds
			default:
				// Sending data at every 0 second from all hosts causes request flooding.
				// To prevent flooding, this loop sleeps for some seconds
				// which is specific to the ID of the host running agent on.
				// The sleep second is up to 60s.
				elapsedSeconds := int(time.Now().Unix() % int64(config.PostMetricsInterval.Seconds()))
				if postDelaySeconds > elapsedSeconds {
					delaySeconds = postDelaySeconds - elapsedSeconds
				}
			}

			// update queueState before sleeping
			if len(postQueue) > 0 {
				qState = queueStateQueued
			} else {
				qState = queueStateDefault
			}

			time.Sleep(time.Duration(delaySeconds) * time.Second)

			err := api.PostMetricsValues(postValues)
			if err != nil {
				// XXX care client error or server error
				logger.Errorf("Failed to post metrics value (will retry): %s", err.Error())
				qState = queueStateRetrying
				go func() {
					for _, v := range orgPostValues {
						postQueue <- v
					}
				}()
				continue
			}
			logger.Debugf("Posting metrics succeeded.")
		}
	}()

	// Periodically update host specs.
	go func() {
		for {
			time.Sleep(specsUpdateInterval)
			UpdateHostSpecs(conf, api, host)
		}
	}()

	for {
		select {
		case result := <-metricsResult:
			created := float64(result.Created.Unix())
			creatingValues := [](*mackerel.CreatingMetricsValue){}
			for name, value := range (map[string]float64)(result.Values) {
				creatingValues = append(
					creatingValues,
					&mackerel.CreatingMetricsValue{host.Id, name, created, value},
				)
			}
			logger.Debugf("Enqueuing task to post metrics.")
			postQueue <- creatingValues
		}
	}
}

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

	meta := spec.Collect(specGenerators())

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

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

	return hostname, meta, interfaces, nil
}

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

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

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

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

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

	return api, host, nil
}

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

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

	loop(ag, conf, api, host)
}

func pluginGenerators(conf *config.Config) []metrics.PluginGenerator {
	generators := []metrics.PluginGenerator{}

	for _, pluginConfig := range conf.Plugin["metrics"] {
		generators = append(generators, metrics.NewPluginGenerator(pluginConfig))
	}

	return generators
}
<commit_msg>refactor<commit_after>package command

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

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

var logger = logging.GetLogger("command")

const idFileName = "id"

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

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

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

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

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

	return nil
}

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

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

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

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

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

	return result, nil
}

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

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

type queueState int

const (
	queueStateFirst queueState = iota
	queueStateDefault
	queueStateQueued
	queueStateRetrying
)

func loop(ag *agent.Agent, conf *config.Config, api *mackerel.API, host *mackerel.Host) {
	metricsResult := ag.Watch()

	postQueue := make(chan []*mackerel.CreatingMetricsValue, conf.Connection.Post_Metrics_Buffer_Size)
	go func() {
		postDelaySeconds := delayByHost(host)
		qState := queueStateFirst
		for values := range postQueue {
			origPostValues := [][](*mackerel.CreatingMetricsValue){values}
			if len(postQueue) > 0 {
				// Bulk posting. However at most "two" metrics are to be posted, so postQueue isn't always empty yet.
				logger.Debugf("Merging datapoints with next queued ones")
				nextValues := <-postQueue
				origPostValues = append(origPostValues, nextValues)
			}

			delaySeconds := 0
			switch qState {
			case queueStateFirst: // request immediately to create graph defs of host
				// nop
			case queueStateQueued:
				delaySeconds = conf.Connection.Post_Metrics_Dequeue_Delay_Seconds
			case queueStateRetrying:
				delaySeconds = conf.Connection.Post_Metrics_Retry_Delay_Seconds
			default:
				// Sending data at every 0 second from all hosts causes request flooding.
				// To prevent flooding, this loop sleeps for some seconds
				// which is specific to the ID of the host running agent on.
				// The sleep second is up to 60s.
				elapsedSeconds := int(time.Now().Unix() % int64(config.PostMetricsInterval.Seconds()))
				if postDelaySeconds > elapsedSeconds {
					delaySeconds = postDelaySeconds - elapsedSeconds
				}
			}

			// update queueState before sleeping
			if len(postQueue) > 0 {
				qState = queueStateQueued
			} else {
				qState = queueStateDefault
			}

			time.Sleep(time.Duration(delaySeconds) * time.Second)

			postValues := [](*mackerel.CreatingMetricsValue){}
			for _, v := range origPostValues {
				postValues = append(postValues, v...)
			}
			err := api.PostMetricsValues(postValues)
			if err != nil {
				// XXX care client error or server error
				logger.Errorf("Failed to post metrics value (will retry): %s", err.Error())
				qState = queueStateRetrying
				go func() {
					for _, v := range origPostValues {
						postQueue <- v
					}
				}()
				continue
			}
			logger.Debugf("Posting metrics succeeded.")
		}
	}()

	// Periodically update host specs.
	go func() {
		for {
			time.Sleep(specsUpdateInterval)
			UpdateHostSpecs(conf, api, host)
		}
	}()

	for {
		select {
		case result := <-metricsResult:
			created := float64(result.Created.Unix())
			creatingValues := [](*mackerel.CreatingMetricsValue){}
			for name, value := range (map[string]float64)(result.Values) {
				creatingValues = append(
					creatingValues,
					&mackerel.CreatingMetricsValue{host.Id, name, created, value},
				)
			}
			logger.Debugf("Enqueuing task to post metrics.")
			postQueue <- creatingValues
		}
	}
}

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

	meta := spec.Collect(specGenerators())

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

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

	return hostname, meta, interfaces, nil
}

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

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

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

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

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

	return api, host, nil
}

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

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

	loop(ag, conf, api, host)
}

func pluginGenerators(conf *config.Config) []metrics.PluginGenerator {
	generators := []metrics.PluginGenerator{}

	for _, pluginConfig := range conf.Plugin["metrics"] {
		generators = append(generators, metrics.NewPluginGenerator(pluginConfig))
	}

	return generators
}
<|endoftext|>


<commit_before>package command

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

	"github.com/hashicorp/otto/appfile"
	"github.com/hashicorp/otto/appfile/detect"
	"github.com/hashicorp/otto/ui"
)

// CompileCommand is the command that is responsible for "compiling" the
// Appfile into a set of data that is used by the other commands for
// execution.
type CompileCommand struct {
	Meta

	Detectors []*detect.Detector
}

func (c *CompileCommand) Run(args []string) int {
	var flagAppfile string
	fs := c.FlagSet("compile", FlagSetNone)
	fs.Usage = func() { c.Ui.Error(c.Help()) }
	fs.StringVar(&flagAppfile, "appfile", "", "")
	if err := fs.Parse(args); err != nil {
		return 1
	}

	// Load a UI
	ui := c.OttoUi()
	ui.Header("Loading Appfile...")

	// Determine all the Appfile paths
	//
	// First, if an Appfile was specified on the command-line, it must
	// exist so we validate that it exists.
	if flagAppfile != "" {
		fi, err := os.Stat(flagAppfile)
		if err != nil {
			c.Ui.Error(fmt.Sprintf(
				"Error loading Appfile: %s", err))
			return 1
		}

		if fi.IsDir() {
			flagAppfile = filepath.Join(flagAppfile, DefaultAppfile)
		}
	}

	// If the Appfile is still blank, just use our current directory
	if flagAppfile == "" {
		var err error
		flagAppfile, err = os.Getwd()
		if err != nil {
			c.Ui.Error(fmt.Sprintf(
				"Error loading working directory: %s", err))
			return 1
		}

		flagAppfile = filepath.Join(flagAppfile, DefaultAppfile)
	}

	// If we have the Appfile, then make sure it is an absoute path
	if flagAppfile != "" {
		var err error
		flagAppfile, err = filepath.Abs(flagAppfile)
		if err != nil {
			c.Ui.Error(fmt.Sprintf(
				"Error getting Appfile path: %s", err))
			return 1
		}
	}

	// Load the appfile. This is the only time we ever load the
	// raw Appfile. All other commands load the compiled Appfile.
	var app *appfile.File
	if fi, err := os.Stat(flagAppfile); err == nil && !fi.IsDir() {
		app, err = appfile.ParseFile(flagAppfile)
		if err != nil {
			c.Ui.Error(err.Error())
			return 1
		}
	}

	// Tell the user what is happening if they have no Appfile
	if app == nil {
		ui.Header("No Appfile found! Detecting project information...")
		ui.Message(fmt.Sprintf(
			"No Appfile was found. If there is no Appfile, Otto will do its best\n" +
				"to detect the type of application this is and set reasonable defaults.\n" +
				"This is a good way to get started with Otto, but over time we recommend\n" +
				"writing a real Appfile since this will allow more complex customizations,\n" +
				"the ability to reference dependencies, versioning, and more."))
	}

	// Parse the detectors
	dataDir, err := c.DataDir()
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}
	detectorDir := filepath.Join(dataDir, DefaultLocalDataDetectorDir)
	log.Printf("[DEBUG] loading detectors from: %s", detectorDir)
	detectConfig, err := detect.ParseDir(detectorDir)
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}
	err = detectConfig.Merge(&detect.Config{Detectors: c.Detectors})
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	// Load the default Appfile so we can merge in any defaults into
	// the loaded Appfile (if there is one).
	appDef, err := appfile.Default(filepath.Dir(flagAppfile), detectConfig)
	if err != nil {
		c.Ui.Error(fmt.Sprintf(
			"Error loading Appfile: %s", err))
		return 1
	}

	// If there was no loaded Appfile and we don't have an application
	// type then we weren't able to detect the type. Error.
	if app == nil && appDef.Application.Type == "" {
		c.Ui.Error(strings.TrimSpace(errCantDetectType))
		return 1
	}

	// Merge the appfiles
	if app != nil {
		if err := appDef.Merge(app); err != nil {
			c.Ui.Error(fmt.Sprintf(
				"Error loading Appfile: %s", err))
			return 1
		}
	}
	app = appDef

	// Compile the Appfile
	ui.Header("Fetching all Appfile dependencies...")
	capp, err := appfile.Compile(app, &appfile.CompileOpts{
		Dir: filepath.Join(
			filepath.Dir(app.Path), DefaultOutputDir, DefaultOutputDirCompiledAppfile),
		Callback: c.compileCallback(ui),
	})
	if err != nil {
		c.Ui.Error(fmt.Sprintf(
			"Error compiling Appfile: %s", err))
		return 1
	}

	// Get a core
	core, err := c.Core(capp)
	if err != nil {
		c.Ui.Error(fmt.Sprintf(
			"Error loading core: %s", err))
		return 1
	}

	// Get the active infrastructure just for UI reasons
	infra := app.ActiveInfrastructure()

	// Before the compilation, output to the user what is going on
	ui.Header("Compiling...")
	ui.Message(fmt.Sprintf(
		"Application:    %s (%s)",
		app.Application.Name,
		app.Application.Type))
	ui.Message(fmt.Sprintf("Project:        %s", app.Project.Name))
	ui.Message(fmt.Sprintf(
		"Infrastructure: %s (%s)",
		infra.Type,
		infra.Flavor))
	ui.Message("")

	// Compile!
	if err := core.Compile(); err != nil {
		c.Ui.Error(fmt.Sprintf(
			"Error compiling: %s", err))
		return 1
	}

	// Success!
	ui.Header("Compilation success!")
	ui.Message(fmt.Sprintf(
		"This means that Otto is now ready to start a development environment,\n" +
			"deploy this application, build the supporting infastructure, and\n" +
			"more. See the help for more information.\n\n" +
			"Supporting files to enable Otto to manage your application from\n" +
			"development to deployment have been placed in the output directory.\n" +
			"These files can be manually inspected to determine what Otto will do."))

	return 0
}

func (c *CompileCommand) Synopsis() string {
	return "Prepares your project for being run."
}

func (c *CompileCommand) Help() string {
	helpText := `
Usage: otto [options] [path]

  Compiles the Appfile into the set of supporting files used for
  development, deploy, etc. If path is not specified, the current directory
  is assumed.

`

	return strings.TrimSpace(helpText)
}

func (c *CompileCommand) compileCallback(ui ui.Ui) func(appfile.CompileEvent) {
	return func(raw appfile.CompileEvent) {
		switch e := raw.(type) {
		case *appfile.CompileEventDep:
			ui.Message(fmt.Sprintf(
				"Fetching dependency: %s", e.Source))
		case *appfile.CompileEventImport:
			ui.Message(fmt.Sprintf(
				"Fetching import: %s", e.Source))
		}
	}
}

const errCantDetectType = `
No Appfile is present and Otto couldn't detect the project type automatically.
Otto does its best without an Appfile to detect what kind of project this is
automatically, but sometimes this fails if the project is in a structure
Otto doesn't recognize or its a project type that Otto doesn't yet support.

Please create an Appfile and specify at a minimum the project type. Below
is an example minimal Appfile specifying the "go" project type:

    application {
        type = "go"
    }

If you believe Otto should've been able to automatically detect your
project type, then please open an issue with the Otto project.
`
<commit_msg>command/compile: handle nil detect<commit_after>package command

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

	"github.com/hashicorp/otto/appfile"
	"github.com/hashicorp/otto/appfile/detect"
	"github.com/hashicorp/otto/ui"
)

// CompileCommand is the command that is responsible for "compiling" the
// Appfile into a set of data that is used by the other commands for
// execution.
type CompileCommand struct {
	Meta

	Detectors []*detect.Detector
}

func (c *CompileCommand) Run(args []string) int {
	var flagAppfile string
	fs := c.FlagSet("compile", FlagSetNone)
	fs.Usage = func() { c.Ui.Error(c.Help()) }
	fs.StringVar(&flagAppfile, "appfile", "", "")
	if err := fs.Parse(args); err != nil {
		return 1
	}

	// Load a UI
	ui := c.OttoUi()
	ui.Header("Loading Appfile...")

	// Determine all the Appfile paths
	//
	// First, if an Appfile was specified on the command-line, it must
	// exist so we validate that it exists.
	if flagAppfile != "" {
		fi, err := os.Stat(flagAppfile)
		if err != nil {
			c.Ui.Error(fmt.Sprintf(
				"Error loading Appfile: %s", err))
			return 1
		}

		if fi.IsDir() {
			flagAppfile = filepath.Join(flagAppfile, DefaultAppfile)
		}
	}

	// If the Appfile is still blank, just use our current directory
	if flagAppfile == "" {
		var err error
		flagAppfile, err = os.Getwd()
		if err != nil {
			c.Ui.Error(fmt.Sprintf(
				"Error loading working directory: %s", err))
			return 1
		}

		flagAppfile = filepath.Join(flagAppfile, DefaultAppfile)
	}

	// If we have the Appfile, then make sure it is an absoute path
	if flagAppfile != "" {
		var err error
		flagAppfile, err = filepath.Abs(flagAppfile)
		if err != nil {
			c.Ui.Error(fmt.Sprintf(
				"Error getting Appfile path: %s", err))
			return 1
		}
	}

	// Load the appfile. This is the only time we ever load the
	// raw Appfile. All other commands load the compiled Appfile.
	var app *appfile.File
	if fi, err := os.Stat(flagAppfile); err == nil && !fi.IsDir() {
		app, err = appfile.ParseFile(flagAppfile)
		if err != nil {
			c.Ui.Error(err.Error())
			return 1
		}
	}

	// Tell the user what is happening if they have no Appfile
	if app == nil {
		ui.Header("No Appfile found! Detecting project information...")
		ui.Message(fmt.Sprintf(
			"No Appfile was found. If there is no Appfile, Otto will do its best\n" +
				"to detect the type of application this is and set reasonable defaults.\n" +
				"This is a good way to get started with Otto, but over time we recommend\n" +
				"writing a real Appfile since this will allow more complex customizations,\n" +
				"the ability to reference dependencies, versioning, and more."))
	}

	// Parse the detectors
	dataDir, err := c.DataDir()
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}
	detectorDir := filepath.Join(dataDir, DefaultLocalDataDetectorDir)
	log.Printf("[DEBUG] loading detectors from: %s", detectorDir)
	detectConfig, err := detect.ParseDir(detectorDir)
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}
	if detectConfig == nil {
		detectConfig = &detect.Config{}
	}
	err = detectConfig.Merge(&detect.Config{Detectors: c.Detectors})
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	// Load the default Appfile so we can merge in any defaults into
	// the loaded Appfile (if there is one).
	appDef, err := appfile.Default(filepath.Dir(flagAppfile), detectConfig)
	if err != nil {
		c.Ui.Error(fmt.Sprintf(
			"Error loading Appfile: %s", err))
		return 1
	}

	// If there was no loaded Appfile and we don't have an application
	// type then we weren't able to detect the type. Error.
	if app == nil && appDef.Application.Type == "" {
		c.Ui.Error(strings.TrimSpace(errCantDetectType))
		return 1
	}

	// Merge the appfiles
	if app != nil {
		if err := appDef.Merge(app); err != nil {
			c.Ui.Error(fmt.Sprintf(
				"Error loading Appfile: %s", err))
			return 1
		}
	}
	app = appDef

	// Compile the Appfile
	ui.Header("Fetching all Appfile dependencies...")
	capp, err := appfile.Compile(app, &appfile.CompileOpts{
		Dir: filepath.Join(
			filepath.Dir(app.Path), DefaultOutputDir, DefaultOutputDirCompiledAppfile),
		Callback: c.compileCallback(ui),
	})
	if err != nil {
		c.Ui.Error(fmt.Sprintf(
			"Error compiling Appfile: %s", err))
		return 1
	}

	// Get a core
	core, err := c.Core(capp)
	if err != nil {
		c.Ui.Error(fmt.Sprintf(
			"Error loading core: %s", err))
		return 1
	}

	// Get the active infrastructure just for UI reasons
	infra := app.ActiveInfrastructure()

	// Before the compilation, output to the user what is going on
	ui.Header("Compiling...")
	ui.Message(fmt.Sprintf(
		"Application:    %s (%s)",
		app.Application.Name,
		app.Application.Type))
	ui.Message(fmt.Sprintf("Project:        %s", app.Project.Name))
	ui.Message(fmt.Sprintf(
		"Infrastructure: %s (%s)",
		infra.Type,
		infra.Flavor))
	ui.Message("")

	// Compile!
	if err := core.Compile(); err != nil {
		c.Ui.Error(fmt.Sprintf(
			"Error compiling: %s", err))
		return 1
	}

	// Success!
	ui.Header("Compilation success!")
	ui.Message(fmt.Sprintf(
		"This means that Otto is now ready to start a development environment,\n" +
			"deploy this application, build the supporting infastructure, and\n" +
			"more. See the help for more information.\n\n" +
			"Supporting files to enable Otto to manage your application from\n" +
			"development to deployment have been placed in the output directory.\n" +
			"These files can be manually inspected to determine what Otto will do."))

	return 0
}

func (c *CompileCommand) Synopsis() string {
	return "Prepares your project for being run."
}

func (c *CompileCommand) Help() string {
	helpText := `
Usage: otto [options] [path]

  Compiles the Appfile into the set of supporting files used for
  development, deploy, etc. If path is not specified, the current directory
  is assumed.

`

	return strings.TrimSpace(helpText)
}

func (c *CompileCommand) compileCallback(ui ui.Ui) func(appfile.CompileEvent) {
	return func(raw appfile.CompileEvent) {
		switch e := raw.(type) {
		case *appfile.CompileEventDep:
			ui.Message(fmt.Sprintf(
				"Fetching dependency: %s", e.Source))
		case *appfile.CompileEventImport:
			ui.Message(fmt.Sprintf(
				"Fetching import: %s", e.Source))
		}
	}
}

const errCantDetectType = `
No Appfile is present and Otto couldn't detect the project type automatically.
Otto does its best without an Appfile to detect what kind of project this is
automatically, but sometimes this fails if the project is in a structure
Otto doesn't recognize or its a project type that Otto doesn't yet support.

Please create an Appfile and specify at a minimum the project type. Below
is an example minimal Appfile specifying the "go" project type:

    application {
        type = "go"
    }

If you believe Otto should've been able to automatically detect your
project type, then please open an issue with the Otto project.
`
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"strings"

	"github.com/hashicorp/consul/command/agent"
	"github.com/mitchellh/cli"
)

// KeyringCommand is a Command implementation that handles querying, installing,
// and removing gossip encryption keys from a keyring.
type KeyringCommand struct {
	Ui cli.Ui
}

func (c *KeyringCommand) Run(args []string) int {
	var installKey, useKey, removeKey string
	var listKeys bool

	cmdFlags := flag.NewFlagSet("keys", flag.ContinueOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }

	cmdFlags.StringVar(&installKey, "install", "", "install key")
	cmdFlags.StringVar(&useKey, "use", "", "use key")
	cmdFlags.StringVar(&removeKey, "remove", "", "remove key")
	cmdFlags.BoolVar(&listKeys, "list", false, "list keys")

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

	c.Ui = &cli.PrefixedUi{
		OutputPrefix: "",
		InfoPrefix:   "==> ",
		ErrorPrefix:  "",
		Ui:           c.Ui,
	}

	// Only accept a single argument
	found := listKeys
	for _, arg := range []string{installKey, useKey, removeKey} {
		if found && len(arg) > 0 {
			c.Ui.Error("Only a single action is allowed")
			return 1
		}
		found = found || len(arg) > 0
	}

	// Fail fast if no actionable args were passed
	if !found {
		c.Ui.Error(c.Help())
		return 1
	}

	// All other operations will require a client connection
	client, err := RPCClient(*rpcAddr)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error connecting to Consul agent: %s", err))
		return 1
	}
	defer client.Close()

	if listKeys {
		c.Ui.Info("Gathering installed encryption keys...")
		r, err := client.ListKeys()
		if err != nil {
			c.Ui.Error(fmt.Sprintf("error: %s", err))
			return 1
		}
		if rval := c.handleResponse(r.Info, r.Messages); rval != 0 {
			return rval
		}
		c.handleList(r.Info, r.Keys)
		return 0
	}

	if installKey != "" {
		c.Ui.Info("Installing new gossip encryption key...")
		r, err := client.InstallKey(installKey)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("error: %s", err))
			return 1
		}
		return c.handleResponse(r.Info, r.Messages)
	}

	if useKey != "" {
		c.Ui.Info("Changing primary gossip encryption key...")
		r, err := client.UseKey(useKey)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("error: %s", err))
			return 1
		}
		return c.handleResponse(r.Info, r.Messages)
	}

	if removeKey != "" {
		c.Ui.Info("Removing gossip encryption key...")
		r, err := client.RemoveKey(removeKey)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("error: %s", err))
			return 1
		}
		return c.handleResponse(r.Info, r.Messages)
	}

	// Should never make it here
	return 0
}

func (c *KeyringCommand) handleResponse(
	info []agent.KeyringInfo,
	messages []agent.KeyringMessage) int {

	var rval int

	for _, i := range info {
		if i.Error != "" {
			pool := i.Pool
			if pool != "WAN" {
				pool = i.Datacenter + " (LAN)"
			}

			c.Ui.Error("")
			c.Ui.Error(fmt.Sprintf("%s error: %s", pool, i.Error))

			for _, msg := range messages {
				if msg.Datacenter != i.Datacenter || msg.Pool != i.Pool {
					continue
				}
				c.Ui.Error(fmt.Sprintf("  %s: %s", msg.Node, msg.Message))
			}
			rval = 1
		}
	}

	if rval == 0 {
		c.Ui.Info("Done!")
	}

	return rval
}

func (c *KeyringCommand) handleList(
	info []agent.KeyringInfo,
	keys []agent.KeyringEntry) {

	installed := make(map[string]map[string][]int)
	for _, key := range keys {
		var nodes int
		for _, i := range info {
			if i.Datacenter == key.Datacenter && i.Pool == key.Pool {
				nodes = i.NumNodes
			}
		}

		pool := key.Pool
		if pool != "WAN" {
			pool = key.Datacenter + " (LAN)"
		}

		if _, ok := installed[pool]; !ok {
			installed[pool] = map[string][]int{key.Key: []int{key.Count, nodes}}
		} else {
			installed[pool][key.Key] = []int{key.Count, nodes}
		}
	}

	for pool, keys := range installed {
		c.Ui.Output("")
		c.Ui.Output(pool + ":")
		for key, num := range keys {
			c.Ui.Output(fmt.Sprintf("  %s [%d/%d]", key, num[0], num[1]))
		}
	}
}

func (c *KeyringCommand) Help() string {
	helpText := `
Usage: consul keyring [options]

  Manages encryption keys used for gossip messages. Gossip encryption is
  optional. When enabled, this command may be used to examine active encryption
  keys in the cluster, add new keys, and remove old ones. When combined, this
  functionality provides the ability to perform key rotation cluster-wide,
  without disrupting the cluster.

  With the exception of the -init argument, all operations performed by this
  command can only be run against server nodes, and affect both the LAN and
  WAN keyrings in lock-step.

Options:

  -install=<key>            Install a new encryption key. This will broadcast
                            the new key to all members in the cluster.
  -use=<key>                Change the primary encryption key, which is used to
                            encrypt messages. The key must already be installed
                            before this operation can succeed.
  -remove=<key>             Remove the given key from the cluster. This
                            operation may only be performed on keys which are
                            not currently the primary key.
  -list                     List all keys currently in use within the cluster.
  -rpc-addr=127.0.0.1:8400  RPC address of the Consul agent.
`
	return strings.TrimSpace(helpText)
}

func (c *KeyringCommand) Synopsis() string {
	return "Manages gossip layer encryption keys"
}
<commit_msg>command/keyring: adjust command help<commit_after>package command

import (
	"flag"
	"fmt"
	"strings"

	"github.com/hashicorp/consul/command/agent"
	"github.com/mitchellh/cli"
)

// KeyringCommand is a Command implementation that handles querying, installing,
// and removing gossip encryption keys from a keyring.
type KeyringCommand struct {
	Ui cli.Ui
}

func (c *KeyringCommand) Run(args []string) int {
	var installKey, useKey, removeKey string
	var listKeys bool

	cmdFlags := flag.NewFlagSet("keys", flag.ContinueOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }

	cmdFlags.StringVar(&installKey, "install", "", "install key")
	cmdFlags.StringVar(&useKey, "use", "", "use key")
	cmdFlags.StringVar(&removeKey, "remove", "", "remove key")
	cmdFlags.BoolVar(&listKeys, "list", false, "list keys")

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

	c.Ui = &cli.PrefixedUi{
		OutputPrefix: "",
		InfoPrefix:   "==> ",
		ErrorPrefix:  "",
		Ui:           c.Ui,
	}

	// Only accept a single argument
	found := listKeys
	for _, arg := range []string{installKey, useKey, removeKey} {
		if found && len(arg) > 0 {
			c.Ui.Error("Only a single action is allowed")
			return 1
		}
		found = found || len(arg) > 0
	}

	// Fail fast if no actionable args were passed
	if !found {
		c.Ui.Error(c.Help())
		return 1
	}

	// All other operations will require a client connection
	client, err := RPCClient(*rpcAddr)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error connecting to Consul agent: %s", err))
		return 1
	}
	defer client.Close()

	if listKeys {
		c.Ui.Info("Gathering installed encryption keys...")
		r, err := client.ListKeys()
		if err != nil {
			c.Ui.Error(fmt.Sprintf("error: %s", err))
			return 1
		}
		if rval := c.handleResponse(r.Info, r.Messages); rval != 0 {
			return rval
		}
		c.handleList(r.Info, r.Keys)
		return 0
	}

	if installKey != "" {
		c.Ui.Info("Installing new gossip encryption key...")
		r, err := client.InstallKey(installKey)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("error: %s", err))
			return 1
		}
		return c.handleResponse(r.Info, r.Messages)
	}

	if useKey != "" {
		c.Ui.Info("Changing primary gossip encryption key...")
		r, err := client.UseKey(useKey)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("error: %s", err))
			return 1
		}
		return c.handleResponse(r.Info, r.Messages)
	}

	if removeKey != "" {
		c.Ui.Info("Removing gossip encryption key...")
		r, err := client.RemoveKey(removeKey)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("error: %s", err))
			return 1
		}
		return c.handleResponse(r.Info, r.Messages)
	}

	// Should never make it here
	return 0
}

func (c *KeyringCommand) handleResponse(
	info []agent.KeyringInfo,
	messages []agent.KeyringMessage) int {

	var rval int

	for _, i := range info {
		if i.Error != "" {
			pool := i.Pool
			if pool != "WAN" {
				pool = i.Datacenter + " (LAN)"
			}

			c.Ui.Error("")
			c.Ui.Error(fmt.Sprintf("%s error: %s", pool, i.Error))

			for _, msg := range messages {
				if msg.Datacenter != i.Datacenter || msg.Pool != i.Pool {
					continue
				}
				c.Ui.Error(fmt.Sprintf("  %s: %s", msg.Node, msg.Message))
			}
			rval = 1
		}
	}

	if rval == 0 {
		c.Ui.Info("Done!")
	}

	return rval
}

func (c *KeyringCommand) handleList(
	info []agent.KeyringInfo,
	keys []agent.KeyringEntry) {

	installed := make(map[string]map[string][]int)
	for _, key := range keys {
		var nodes int
		for _, i := range info {
			if i.Datacenter == key.Datacenter && i.Pool == key.Pool {
				nodes = i.NumNodes
			}
		}

		pool := key.Pool
		if pool != "WAN" {
			pool = key.Datacenter + " (LAN)"
		}

		if _, ok := installed[pool]; !ok {
			installed[pool] = map[string][]int{key.Key: []int{key.Count, nodes}}
		} else {
			installed[pool][key.Key] = []int{key.Count, nodes}
		}
	}

	for pool, keys := range installed {
		c.Ui.Output("")
		c.Ui.Output(pool + ":")
		for key, num := range keys {
			c.Ui.Output(fmt.Sprintf("  %s [%d/%d]", key, num[0], num[1]))
		}
	}
}

func (c *KeyringCommand) Help() string {
	helpText := `
Usage: consul keyring [options]

  Manages encryption keys used for gossip messages. Gossip encryption is
  optional. When enabled, this command may be used to examine active encryption
  keys in the cluster, add new keys, and remove old ones. When combined, this
  functionality provides the ability to perform key rotation cluster-wide,
  without disrupting the cluster.

  All operations performed by this command can only be run against server nodes,
  and affect both the LAN and WAN keyrings in lock-step.

  All variations of the keyring command return 0 if all nodes reply and there
  are no errors. If any node fails to reply or reports failure, the exit code
  will be 1.

Options:

  -install=<key>            Install a new encryption key. This will broadcast
                            the new key to all members in the cluster.
  -use=<key>                Change the primary encryption key, which is used to
                            encrypt messages. The key must already be installed
                            before this operation can succeed.
  -remove=<key>             Remove the given key from the cluster. This
                            operation may only be performed on keys which are
                            not currently the primary key.
  -list                     List all keys currently in use within the cluster.
  -rpc-addr=127.0.0.1:8400  RPC address of the Consul agent.
`
	return strings.TrimSpace(helpText)
}

func (c *KeyringCommand) Synopsis() string {
	return "Manages gossip layer encryption keys"
}
<|endoftext|>