<commit_before>package alerting

import (
	"errors"

	"fmt"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/log"
	m "github.com/grafana/grafana/pkg/models"
)

type DashAlertExtractor struct {
	Dash  *m.Dashboard
	OrgId int64
	log   log.Logger
}

func NewDashAlertExtractor(dash *m.Dashboard, orgId int64) *DashAlertExtractor {
	return &DashAlertExtractor{
		Dash:  dash,
		OrgId: orgId,
		log:   log.New("alerting.extractor"),
	}
}

func (e *DashAlertExtractor) lookupDatasourceId(dsName string) (*m.DataSource, error) {
	if dsName == "" {
		query := &m.GetDataSourcesQuery{OrgId: e.OrgId}
		if err := bus.Dispatch(query); err != nil {
			return nil, err
		} else {
			for _, ds := range query.Result {
				if ds.IsDefault {
					return ds, nil
				}
			}
		}
	} else {
		query := &m.GetDataSourceByNameQuery{Name: dsName, OrgId: e.OrgId}
		if err := bus.Dispatch(query); err != nil {
			return nil, err
		} else {
			return query.Result, nil
		}
	}

	return nil, errors.New("Could not find datasource id for " + dsName)
}

func findPanelQueryByRefId(panel *simplejson.Json, refId string) *simplejson.Json {
	for _, targetsObj := range panel.Get("targets").MustArray() {
		target := simplejson.NewFromAny(targetsObj)

		if target.Get("refId").MustString() == refId {
			return target
		}
	}
	return nil
}

func (e *DashAlertExtractor) GetAlerts() ([]*m.Alert, error) {
	e.log.Debug("GetAlerts")

	alerts := make([]*m.Alert, 0)

	for _, rowObj := range e.Dash.Data.Get("rows").MustArray() {
		row := simplejson.NewFromAny(rowObj)

		for _, panelObj := range row.Get("panels").MustArray() {
			panel := simplejson.NewFromAny(panelObj)
			jsonAlert, hasAlert := panel.CheckGet("alert")

			if !hasAlert {
				continue
			}

			// backward compatability check, can be removed later
			enabled, hasEnabled := jsonAlert.CheckGet("enabled")
			if hasEnabled && enabled.MustBool() == false {
				continue
			}

			frequency, err := getTimeDurationStringToSeconds(jsonAlert.Get("frequency").MustString())
			if err != nil {
				return nil, ValidationError{Reason: "Could not parse frequency"}
			}

			alert := &m.Alert{
				DashboardId: e.Dash.Id,
				OrgId:       e.OrgId,
				PanelId:     panel.Get("id").MustInt64(),
				Id:          jsonAlert.Get("id").MustInt64(),
				Name:        jsonAlert.Get("name").MustString(),
				Handler:     jsonAlert.Get("handler").MustInt64(),
				Message:     jsonAlert.Get("message").MustString(),
				Frequency:   frequency,
			}

			for _, condition := range jsonAlert.Get("conditions").MustArray() {
				jsonCondition := simplejson.NewFromAny(condition)

				jsonQuery := jsonCondition.Get("query")
				queryRefId := jsonQuery.Get("params").MustArray()[0].(string)
				panelQuery := findPanelQueryByRefId(panel, queryRefId)

				if panelQuery == nil {
					reason := fmt.Sprintf("Alert on PanelId: %v refers to query(%s) that cannot be found", alert.PanelId, queryRefId)
					return nil, ValidationError{Reason: reason}
				}

				dsName := ""
				if panelQuery.Get("datasource").MustString() != "" {
					dsName = panelQuery.Get("datasource").MustString()
				} else if panel.Get("datasource").MustString() != "" {
					dsName = panel.Get("datasource").MustString()
				}

				if datasource, err := e.lookupDatasourceId(dsName); err != nil {
					return nil, err
				} else {
					jsonQuery.SetPath([]string{"datasourceId"}, datasource.Id)
				}

				if interval, err := panel.Get("interval").String(); err == nil {
					panelQuery.Set("interval", interval)
				}

				jsonQuery.Set("model", panelQuery.Interface())
			}

			alert.Settings = jsonAlert

			// validate
			_, err = NewRuleFromDBAlert(alert)
			if err == nil && alert.ValidToSave() {
				alerts = append(alerts, alert)
			} else {
				e.log.Error("Failed to extract alerts from dashboard", "error", err)
				return nil, err
			}
		}
	}

	e.log.Debug("Extracted alerts from dashboard", "alertCount", len(alerts))
	return alerts, nil
}
<commit_msg>chore(alerting): removes double logging<commit_after>package alerting

import (
	"errors"

	"fmt"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/log"
	m "github.com/grafana/grafana/pkg/models"
)

type DashAlertExtractor struct {
	Dash  *m.Dashboard
	OrgId int64
	log   log.Logger
}

func NewDashAlertExtractor(dash *m.Dashboard, orgId int64) *DashAlertExtractor {
	return &DashAlertExtractor{
		Dash:  dash,
		OrgId: orgId,
		log:   log.New("alerting.extractor"),
	}
}

func (e *DashAlertExtractor) lookupDatasourceId(dsName string) (*m.DataSource, error) {
	if dsName == "" {
		query := &m.GetDataSourcesQuery{OrgId: e.OrgId}
		if err := bus.Dispatch(query); err != nil {
			return nil, err
		} else {
			for _, ds := range query.Result {
				if ds.IsDefault {
					return ds, nil
				}
			}
		}
	} else {
		query := &m.GetDataSourceByNameQuery{Name: dsName, OrgId: e.OrgId}
		if err := bus.Dispatch(query); err != nil {
			return nil, err
		} else {
			return query.Result, nil
		}
	}

	return nil, errors.New("Could not find datasource id for " + dsName)
}

func findPanelQueryByRefId(panel *simplejson.Json, refId string) *simplejson.Json {
	for _, targetsObj := range panel.Get("targets").MustArray() {
		target := simplejson.NewFromAny(targetsObj)

		if target.Get("refId").MustString() == refId {
			return target
		}
	}
	return nil
}

func (e *DashAlertExtractor) GetAlerts() ([]*m.Alert, error) {
	e.log.Debug("GetAlerts")

	alerts := make([]*m.Alert, 0)

	for _, rowObj := range e.Dash.Data.Get("rows").MustArray() {
		row := simplejson.NewFromAny(rowObj)

		for _, panelObj := range row.Get("panels").MustArray() {
			panel := simplejson.NewFromAny(panelObj)
			jsonAlert, hasAlert := panel.CheckGet("alert")

			if !hasAlert {
				continue
			}

			// backward compatability check, can be removed later
			enabled, hasEnabled := jsonAlert.CheckGet("enabled")
			if hasEnabled && enabled.MustBool() == false {
				continue
			}

			frequency, err := getTimeDurationStringToSeconds(jsonAlert.Get("frequency").MustString())
			if err != nil {
				return nil, ValidationError{Reason: "Could not parse frequency"}
			}

			alert := &m.Alert{
				DashboardId: e.Dash.Id,
				OrgId:       e.OrgId,
				PanelId:     panel.Get("id").MustInt64(),
				Id:          jsonAlert.Get("id").MustInt64(),
				Name:        jsonAlert.Get("name").MustString(),
				Handler:     jsonAlert.Get("handler").MustInt64(),
				Message:     jsonAlert.Get("message").MustString(),
				Frequency:   frequency,
			}

			for _, condition := range jsonAlert.Get("conditions").MustArray() {
				jsonCondition := simplejson.NewFromAny(condition)

				jsonQuery := jsonCondition.Get("query")
				queryRefId := jsonQuery.Get("params").MustArray()[0].(string)
				panelQuery := findPanelQueryByRefId(panel, queryRefId)

				if panelQuery == nil {
					reason := fmt.Sprintf("Alert on PanelId: %v refers to query(%s) that cannot be found", alert.PanelId, queryRefId)
					return nil, ValidationError{Reason: reason}
				}

				dsName := ""
				if panelQuery.Get("datasource").MustString() != "" {
					dsName = panelQuery.Get("datasource").MustString()
				} else if panel.Get("datasource").MustString() != "" {
					dsName = panel.Get("datasource").MustString()
				}

				if datasource, err := e.lookupDatasourceId(dsName); err != nil {
					return nil, err
				} else {
					jsonQuery.SetPath([]string{"datasourceId"}, datasource.Id)
				}

				if interval, err := panel.Get("interval").String(); err == nil {
					panelQuery.Set("interval", interval)
				}

				jsonQuery.Set("model", panelQuery.Interface())
			}

			alert.Settings = jsonAlert

			// validate
			_, err = NewRuleFromDBAlert(alert)
			if err == nil && alert.ValidToSave() {
				alerts = append(alerts, alert)
			} else {
				return nil, err
			}
		}
	}

	e.log.Debug("Extracted alerts from dashboard", "alertCount", len(alerts))
	return alerts, nil
}
<|endoftext|>


<commit_before>package tests

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"testing"
	"time"

	"github.com/grafana/grafana/pkg/api/routing"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/services/encryption/ossencryption"
	"github.com/grafana/grafana/pkg/services/ngalert"
	apimodels "github.com/grafana/grafana/pkg/services/ngalert/api/tooling/definitions"
	"github.com/grafana/grafana/pkg/services/ngalert/metrics"
	"github.com/grafana/grafana/pkg/services/ngalert/models"
	"github.com/grafana/grafana/pkg/services/ngalert/store"
	"github.com/grafana/grafana/pkg/services/sqlstore"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/stretchr/testify/require"
)

// SetupTestEnv initializes a store to used by the tests.
func SetupTestEnv(t *testing.T, baseInterval time.Duration) (*ngalert.AlertNG, *store.DBstore) {
	t.Helper()

	cfg := setting.NewCfg()
	cfg.AlertingBaseInterval = baseInterval
	// AlertNG database migrations run and the relative database tables are created only when it's enabled
	cfg.UnifiedAlerting.Enabled = true

	m := metrics.NewNGAlert(prometheus.NewRegistry())
	ng, err := ngalert.ProvideService(
		cfg, nil, routing.NewRouteRegister(), sqlstore.InitTestDB(t),
		nil, nil, nil, nil, ossencryption.ProvideService(), m,
	)
	require.NoError(t, err)
	return ng, &store.DBstore{
		SQLStore:     ng.SQLStore,
		BaseInterval: baseInterval * time.Second,
		Logger:       log.New("ngalert-test"),
	}
}

// CreateTestAlertRule creates a dummy alert definition to be used by the tests.
func CreateTestAlertRule(t *testing.T, dbstore *store.DBstore, intervalSeconds int64, orgID int64) *models.AlertRule {
	d := rand.Intn(1000)
	ruleGroup := fmt.Sprintf("ruleGroup-%d", d)
	err := dbstore.UpdateRuleGroup(store.UpdateRuleGroupCmd{
		OrgID:        orgID,
		NamespaceUID: "namespace",
		RuleGroupConfig: apimodels.PostableRuleGroupConfig{
			Name:     ruleGroup,
			Interval: model.Duration(time.Duration(intervalSeconds) * time.Second),
			Rules: []apimodels.PostableExtendedRuleNode{
				{
					ApiRuleNode: &apimodels.ApiRuleNode{
						Annotations: map[string]string{"testAnnoKey": "testAnnoValue"},
					},
					GrafanaManagedAlert: &apimodels.PostableGrafanaRule{
						Title:     fmt.Sprintf("an alert definition %d", d),
						Condition: "A",
						Data: []models.AlertQuery{
							{
								Model: json.RawMessage(`{
										"datasourceUid": "-100",
										"type":"math",
										"expression":"2 + 2 > 1"
									}`),
								RelativeTimeRange: models.RelativeTimeRange{
									From: models.Duration(5 * time.Hour),
									To:   models.Duration(3 * time.Hour),
								},
								RefID: "A",
							},
						},
					},
				},
			},
		},
	})
	require.NoError(t, err)

	q := models.ListRuleGroupAlertRulesQuery{
		OrgID:        orgID,
		NamespaceUID: "namespace",
		RuleGroup:    ruleGroup,
	}
	err = dbstore.GetRuleGroupAlertRules(&q)
	require.NoError(t, err)
	require.NotEmpty(t, q.Result)

	rule := q.Result[0]
	t.Logf("alert definition: %v with interval: %d created", rule.GetKey(), rule.IntervalSeconds)
	return rule
}

// updateTestAlertRule update a dummy alert definition to be used by the tests.
func UpdateTestAlertRuleIntervalSeconds(t *testing.T, dbstore *store.DBstore, existingRule *models.AlertRule, intervalSeconds int64) *models.AlertRule {
	cmd := store.UpdateRuleGroupCmd{
		OrgID:        1,
		NamespaceUID: "namespace",
		RuleGroupConfig: apimodels.PostableRuleGroupConfig{
			Name:     existingRule.RuleGroup,
			Interval: model.Duration(time.Duration(intervalSeconds) * time.Second),
			Rules: []apimodels.PostableExtendedRuleNode{
				{
					GrafanaManagedAlert: &apimodels.PostableGrafanaRule{
						UID: existingRule.UID,
					},
				},
			},
		},
	}

	err := dbstore.UpdateRuleGroup(cmd)
	require.NoError(t, err)

	q := models.ListRuleGroupAlertRulesQuery{
		OrgID:        1,
		NamespaceUID: "namespace",
		RuleGroup:    existingRule.RuleGroup,
	}
	err = dbstore.GetRuleGroupAlertRules(&q)
	require.NoError(t, err)
	require.NotEmpty(t, q.Result)

	rule := q.Result[0]
	t.Logf("alert definition: %v with interval: %d created", rule.GetKey(), rule.IntervalSeconds)
	return rule
}
<commit_msg>Fix test (#41287)<commit_after>package tests

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"testing"
	"time"

	"github.com/grafana/grafana/pkg/api/routing"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/services/encryption/ossencryption"
	"github.com/grafana/grafana/pkg/services/ngalert"
	apimodels "github.com/grafana/grafana/pkg/services/ngalert/api/tooling/definitions"
	"github.com/grafana/grafana/pkg/services/ngalert/metrics"
	"github.com/grafana/grafana/pkg/services/ngalert/models"
	"github.com/grafana/grafana/pkg/services/ngalert/store"
	"github.com/grafana/grafana/pkg/services/sqlstore"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/stretchr/testify/require"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

// SetupTestEnv initializes a store to used by the tests.
func SetupTestEnv(t *testing.T, baseInterval time.Duration) (*ngalert.AlertNG, *store.DBstore) {
	t.Helper()

	cfg := setting.NewCfg()
	cfg.AlertingBaseInterval = baseInterval
	// AlertNG database migrations run and the relative database tables are created only when it's enabled
	cfg.UnifiedAlerting.Enabled = true

	m := metrics.NewNGAlert(prometheus.NewRegistry())
	ng, err := ngalert.ProvideService(
		cfg, nil, routing.NewRouteRegister(), sqlstore.InitTestDB(t),
		nil, nil, nil, nil, ossencryption.ProvideService(), m,
	)
	require.NoError(t, err)
	return ng, &store.DBstore{
		SQLStore:     ng.SQLStore,
		BaseInterval: baseInterval * time.Second,
		Logger:       log.New("ngalert-test"),
	}
}

// CreateTestAlertRule creates a dummy alert definition to be used by the tests.
func CreateTestAlertRule(t *testing.T, dbstore *store.DBstore, intervalSeconds int64, orgID int64) *models.AlertRule {
	d := rand.Intn(1000)
	ruleGroup := fmt.Sprintf("ruleGroup-%d", d)
	err := dbstore.UpdateRuleGroup(store.UpdateRuleGroupCmd{
		OrgID:        orgID,
		NamespaceUID: "namespace",
		RuleGroupConfig: apimodels.PostableRuleGroupConfig{
			Name:     ruleGroup,
			Interval: model.Duration(time.Duration(intervalSeconds) * time.Second),
			Rules: []apimodels.PostableExtendedRuleNode{
				{
					ApiRuleNode: &apimodels.ApiRuleNode{
						Annotations: map[string]string{"testAnnoKey": "testAnnoValue"},
					},
					GrafanaManagedAlert: &apimodels.PostableGrafanaRule{
						Title:     fmt.Sprintf("an alert definition %d", d),
						Condition: "A",
						Data: []models.AlertQuery{
							{
								Model: json.RawMessage(`{
										"datasourceUid": "-100",
										"type":"math",
										"expression":"2 + 2 > 1"
									}`),
								RelativeTimeRange: models.RelativeTimeRange{
									From: models.Duration(5 * time.Hour),
									To:   models.Duration(3 * time.Hour),
								},
								RefID: "A",
							},
						},
					},
				},
			},
		},
	})
	require.NoError(t, err)

	q := models.ListRuleGroupAlertRulesQuery{
		OrgID:        orgID,
		NamespaceUID: "namespace",
		RuleGroup:    ruleGroup,
	}
	err = dbstore.GetRuleGroupAlertRules(&q)
	require.NoError(t, err)
	require.NotEmpty(t, q.Result)

	rule := q.Result[0]
	t.Logf("alert definition: %v with title: %q interval: %d created", rule.GetKey(), rule.Title, rule.IntervalSeconds)
	return rule
}

// updateTestAlertRule update a dummy alert definition to be used by the tests.
func UpdateTestAlertRuleIntervalSeconds(t *testing.T, dbstore *store.DBstore, existingRule *models.AlertRule, intervalSeconds int64) *models.AlertRule {
	cmd := store.UpdateRuleGroupCmd{
		OrgID:        1,
		NamespaceUID: "namespace",
		RuleGroupConfig: apimodels.PostableRuleGroupConfig{
			Name:     existingRule.RuleGroup,
			Interval: model.Duration(time.Duration(intervalSeconds) * time.Second),
			Rules: []apimodels.PostableExtendedRuleNode{
				{
					GrafanaManagedAlert: &apimodels.PostableGrafanaRule{
						UID: existingRule.UID,
					},
				},
			},
		},
	}

	err := dbstore.UpdateRuleGroup(cmd)
	require.NoError(t, err)

	q := models.ListRuleGroupAlertRulesQuery{
		OrgID:        1,
		NamespaceUID: "namespace",
		RuleGroup:    existingRule.RuleGroup,
	}
	err = dbstore.GetRuleGroupAlertRules(&q)
	require.NoError(t, err)
	require.NotEmpty(t, q.Result)

	rule := q.Result[0]
	t.Logf("alert definition: %v with title: %s and interval: %d created", rule.GetKey(), rule.Title, rule.IntervalSeconds)
	return rule
}
<|endoftext|>


<commit_before>package snapshotcontroller

import (
	"fmt"
	"reflect"
	"sync"
	"time"

	"github.com/kubernetes-incubator/external-storage/lib/controller"
	"github.com/kubernetes-incubator/external-storage/snapshot/pkg/client"
	snapshotcontroller "github.com/kubernetes-incubator/external-storage/snapshot/pkg/controller/snapshot-controller"
	snapshotvolume "github.com/kubernetes-incubator/external-storage/snapshot/pkg/volume"
	"github.com/libopenstorage/stork/drivers/volume"
	stork "github.com/libopenstorage/stork/pkg/apis/stork"
	storkapi "github.com/libopenstorage/stork/pkg/apis/stork/v1alpha1"
	"github.com/libopenstorage/stork/pkg/snapshot/rule"
	"github.com/portworx/sched-ops/k8s"
	log "github.com/sirupsen/logrus"
	apiextensionsv1beta1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	apiextensionsclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

const (
	snapshotProvisionerName               = "stork-snapshot"
	snapshotProvisionerID                 = "stork"
	provisionerIDAnn                      = "snapshotProvisionerIdentity"
	defaultSyncDuration     time.Duration = 60 * time.Second
)

// SnapshotController Snapshot Controller
type SnapshotController struct {
	Driver      volume.Driver
	lock        sync.Mutex
	started     bool
	stopChannel chan struct{}
}

// GetProvisionerName Gets the name of the provisioner
func GetProvisionerName() string {
	return snapshotProvisionerName
}

// Start Starts the snapshot controller
func (s *SnapshotController) Start() error {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.started {
		return fmt.Errorf("Extender has already been started")
	}
	config, err := rest.InClusterConfig()
	if err != nil {
		return err
	}

	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		return err
	}

	if clientset == nil {
		return k8s.ErrK8SApiAccountNotSet
	}

	aeclientset, err := apiextensionsclient.NewForConfig(config)
	if err != nil {
		return err
	}

	snapshotClient, snapshotScheme, err := client.NewClient(config)
	if err != nil {
		return err
	}

	log.Infof("Registering CRDs")
	err = client.CreateCRD(aeclientset)
	if err != nil {
		return err
	}

	storkRuleResource := k8s.CustomResource{
		Name:    "rule",
		Plural:  "rules",
		Group:   stork.GroupName,
		Version: stork.Version,
		Scope:   apiextensionsv1beta1.NamespaceScoped,
		Kind:    reflect.TypeOf(storkapi.Rule{}).Name(),
	}

	err = k8s.Instance().CreateCRD(storkRuleResource)
	if err != nil {
		if !errors.IsAlreadyExists(err) {
			return fmt.Errorf("Failed to create CRD due to: %v", err)
		}
	}

	err = k8s.Instance().ValidateCRD(storkRuleResource)
	if err != nil {
		return fmt.Errorf("Failed to validate stork rules CRD due to: %v", err)
	}

	err = client.WaitForSnapshotResource(snapshotClient)
	if err != nil {
		return err
	}

	plugins := make(map[string]snapshotvolume.Plugin)
	plugins[s.Driver.String()] = s.Driver.GetSnapshotPlugin()

	snapController := snapshotcontroller.NewSnapshotController(snapshotClient, snapshotScheme,
		clientset, &plugins, defaultSyncDuration)

	s.stopChannel = make(chan struct{})

	snapController.Run(s.stopChannel)

	serverVersion, err := clientset.Discovery().ServerVersion()
	if err != nil {
		log.Fatalf("Error getting server version: %v", err)
	}

	snapProvisioner := newSnapshotProvisioner(clientset, snapshotClient, plugins, snapshotProvisionerID)

	provisioner := controller.NewProvisionController(
		clientset,
		snapshotProvisionerName,
		snapProvisioner,
		serverVersion.GitVersion,
	)
	provisioner.Run(s.stopChannel)

	if err := rule.PerformRuleRecovery(); err != nil {
		log.Errorf("failed to perform recovery for snapshot rules due to: %v", err)
		return err
	}

	s.started = true
	return nil
}

// Stop Stops the snapshot controller
func (s *SnapshotController) Stop() error {
	s.lock.Lock()
	defer s.lock.Unlock()

	if !s.started {
		return fmt.Errorf("Extender has not been started")
	}

	close(s.stopChannel)

	s.started = false
	return nil
}
<commit_msg>Start snapshot provisioner in background since it blocks<commit_after>package snapshotcontroller

import (
	"fmt"
	"reflect"
	"sync"
	"time"

	"github.com/kubernetes-incubator/external-storage/lib/controller"
	"github.com/kubernetes-incubator/external-storage/snapshot/pkg/client"
	snapshotcontroller "github.com/kubernetes-incubator/external-storage/snapshot/pkg/controller/snapshot-controller"
	snapshotvolume "github.com/kubernetes-incubator/external-storage/snapshot/pkg/volume"
	"github.com/libopenstorage/stork/drivers/volume"
	stork "github.com/libopenstorage/stork/pkg/apis/stork"
	storkapi "github.com/libopenstorage/stork/pkg/apis/stork/v1alpha1"
	"github.com/libopenstorage/stork/pkg/snapshot/rule"
	"github.com/portworx/sched-ops/k8s"
	log "github.com/sirupsen/logrus"
	apiextensionsv1beta1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	apiextensionsclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

const (
	snapshotProvisionerName               = "stork-snapshot"
	snapshotProvisionerID                 = "stork"
	provisionerIDAnn                      = "snapshotProvisionerIdentity"
	defaultSyncDuration     time.Duration = 60 * time.Second
)

// SnapshotController Snapshot Controller
type SnapshotController struct {
	Driver      volume.Driver
	lock        sync.Mutex
	started     bool
	stopChannel chan struct{}
}

// GetProvisionerName Gets the name of the provisioner
func GetProvisionerName() string {
	return snapshotProvisionerName
}

// Start Starts the snapshot controller
func (s *SnapshotController) Start() error {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.started {
		return fmt.Errorf("Extender has already been started")
	}
	config, err := rest.InClusterConfig()
	if err != nil {
		return err
	}

	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		return err
	}

	if clientset == nil {
		return k8s.ErrK8SApiAccountNotSet
	}

	aeclientset, err := apiextensionsclient.NewForConfig(config)
	if err != nil {
		return err
	}

	snapshotClient, snapshotScheme, err := client.NewClient(config)
	if err != nil {
		return err
	}

	log.Infof("Registering CRDs")
	err = client.CreateCRD(aeclientset)
	if err != nil {
		return err
	}

	storkRuleResource := k8s.CustomResource{
		Name:    "rule",
		Plural:  "rules",
		Group:   stork.GroupName,
		Version: stork.Version,
		Scope:   apiextensionsv1beta1.NamespaceScoped,
		Kind:    reflect.TypeOf(storkapi.Rule{}).Name(),
	}

	err = k8s.Instance().CreateCRD(storkRuleResource)
	if err != nil {
		if !errors.IsAlreadyExists(err) {
			return fmt.Errorf("Failed to create CRD due to: %v", err)
		}
	}

	err = k8s.Instance().ValidateCRD(storkRuleResource)
	if err != nil {
		return fmt.Errorf("Failed to validate stork rules CRD due to: %v", err)
	}

	err = client.WaitForSnapshotResource(snapshotClient)
	if err != nil {
		return err
	}

	plugins := make(map[string]snapshotvolume.Plugin)
	plugins[s.Driver.String()] = s.Driver.GetSnapshotPlugin()

	snapController := snapshotcontroller.NewSnapshotController(snapshotClient, snapshotScheme,
		clientset, &plugins, defaultSyncDuration)

	s.stopChannel = make(chan struct{})

	snapController.Run(s.stopChannel)

	serverVersion, err := clientset.Discovery().ServerVersion()
	if err != nil {
		log.Fatalf("Error getting server version: %v", err)
	}

	snapProvisioner := newSnapshotProvisioner(clientset, snapshotClient, plugins, snapshotProvisionerID)

	provisioner := controller.NewProvisionController(
		clientset,
		snapshotProvisionerName,
		snapProvisioner,
		serverVersion.GitVersion,
	)
	go provisioner.Run(s.stopChannel)

	if err := rule.PerformRuleRecovery(); err != nil {
		log.Errorf("failed to perform recovery for snapshot rules due to: %v", err)
		return err
	}

	s.started = true
	return nil
}

// Stop Stops the snapshot controller
func (s *SnapshotController) Stop() error {
	s.lock.Lock()
	defer s.lock.Unlock()

	if !s.started {
		return fmt.Errorf("Extender has not been started")
	}

	close(s.stopChannel)

	s.started = false
	return nil
}
<|endoftext|>


<commit_before>package graphite

import (
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

func TestGraphiteFunctions(t *testing.T) {
	Convey("Testing Graphite Functions", t, func() {

		Convey("formatting time range for now", func() {

			timeRange := formatTimeRange("now")
			So(timeRange, ShouldEqual, "now")

		})

		Convey("formatting time range for now-1m", func() {

			timeRange := formatTimeRange("now-1m")
			So(timeRange, ShouldEqual, "now-1min")

		})

		Convey("formatting time range for now-1M", func() {

			timeRange := formatTimeRange("now-1M")
			So(timeRange, ShouldEqual, "now-1mon")

		})

		Convey("fix interval format in query for 1m", func() {

			timeRange := formatTimeRange("aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1m'), 4)")
			So(timeRange, ShouldEqual, "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1min'), 4)")

		})

		Convey("fix interval format in query for 1M", func() {

			timeRange := formatTimeRange("aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1M'), 4)")
			So(timeRange, ShouldEqual, "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1mon'), 4)")

		})

	})
}
<commit_msg>Commented strange behavior of tests<commit_after>package graphite

import (
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

func TestGraphiteFunctions(t *testing.T) {
	Convey("Testing Graphite Functions", t, func() {

		Convey("formatting time range for now", func() {

			timeRange := formatTimeRange("now")
			So(timeRange, ShouldEqual, "now")

		})

		Convey("formatting time range for now-1m", func() {

			timeRange := formatTimeRange("now-1m")
			So(timeRange, ShouldEqual, "now-1min")

		})

		Convey("formatting time range for now-1M", func() {

			timeRange := formatTimeRange("now-1M")
			So(timeRange, ShouldEqual, "now-1mon")

		})

		Convey("fix interval format in query for 1m", func() {

			timeRange := formatTimeRange("aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1m'), 4)")
			So(timeRange, ShouldEqual, "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1min'), 4)")

		})

		Convey("fix interval format in query for 1M", func() {

			timeRange := formatTimeRange("aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1M'), 4)")
			So(timeRange, ShouldEqual, "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1mon'), 4)")

		})

		/*
		   		Convey("should not override query", func() {

		       	timeRange := formatTimeRange("app.grafana.*.dashboards.views.1M.count")
		       	So(timeRange, ShouldEqual, "app.grafana.*.dashboards.views.1M.count")

		   		})

		   		Convey("should not override query", func() {

		       	timeRange := formatTimeRange("app.grafana.*.dashboards.views.1m.count")
		       	So(timeRange, ShouldEqual, "app.grafana.*.dashboards.views.1m.count")

		   		})

		*/
	})
}
<|endoftext|>


<commit_before>package graphite

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestGraphiteFunctions(t *testing.T) {
	Convey("Testing Graphite Functions", t, func() {
		Convey("formatting time range for now", func() {
			timeRange := formatTimeRange("now")
			So(timeRange, ShouldEqual, "now")
		})

		Convey("formatting time range for now-1m", func() {
			timeRange := formatTimeRange("now-1m")
			So(timeRange, ShouldEqual, "-1min")
		})

		Convey("formatting time range for now-1M", func() {
			timeRange := formatTimeRange("now-1M")
			So(timeRange, ShouldEqual, "-1mon")
		})

		Convey("fix interval format in query for 1m", func() {
			timeRange := fixIntervalFormat("aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1m'), 4)")
			So(timeRange, ShouldEqual, "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1min'), 4)")
		})

		Convey("fix interval format in query for 1M", func() {
			timeRange := fixIntervalFormat("aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1M'), 4)")
			So(timeRange, ShouldEqual, "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1mon'), 4)")
		})

		Convey("should not override query for 1M", func() {
			timeRange := fixIntervalFormat("app.grafana.*.dashboards.views.1M.count")
			So(timeRange, ShouldEqual, "app.grafana.*.dashboards.views.1M.count")
		})

		Convey("should not override query for 1m", func() {
			timeRange := fixIntervalFormat("app.grafana.*.dashboards.views.1m.count")
			So(timeRange, ShouldEqual, "app.grafana.*.dashboards.views.1m.count")
		})
	})
}
<commit_msg>Chore: Rewrite tsdb graphite test to standard library (#30088)<commit_after>package graphite

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestFormatTimeRange(t *testing.T) {
	testCases := []struct {
		input    string
		expected string
	}{
		{"now", "now"},
		{"now-1m", "-1min"},
		{"now-1M", "-1mon"},
	}

	for _, tc := range testCases {
		t.Run(tc.input, func(t *testing.T) {
			tr := formatTimeRange(tc.input)
			assert.Equal(t, tc.expected, tr)
		})
	}
}

func TestFixIntervalFormat(t *testing.T) {
	testCases := []struct {
		name     string
		target   string
		expected string
	}{
		{
			name:     "should transform 1m to graphite unit (1min) when used as interval string",
			target:   "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1m'), 4)",
			expected: "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1min'), 4)",
		},
		{
			name:     "should transform 1M to graphite unit (1mon) when used as interval string",
			target:   "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1M'), 4)",
			expected: "aliasByNode(hitcount(averageSeries(app.grafana.*.dashboards.views.count), '1mon'), 4)",
		},
		{
			name:     "should not transform 1m when not used as interval string",
			target:   "app.grafana.*.dashboards.views.1m.count",
			expected: "app.grafana.*.dashboards.views.1m.count",
		},
		{
			name:     "should not transform 1M when not used as interval string",
			target:   "app.grafana.*.dashboards.views.1M.count",
			expected: "app.grafana.*.dashboards.views.1M.count",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			tr := fixIntervalFormat(tc.target)
			assert.Equal(t, tc.expected, tr)
		})
	}
}
<|endoftext|>


<commit_before>/*
Package loom implements a set of functions to interact with remote servers using SSH.
It is based on the Python fabric library.
*/
package loom

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"

	"code.google.com/p/go.crypto/ssh"
)

// Config contains ssh and other configuration data needed for all the public functions in loom.
type Config struct {
	// The user name used in SSH connections. If not specified, the current user is assumed.
	User string

	// Password for SSH connections. This is optional. If the user has an ~/.ssh/id_rsa keyfile,
	// that will also be tried. In addition, other key files can be specified.
	Password string

	// The machine:port to connect to.
	Host string

	// The file names of additional key files to use for authentication (~/.ssh/id_rsa is defaulted).
	// RSA (PKCS#1), DSA (OpenSSL), and ECDSA private keys are supported.
	KeyFiles []string

	// If true, send command output to stdout.
	DisplayOutput bool

	// If true, errors are fatal and will abort immediately.
	AbortOnError bool
}

// parsekey is a private function that reads in a keyfile containing a private key and parses it.
func parsekey(file string) (ssh.Signer, error) {
	var private ssh.Signer
	privateBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return private, err
	}

	private, err = ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		return private, nil
	}
	return private, nil
}

// connect is a private function to set up the ssh connection. It is called at the beginning of every public
// function.
func (config *Config) connect() (*ssh.Session, error) {

	sshconfig := &ssh.ClientConfig{
		User: config.User,
	}

	if config.User == "" {
		u, err := user.Current()
		if err != nil {
			return nil, err
		}
		sshconfig.User = u.Username
	}

	if config.Password != "" {
		sshconfig.Auth = append(sshconfig.Auth, ssh.Password(config.Password))
	}

	// By default, we try to include ~/.ssh/id_rsa. It is not an error if this file
	// doesn't exist.
	keyfile := os.Getenv("HOME") + "/.ssh/id_rsa"
	pkey, err := parsekey(keyfile)
	if err == nil {
		sshconfig.Auth = append(sshconfig.Auth, ssh.PublicKeys(pkey))
	}

	// Include any additional key files
	for _, keyfile = range config.KeyFiles {
		pkey, err = parsekey(keyfile)
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return nil, err
		}
		sshconfig.Auth = append(sshconfig.Auth, ssh.PublicKeys(pkey))
	}

	host := config.Host
	if strings.Contains(host, ":") == false {
		host = host + ":22"
	}
	client, err := ssh.Dial("tcp", host, sshconfig)
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return nil, err
	}

	session, err := client.NewSession()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return nil, err
	}
	return session, err
}

// doRun is called by both Run() and Sudo() to execute a command.
func (config *Config) doRun(cmd string, sudo bool) (string, error) {

	session, err := config.connect()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return "", err
	}
	defer session.Close()

	// Set up terminal modes
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}
	// Request pseudo terminal
	if err := session.RequestPty("xterm", 80, 40, modes); err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return "", err
	}

	if sudo == true {
		cmd = fmt.Sprintf("/usr/bin/sudo bash <<CMD\nexport PATH=/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin:/root/bin\n%s\nCMD", cmd)
	}

	// TODO: use pipes instead of CombinedOutput so that we can show the output of commands more interactively, instead
	// of now, which is after they're completely done executing.
	output, err := session.CombinedOutput(cmd)
	if err != nil {
		if config.DisplayOutput == true && len(output) > 0 {
			fmt.Printf("%s", string(output))
		}
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return "", err
	}
	session.SendRequest("close", false, nil)
	if config.DisplayOutput == true {
		fmt.Printf("%s", string(output))
	}
	return string(output), nil
}

// Run takes a command and runs it on the remote host, using ssh.
func (config *Config) Run(cmd string) (string, error) {
	if config.DisplayOutput == true {
		fmt.Printf("run: %s\n", cmd)
	}
	return config.doRun(cmd, false)
}

// Sudo takes a command and runs it as root on the remote host, using sudo over ssh.
func (config *Config) Sudo(cmd string) (string, error) {
	if config.DisplayOutput == true {
		fmt.Printf("sudo: %s\n", cmd)
	}
	return config.doRun(cmd, true)
}

// Put copies one or more local files to the remote host, using scp. localfiles can
// contain wildcards, and remotefile can be either a directory or a file.
func (config *Config) Put(localfiles string, remotefile string) error {

	files, err := filepath.Glob(localfiles)
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	if len(files) == 0 {
		err = fmt.Errorf("No files match %s", localfiles)
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	for _, localfile := range files {

		fmt.Printf("put: %s %s\n", localfile, remotefile)
		contents, err := ioutil.ReadFile(localfile)
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return err
		}

		// get the local file mode bits
		fi, err := os.Stat(localfile)
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return err
		}
		// the file mode bits are the 9 least significant bits of Mode()
		mode := fi.Mode() & 1023

		session, err := config.connect()
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return err
		}
		var stdoutBuf bytes.Buffer
		var stderrBuf bytes.Buffer
		session.Stdout = &stdoutBuf
		session.Stderr = &stderrBuf

		w, _ := session.StdinPipe()

		_, lfile := filepath.Split(localfile)
		err = session.Start("/usr/bin/scp -qrt " + remotefile)
		if err != nil {
			w.Close()
			session.Close()
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return err
		}
		fmt.Fprintf(w, "C%04o %d %s\n", mode, len(contents), lfile /*remotefile*/)
		w.Write(contents)
		fmt.Fprint(w, "\x00")
		w.Close()

		err = session.Wait()
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			session.Close()
			return err
		}

		if config.DisplayOutput == true {
			stdout := stdoutBuf.String()
			stderr := stderrBuf.String()
			fmt.Printf("%s%s", stderr, stdout)
		}
		session.Close()

	}

	return nil
}

// PutString generates a new file on the remote host containing data. The file is created with mode 0644.
func (config *Config) PutString(data string, remotefile string) error {

	if config.DisplayOutput == true {
		fmt.Printf("putstring: %s\n", remotefile)
	}
	session, err := config.connect()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	var stdoutBuf bytes.Buffer
	var stderrBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	session.Stderr = &stderrBuf

	w, _ := session.StdinPipe()

	_, rfile := filepath.Split(remotefile)
	err = session.Start("/usr/bin/scp -qrt " + remotefile)
	if err != nil {
		w.Close()
		session.Close()
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	fmt.Fprintf(w, "C0644 %d %s\n", len(data), rfile)
	w.Write([]byte(data))
	fmt.Fprint(w, "\x00")
	w.Close()

	err = session.Wait()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		session.Close()
		return err
	}

	if config.DisplayOutput == true {
		stdout := stdoutBuf.String()
		stderr := stderrBuf.String()
		fmt.Printf("%s%s", stderr, stdout)
	}
	session.Close()

	return nil
}

// Get copies the file from the remote host to the local host, using scp. Wildcards are not currently supported.
func (config *Config) Get(remotefile string, localfile string) error {

	if config.DisplayOutput == true {
		fmt.Printf("get: %s %s\n", remotefile, localfile)
	}

	session, err := config.connect()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	defer session.Close()

	// TODO: Handle wildcards in remotefile

	var stdoutBuf bytes.Buffer
	var stderrBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	session.Stderr = &stderrBuf
	w, _ := session.StdinPipe()

	err = session.Start("/usr/bin/scp -qrf " + remotefile)
	if err != nil {
		w.Close()
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	// TODO: better error checking than just firing and forgetting these nulls.
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")

	err = session.Wait()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}

	stdout := stdoutBuf.String()
	//stderr := stderrBuf.String()

	// first line of stdout contains file information
	fields := strings.SplitN(stdout, "\n", 2)
	mode, _ := strconv.ParseInt(fields[0][1:5], 8, 32)

	// need to generate final local file name
	// localfile could be a directory or a filename
	// if it's a directory, we need to append the remotefile filename
	// if it doesn't exist, we assume file
	var lfile string
	_, rfile := filepath.Split(remotefile)
	l := len(localfile)
	if localfile[l-1] == '/' {
		localfile = localfile[:l-1]
	}
	fi, err := os.Stat(localfile)
	if err != nil || fi.IsDir() == false {
		lfile = localfile
	} else if fi.IsDir() == true {
		lfile = localfile + "/" + rfile
	}
	// there's a trailing 0 in the file that we need to nuke
	l = len(fields[1])
	err = ioutil.WriteFile(lfile, []byte(fields[1][:l-1]), os.FileMode(mode))
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	return nil
}

// Local executes a command on the local host.
func (config *Config) Local(cmd string) (string, error) {
	if config.DisplayOutput == true {
		fmt.Printf("local: %s\n", cmd)
	}
	fields := strings.Split(cmd, " ")
	var command *exec.Cmd
	if len(fields) == 1 {
		command = exec.Command(fields[0])
	} else {
		command = exec.Command(fields[0], fields[1:]...)
	}
	var out bytes.Buffer
	command.Stdout = &out
	err := command.Run()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return "", err
	}
	if config.DisplayOutput == true {
		fmt.Printf("%s", out.String())
	}
	return out.String(), nil
}
<commit_msg>Changed Local() to execute the command under /bin/sh. Previously, it would try to execute the command directly. I think having access to shell commands more closely aligns to how people expect to use Local(), and matches how Run/Sudo() work remotely.<commit_after>/*
Package loom implements a set of functions to interact with remote servers using SSH.
It is based on the Python fabric library.
*/
package loom

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"

	"code.google.com/p/go.crypto/ssh"
)

// Config contains ssh and other configuration data needed for all the public functions in loom.
type Config struct {
	// The user name used in SSH connections. If not specified, the current user is assumed.
	User string

	// Password for SSH connections. This is optional. If the user has an ~/.ssh/id_rsa keyfile,
	// that will also be tried. In addition, other key files can be specified.
	Password string

	// The machine:port to connect to.
	Host string

	// The file names of additional key files to use for authentication (~/.ssh/id_rsa is defaulted).
	// RSA (PKCS#1), DSA (OpenSSL), and ECDSA private keys are supported.
	KeyFiles []string

	// If true, send command output to stdout.
	DisplayOutput bool

	// If true, errors are fatal and will abort immediately.
	AbortOnError bool
}

// parsekey is a private function that reads in a keyfile containing a private key and parses it.
func parsekey(file string) (ssh.Signer, error) {
	var private ssh.Signer
	privateBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return private, err
	}

	private, err = ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		return private, nil
	}
	return private, nil
}

// connect is a private function to set up the ssh connection. It is called at the beginning of every public
// function.
func (config *Config) connect() (*ssh.Session, error) {

	sshconfig := &ssh.ClientConfig{
		User: config.User,
	}

	if config.User == "" {
		u, err := user.Current()
		if err != nil {
			return nil, err
		}
		sshconfig.User = u.Username
	}

	if config.Password != "" {
		sshconfig.Auth = append(sshconfig.Auth, ssh.Password(config.Password))
	}

	// By default, we try to include ~/.ssh/id_rsa. It is not an error if this file
	// doesn't exist.
	keyfile := os.Getenv("HOME") + "/.ssh/id_rsa"
	pkey, err := parsekey(keyfile)
	if err == nil {
		sshconfig.Auth = append(sshconfig.Auth, ssh.PublicKeys(pkey))
	}

	// Include any additional key files
	for _, keyfile = range config.KeyFiles {
		pkey, err = parsekey(keyfile)
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return nil, err
		}
		sshconfig.Auth = append(sshconfig.Auth, ssh.PublicKeys(pkey))
	}

	host := config.Host
	if strings.Contains(host, ":") == false {
		host = host + ":22"
	}
	client, err := ssh.Dial("tcp", host, sshconfig)
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return nil, err
	}

	session, err := client.NewSession()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return nil, err
	}
	return session, err
}

// doRun is called by both Run() and Sudo() to execute a command.
func (config *Config) doRun(cmd string, sudo bool) (string, error) {

	session, err := config.connect()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return "", err
	}
	defer session.Close()

	// Set up terminal modes
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}
	// Request pseudo terminal
	if err := session.RequestPty("xterm", 80, 40, modes); err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return "", err
	}

	if sudo == true {
		cmd = fmt.Sprintf("/usr/bin/sudo bash <<CMD\nexport PATH=/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin:/root/bin\n%s\nCMD", cmd)
	}

	// TODO: use pipes instead of CombinedOutput so that we can show the output of commands more interactively, instead
	// of now, which is after they're completely done executing.
	output, err := session.CombinedOutput(cmd)
	if err != nil {
		if config.DisplayOutput == true && len(output) > 0 {
			fmt.Printf("%s", string(output))
		}
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return "", err
	}
	session.SendRequest("close", false, nil)
	if config.DisplayOutput == true {
		fmt.Printf("%s", string(output))
	}
	return string(output), nil
}

// Run takes a command and runs it on the remote host, using ssh.
func (config *Config) Run(cmd string) (string, error) {
	if config.DisplayOutput == true {
		fmt.Printf("run: %s\n", cmd)
	}
	return config.doRun(cmd, false)
}

// Sudo takes a command and runs it as root on the remote host, using sudo over ssh.
func (config *Config) Sudo(cmd string) (string, error) {
	if config.DisplayOutput == true {
		fmt.Printf("sudo: %s\n", cmd)
	}
	return config.doRun(cmd, true)
}

// Put copies one or more local files to the remote host, using scp. localfiles can
// contain wildcards, and remotefile can be either a directory or a file.
func (config *Config) Put(localfiles string, remotefile string) error {

	files, err := filepath.Glob(localfiles)
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	if len(files) == 0 {
		err = fmt.Errorf("No files match %s", localfiles)
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	for _, localfile := range files {

		fmt.Printf("put: %s %s\n", localfile, remotefile)
		contents, err := ioutil.ReadFile(localfile)
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return err
		}

		// get the local file mode bits
		fi, err := os.Stat(localfile)
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return err
		}
		// the file mode bits are the 9 least significant bits of Mode()
		mode := fi.Mode() & 1023

		session, err := config.connect()
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return err
		}
		var stdoutBuf bytes.Buffer
		var stderrBuf bytes.Buffer
		session.Stdout = &stdoutBuf
		session.Stderr = &stderrBuf

		w, _ := session.StdinPipe()

		_, lfile := filepath.Split(localfile)
		err = session.Start("/usr/bin/scp -qrt " + remotefile)
		if err != nil {
			w.Close()
			session.Close()
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			return err
		}
		fmt.Fprintf(w, "C%04o %d %s\n", mode, len(contents), lfile /*remotefile*/)
		w.Write(contents)
		fmt.Fprint(w, "\x00")
		w.Close()

		err = session.Wait()
		if err != nil {
			if config.AbortOnError == true {
				log.Fatalf("%s", err)
			}
			session.Close()
			return err
		}

		if config.DisplayOutput == true {
			stdout := stdoutBuf.String()
			stderr := stderrBuf.String()
			fmt.Printf("%s%s", stderr, stdout)
		}
		session.Close()

	}

	return nil
}

// PutString generates a new file on the remote host containing data. The file is created with mode 0644.
func (config *Config) PutString(data string, remotefile string) error {

	if config.DisplayOutput == true {
		fmt.Printf("putstring: %s\n", remotefile)
	}
	session, err := config.connect()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	var stdoutBuf bytes.Buffer
	var stderrBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	session.Stderr = &stderrBuf

	w, _ := session.StdinPipe()

	_, rfile := filepath.Split(remotefile)
	err = session.Start("/usr/bin/scp -qrt " + remotefile)
	if err != nil {
		w.Close()
		session.Close()
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	fmt.Fprintf(w, "C0644 %d %s\n", len(data), rfile)
	w.Write([]byte(data))
	fmt.Fprint(w, "\x00")
	w.Close()

	err = session.Wait()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		session.Close()
		return err
	}

	if config.DisplayOutput == true {
		stdout := stdoutBuf.String()
		stderr := stderrBuf.String()
		fmt.Printf("%s%s", stderr, stdout)
	}
	session.Close()

	return nil
}

// Get copies the file from the remote host to the local host, using scp. Wildcards are not currently supported.
func (config *Config) Get(remotefile string, localfile string) error {

	if config.DisplayOutput == true {
		fmt.Printf("get: %s %s\n", remotefile, localfile)
	}

	session, err := config.connect()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	defer session.Close()

	// TODO: Handle wildcards in remotefile

	var stdoutBuf bytes.Buffer
	var stderrBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	session.Stderr = &stderrBuf
	w, _ := session.StdinPipe()

	err = session.Start("/usr/bin/scp -qrf " + remotefile)
	if err != nil {
		w.Close()
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	// TODO: better error checking than just firing and forgetting these nulls.
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")
	fmt.Fprintf(w, "\x00")

	err = session.Wait()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}

	stdout := stdoutBuf.String()
	//stderr := stderrBuf.String()

	// first line of stdout contains file information
	fields := strings.SplitN(stdout, "\n", 2)
	mode, _ := strconv.ParseInt(fields[0][1:5], 8, 32)

	// need to generate final local file name
	// localfile could be a directory or a filename
	// if it's a directory, we need to append the remotefile filename
	// if it doesn't exist, we assume file
	var lfile string
	_, rfile := filepath.Split(remotefile)
	l := len(localfile)
	if localfile[l-1] == '/' {
		localfile = localfile[:l-1]
	}
	fi, err := os.Stat(localfile)
	if err != nil || fi.IsDir() == false {
		lfile = localfile
	} else if fi.IsDir() == true {
		lfile = localfile + "/" + rfile
	}
	// there's a trailing 0 in the file that we need to nuke
	l = len(fields[1])
	err = ioutil.WriteFile(lfile, []byte(fields[1][:l-1]), os.FileMode(mode))
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return err
	}
	return nil
}

// Local executes a command on the local host.
func (config *Config) Local(cmd string) (string, error) {
	if config.DisplayOutput == true {
		fmt.Printf("local: %s\n", cmd)
	}
	var command *exec.Cmd
	command = exec.Command("/bin/sh", "-c", cmd)
	var out bytes.Buffer
	command.Stdout = &out
	err := command.Run()
	if err != nil {
		if config.AbortOnError == true {
			log.Fatalf("%s", err)
		}
		return "", err
	}
	if config.DisplayOutput == true {
		fmt.Printf("%s", out.String())
	}
	return out.String(), nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Square Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package aggregate

// This file culminates in the definition of `aggregateBy`, which takes a SeriesList and an Aggregator and a list of tags,
// and produces an aggregated SeriesList with one list per group, each group having been aggregated into it.

import (
	"math"

	"github.com/square/metrics/api"
)

type group struct {
	List   []api.Timeseries
	TagSet api.TagSet
}

// If the given group will accept this given series (since it belongs to this group)
// then groupAccept will return true.
func groupAccepts(left api.TagSet, right api.TagSet, tags []string) bool {
	for _, tag := range tags {
		if left[tag] != right[tag] {
			return false
		}
	}
	return true
}

// Adds the series to the corresponding bucket, possibly modifying the input `rows` and returning a new list.
func addToGroup(rows []group, series api.Timeseries, tags []string) []group {
	// First we delete all tags with names other than those found in 'tags'
	newTags := api.NewTagSet()
	for _, tag := range tags {
		newTags[tag] = series.TagSet[tag]
	}
	// replace series' TagSet with newTags
	series.TagSet = newTags

	// Next, find the best bucket for this series:
	for i, row := range rows {
		if groupAccepts(row.TagSet, series.TagSet, tags) {
			rows[i].List = append(rows[i].List, series)
			return rows
		}
	}
	// Otherwise, no bucket yet exists
	return append(rows, group{
		[]api.Timeseries{series},
		newTags,
	})
}

// Groups the given SeriesList by tags, producing a list of lists (of type groupResult)
func groupBy(list api.SeriesList, tags []string) []group {
	result := []group{}
	for _, series := range list.Series {
		result = addToGroup(result, series, tags)
	}
	return result
}

var aggregateMap = map[string]func([]float64) float64{
	"aggregate.sum":
	// The aggregatefor sum finds the sum of the given array.
	func(array []float64) float64 {
		sum := 0.0
		for _, v := range array {
			sum += v
		}
		return sum
	},
	"aggregate.mean":
	// The mean aggregator returns the mean of the given array
	func(array []float64) float64 {
		if len(array) == 0 {
			// The mean of an empty list is not well-defined
			return math.NaN()
		}
		sum := 0.0
		for _, v := range array {
			sum += v
		}
		return sum / float64(len(array))
	},
	"aggregate.min":
	// The minimum aggregator returns the minimum
	func(array []float64) float64 {
		if len(array) == 0 {
			// The minimum of an empty list is not well-defined
			return math.NaN()
		}
		min := array[0]
		for _, v := range array {
			min = math.Min(min, v)
		}
		return min
	},
	"aggregate.max":
	// The maximum aggregator returns the maximum
	func(array []float64) float64 {
		if len(array) == 0 {
			// The maximum of an empty list is not well-defined
			return math.NaN()
		}
		max := array[0]
		for _, v := range array {
			max = math.Max(max, v)
		}
		return max
	},
}

// GetAggregate gives the aggregate of the given name (and false if it doesn't exist).
func GetAggregate(name string) (func([]float64) float64, bool) {
	aggregate, ok := aggregateMap[name]
	return aggregate, ok
}

// AddAggregate adds an aggregate of a given name to the aggregate map.
func NewAggregate(name string, aggregate func([]float64) float64) {
	aggregateMap[name] = aggregate
}

// applyAggregation takes an aggregation function ( [float64] => float64 ) and applies it to a given list of Timeseries
// the list must be non-empty, or an error is returned
func applyAggregation(group group, aggregator func([]float64) float64) api.Timeseries {
	list := group.List
	tagSet := group.TagSet

	if len(list) == 0 {
		// This case should not actually occur, provided the rest of the code has been implemented correctly.
		// So when it does, issue a panic:
		panic("applyAggregation given empty group for tagset")
	}

	series := api.Timeseries{
		Values: make([]float64, len(list[0].Values)), // The first Series in the given list is used to determine this length
		TagSet: tagSet,                               // The tagset is supplied by an argument (it will be the values grouped on)
	}

	// Make a slice of time to reuse.
	// Each entry corresponds to a particular Series, all having the same index within their corresponding Series.
	timeSlice := make([]float64, len(list))

	for i := range series.Values {
		// We need to determine each value in turn.
		for j := range timeSlice {
			timeSlice[j] = list[j].Values[i]
		}
		// Find the aggregated value:
		series.Values[i] = aggregator(timeSlice)
	}

	return series
}

// This function is the culmination of all others.
// `aggregateBy` takes a series list, an aggregator, and a set of tags.
// It produces a SeriesList which is the result of grouping by the tags and then aggregating each group
// into a single Series.
func AggregateBy(list api.SeriesList, aggregator func([]float64) float64, tags []string) api.SeriesList {
	// Begin by grouping the input:
	groups := groupBy(list, tags)

	result := api.SeriesList{
		Series:    make([]api.Timeseries, len(groups)),
		Timerange: list.Timerange,
		Name:      list.Name,
	}

	for i, group := range groups {
		// The group contains a list of Series and a TagSet.
		result.Series[i] = applyAggregation(group, aggregator)
	}
	return result
}
<commit_msg>add panic for redeclaration of aggregates<commit_after>// Copyright 2015 Square Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package aggregate

// This file culminates in the definition of `aggregateBy`, which takes a SeriesList and an Aggregator and a list of tags,
// and produces an aggregated SeriesList with one list per group, each group having been aggregated into it.

import (
	"fmt"
	"math"

	"github.com/square/metrics/api"
)

type group struct {
	List   []api.Timeseries
	TagSet api.TagSet
}

// If the given group will accept this given series (since it belongs to this group)
// then groupAccept will return true.
func groupAccepts(left api.TagSet, right api.TagSet, tags []string) bool {
	for _, tag := range tags {
		if left[tag] != right[tag] {
			return false
		}
	}
	return true
}

// Adds the series to the corresponding bucket, possibly modifying the input `rows` and returning a new list.
func addToGroup(rows []group, series api.Timeseries, tags []string) []group {
	// First we delete all tags with names other than those found in 'tags'
	newTags := api.NewTagSet()
	for _, tag := range tags {
		newTags[tag] = series.TagSet[tag]
	}
	// replace series' TagSet with newTags
	series.TagSet = newTags

	// Next, find the best bucket for this series:
	for i, row := range rows {
		if groupAccepts(row.TagSet, series.TagSet, tags) {
			rows[i].List = append(rows[i].List, series)
			return rows
		}
	}
	// Otherwise, no bucket yet exists
	return append(rows, group{
		[]api.Timeseries{series},
		newTags,
	})
}

// Groups the given SeriesList by tags, producing a list of lists (of type groupResult)
func groupBy(list api.SeriesList, tags []string) []group {
	result := []group{}
	for _, series := range list.Series {
		result = addToGroup(result, series, tags)
	}
	return result
}

var aggregateMap = map[string]func([]float64) float64{
	"aggregate.sum":
	// The aggregatefor sum finds the sum of the given array.
	func(array []float64) float64 {
		sum := 0.0
		for _, v := range array {
			sum += v
		}
		return sum
	},
	"aggregate.mean":
	// The mean aggregator returns the mean of the given array
	func(array []float64) float64 {
		if len(array) == 0 {
			// The mean of an empty list is not well-defined
			return math.NaN()
		}
		sum := 0.0
		for _, v := range array {
			sum += v
		}
		return sum / float64(len(array))
	},
	"aggregate.min":
	// The minimum aggregator returns the minimum
	func(array []float64) float64 {
		if len(array) == 0 {
			// The minimum of an empty list is not well-defined
			return math.NaN()
		}
		min := array[0]
		for _, v := range array {
			min = math.Min(min, v)
		}
		return min
	},
	"aggregate.max":
	// The maximum aggregator returns the maximum
	func(array []float64) float64 {
		if len(array) == 0 {
			// The maximum of an empty list is not well-defined
			return math.NaN()
		}
		max := array[0]
		for _, v := range array {
			max = math.Max(max, v)
		}
		return max
	},
}

// GetAggregate gives the aggregate of the given name (and false if it doesn't exist).
func GetAggregate(name string) (func([]float64) float64, bool) {
	aggregate, ok := aggregateMap[name]
	return aggregate, ok
}

// AddAggregate adds an aggregate of a given name to the aggregate map.
func NewAggregate(name string, aggregate func([]float64) float64) {
	if _, ok := aggregateMap[name]; ok {
		panic(fmt.Sprintf("aggregate %s has already been declared", name))
	}
	aggregateMap[name] = aggregate
}

// applyAggregation takes an aggregation function ( [float64] => float64 ) and applies it to a given list of Timeseries
// the list must be non-empty, or an error is returned
func applyAggregation(group group, aggregator func([]float64) float64) api.Timeseries {
	list := group.List
	tagSet := group.TagSet

	if len(list) == 0 {
		// This case should not actually occur, provided the rest of the code has been implemented correctly.
		// So when it does, issue a panic:
		panic("applyAggregation given empty group for tagset")
	}

	series := api.Timeseries{
		Values: make([]float64, len(list[0].Values)), // The first Series in the given list is used to determine this length
		TagSet: tagSet,                               // The tagset is supplied by an argument (it will be the values grouped on)
	}

	// Make a slice of time to reuse.
	// Each entry corresponds to a particular Series, all having the same index within their corresponding Series.
	timeSlice := make([]float64, len(list))

	for i := range series.Values {
		// We need to determine each value in turn.
		for j := range timeSlice {
			timeSlice[j] = list[j].Values[i]
		}
		// Find the aggregated value:
		series.Values[i] = aggregator(timeSlice)
	}

	return series
}

// This function is the culmination of all others.
// `aggregateBy` takes a series list, an aggregator, and a set of tags.
// It produces a SeriesList which is the result of grouping by the tags and then aggregating each group
// into a single Series.
func AggregateBy(list api.SeriesList, aggregator func([]float64) float64, tags []string) api.SeriesList {
	// Begin by grouping the input:
	groups := groupBy(list, tags)

	result := api.SeriesList{
		Series:    make([]api.Timeseries, len(groups)),
		Timerange: list.Timerange,
		Name:      list.Name,
	}

	for i, group := range groups {
		// The group contains a list of Series and a TagSet.
		result.Series[i] = applyAggregation(group, aggregator)
	}
	return result
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strings"
	"time"
	"unicode"

	"github.com/spiffe/go-spiffe/v2/spiffetls/tlsconfig"
	"github.com/spiffe/go-spiffe/v2/workloadapi"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

var (
	flAddr          string
	flService       string
	flUserAgent     string
	flConnTimeout   time.Duration
	flRPCHeaders    = rpcHeaders{MD: make(metadata.MD)}
	flRPCTimeout    time.Duration
	flTLS           bool
	flTLSNoVerify   bool
	flTLSCACert     string
	flTLSClientCert string
	flTLSClientKey  string
	flTLSServerName string
	flVerbose       bool
	flGZIP          bool
	flSPIFFE        bool
)

const (
	// StatusInvalidArguments indicates specified invalid arguments.
	StatusInvalidArguments = 1
	// StatusConnectionFailure indicates connection failed.
	StatusConnectionFailure = 2
	// StatusRPCFailure indicates rpc failed.
	StatusRPCFailure = 3
	// StatusUnhealthy indicates rpc succeeded but indicates unhealthy service.
	StatusUnhealthy = 4
	// StatusSpiffeFailed indicates failure to retrieve credentials using spiffe workload API
	StatusSpiffeFailed = 20
)

func init() {
	flagSet := flag.NewFlagSet("", flag.ContinueOnError)
	log.SetFlags(0)
	flagSet.StringVar(&flAddr, "addr", "", "(required) tcp host:port to connect")
	flagSet.StringVar(&flService, "service", "", "service name to check (default: \"\")")
	flagSet.StringVar(&flUserAgent, "user-agent", "grpc_health_probe", "user-agent header value of health check requests")
	// timeouts
	flagSet.DurationVar(&flConnTimeout, "connect-timeout", time.Second, "timeout for establishing connection")
	flagSet.Var(&flRPCHeaders, "rpc-header", "additional RPC headers in 'name: value' format. May specify more than one via multiple flags.")
	flagSet.DurationVar(&flRPCTimeout, "rpc-timeout", time.Second, "timeout for health check rpc")
	// tls settings
	flagSet.BoolVar(&flTLS, "tls", false, "use TLS (default: false, INSECURE plaintext transport)")
	flagSet.BoolVar(&flTLSNoVerify, "tls-no-verify", false, "(with -tls) don't verify the certificate (INSECURE) presented by the server (default: false)")
	flagSet.StringVar(&flTLSCACert, "tls-ca-cert", "", "(with -tls, optional) file containing trusted certificates for verifying server")
	flagSet.StringVar(&flTLSClientCert, "tls-client-cert", "", "(with -tls, optional) client certificate for authenticating to the server (requires -tls-client-key)")
	flagSet.StringVar(&flTLSClientKey, "tls-client-key", "", "(with -tls) client private key for authenticating to the server (requires -tls-client-cert)")
	flagSet.StringVar(&flTLSServerName, "tls-server-name", "", "(with -tls) override the hostname used to verify the server certificate")
	flagSet.BoolVar(&flVerbose, "v", false, "verbose logs")
	flagSet.BoolVar(&flGZIP, "gzip", false, "use GZIPCompressor for requests and GZIPDecompressor for response (default: false)")
	flagSet.BoolVar(&flSPIFFE, "spiffe", false, "use SPIFFE to obtain mTLS credentials")

	err := flagSet.Parse(os.Args[1:])
	if err != nil {
		os.Exit(StatusInvalidArguments)
	}

	argError := func(s string, v ...interface{}) {
		log.Printf("error: "+s, v...)
		os.Exit(StatusInvalidArguments)
	}

	if flAddr == "" {
		argError("-addr not specified")
	}
	if flConnTimeout <= 0 {
		argError("-connect-timeout must be greater than zero (specified: %v)", flConnTimeout)
	}
	if flRPCTimeout <= 0 {
		argError("-rpc-timeout must be greater than zero (specified: %v)", flRPCTimeout)
	}
	if !flTLS && flTLSNoVerify {
		argError("specified -tls-no-verify without specifying -tls")
	}
	if !flTLS && flTLSCACert != "" {
		argError("specified -tls-ca-cert without specifying -tls")
	}
	if !flTLS && flTLSClientCert != "" {
		argError("specified -tls-client-cert without specifying -tls")
	}
	if !flTLS && flTLSServerName != "" {
		argError("specified -tls-server-name without specifying -tls")
	}
	if flTLSClientCert != "" && flTLSClientKey == "" {
		argError("specified -tls-client-cert without specifying -tls-client-key")
	}
	if flTLSClientCert == "" && flTLSClientKey != "" {
		argError("specified -tls-client-key without specifying -tls-client-cert")
	}
	if flTLSNoVerify && flTLSCACert != "" {
		argError("cannot specify -tls-ca-cert with -tls-no-verify (CA cert would not be used)")
	}
	if flTLSNoVerify && flTLSServerName != "" {
		argError("cannot specify -tls-server-name with -tls-no-verify (server name would not be used)")
	}

	if flVerbose {
		log.Printf("parsed options:")
		log.Printf("> addr=%s conn_timeout=%v rpc_timeout=%v", flAddr, flConnTimeout, flRPCTimeout)
		if flRPCHeaders.Len() > 0 {
			log.Printf("> headers: %s", flRPCHeaders)
		}
		log.Printf("> tls=%v", flTLS)
		if flTLS {
			log.Printf("  > no-verify=%v ", flTLSNoVerify)
			log.Printf("  > ca-cert=%s", flTLSCACert)
			log.Printf("  > client-cert=%s", flTLSClientCert)
			log.Printf("  > client-key=%s", flTLSClientKey)
			log.Printf("  > server-name=%s", flTLSServerName)
		}
		log.Printf("> spiffe=%v", flSPIFFE)
	}
}

type rpcHeaders struct{ metadata.MD }

func (s *rpcHeaders) String() string { return fmt.Sprintf("%v", s.MD) }

func (s *rpcHeaders) Set(value string) error {
	parts := strings.SplitN(value, ":", 2)
	if len(parts) != 2 {
		return fmt.Errorf("invalid RPC header, expected 'key: value', got %q", value)
	}
	trimmed := strings.TrimLeftFunc(parts[1], unicode.IsSpace)
	s.Append(parts[0], trimmed)
	return nil
}

func buildCredentials(skipVerify bool, caCerts, clientCert, clientKey, serverName string) (credentials.TransportCredentials, error) {
	var cfg tls.Config

	if clientCert != "" && clientKey != "" {
		keyPair, err := tls.LoadX509KeyPair(clientCert, clientKey)
		if err != nil {
			return nil, fmt.Errorf("failed to load tls client cert/key pair. error=%v", err)
		}
		cfg.Certificates = []tls.Certificate{keyPair}
	}

	if skipVerify {
		cfg.InsecureSkipVerify = true
	} else if caCerts != "" {
		// override system roots
		rootCAs := x509.NewCertPool()
		pem, err := ioutil.ReadFile(caCerts)
		if err != nil {
			return nil, fmt.Errorf("failed to load root CA certificates from file (%s) error=%v", caCerts, err)
		}
		if !rootCAs.AppendCertsFromPEM(pem) {
			return nil, fmt.Errorf("no root CA certs parsed from file %s", caCerts)
		}
		cfg.RootCAs = rootCAs
	}
	if serverName != "" {
		cfg.ServerName = serverName
	}
	return credentials.NewTLS(&cfg), nil
}

func main() {
	retcode := 0
	defer func() { os.Exit(retcode) }()

	ctx, cancel := context.WithCancel(context.Background())

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		sig := <-c
		if sig == os.Interrupt {
			log.Printf("cancellation received")
			cancel()
			return
		}
	}()

	opts := []grpc.DialOption{
		grpc.WithUserAgent(flUserAgent),
		grpc.WithBlock(),
	}
	if flTLS && flSPIFFE {
		log.Printf("-tls and -spiffe are mutually incompatible")
		retcode = StatusInvalidArguments
		return
	}
	if flTLS {
		creds, err := buildCredentials(flTLSNoVerify, flTLSCACert, flTLSClientCert, flTLSClientKey, flTLSServerName)
		if err != nil {
			log.Printf("failed to initialize tls credentials. error=%v", err)
			retcode = StatusInvalidArguments
			return
		}
		opts = append(opts, grpc.WithTransportCredentials(creds))
	} else if flSPIFFE {
		spiffeCtx, _ := context.WithTimeout(ctx, flRPCTimeout)
		source, err := workloadapi.NewX509Source(spiffeCtx)
		if err != nil {
			log.Printf("failed to initialize tls credentials with spiffe. error=%v", err)
			retcode = StatusSpiffeFailed
			return
		}
		if flVerbose {
			svid, err := source.GetX509SVID()
			if err != nil {
				log.Fatalf("error getting x509 svid: %+v", err)
			}
			log.Printf("SPIFFE Verifiable Identity Document (SVID): %q", svid.ID)
		}
		creds := credentials.NewTLS(tlsconfig.MTLSClientConfig(source, source, tlsconfig.AuthorizeAny()))
		opts = append(opts, grpc.WithTransportCredentials(creds))
	} else {
		opts = append(opts, grpc.WithInsecure())
	}

	if flGZIP {
		opts = append(opts,
			grpc.WithCompressor(grpc.NewGZIPCompressor()),
			grpc.WithDecompressor(grpc.NewGZIPDecompressor()),
		)
	}

	if flVerbose {
		log.Print("establishing connection")
	}
	connStart := time.Now()
	dialCtx, dialCancel := context.WithTimeout(ctx, flConnTimeout)
	defer dialCancel()
	conn, err := grpc.DialContext(dialCtx, flAddr, opts...)
	if err != nil {
		if err == context.DeadlineExceeded {
			log.Printf("timeout: failed to connect service %q within %v", flAddr, flConnTimeout)
		} else {
			log.Printf("error: failed to connect service at %q: %+v", flAddr, err)
		}
		retcode = StatusConnectionFailure
		return
	}
	connDuration := time.Since(connStart)
	defer conn.Close()
	if flVerbose {
		log.Printf("connection established (took %v)", connDuration)
	}

	rpcStart := time.Now()
	rpcCtx, rpcCancel := context.WithTimeout(ctx, flRPCTimeout)
	defer rpcCancel()
	rpcCtx = metadata.NewOutgoingContext(rpcCtx, flRPCHeaders.MD)
	resp, err := healthpb.NewHealthClient(conn).Check(rpcCtx,
		&healthpb.HealthCheckRequest{
			Service: flService})
	if err != nil {
		if stat, ok := status.FromError(err); ok && stat.Code() == codes.Unimplemented {
			log.Printf("error: this server does not implement the grpc health protocol (grpc.health.v1.Health): %s", stat.Message())
		} else if stat, ok := status.FromError(err); ok && stat.Code() == codes.DeadlineExceeded {
			log.Printf("timeout: health rpc did not complete within %v", flRPCTimeout)
		} else {
			log.Printf("error: health rpc failed: %+v", err)
		}
		retcode = StatusRPCFailure
		return
	}
	rpcDuration := time.Since(rpcStart)

	if resp.GetStatus() != healthpb.HealthCheckResponse_SERVING {
		log.Printf("service unhealthy (responded with %q)", resp.GetStatus().String())
		retcode = StatusUnhealthy
		return
	}
	if flVerbose {
		log.Printf("time elapsed: connect=%v rpc=%v", connDuration, rpcDuration)
	}
	log.Printf("status: %v", resp.GetStatus().String())
}
<commit_msg>add support for ALTS transport (#106)<commit_after>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strings"
	"time"
	"unicode"

	"github.com/spiffe/go-spiffe/v2/spiffetls/tlsconfig"
	"github.com/spiffe/go-spiffe/v2/workloadapi"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/alts"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

var (
	flAddr          string
	flService       string
	flUserAgent     string
	flConnTimeout   time.Duration
	flRPCHeaders    = rpcHeaders{MD: make(metadata.MD)}
	flRPCTimeout    time.Duration
	flTLS           bool
	flTLSNoVerify   bool
	flTLSCACert     string
	flTLSClientCert string
	flTLSClientKey  string
	flTLSServerName string
	flALTS          bool
	flVerbose       bool
	flGZIP          bool
	flSPIFFE        bool
)

const (
	// StatusInvalidArguments indicates specified invalid arguments.
	StatusInvalidArguments = 1
	// StatusConnectionFailure indicates connection failed.
	StatusConnectionFailure = 2
	// StatusRPCFailure indicates rpc failed.
	StatusRPCFailure = 3
	// StatusUnhealthy indicates rpc succeeded but indicates unhealthy service.
	StatusUnhealthy = 4
	// StatusSpiffeFailed indicates failure to retrieve credentials using spiffe workload API
	StatusSpiffeFailed = 20
)

func init() {
	flagSet := flag.NewFlagSet("", flag.ContinueOnError)
	log.SetFlags(0)
	flagSet.StringVar(&flAddr, "addr", "", "(required) tcp host:port to connect")
	flagSet.StringVar(&flService, "service", "", "service name to check (default: \"\")")
	flagSet.StringVar(&flUserAgent, "user-agent", "grpc_health_probe", "user-agent header value of health check requests")
	// timeouts
	flagSet.DurationVar(&flConnTimeout, "connect-timeout", time.Second, "timeout for establishing connection")
	flagSet.Var(&flRPCHeaders, "rpc-header", "additional RPC headers in 'name: value' format. May specify more than one via multiple flags.")
	flagSet.DurationVar(&flRPCTimeout, "rpc-timeout", time.Second, "timeout for health check rpc")
	// tls settings
	flagSet.BoolVar(&flTLS, "tls", false, "use TLS (default: false, INSECURE plaintext transport)")
	flagSet.BoolVar(&flTLSNoVerify, "tls-no-verify", false, "(with -tls) don't verify the certificate (INSECURE) presented by the server (default: false)")
	flagSet.StringVar(&flTLSCACert, "tls-ca-cert", "", "(with -tls, optional) file containing trusted certificates for verifying server")
	flagSet.StringVar(&flTLSClientCert, "tls-client-cert", "", "(with -tls, optional) client certificate for authenticating to the server (requires -tls-client-key)")
	flagSet.StringVar(&flTLSClientKey, "tls-client-key", "", "(with -tls) client private key for authenticating to the server (requires -tls-client-cert)")
	flagSet.StringVar(&flTLSServerName, "tls-server-name", "", "(with -tls) override the hostname used to verify the server certificate")
	flagSet.BoolVar(&flALTS, "alts", false, "use ALTS (default: false, INSECURE plaintext transport)")
	flagSet.BoolVar(&flVerbose, "v", false, "verbose logs")
	flagSet.BoolVar(&flGZIP, "gzip", false, "use GZIPCompressor for requests and GZIPDecompressor for response (default: false)")
	flagSet.BoolVar(&flSPIFFE, "spiffe", false, "use SPIFFE to obtain mTLS credentials")

	err := flagSet.Parse(os.Args[1:])
	if err != nil {
		os.Exit(StatusInvalidArguments)
	}

	argError := func(s string, v ...interface{}) {
		log.Printf("error: "+s, v...)
		os.Exit(StatusInvalidArguments)
	}

	if flAddr == "" {
		argError("-addr not specified")
	}
	if flConnTimeout <= 0 {
		argError("-connect-timeout must be greater than zero (specified: %v)", flConnTimeout)
	}
	if flRPCTimeout <= 0 {
		argError("-rpc-timeout must be greater than zero (specified: %v)", flRPCTimeout)
	}
	if flALTS && flSPIFFE {
		argError("-alts and -spiffe are mutually incompatible")
	}
	if flTLS && flALTS {
		argError("cannot specify -tls with -alts")
	}
	if !flTLS && flTLSNoVerify {
		argError("specified -tls-no-verify without specifying -tls")
	}
	if !flTLS && flTLSCACert != "" {
		argError("specified -tls-ca-cert without specifying -tls")
	}
	if !flTLS && flTLSClientCert != "" {
		argError("specified -tls-client-cert without specifying -tls")
	}
	if !flTLS && flTLSServerName != "" {
		argError("specified -tls-server-name without specifying -tls")
	}
	if flTLSClientCert != "" && flTLSClientKey == "" {
		argError("specified -tls-client-cert without specifying -tls-client-key")
	}
	if flTLSClientCert == "" && flTLSClientKey != "" {
		argError("specified -tls-client-key without specifying -tls-client-cert")
	}
	if flTLSNoVerify && flTLSCACert != "" {
		argError("cannot specify -tls-ca-cert with -tls-no-verify (CA cert would not be used)")
	}
	if flTLSNoVerify && flTLSServerName != "" {
		argError("cannot specify -tls-server-name with -tls-no-verify (server name would not be used)")
	}

	if flVerbose {
		log.Printf("parsed options:")
		log.Printf("> addr=%s conn_timeout=%v rpc_timeout=%v", flAddr, flConnTimeout, flRPCTimeout)
		if flRPCHeaders.Len() > 0 {
			log.Printf("> headers: %s", flRPCHeaders)
		}
		log.Printf("> tls=%v", flTLS)
		if flTLS {
			log.Printf("  > no-verify=%v ", flTLSNoVerify)
			log.Printf("  > ca-cert=%s", flTLSCACert)
			log.Printf("  > client-cert=%s", flTLSClientCert)
			log.Printf("  > client-key=%s", flTLSClientKey)
			log.Printf("  > server-name=%s", flTLSServerName)
		}
		log.Printf("> alts=%v", flALTS)
		log.Printf("> spiffe=%v", flSPIFFE)
	}
}

type rpcHeaders struct{ metadata.MD }

func (s *rpcHeaders) String() string { return fmt.Sprintf("%v", s.MD) }

func (s *rpcHeaders) Set(value string) error {
	parts := strings.SplitN(value, ":", 2)
	if len(parts) != 2 {
		return fmt.Errorf("invalid RPC header, expected 'key: value', got %q", value)
	}
	trimmed := strings.TrimLeftFunc(parts[1], unicode.IsSpace)
	s.Append(parts[0], trimmed)
	return nil
}

func buildCredentials(skipVerify bool, caCerts, clientCert, clientKey, serverName string) (credentials.TransportCredentials, error) {
	var cfg tls.Config

	if clientCert != "" && clientKey != "" {
		keyPair, err := tls.LoadX509KeyPair(clientCert, clientKey)
		if err != nil {
			return nil, fmt.Errorf("failed to load tls client cert/key pair. error=%v", err)
		}
		cfg.Certificates = []tls.Certificate{keyPair}
	}

	if skipVerify {
		cfg.InsecureSkipVerify = true
	} else if caCerts != "" {
		// override system roots
		rootCAs := x509.NewCertPool()
		pem, err := ioutil.ReadFile(caCerts)
		if err != nil {
			return nil, fmt.Errorf("failed to load root CA certificates from file (%s) error=%v", caCerts, err)
		}
		if !rootCAs.AppendCertsFromPEM(pem) {
			return nil, fmt.Errorf("no root CA certs parsed from file %s", caCerts)
		}
		cfg.RootCAs = rootCAs
	}
	if serverName != "" {
		cfg.ServerName = serverName
	}
	return credentials.NewTLS(&cfg), nil
}

func main() {
	retcode := 0
	defer func() { os.Exit(retcode) }()

	ctx, cancel := context.WithCancel(context.Background())

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		sig := <-c
		if sig == os.Interrupt {
			log.Printf("cancellation received")
			cancel()
			return
		}
	}()

	opts := []grpc.DialOption{
		grpc.WithUserAgent(flUserAgent),
		grpc.WithBlock(),
	}
	if flTLS && flSPIFFE {
		log.Printf("-tls and -spiffe are mutually incompatible")
		retcode = StatusInvalidArguments
		return
	}
	if flTLS {
		creds, err := buildCredentials(flTLSNoVerify, flTLSCACert, flTLSClientCert, flTLSClientKey, flTLSServerName)
		if err != nil {
			log.Printf("failed to initialize tls credentials. error=%v", err)
			retcode = StatusInvalidArguments
			return
		}
		opts = append(opts, grpc.WithTransportCredentials(creds))
	} else if flALTS {
		creds := alts.NewServerCreds(alts.DefaultServerOptions())
		opts = append(opts, grpc.WithTransportCredentials(creds))
	} else if flSPIFFE {
		spiffeCtx, _ := context.WithTimeout(ctx, flRPCTimeout)
		source, err := workloadapi.NewX509Source(spiffeCtx)
		if err != nil {
			log.Printf("failed to initialize tls credentials with spiffe. error=%v", err)
			retcode = StatusSpiffeFailed
			return
		}
		if flVerbose {
			svid, err := source.GetX509SVID()
			if err != nil {
				log.Fatalf("error getting x509 svid: %+v", err)
			}
			log.Printf("SPIFFE Verifiable Identity Document (SVID): %q", svid.ID)
		}
		creds := credentials.NewTLS(tlsconfig.MTLSClientConfig(source, source, tlsconfig.AuthorizeAny()))
		opts = append(opts, grpc.WithTransportCredentials(creds))
	} else {
		opts = append(opts, grpc.WithInsecure())
	}

	if flGZIP {
		opts = append(opts,
			grpc.WithCompressor(grpc.NewGZIPCompressor()),
			grpc.WithDecompressor(grpc.NewGZIPDecompressor()),
		)
	}

	if flVerbose {
		log.Print("establishing connection")
	}
	connStart := time.Now()
	dialCtx, dialCancel := context.WithTimeout(ctx, flConnTimeout)
	defer dialCancel()
	conn, err := grpc.DialContext(dialCtx, flAddr, opts...)
	if err != nil {
		if err == context.DeadlineExceeded {
			log.Printf("timeout: failed to connect service %q within %v", flAddr, flConnTimeout)
		} else {
			log.Printf("error: failed to connect service at %q: %+v", flAddr, err)
		}
		retcode = StatusConnectionFailure
		return
	}
	connDuration := time.Since(connStart)
	defer conn.Close()
	if flVerbose {
		log.Printf("connection established (took %v)", connDuration)
	}

	rpcStart := time.Now()
	rpcCtx, rpcCancel := context.WithTimeout(ctx, flRPCTimeout)
	defer rpcCancel()
	rpcCtx = metadata.NewOutgoingContext(rpcCtx, flRPCHeaders.MD)
	resp, err := healthpb.NewHealthClient(conn).Check(rpcCtx,
		&healthpb.HealthCheckRequest{
			Service: flService})
	if err != nil {
		if stat, ok := status.FromError(err); ok && stat.Code() == codes.Unimplemented {
			log.Printf("error: this server does not implement the grpc health protocol (grpc.health.v1.Health): %s", stat.Message())
		} else if stat, ok := status.FromError(err); ok && stat.Code() == codes.DeadlineExceeded {
			log.Printf("timeout: health rpc did not complete within %v", flRPCTimeout)
		} else {
			log.Printf("error: health rpc failed: %+v", err)
		}
		retcode = StatusRPCFailure
		return
	}
	rpcDuration := time.Since(rpcStart)

	if resp.GetStatus() != healthpb.HealthCheckResponse_SERVING {
		log.Printf("service unhealthy (responded with %q)", resp.GetStatus().String())
		retcode = StatusUnhealthy
		return
	}
	if flVerbose {
		log.Printf("time elapsed: connect=%v rpc=%v", connDuration, rpcDuration)
	}
	log.Printf("status: %v", resp.GetStatus().String())
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/bwmarrin/discordgo"
	"github.com/spf13/viper"
)

const botTestChannel = "276738368749043712"

var token string

func init() {
}

func main() {
	viper.SetConfigType("yaml")
	viper.SetConfigName("config")
	viper.AddConfigPath("./")

	err := viper.ReadInConfig()
	if err != nil {
		fmt.Println(err)
	}

	token = viper.GetString("discordToken")

	if token == "" {
		fmt.Println("No token")
		os.Exit(1)
	}
	Session, err := discordgo.New("Bot " + token)
	if err != nil {
		fmt.Println("Can't create session", err)
	}

	Session.State.User, err = Session.User("@me")
	if err != nil {
		fmt.Println("Error fetching user info", err)
	}

	err = Session.Open()
	if err != nil {
		fmt.Println("Error connecting to discord", err)
	}

	//Session.ChannelMessageSend("274304138953752578", "Hallo Welt")
	channels, _ := Session.GuildChannels("208374693185585152")
	for _, channel := range channels {
		if channel.Type == "text" {
			//Session.ChannelMessageSend(channel.ID, fmt.Sprintf("Dies ist der Channel %s.", channel.Name))
		}
	}

	Session.AddHandler(messageCreateHandler)

	log.Printf(`Now running. Press CTRL-C to exit.`)
	sc := make(chan os.Signal, 1)
	signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
	<-sc

	// Clean up
	Session.Close()

}
<commit_msg>Add guildWars2Token<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/bwmarrin/discordgo"
	"github.com/spf13/viper"
	"github.com/yasvisu/gw2api"
)

var botTestChannel string
var discordToken string
var guildWars2Token string

func main() {
	viper.SetConfigType("yaml")
	viper.SetConfigName("config")
	viper.AddConfigPath("./")

	err := viper.ReadInConfig()
	if err != nil {
		fmt.Println(err)
	}

	session, err := createDiscordSession()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	api, err := createGuildWarsSession()
	if err != nil {
		fmt.Println(err)
		os.Exit(2)
	}

	_ = api

	log.Printf(`Now running. Press CTRL-C to exit.`)
	sc := make(chan os.Signal, 1)
	signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
	<-sc

	// Clean up
	session.Close()

}

func createDiscordSession() (*discordgo.Session, error) {
	discordToken = viper.GetString("discordToken")
	if discordToken == "" {
		return nil, fmt.Errorf("No discordToken")
	}
	Session, err := discordgo.New("Bot " + discordToken)
	if err != nil {
		fmt.Println("Can't create session", err)
		return nil, err
	}

	Session.State.User, err = Session.User("@me")
	if err != nil {
		fmt.Println("Error fetching user info", err)
		return nil, err
	}

	err = Session.Open()
	if err != nil {
		fmt.Println("Error connecting to discord", err)
		return nil, err
	}

	Session.AddHandler(messageCreateHandler)
	return Session, nil
}

func createGuildWarsSession() (*gw2api.GW2Api, error) {
	guildWars2Token = viper.GetString("guildWars2Token")
	api, err := gw2api.NewAuthenticatedGW2Api(guildWars2Token)
	if err != nil {
		fmt.Println("Can't create gw2api")
		return nil, err
	}

	return api, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"time"
)

func apiFakeDataProvider() []byte {
	return []byte(`<?xml version="1.0" encoding="UTF-8"?>
    <bicing_stations>
     <updatetime><![CDATA[1415996588]]></updatetime>
      <station>
        <id>1</id>
          <type>BIKE</type>
            <lat>41.397952</lat>
              <long>2.180042</long>
                <street><![CDATA[Gran Via Corts Catalanes]]></street>
                  <height>21</height>
                    <streetNumber>760</streetNumber>
                      <nearbyStationList>24, 369, 387, 426</nearbyStationList>
                        <status>OPN</status>
                          <slots>0</slots>
                            <bikes>24</bikes>
                             </station>
                             </bicing_stations>`)
}

func doAPIRequest() []byte {
	response, err := http.Get("http://wservice.viabicing.cat/v1/getstations.php?v=1")
	if err != nil {
		fmt.Printf("Error with the request %s", err)
		os.Exit(1)
	}

	defer response.Body.Close()
	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		fmt.Printf("Error with the request %s", err)
		os.Exit(1)
	}

	return contents
}

func main() {
	ticker := time.NewTicker(2 * time.Second)
	quit := make(chan struct{})
	go func() {
		for {
			select {
			case <-ticker.C:
				data := obtainAPIData()
				persistCollection(data)
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()

	<-quit
}

func obtainAPIData() stationStateCollection {
	startTime := time.Now()
	apiData := doAPIRequest()
	requestEndTime := time.Now()

	var stationCollection stationStateCollection

	err := xml.Unmarshal(apiData, &stationCollection)
	if err != nil {
		fmt.Printf("Unmarshal error: %v, structure :%v", err, apiData)
		return stationCollection
	}

	fmt.Printf("Data successfully received, request time: %v, unmarshalling time: %v\n", requestEndTime.Sub(startTime), time.Since(requestEndTime))
	return stationCollection
}

type stationStateCollection struct {
	StationStates []stationState `xml:"station"`
	Updatetime    int            `xml:"updatetime"`
}

func (s stationStateCollection) Print() {
	for i := 0; i < len(s.StationStates); i++ {
		s.StationStates[i].Print()
	}
}

type stationState struct {
	// TODO review which of these fields need to be parsed and which not (we could potentially have different queries for the station state and the station data, as the second will change less frequently or may even not change at all)
	ID                int     `xml:"id"`
	Type              string  `xml:"type"`
	Latitude          float64 `xml:"lat"`
	Longitude         float64 `xml:"long"`
	Street            string  `xml:"street"`
	Height            int     `xml:"height"`
	StreetNumber      string  `xml:"streetNumber"` // Temporary, sometimes it is not set
	NearbyStationList string  `xml:"nearbyStationList"`
	Status            string  `xml:"status"`
	FreeSlots         int     `xml:"slots"`
	Bikes             int     `xml:"bikes"`
}

func (s stationState) Print() {
	fmt.Printf("Id : %v\n", s.ID)
	fmt.Printf("Type : %v\n", s.Type)
	fmt.Printf("Latitude : %v\n", s.Latitude)
	fmt.Printf("Longitude : %v\n", s.Longitude)
	fmt.Printf("Street : %v\n", s.Street)
	fmt.Printf("Height : %v\n", s.Height)
	fmt.Printf("StreetNumber : %v\n", s.StreetNumber)
	fmt.Printf("NearbyStationList : %v\n", s.NearbyStationList)
	fmt.Printf("Status : %v\n", s.Status)
	fmt.Printf("FreeSlots : %v\n", s.FreeSlots)
	fmt.Printf("Bikes : %v\n", s.Bikes)
}

func persistCollection(s stationStateCollection) {
	fmt.Println("Calling persistCollection", s.Updatetime)
}
<commit_msg>Successfully created station state storage<commit_after>package main

import (
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"time"
)

func apiFakeDataProvider() []byte {
	return []byte(`<?xml version="1.0" encoding="UTF-8"?>
    <bicing_stations>
     <updatetime><![CDATA[1415996588]]></updatetime>
      <station>
        <id>1</id>
          <type>BIKE</type>
            <lat>41.397952</lat>
              <long>2.180042</long>
                <street><![CDATA[Gran Via Corts Catalanes]]></street>
                  <height>21</height>
                    <streetNumber>760</streetNumber>
                      <nearbyStationList>24, 369, 387, 426</nearbyStationList>
                        <status>OPN</status>
                          <slots>0</slots>
                            <bikes>24</bikes>
                             </station>
                             </bicing_stations>`)
}

func doAPIRequest() []byte {
	response, err := http.Get("http://wservice.viabicing.cat/v1/getstations.php?v=1")
	if err != nil {
		fmt.Printf("Error with the request %s", err)
		os.Exit(1)
	}

	defer response.Body.Close()
	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		fmt.Printf("Error with the request %s", err)
		os.Exit(1)
	}

	return contents
}

func main() {
	ticker := time.NewTicker(2 * time.Second)
	quit := make(chan struct{})

	go func(s *StationStateStorage) {
		for {
			select {
			case <-ticker.C:
				data := obtainAPIData()
				s.PersistCollection(data)
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}(NewStorage())

	<-quit
}

func obtainAPIData() StationStateCollection {
	startTime := time.Now()
	apiData := doAPIRequest()
	requestEndTime := time.Now()

	var stationCollection StationStateCollection

	err := xml.Unmarshal(apiData, &stationCollection)
	if err != nil {
		fmt.Printf("Unmarshal error: %v, structure :%v", err, apiData)
		return stationCollection
	}

	fmt.Printf("Data successfully received, request time: %v, unmarshalling time: %v\n", requestEndTime.Sub(startTime), time.Since(requestEndTime))
	return stationCollection
}

type StationStateCollection struct {
	StationStates []stationState `xml:"station"`
	Updatetime    int            `xml:"updatetime"`
}

func (s StationStateCollection) Print() {
	for i := 0; i < len(s.StationStates); i++ {
		s.StationStates[i].Print()
	}
}

type stationState struct {
	// TODO review which of these fields need to be parsed and which not (we could potentially have different queries for the station state and the station data, as the second will change less frequently or may even not change at all)
	ID                int     `xml:"id"`
	Type              string  `xml:"type"`
	Latitude          float64 `xml:"lat"`
	Longitude         float64 `xml:"long"`
	Street            string  `xml:"street"`
	Height            int     `xml:"height"`
	StreetNumber      string  `xml:"streetNumber"` // Temporary, sometimes it is not set
	NearbyStationList string  `xml:"nearbyStationList"`
	Status            string  `xml:"status"`
	FreeSlots         int     `xml:"slots"`
	Bikes             int     `xml:"bikes"`
}

func (s stationState) Print() {
	fmt.Printf("Id : %v\n", s.ID)
	fmt.Printf("Type : %v\n", s.Type)
	fmt.Printf("Latitude : %v\n", s.Latitude)
	fmt.Printf("Longitude : %v\n", s.Longitude)
	fmt.Printf("Street : %v\n", s.Street)
	fmt.Printf("Height : %v\n", s.Height)
	fmt.Printf("StreetNumber : %v\n", s.StreetNumber)
	fmt.Printf("NearbyStationList : %v\n", s.NearbyStationList)
	fmt.Printf("Status : %v\n", s.Status)
	fmt.Printf("FreeSlots : %v\n", s.FreeSlots)
	fmt.Printf("Bikes : %v\n", s.Bikes)
}

func persistCollection(s StationStateCollection) {
	fmt.Println("Calling persistCollection", s.Updatetime)
}

func NewStorage() *StationStateStorage {
	return &StationStateStorage{}
}

type StationStateStorage struct{}

func (storage *StationStateStorage) PersistCollection(collection StationStateCollection) {

}
<|endoftext|>


<commit_before>package main

import (
    "log"
    "math/rand"
    "time"
)

func main() {
    log.Print("started.")

    // チャネル
    sleep1_finished := make(chan bool)

    go func() {
        // 0.2秒かかるコマンド
        log.Print("sleep1 started.")
        time.Sleep(200 * 1000 * 1000 * time.Nanosecond)
        log.Print("sleep1 finished.")
        sleep1_finished <- true
    }()

    // 終わるまで待つ
    <- sleep1_finished

    // 136枚の中から14枚返したい
    list := sub()
    for j := 0; j < 14; j++ {
        log.Println(list[j])
    }
}

// http://d.hatena.ne.jp/hake/20150930/p1
func shuffle(list []int){
	for i := len(list); i > 1; i-- {
		j := rand.Intn(i)          // 0　.. i-1 の乱数発生
		list[i - 1], list[j] = list[j], list[i - 1]
	}
}

func sub()([]int) {
	rand.Seed(time.Now().UnixNano())

	// データ要素数指定、および初期データ作成
	size := 136
	list := make([]int, size, size)
	for i := 0; i < size; i++ { list[i] = i }

	// シャッフル
	shuffle(list)

  return list
}
<commit_msg>Refactor<commit_after>package main

import (
    "log"
    "math/rand"
    "time"
)

func main() {
    log.Print("started.")

    // チャネル
    sleep1_finished := make(chan bool)

    go func() {
        // 0.2秒かかるコマンド
        log.Print("sleep1 started.")
        time.Sleep(200 * 1000 * 1000 * time.Nanosecond)
        log.Print("sleep1 finished.")
        sleep1_finished <- true
    }()

    // 終わるまで待つ
    <- sleep1_finished

    // 136枚の中から14枚返したい
    list := shuffled_cards()
    for j := 0; j < 14; j++ {
        log.Println(list[j])
    }
}

// http://d.hatena.ne.jp/hake/20150930/p1
func shuffle(list []int){
	for i := len(list); i > 1; i-- {
		j := rand.Intn(i)          // 0　.. i-1 の乱数発生
		list[i - 1], list[j] = list[j], list[i - 1]
	}
}

func shuffled_cards()([]int) {
	rand.Seed(time.Now().UnixNano())

	// データ要素数指定、および初期データ作成
	size := 136
	list := make([]int, size, size)
	for i := 0; i < size; i++ { list[i] = i }

	// シャッフル
	shuffle(list)

  return list
}
<|endoftext|>