<commit_before>package util

import (
	"encoding/base64"
	"fmt"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"time"

	"github.com/pborman/uuid"

	kerrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/wait"
	restclient "k8s.io/client-go/rest"
	kapi "k8s.io/kubernetes/pkg/apis/core"
	kclientset "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset"
	kcoreclient "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/typed/core/internalversion"
	"k8s.io/kubernetes/pkg/quota"

	configapi "github.com/openshift/origin/pkg/cmd/server/apis/config"
	cmdutil "github.com/openshift/origin/pkg/cmd/util"
	oauthapi "github.com/openshift/origin/pkg/oauth/apis/oauth"
	oauthclient "github.com/openshift/origin/pkg/oauth/generated/internalclientset"
	"github.com/openshift/origin/pkg/serviceaccounts"
	userapi "github.com/openshift/origin/pkg/user/apis/user"
	userclient "github.com/openshift/origin/pkg/user/generated/internalclientset"
)

// GetBaseDir returns the base directory used for test.
func GetBaseDir() string {
	return cmdutil.Env("BASETMPDIR", path.Join(os.TempDir(), "openshift-"+Namespace()))
}

func KubeConfigPath() string {
	return filepath.Join(GetBaseDir(), "openshift.local.config", "master", "admin.kubeconfig")
}

func GetClusterAdminKubeClient(adminKubeConfigFile string) (kclientset.Interface, error) {
	clientConfig, err := GetClusterAdminClientConfig(adminKubeConfigFile)
	if err != nil {
		return nil, err
	}

	return kclientset.NewForConfig(clientConfig)
}

func GetClusterAdminClientConfig(adminKubeConfigFile string) (*restclient.Config, error) {
	return configapi.GetClientConfig(adminKubeConfigFile, nil)
}

// GetClusterAdminClientConfigOrDie returns a REST config for the cluster admin
// user or panic.
func GetClusterAdminClientConfigOrDie(adminKubeConfigFile string) *restclient.Config {
	conf, err := GetClusterAdminClientConfig(adminKubeConfigFile)
	if err != nil {
		panic(err)
	}
	return conf
}

func GetClientForUser(clusterAdminConfig *restclient.Config, username string) (kclientset.Interface, *restclient.Config, error) {
	userClient, err := userclient.NewForConfig(clusterAdminConfig)
	if err != nil {
		return nil, nil, err
	}

	user, err := userClient.User().Users().Get(username, metav1.GetOptions{})
	if err != nil {
		user = &userapi.User{
			ObjectMeta: metav1.ObjectMeta{Name: username},
		}
		user, err = userClient.User().Users().Create(user)
		if err != nil {
			return nil, nil, err
		}
	}

	oauthClient, err := oauthclient.NewForConfig(clusterAdminConfig)
	if err != nil {
		return nil, nil, err
	}

	oauthClientObj := &oauthapi.OAuthClient{
		ObjectMeta: metav1.ObjectMeta{Name: "test-integration-client"},
	}
	if _, err := oauthClient.Oauth().OAuthClients().Create(oauthClientObj); err != nil && !kerrs.IsAlreadyExists(err) {
		return nil, nil, err
	}

	randomToken := uuid.NewRandom()
	accesstoken := base64.RawURLEncoding.EncodeToString([]byte(randomToken))
	for i := len(accesstoken); i < 32; i++ {
		accesstoken = accesstoken + "A"
	}
	token := &oauthapi.OAuthAccessToken{
		ObjectMeta: metav1.ObjectMeta{Name: accesstoken},
		ClientName: oauthClientObj.Name,
		UserName:   username,
		UserUID:    string(user.UID),
	}
	if _, err := oauthClient.Oauth().OAuthAccessTokens().Create(token); err != nil {
		return nil, nil, err
	}

	userClientConfig := restclient.AnonymousClientConfig(clusterAdminConfig)
	userClientConfig.BearerToken = token.Name

	kubeClientset, err := kclientset.NewForConfig(userClientConfig)
	if err != nil {
		return nil, nil, err
	}

	return kubeClientset, userClientConfig, nil
}

func GetScopedClientForUser(clusterAdminClientConfig *restclient.Config, username string, scopes []string) (kclientset.Interface, *restclient.Config, error) {
	// make sure the user exists
	if _, _, err := GetClientForUser(clusterAdminClientConfig, username); err != nil {
		return nil, nil, err
	}
	user, err := userclient.NewForConfigOrDie(clusterAdminClientConfig).User().Users().Get(username, metav1.GetOptions{})
	if err != nil {
		return nil, nil, err
	}

	token := &oauthapi.OAuthAccessToken{
		ObjectMeta:  metav1.ObjectMeta{Name: fmt.Sprintf("%s-token-plus-some-padding-here-to-make-the-limit-%d", username, rand.Int())},
		ClientName:  "openshift-challenging-client",
		ExpiresIn:   86400,
		Scopes:      scopes,
		RedirectURI: "https://127.0.0.1:12000/oauth/token/implicit",
		UserName:    user.Name,
		UserUID:     string(user.UID),
	}
	if _, err := oauthclient.NewForConfigOrDie(clusterAdminClientConfig).Oauth().OAuthAccessTokens().Create(token); err != nil {
		return nil, nil, err
	}

	scopedConfig := restclient.AnonymousClientConfig(clusterAdminClientConfig)
	scopedConfig.BearerToken = token.Name
	kubeClient, err := kclientset.NewForConfig(scopedConfig)
	if err != nil {
		return nil, nil, err
	}
	return kubeClient, scopedConfig, nil
}

func GetClientForServiceAccount(adminClient kclientset.Interface, clientConfig restclient.Config, namespace, name string) (*kclientset.Clientset, *restclient.Config, error) {
	_, err := adminClient.Core().Namespaces().Create(&kapi.Namespace{ObjectMeta: metav1.ObjectMeta{Name: namespace}})
	if err != nil && !kerrs.IsAlreadyExists(err) {
		return nil, nil, err
	}

	sa, err := adminClient.Core().ServiceAccounts(namespace).Create(&kapi.ServiceAccount{ObjectMeta: metav1.ObjectMeta{Name: name}})
	if kerrs.IsAlreadyExists(err) {
		sa, err = adminClient.Core().ServiceAccounts(namespace).Get(name, metav1.GetOptions{})
	}
	if err != nil {
		return nil, nil, err
	}

	token := ""
	err = wait.Poll(time.Second, 30*time.Second, func() (bool, error) {
		selector := fields.OneTermEqualSelector(kapi.SecretTypeField, string(kapi.SecretTypeServiceAccountToken))
		secrets, err := adminClient.Core().Secrets(namespace).List(metav1.ListOptions{FieldSelector: selector.String()})
		if err != nil {
			return false, err
		}
		for _, secret := range secrets.Items {
			if serviceaccounts.IsValidServiceAccountToken(sa, &secret) {
				token = string(secret.Data[kapi.ServiceAccountTokenKey])
				return true, nil
			}
		}
		return false, nil
	})
	if err != nil {
		return nil, nil, err
	}

	saClientConfig := restclient.AnonymousClientConfig(&clientConfig)
	saClientConfig.BearerToken = token

	kubeClientset, err := kclientset.NewForConfig(saClientConfig)
	if err != nil {
		return nil, nil, err
	}

	return kubeClientset, saClientConfig, nil
}

// WaitForResourceQuotaSync watches given resource quota until its hard limit is updated to match the desired
// spec or timeout occurs.
func WaitForResourceQuotaLimitSync(
	client kcoreclient.ResourceQuotaInterface,
	name string,
	hardLimit kapi.ResourceList,
	timeout time.Duration,
) error {

	startTime := time.Now()
	endTime := startTime.Add(timeout)

	expectedResourceNames := quota.ResourceNames(hardLimit)

	list, err := client.List(metav1.ListOptions{FieldSelector: fields.Set{"metadata.name": name}.AsSelector().String()})
	if err != nil {
		return err
	}

	for i := range list.Items {
		used := quota.Mask(list.Items[i].Status.Hard, expectedResourceNames)
		if isLimitSynced(used, hardLimit) {
			return nil
		}
	}

	rv := list.ResourceVersion
	w, err := client.Watch(metav1.ListOptions{FieldSelector: fields.Set{"metadata.name": name}.AsSelector().String(), ResourceVersion: rv})
	if err != nil {
		return err
	}
	defer w.Stop()

	for time.Now().Before(endTime) {
		select {
		case val, ok := <-w.ResultChan():
			if !ok {
				// reget and re-watch
				continue
			}
			if rq, ok := val.Object.(*kapi.ResourceQuota); ok {
				used := quota.Mask(rq.Status.Hard, expectedResourceNames)
				if isLimitSynced(used, hardLimit) {
					return nil
				}
			}
		case <-time.After(endTime.Sub(time.Now())):
			return wait.ErrWaitTimeout
		}
	}
	return wait.ErrWaitTimeout
}

func isLimitSynced(received, expected kapi.ResourceList) bool {
	resourceNames := quota.ResourceNames(expected)
	masked := quota.Mask(received, resourceNames)
	if len(masked) != len(expected) {
		return false
	}
	if le, _ := quota.LessThanOrEqual(masked, expected); !le {
		return false
	}
	if le, _ := quota.LessThanOrEqual(expected, masked); !le {
		return false
	}
	return true
}
<commit_msg>Turn off client rate limiting in integration tests<commit_after>package util

import (
	"encoding/base64"
	"fmt"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"time"

	"github.com/pborman/uuid"

	kerrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/wait"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/util/flowcontrol"
	kapi "k8s.io/kubernetes/pkg/apis/core"
	kclientset "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset"
	kcoreclient "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/typed/core/internalversion"
	"k8s.io/kubernetes/pkg/quota"

	configapi "github.com/openshift/origin/pkg/cmd/server/apis/config"
	cmdutil "github.com/openshift/origin/pkg/cmd/util"
	oauthapi "github.com/openshift/origin/pkg/oauth/apis/oauth"
	oauthclient "github.com/openshift/origin/pkg/oauth/generated/internalclientset"
	"github.com/openshift/origin/pkg/serviceaccounts"
	userapi "github.com/openshift/origin/pkg/user/apis/user"
	userclient "github.com/openshift/origin/pkg/user/generated/internalclientset"
)

// GetBaseDir returns the base directory used for test.
func GetBaseDir() string {
	return cmdutil.Env("BASETMPDIR", path.Join(os.TempDir(), "openshift-"+Namespace()))
}

func KubeConfigPath() string {
	return filepath.Join(GetBaseDir(), "openshift.local.config", "master", "admin.kubeconfig")
}

func GetClusterAdminKubeClient(adminKubeConfigFile string) (kclientset.Interface, error) {
	clientConfig, err := GetClusterAdminClientConfig(adminKubeConfigFile)
	if err != nil {
		return nil, err
	}

	return kclientset.NewForConfig(clientConfig)
}

func GetClusterAdminClientConfig(adminKubeConfigFile string) (*restclient.Config, error) {
	conf, err := configapi.GetClientConfig(adminKubeConfigFile, nil)
	if err != nil {
		return nil, err
	}
	return turnOffRateLimiting(conf), nil
}

// GetClusterAdminClientConfigOrDie returns a REST config for the cluster admin
// user or panic.
func GetClusterAdminClientConfigOrDie(adminKubeConfigFile string) *restclient.Config {
	conf, err := GetClusterAdminClientConfig(adminKubeConfigFile)
	if err != nil {
		panic(err)
	}
	return conf
}

func GetClientForUser(clusterAdminConfig *restclient.Config, username string) (kclientset.Interface, *restclient.Config, error) {
	userClient, err := userclient.NewForConfig(clusterAdminConfig)
	if err != nil {
		return nil, nil, err
	}

	user, err := userClient.User().Users().Get(username, metav1.GetOptions{})
	if err != nil && !kerrs.IsNotFound(err) {
		return nil, nil, err
	}
	if err != nil {
		user = &userapi.User{
			ObjectMeta: metav1.ObjectMeta{Name: username},
		}
		user, err = userClient.User().Users().Create(user)
		if err != nil {
			return nil, nil, err
		}
	}

	oauthClient, err := oauthclient.NewForConfig(clusterAdminConfig)
	if err != nil {
		return nil, nil, err
	}

	oauthClientObj := &oauthapi.OAuthClient{
		ObjectMeta: metav1.ObjectMeta{Name: "test-integration-client"},
	}
	if _, err := oauthClient.Oauth().OAuthClients().Create(oauthClientObj); err != nil && !kerrs.IsAlreadyExists(err) {
		return nil, nil, err
	}

	randomToken := uuid.NewRandom()
	accesstoken := base64.RawURLEncoding.EncodeToString([]byte(randomToken))
	// make sure the token is long enough to pass validation
	for i := len(accesstoken); i < 32; i++ {
		accesstoken += "A"
	}
	token := &oauthapi.OAuthAccessToken{
		ObjectMeta: metav1.ObjectMeta{Name: accesstoken},
		ClientName: oauthClientObj.Name,
		UserName:   username,
		UserUID:    string(user.UID),
	}
	if _, err := oauthClient.Oauth().OAuthAccessTokens().Create(token); err != nil {
		return nil, nil, err
	}

	userClientConfig := restclient.AnonymousClientConfig(turnOffRateLimiting(clusterAdminConfig))
	userClientConfig.BearerToken = token.Name

	kubeClientset, err := kclientset.NewForConfig(userClientConfig)
	if err != nil {
		return nil, nil, err
	}

	return kubeClientset, userClientConfig, nil
}

func GetScopedClientForUser(clusterAdminClientConfig *restclient.Config, username string, scopes []string) (kclientset.Interface, *restclient.Config, error) {
	// make sure the user exists
	if _, _, err := GetClientForUser(clusterAdminClientConfig, username); err != nil {
		return nil, nil, err
	}
	user, err := userclient.NewForConfigOrDie(clusterAdminClientConfig).User().Users().Get(username, metav1.GetOptions{})
	if err != nil {
		return nil, nil, err
	}

	token := &oauthapi.OAuthAccessToken{
		ObjectMeta:  metav1.ObjectMeta{Name: fmt.Sprintf("%s-token-plus-some-padding-here-to-make-the-limit-%d", username, rand.Int())},
		ClientName:  "openshift-challenging-client",
		ExpiresIn:   86400,
		Scopes:      scopes,
		RedirectURI: "https://127.0.0.1:12000/oauth/token/implicit",
		UserName:    user.Name,
		UserUID:     string(user.UID),
	}
	if _, err := oauthclient.NewForConfigOrDie(clusterAdminClientConfig).Oauth().OAuthAccessTokens().Create(token); err != nil {
		return nil, nil, err
	}

	scopedConfig := restclient.AnonymousClientConfig(turnOffRateLimiting(clusterAdminClientConfig))
	scopedConfig.BearerToken = token.Name
	kubeClient, err := kclientset.NewForConfig(scopedConfig)
	if err != nil {
		return nil, nil, err
	}
	return kubeClient, scopedConfig, nil
}

func GetClientForServiceAccount(adminClient kclientset.Interface, clientConfig restclient.Config, namespace, name string) (*kclientset.Clientset, *restclient.Config, error) {
	_, err := adminClient.Core().Namespaces().Create(&kapi.Namespace{ObjectMeta: metav1.ObjectMeta{Name: namespace}})
	if err != nil && !kerrs.IsAlreadyExists(err) {
		return nil, nil, err
	}

	sa, err := adminClient.Core().ServiceAccounts(namespace).Create(&kapi.ServiceAccount{ObjectMeta: metav1.ObjectMeta{Name: name}})
	if kerrs.IsAlreadyExists(err) {
		sa, err = adminClient.Core().ServiceAccounts(namespace).Get(name, metav1.GetOptions{})
	}
	if err != nil {
		return nil, nil, err
	}

	token := ""
	err = wait.Poll(time.Second, 30*time.Second, func() (bool, error) {
		selector := fields.OneTermEqualSelector(kapi.SecretTypeField, string(kapi.SecretTypeServiceAccountToken))
		secrets, err := adminClient.Core().Secrets(namespace).List(metav1.ListOptions{FieldSelector: selector.String()})
		if err != nil {
			return false, err
		}
		for _, secret := range secrets.Items {
			if serviceaccounts.IsValidServiceAccountToken(sa, &secret) {
				token = string(secret.Data[kapi.ServiceAccountTokenKey])
				return true, nil
			}
		}
		return false, nil
	})
	if err != nil {
		return nil, nil, err
	}

	saClientConfig := restclient.AnonymousClientConfig(turnOffRateLimiting(&clientConfig))
	saClientConfig.BearerToken = token

	kubeClientset, err := kclientset.NewForConfig(saClientConfig)
	if err != nil {
		return nil, nil, err
	}

	return kubeClientset, saClientConfig, nil
}

// WaitForResourceQuotaLimitSync watches given resource quota until its hard limit is updated to match the desired
// spec or timeout occurs.
func WaitForResourceQuotaLimitSync(
	client kcoreclient.ResourceQuotaInterface,
	name string,
	hardLimit kapi.ResourceList,
	timeout time.Duration,
) error {
	startTime := time.Now()
	endTime := startTime.Add(timeout)

	expectedResourceNames := quota.ResourceNames(hardLimit)

	list, err := client.List(metav1.ListOptions{FieldSelector: fields.Set{"metadata.name": name}.AsSelector().String()})
	if err != nil {
		return err
	}

	for i := range list.Items {
		used := quota.Mask(list.Items[i].Status.Hard, expectedResourceNames)
		if isLimitSynced(used, hardLimit) {
			return nil
		}
	}

	rv := list.ResourceVersion
	w, err := client.Watch(metav1.ListOptions{FieldSelector: fields.Set{"metadata.name": name}.AsSelector().String(), ResourceVersion: rv})
	if err != nil {
		return err
	}
	defer w.Stop()

	for time.Now().Before(endTime) {
		select {
		case val, ok := <-w.ResultChan():
			if !ok {
				// reget and re-watch
				continue
			}
			if rq, ok := val.Object.(*kapi.ResourceQuota); ok {
				used := quota.Mask(rq.Status.Hard, expectedResourceNames)
				if isLimitSynced(used, hardLimit) {
					return nil
				}
			}
		case <-time.After(endTime.Sub(time.Now())):
			return wait.ErrWaitTimeout
		}
	}
	return wait.ErrWaitTimeout
}

func isLimitSynced(received, expected kapi.ResourceList) bool {
	resourceNames := quota.ResourceNames(expected)
	masked := quota.Mask(received, resourceNames)
	if len(masked) != len(expected) {
		return false
	}
	if le, _ := quota.LessThanOrEqual(masked, expected); !le {
		return false
	}
	if le, _ := quota.LessThanOrEqual(expected, masked); !le {
		return false
	}
	return true
}

// turnOffRateLimiting reduces the chance that a flaky test can be written while using this package
func turnOffRateLimiting(config *restclient.Config) *restclient.Config {
	configCopy := *config
	configCopy.QPS = 10000
	configCopy.Burst = 10000
	configCopy.RateLimiter = flowcontrol.NewFakeAlwaysRateLimiter()
	// We do not set a timeout because that will cause watches to fail
	// Integration tests are already limited to 5 minutes
	// configCopy.Timeout = time.Minute
	return &configCopy
}
<|endoftext|>


<commit_before>package assert

import (
	"reflect"
	"testing"
)

var FailFunc func(format string, args ...interface{})

func getFailFunc(t *testing.T) func(format string, args ...interface{}) {
	if FailFunc == nil {
		return t.Errorf
	}
	return FailFunc
}

func True(t *testing.T, val bool) bool {
	if !val {
		getFailFunc(t)("Expected: [true] got: [%v]\n%s", val, SourceInfo(2))
		return true
	}
	return false
}

func False(t *testing.T, val bool) bool {
	if val {
		getFailFunc(t)("Expected: [false] got: [%v]\n%s", val, SourceInfo(2))
		return true
	}
	return true
}

func Equal(t *testing.T, val, expect interface{}) bool {
	if !reflect.DeepEqual(val, expect) {
		getFailFunc(t)("Expected: [%v] got: [%v]\n%s", expect, val, SourceInfo(2))
		return true
	}
	return false
}

func NotEqual(t *testing.T, val, expect interface{}) bool {
	if reflect.DeepEqual(val, expect) {
		getFailFunc(t)("Not Expecting: [%v] got: [%v]\n%s", expect, val, SourceInfo(2))
		return true
	}
	return false
}

func Nil(t *testing.T, val interface{}) bool {
	if val != nil && !reflect.ValueOf(val).IsNil() {
		getFailFunc(t)("Expecting: [nil] got: [%v]\n%s", val, SourceInfo(2))
		return true
	}
	return false
}

func NotNil(t *testing.T, val interface{}) bool {
	if val == nil || reflect.ValueOf(val).IsNil() {
		getFailFunc(t)("Expecting: [not nil] got: [%v]\n%s", val, SourceInfo(2))
		return true
	}
	return false
}

type TestFunc func(t *testing.T)

func MustPanic(t *testing.T, f TestFunc) {
	defer func() {
		if r := recover(); r != nil {
		}
	}()
	f(t)
	ci, _ := GetCallerInfo(1)
	getFailFunc(t)("Expecting Panic:\n%s:[%d]\n%s", ci.filename, ci.lineNum, ci.lineSrc)
}

func MustNotPanic(t *testing.T, f TestFunc) {
	defer func() {
		if r := recover(); r != nil {
			ci, _ := GetCallerInfo(1)
			getFailFunc(t)("Not Expecting Panic:\n%s:[%d]\n%s", ci.filename, ci.lineNum, ci.lineSrc)
		}
	}()
	f(t)
}
<commit_msg>changed wasy fail fun is set<commit_after>package assert

import (
	"reflect"
	"testing"
)

type FailFunc func(format string, args ...interface{})

var GetFailFunc func(t *testing.T) FailFunc

func init() {
	GetFailFunc = func(t *testing.T) FailFunc {
		return t.Errorf
	}
}

func True(t *testing.T, val bool) {
	if !val {
		GetFailFunc(t)("Expected: [true] got: [%v]\n%s", val, SourceInfo(2))
	}
}

func False(t *testing.T, val bool) {
	if val {
		GetFailFunc(t)("Expected: [false] got: [%v]\n%s", val, SourceInfo(2))
	}
}

func Equal(t *testing.T, val, expect interface{}) {
	if !reflect.DeepEqual(val, expect) {
		GetFailFunc(t)("Expected: [%v] got: [%v]\n%s", expect, val, SourceInfo(2))
	}
}

func NotEqual(t *testing.T, val, expect interface{}) {
	if reflect.DeepEqual(val, expect) {
		GetFailFunc(t)("Not Expecting: [%v] got: [%v]\n%s", expect, val, SourceInfo(2))
	}
}

func Nil(t *testing.T, val interface{}) {
	if val != nil && !reflect.ValueOf(val).IsNil() {
		GetFailFunc(t)("Expecting: [nil] got: [%v]\n%s", val, SourceInfo(2))
	}
}

func NotNil(t *testing.T, val interface{}) {
	if val == nil || reflect.ValueOf(val).IsNil() {
		GetFailFunc(t)("Expecting: [not nil] got: [%v]\n%s", val, SourceInfo(2))
	}
}

type TestFunc func(t *testing.T)

func MustPanic(t *testing.T, f TestFunc) {
	defer func() {
		if r := recover(); r != nil {
		}
	}()
	f(t)
	ci, _ := GetCallerInfo(1)
	GetFailFunc(t)("Expecting Panic:\n%s:[%d]\n%s", ci.filename, ci.lineNum, ci.lineSrc)
}

func MustNotPanic(t *testing.T, f TestFunc) {
	defer func() {
		if r := recover(); r != nil {
			ci, _ := GetCallerInfo(1)
			GetFailFunc(t)("Not Expecting Panic:\n%s:[%d]\n%s", ci.filename, ci.lineNum, ci.lineSrc)
		}
	}()
	f(t)
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Cisco and/or its affiliates.
//
// 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 cryptodata

import (
	"github.com/ligato/cn-infra/db/keyval"
	"github.com/golang/protobuf/proto"
)

// KvBytesPluginWrapper wraps keyval.KvBytesPlugin with additional support of reading encrypted data
type KvBytesPluginWrapper struct {
	// Wrapped KvBytesPlugin
	keyval.KvBytesPlugin
	// Function used for decrypting arbitrary data later
	decryptFunc DecryptFunc
	// ArbitraryDecrypter is used to decrypt data
	decrypter ArbitraryDecrypter
}

// BytesBrokerWrapper wraps keyval.BytesBroker with additional support of reading encrypted data
type BytesBrokerWrapper struct {
	// Wrapped BytesBroker
	keyval.BytesBroker
	// Function used for decrypting arbitrary data later
	decryptFunc DecryptFunc
	// ArbitraryDecrypter is used to decrypt data
	decrypter ArbitraryDecrypter
}

// BytesWatcherWrapper wraps keyval.BytesWatcher with additional support of reading encrypted data
type BytesWatcherWrapper struct {
	// Wrapped BytesWatcher
	keyval.BytesWatcher
	// Function used for decrypting arbitrary data later
	decryptFunc DecryptFunc
	// ArbitraryDecrypter is used to decrypt data
	decrypter ArbitraryDecrypter
}

// BytesWatchRespWrapper wraps keyval.BytesWatchResp with additional support of reading encrypted data
type BytesWatchRespWrapper struct {
	// Wrapped BytesWatchResp
	keyval.BytesWatchResp
	// Function used for decrypting arbitrary data later
	decryptFunc DecryptFunc
	// ArbitraryDecrypter is used to decrypt data
	decrypter ArbitraryDecrypter
}

// NewKvBytesPluginWrapper creates wrapper for provided CoreBrokerWatcher, adding support for decrypting encrypted
// data
func NewKvBytesPluginWrapper(cbw keyval.KvBytesPlugin, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *KvBytesPluginWrapper {
	return &KvBytesPluginWrapper{
		KvBytesPlugin: cbw,
		decryptFunc:   decryptFunc,
		decrypter:     decrypter,
	}
}

// NewBytesBrokerWrapper creates wrapper for provided BytesBroker, adding support for decrypting encrypted data
func NewBytesBrokerWrapper(pb keyval.BytesBroker, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *BytesBrokerWrapper {
	return &BytesBrokerWrapper{
		BytesBroker: pb,
		decryptFunc: decryptFunc,
		decrypter:   decrypter,
	}
}

// NewBytesWatcherWrapper creates wrapper for provided BytesWatcher, adding support for decrypting encrypted data
func NewBytesWatcherWrapper(pb keyval.BytesWatcher, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *BytesWatcherWrapper {
	return &BytesWatcherWrapper{
		BytesWatcher: pb,
		decryptFunc:  decryptFunc,
		decrypter:    decrypter,
	}
}

// NewBroker returns a BytesBroker instance with support for decrypting values that prepends given <keyPrefix> to all
// keys in its calls.
// To avoid using a prefix, pass keyval.Root constant as argument.
func (cbw *KvBytesPluginWrapper) NewBroker(prefix string) keyval.BytesBroker {
	return NewBytesBrokerWrapper(cbw.KvBytesPlugin.NewBroker(prefix), cbw.decrypter, cbw.decryptFunc)
}

// NewWatcher returns a BytesWatcher instance with support for decrypting values that prepends given <keyPrefix> to all
// keys during watch subscribe phase.
// The prefix is removed from the key retrieved by GetKey() in BytesWatchResp.
// To avoid using a prefix, pass keyval.Root constant as argument.
func (cbw *KvBytesPluginWrapper) NewWatcher(prefix string) keyval.BytesWatcher {
	return NewBytesWatcherWrapper(cbw.KvBytesPlugin.NewWatcher(prefix), cbw.decrypter, cbw.decryptFunc)
}

// GetValue retrieves and tries to decrypt one item under the provided key.
func (cbb *BytesBrokerWrapper) GetValue(key string) (data []byte, found bool, revision int64, err error) {
	data, found, revision, err = cbb.BytesBroker.GetValue(key)
	if err == nil {
		objData, err := cbb.decrypter.Decrypt(data, cbb.decryptFunc)
		if err != nil {
			return data, found, revision, err
		}

		data = objData.([]byte)
	}
	return
}

// Watch starts subscription for changes associated with the selected keys.
// Watch events will be delivered to callback (not channel) <respChan>.
// Channel <closeChan> can be used to close watching on respective key
func (b *BytesWatcherWrapper) Watch(respChan func(keyval.BytesWatchResp), closeChan chan string, keys ...string) error {
	return b.BytesWatcher.Watch(func(resp keyval.BytesWatchResp) {
		respChan(&BytesWatchRespWrapper{
			BytesWatchResp: resp,
			decrypter:      b.decrypter,
			decryptFunc:    b.decryptFunc,
		})
	}, closeChan, keys...)
}

// GetValue returns the value of the pair.
func (r *BytesWatchRespWrapper) GetValue() []byte {
	data, err := r.decrypter.Decrypt(r.BytesWatchResp.GetValue(), r.decryptFunc)
	if err != nil {
		return nil
	}
	return data.([]byte)
}

// GetPrevValue returns the previous value of the pair.
func (r *BytesWatchRespWrapper) GetPrevValue() []byte {
	data, err := r.decrypter.Decrypt(r.BytesWatchResp.GetPrevValue(), r.decryptFunc)
	if err != nil {
		return nil
	}
	return data.([]byte)
}

// KvProtoPluginWrapper wraps keyval.KvProtoPlugin with additional support of reading encrypted data
type KvProtoPluginWrapper struct {
	// Wrapped KvProtoPlugin
	keyval.KvProtoPlugin
	// Function used for decrypting arbitrary data later
	decryptFunc DecryptFunc
	// ArbitraryDecrypter is used to decrypt data
	decrypter ArbitraryDecrypter
}

// ProtoBrokerWrapper wraps keyval.ProtoBroker with additional support of reading encrypted data
type ProtoBrokerWrapper struct {
	keyval.ProtoBroker
	// Function used for decrypting arbitrary data later
	decryptFunc DecryptFunc
	// ArbitraryDecrypter is used to decrypt data
	decrypter ArbitraryDecrypter
}

// NewKvProtoPluginWrapper creates wrapper for provided KvProtoPlugin, adding support for decrypting encrypted data
func NewKvProtoPluginWrapper(kvp keyval.KvProtoPlugin, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *KvProtoPluginWrapper {
	return &KvProtoPluginWrapper{
		KvProtoPlugin: kvp,
		decryptFunc:   decryptFunc,
		decrypter:     decrypter,
	}
}

// NewProtoBrokerWrapper creates wrapper for provided ProtoBroker, adding support for decrypting encrypted data
func NewProtoBrokerWrapper(pb keyval.ProtoBroker, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *ProtoBrokerWrapper {
	return &ProtoBrokerWrapper{
		ProtoBroker: pb,
		decryptFunc: decryptFunc,
		decrypter:   decrypter,
	}
}

// NewBroker returns a BytesBroker instance with support for decrypting values that prepends given <keyPrefix> to all
// keys in its calls.
// To avoid using a prefix, pass keyval.Root constant as argument.
func (kvp *KvProtoPluginWrapper) NewBroker(prefix string) keyval.ProtoBroker {
	return NewProtoBrokerWrapper(kvp.KvProtoPlugin.NewBroker(prefix), kvp.decrypter, kvp.decryptFunc)
}

// GetValue retrieves one item under the provided <key>. If the item exists,
// it is unmarshaled into the <reqObj> and its fields are decrypted.
func (db *ProtoBrokerWrapper) GetValue(key string, reqObj proto.Message) (bool, int64, error) {
	found, revision, err := db.ProtoBroker.GetValue(key, reqObj)
	if !found || err != nil {
		return found, revision, err
	}

	_, err = db.decrypter.Decrypt(reqObj, db.decryptFunc)
	return found, revision, err
}
<commit_msg>Add support for BytesBroker.ListValues<commit_after>// Copyright (c) 2018 Cisco and/or its affiliates.
//
// 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 cryptodata

import (
	"github.com/ligato/cn-infra/db/keyval"
	"github.com/golang/protobuf/proto"
)

// decryptData holds values required for decrypting
type decryptData struct {
	// Function used for decrypting arbitrary data later
	decryptFunc DecryptFunc
	// ArbitraryDecrypter is used to decrypt data
	decrypter ArbitraryDecrypter
}

// KvBytesPluginWrapper wraps keyval.KvBytesPlugin with additional support of reading encrypted data
type KvBytesPluginWrapper struct {
	keyval.KvBytesPlugin
	decryptData
}

// BytesBrokerWrapper wraps keyval.BytesBroker with additional support of reading encrypted data
type BytesBrokerWrapper struct {
	keyval.BytesBroker
	decryptData
}

// BytesWatcherWrapper wraps keyval.BytesWatcher with additional support of reading encrypted data
type BytesWatcherWrapper struct {
	keyval.BytesWatcher
	decryptData
}

// BytesKeyValWrapper wraps keyval.BytesKeyVal with additional support of reading encrypted data
type BytesKeyValWrapper struct {
	keyval.BytesKeyVal
	decryptData
}

// BytesWatchRespWrapper wraps keyval.BytesWatchResp with additional support of reading encrypted data
type BytesWatchRespWrapper struct {
	keyval.BytesWatchResp
	BytesKeyValWrapper
}

// BytesKeyValIterator wraps keyval.BytesKeyValIterator with additional support of reading encrypted data
type BytesKeyValIteratorWrapper struct {
	keyval.BytesKeyValIterator
	decryptData
}

// NewKvBytesPluginWrapper creates wrapper for provided CoreBrokerWatcher, adding support for decrypting encrypted
// data
func NewKvBytesPluginWrapper(cbw keyval.KvBytesPlugin, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *KvBytesPluginWrapper {
	return &KvBytesPluginWrapper{
		KvBytesPlugin: cbw,
		decryptData: decryptData{
			decryptFunc: decryptFunc,
			decrypter:   decrypter,
		},
	}
}

// NewBytesBrokerWrapper creates wrapper for provided BytesBroker, adding support for decrypting encrypted data
func NewBytesBrokerWrapper(pb keyval.BytesBroker, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *BytesBrokerWrapper {
	return &BytesBrokerWrapper{
		BytesBroker: pb,
		decryptData: decryptData{
			decryptFunc: decryptFunc,
			decrypter:   decrypter,
		},
	}
}

// NewBytesWatcherWrapper creates wrapper for provided BytesWatcher, adding support for decrypting encrypted data
func NewBytesWatcherWrapper(pb keyval.BytesWatcher, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *BytesWatcherWrapper {
	return &BytesWatcherWrapper{
		BytesWatcher: pb,
		decryptData: decryptData{
			decryptFunc: decryptFunc,
			decrypter:   decrypter,
		},
	}
}

// NewBroker returns a BytesBroker instance with support for decrypting values that prepends given <keyPrefix> to all
// keys in its calls.
// To avoid using a prefix, pass keyval.Root constant as argument.
func (cbw *KvBytesPluginWrapper) NewBroker(prefix string) keyval.BytesBroker {
	return NewBytesBrokerWrapper(cbw.KvBytesPlugin.NewBroker(prefix), cbw.decrypter, cbw.decryptFunc)
}

// NewWatcher returns a BytesWatcher instance with support for decrypting values that prepends given <keyPrefix> to all
// keys during watch subscribe phase.
// The prefix is removed from the key retrieved by GetKey() in BytesWatchResp.
// To avoid using a prefix, pass keyval.Root constant as argument.
func (cbw *KvBytesPluginWrapper) NewWatcher(prefix string) keyval.BytesWatcher {
	return NewBytesWatcherWrapper(cbw.KvBytesPlugin.NewWatcher(prefix), cbw.decrypter, cbw.decryptFunc)
}

// GetValue retrieves and tries to decrypt one item under the provided key.
func (cbb *BytesBrokerWrapper) GetValue(key string) (data []byte, found bool, revision int64, err error) {
	data, found, revision, err = cbb.BytesBroker.GetValue(key)
	if err == nil {
		objData, err := cbb.decrypter.Decrypt(data, cbb.decryptFunc)
		if err != nil {
			return data, found, revision, err
		}

		data = objData.([]byte)
	}
	return
}

// ListValues returns an iterator that enables to traverse all items stored
// under the provided <key>.
func (cbb *BytesBrokerWrapper) ListValues(key string) (keyval.BytesKeyValIterator, error) {
	kv, err := cbb.BytesBroker.ListValues(key)
	if err != nil {
		return kv, err
	}
	return &BytesKeyValIteratorWrapper{
		BytesKeyValIterator: kv,
		decryptData:         cbb.decryptData,
	}, nil
}

// Watch starts subscription for changes associated with the selected keys.
// Watch events will be delivered to callback (not channel) <respChan>.
// Channel <closeChan> can be used to close watching on respective key
func (b *BytesWatcherWrapper) Watch(respChan func(keyval.BytesWatchResp), closeChan chan string, keys ...string) error {
	return b.BytesWatcher.Watch(func(resp keyval.BytesWatchResp) {
		respChan(&BytesWatchRespWrapper{
			BytesWatchResp: resp,
			BytesKeyValWrapper: BytesKeyValWrapper{
				BytesKeyVal: resp,
				decryptData: b.decryptData,
			},
		})
	}, closeChan, keys...)
}

// GetValue returns the value of the pair.
func (r *BytesWatchRespWrapper) GetValue() []byte {
	return r.BytesKeyValWrapper.GetValue()
}

// GetPrevValue returns the previous value of the pair.
func (r *BytesWatchRespWrapper) GetPrevValue() []byte {
	return r.BytesKeyValWrapper.GetPrevValue()
}

// GetValue returns the value of the pair.
func (r *BytesKeyValWrapper) GetValue() []byte {
	data, err := r.decrypter.Decrypt(r.BytesKeyVal.GetValue(), r.decryptFunc)
	if err != nil {
		return nil
	}
	return data.([]byte)
}

// GetPrevValue returns the previous value of the pair.
func (r *BytesKeyValWrapper) GetPrevValue() []byte {
	data, err := r.decrypter.Decrypt(r.BytesKeyVal.GetPrevValue(), r.decryptFunc)
	if err != nil {
		return nil
	}
	return data.([]byte)
}

// GetNext retrieves the following item from the context.
// When there are no more items to get, <stop> is returned as *true*
// and <kv> is simply *nil*.
func (r *BytesKeyValIteratorWrapper) GetNext() (kv keyval.BytesKeyVal, stop bool) {
	if stop {
		return kv, stop
	}
	return &BytesKeyValWrapper{
		BytesKeyVal: kv,
		decryptData: r.decryptData,
	}, stop
}

// KvProtoPluginWrapper wraps keyval.KvProtoPlugin with additional support of reading encrypted data
type KvProtoPluginWrapper struct {
	keyval.KvProtoPlugin
	decryptData
}

// ProtoBrokerWrapper wraps keyval.ProtoBroker with additional support of reading encrypted data
type ProtoBrokerWrapper struct {
	keyval.ProtoBroker
	decryptData
}

// NewKvProtoPluginWrapper creates wrapper for provided KvProtoPlugin, adding support for decrypting encrypted data
func NewKvProtoPluginWrapper(kvp keyval.KvProtoPlugin, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *KvProtoPluginWrapper {
	return &KvProtoPluginWrapper{
		KvProtoPlugin: kvp,
		decryptData: decryptData{
			decryptFunc: decryptFunc,
			decrypter:   decrypter,
		},
	}
}

// NewProtoBrokerWrapper creates wrapper for provided ProtoBroker, adding support for decrypting encrypted data
func NewProtoBrokerWrapper(pb keyval.ProtoBroker, decrypter ArbitraryDecrypter, decryptFunc DecryptFunc) *ProtoBrokerWrapper {
	return &ProtoBrokerWrapper{
		ProtoBroker: pb,
		decryptData: decryptData{
			decryptFunc: decryptFunc,
			decrypter:   decrypter,
		},
	}
}

// NewBroker returns a BytesBroker instance with support for decrypting values that prepends given <keyPrefix> to all
// keys in its calls.
// To avoid using a prefix, pass keyval.Root constant as argument.
func (kvp *KvProtoPluginWrapper) NewBroker(prefix string) keyval.ProtoBroker {
	return NewProtoBrokerWrapper(kvp.KvProtoPlugin.NewBroker(prefix), kvp.decrypter, kvp.decryptFunc)
}

// GetValue retrieves one item under the provided <key>. If the item exists,
// it is unmarshaled into the <reqObj> and its fields are decrypted.
func (db *ProtoBrokerWrapper) GetValue(key string, reqObj proto.Message) (bool, int64, error) {
	found, revision, err := db.ProtoBroker.GetValue(key, reqObj)
	if !found || err != nil {
		return found, revision, err
	}

	_, err = db.decrypter.Decrypt(reqObj, db.decryptFunc)
	return found, revision, err
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"strings"
	"time"

	"github.com/sam701/awstools/config"
	"github.com/sam701/awstools/cred"
	"github.com/sam701/awstools/sess"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/urfave/cli"
)

var scriptOutput io.Writer = os.Stdout

func actionAssumeRole(c *cli.Context) error {
	if len(c.Args()) == 2 {
		var account, role string
		account = c.Args().Get(0)
		role = c.Args().Get(1)
		exportScriptPath := c.String("export")
		if exportScriptPath != "" {
			f, err := os.OpenFile(exportScriptPath, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0600)
			if err != nil {
				log.Fatalln("ERROR:", err)
			}
			defer f.Close()
			scriptOutput = f
		}
		assumeRole(account, role)
	} else {
		cli.ShowCommandHelp(c, "assume")
	}
	return nil
}

func assumeRole(account, role string) {
	account = adjustAccountName(account)
	role = adjustRoleName(role)

	err := tryToAssumeRole(account, role)

	if err != nil {
		getMainAccountMfaSessionToken()
		err = tryToAssumeRole(account, role)
		if err != nil {
			log.Fatalln(err)
		}

		needToBeRotatedChan := make(chan bool)
		go isNeedRotateKey(needToBeRotatedChan)
		if <-needToBeRotatedChan {
			rotateMainAccountKey()
		}
	}
}

func isNeedRotateKey(needToBeRotated chan<- bool) {
	session := sess.New(config.Current.Profiles.MainAccountMfaSession)
	cl := iam.New(session)

	keyId := cred.GetMainAccountKeyId(config.Current.Profiles.MainAccount)
	out, err := cl.ListAccessKeys(&iam.ListAccessKeysInput{})
	if err != nil {
		log.Fatalln("ERROR", err)
	}

	var creationTime time.Time
	for _, md := range out.AccessKeyMetadata {
		if keyId == *md.AccessKeyId {
			creationTime = *md.CreateDate
			break
		}
	}
	if creationTime.IsZero() {
		log.Fatalln("Cannot get creation time for key", keyId)
	}

	needToBeRotated <- int(time.Now().Sub(creationTime).Minutes()) >= config.Current.KeyRotationIntervalMinutes
}

func adjustAccountName(account string) string {
	var candidate string
	for k, _ := range config.Current.Accounts {
		if strings.Contains(k, account) {
			if candidate == "" {
				candidate = k
			} else {
				log.Fatalf("Ambiguous account name. Possible matches: %s, %s\n", candidate, k)
			}
		}
	}
	if candidate == "" {
		log.Fatalln("No such account:", account)
	}
	return candidate
}

func adjustRoleName(role string) string {
	if role[0] == 'r' {
		return "ReadOnlyAccess"
	} else if role[0] == 'w' {
		return "PowerUserAccess"
	} else {
		return role
	}
}

var userName string

func getUserName() string {
	if userName == "" {
		client := iam.New(sess.New(config.Current.Profiles.MainAccount))
		data, err := client.GetUser(&iam.GetUserInput{})
		if err != nil {
			log.Fatalln("ERROR:", err)
		}
		userName = *data.User.UserName
	}
	return userName
}

func readMfaToken() string {
	scanner := bufio.NewScanner(os.Stdin)
	fmt.Print("MFA token: ")
	scanner.Scan()
	return scanner.Text()
}

func getMainAccountMfaSessionToken() {
	token := readMfaToken()
	session := sess.New(config.Current.Profiles.MainAccount)
	stsClient := sts.New(session)
	data, err := stsClient.GetSessionToken(&sts.GetSessionTokenInput{
		SerialNumber: aws.String(fmt.Sprintf("arn:aws:iam::%s:mfa/%s",
			accountId(config.Current.Profiles.MainAccount),
			getUserName())),
		TokenCode: aws.String(token),
	})
	if err != nil {
		log.Fatalln("ERROR:", err)
	}

	persistSharedCredentials(data.Credentials, config.Current.Profiles.MainAccountMfaSession)
}

func accountId(accountName string) string {
	accountId := config.Current.Accounts[accountName]
	if accountId == "" {
		log.Fatalln("Unknown account name:", accountName)
	}
	return accountId
}

func tryToAssumeRole(accountName, role string) error {
	session := sess.New(config.Current.Profiles.MainAccountMfaSession)
	accountId := accountId(accountName)

	stsClient := sts.New(session)
	assumeData, err := stsClient.AssumeRole(&sts.AssumeRoleInput{
		RoleArn:         aws.String(fmt.Sprintf("arn:aws:iam::%s:role/%s", accountId, role)),
		RoleSessionName: aws.String(getUserName()),
	})
	if err != nil {
		return err
	}

	profile := fmt.Sprintf("%s %s", accountName, role)
	persistSharedCredentials(assumeData.Credentials, profile)
	printExport(assumeData.Credentials)

	cred.SetProfileRegion(profile, config.Current.DefaultRegion)
	return nil
}

func printExport(cred *sts.Credentials) {
	shell := path.Base(os.Getenv("SHELL"))
	var pattern string
	if shell == "fish" {
		pattern = "set -xg %s \"%s\"\n"
	} else {
		pattern = "export %s=\"%s\"\n"
	}

	exp := func(key, value string) {
		fmt.Fprintf(scriptOutput, pattern, key, value)
	}

	exp("AWS_ACCESS_KEY_ID", *cred.AccessKeyId)
	exp("AWS_SECRET_ACCESS_KEY", *cred.SecretAccessKey)
	exp("AWS_SESSION_TOKEN", *cred.SessionToken)
}

func persistSharedCredentials(c *sts.Credentials, profile string) {
	cred.SaveCredentials(profile, *c.AccessKeyId, *c.SecretAccessKey, *c.SessionToken)
}
<commit_msg>Revert "Rotate access keys using MFA session after assumeRole has been successfully completed"<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"strings"
	"time"

	"github.com/sam701/awstools/config"
	"github.com/sam701/awstools/cred"
	"github.com/sam701/awstools/sess"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/urfave/cli"
)

var scriptOutput io.Writer = os.Stdout

func actionAssumeRole(c *cli.Context) error {
	if len(c.Args()) == 2 {
		var account, role string
		account = c.Args().Get(0)
		role = c.Args().Get(1)
		exportScriptPath := c.String("export")
		if exportScriptPath != "" {
			f, err := os.OpenFile(exportScriptPath, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0600)
			if err != nil {
				log.Fatalln("ERROR:", err)
			}
			defer f.Close()
			scriptOutput = f
		}
		assumeRole(account, role)
	} else {
		cli.ShowCommandHelp(c, "assume")
	}
	return nil
}

func assumeRole(account, role string) {
	account = adjustAccountName(account)
	role = adjustRoleName(role)

	err := tryToAssumeRole(account, role)
	if err != nil {
		needToBeRotatedChan := make(chan bool)
		go isNeedRotateKey(needToBeRotatedChan)

		getMainAccountMfaSessionToken()
		err = tryToAssumeRole(account, role)
		if err != nil {
			log.Fatalln(err)
		}

		if <-needToBeRotatedChan {
			rotateMainAccountKey()
		}
	}
}

func isNeedRotateKey(needToBeRotated chan<- bool) {
	session := sess.New(config.Current.Profiles.MainAccount)
	cl := iam.New(session)

	keyId := cred.GetMainAccountKeyId(config.Current.Profiles.MainAccount)
	out, err := cl.ListAccessKeys(&iam.ListAccessKeysInput{})
	if err != nil {
		log.Fatalln("ERROR", err)
	}

	var creationTime time.Time
	for _, md := range out.AccessKeyMetadata {
		if keyId == *md.AccessKeyId {
			creationTime = *md.CreateDate
			break
		}
	}
	if creationTime.IsZero() {
		log.Fatalln("Cannot get creation time for key", keyId)
	}

	needToBeRotated <- int(time.Now().Sub(creationTime).Minutes()) >= config.Current.KeyRotationIntervalMinutes
}

func adjustAccountName(account string) string {
	var candidate string
	for k, _ := range config.Current.Accounts {
		if strings.Contains(k, account) {
			if candidate == "" {
				candidate = k
			} else {
				log.Fatalf("Ambiguous account name. Possible matches: %s, %s\n", candidate, k)
			}
		}
	}
	if candidate == "" {
		log.Fatalln("No such account:", account)
	}
	return candidate
}

func adjustRoleName(role string) string {
	if role[0] == 'r' {
		return "ReadOnlyAccess"
	} else if role[0] == 'w' {
		return "PowerUserAccess"
	} else {
		return role
	}
}

var userName string

func getUserName() string {
	if userName == "" {
		client := iam.New(sess.New(config.Current.Profiles.MainAccount))
		data, err := client.GetUser(&iam.GetUserInput{})
		if err != nil {
			log.Fatalln("ERROR:", err)
		}
		userName = *data.User.UserName
	}
	return userName
}

func readMfaToken() string {
	scanner := bufio.NewScanner(os.Stdin)
	fmt.Print("MFA token: ")
	scanner.Scan()
	return scanner.Text()
}

func getMainAccountMfaSessionToken() {
	token := readMfaToken()
	session := sess.New(config.Current.Profiles.MainAccount)
	stsClient := sts.New(session)
	data, err := stsClient.GetSessionToken(&sts.GetSessionTokenInput{
		SerialNumber: aws.String(fmt.Sprintf("arn:aws:iam::%s:mfa/%s",
			accountId(config.Current.Profiles.MainAccount),
			getUserName())),
		TokenCode: aws.String(token),
	})
	if err != nil {
		log.Fatalln("ERROR:", err)
	}

	persistSharedCredentials(data.Credentials, config.Current.Profiles.MainAccountMfaSession)
}

func accountId(accountName string) string {
	accountId := config.Current.Accounts[accountName]
	if accountId == "" {
		log.Fatalln("Unknown account name:", accountName)
	}
	return accountId
}

func tryToAssumeRole(accountName, role string) error {
	session := sess.New(config.Current.Profiles.MainAccountMfaSession)
	accountId := accountId(accountName)

	stsClient := sts.New(session)
	assumeData, err := stsClient.AssumeRole(&sts.AssumeRoleInput{
		RoleArn:         aws.String(fmt.Sprintf("arn:aws:iam::%s:role/%s", accountId, role)),
		RoleSessionName: aws.String(getUserName()),
	})
	if err != nil {
		return err
	}

	profile := fmt.Sprintf("%s %s", accountName, role)
	persistSharedCredentials(assumeData.Credentials, profile)
	printExport(assumeData.Credentials)

	cred.SetProfileRegion(profile, config.Current.DefaultRegion)
	return nil
}

func printExport(cred *sts.Credentials) {
	shell := path.Base(os.Getenv("SHELL"))
	var pattern string
	if shell == "fish" {
		pattern = "set -xg %s \"%s\"\n"
	} else {
		pattern = "export %s=\"%s\"\n"
	}

	exp := func(key, value string) {
		fmt.Fprintf(scriptOutput, pattern, key, value)
	}

	exp("AWS_ACCESS_KEY_ID", *cred.AccessKeyId)
	exp("AWS_SECRET_ACCESS_KEY", *cred.SecretAccessKey)
	exp("AWS_SESSION_TOKEN", *cred.SessionToken)
}

func persistSharedCredentials(c *sts.Credentials, profile string) {
	cred.SaveCredentials(profile, *c.AccessKeyId, *c.SecretAccessKey, *c.SessionToken)
}
<|endoftext|>


<commit_before>// Copyright William Schwartz 2014. See the LICENSE file for more information.

package pigosat

// #include "picosat.h"
import "C"
import "unsafe"

// For testing litArrayToSlice in TestlitArrayToSlice since you can't import CGo
// into test files.
var cArray123 = [3]C.int{1, 2, 3}
var cArray1230 = [4]C.int{1, 2, 3, 0}
var cZero C.int = 0

// Assume adds a unit clause (i.e., a clause containing one literal) containing
// just the given literal for just the next call to Solve. You can add arbitrary
// many assumptions before the next Solve. After the next call to Solve, the
// assumptions are removed.
//
// More precisely, assumptions actually remain valid even after the next call to
// Solve() returns. Valid means they remain 'assumed' until a call to
// AddClauses(), Assume(), or another Solve(), following the first Solve().
// They need to stay valid for FailedAssumptions(), etc., to work correctly.
func (p *Pigosat) Assume(lit Literal) {
	defer p.ready(false)()
	// void picosat_assume (PicoSAT *, int lit);
	C.picosat_assume(p.p, C.int(lit))
}

// FailedAssumption returns true if the literal is a failed assumption, meaning
// the literal was added with Assume then used to derive unsatisfiability.
// The return value is only meaningful as long the current assumptions are
// valid. See Assume() for more details.
//
// From the Picosat documentation:
//   This is as accurate as generating core literals, but still of course is an
//   overapproximation of the set of assumptions really necessary. The technique
//   does not need clausal core generation nor tracing to be enabled and thus
//   can be much more effective.
func (p *Pigosat) FailedAssumption(lit Literal) bool {
	defer p.ready(true)()
	// picoast_failed_assumption SIGABRTs if the following conditional is true
	if p.res() != Unsatisfiable || lit == 0 {
		return false
	}
	// int picosat_failed_assumption (PicoSAT *, int lit);
	return C.picosat_failed_assumption(p.p, C.int(lit)) != 0
}

// litArrayToSlice converts a 0-terminated C array of ints (of length at most
// maxLen) to a Go slice of Literals.
func litArrayToSlice(litPtr *C.int, maxLen int) []Literal {
	if litPtr == nil {
		panic("NULL pointer")
	}
	lits := []Literal{}
	for *litPtr != 0 {
		if len(lits) > maxLen {
			panic("Array not zero-terminated")
		}
		lits = append(lits, Literal(*litPtr))
		litPtr = (*C.int)(unsafe.Pointer(
			uintptr(unsafe.Pointer(litPtr)) + uintptr(C.sizeof_int)))
	}
	return lits
}

// FailedAssumptions returns a list of failed assumptions in the last call to
// Solve() or an empty list if the last call to Solve() did not return
// Unsatisfiable. See Assume() and FailedAssumption() for more details.
func (p *Pigosat) FailedAssumptions() []Literal {
	defer p.ready(true)()
	if p.res() != Unsatisfiable {
		return []Literal{}
	}

	// const int * picosat_failed_assumptions (PicoSAT *);
	litPtr := C.picosat_failed_assumptions(p.p)
	return litArrayToSlice(litPtr, int(C.picosat_variables(p.p)))
}

// MaxSatisfiableAssumptions computes a maximal subset of satisfiable
// assumptions. See Assume() for more details. You need to set the assumptions
// and call Solve() before calling this function. The result is a list of
// assumptions that consistently can be asserted at the same time.
//
// It could be beneficial to set the default phase of assumptions
// to true (positive). This can speed up the computation.
func (p *Pigosat) MaxSatisfiableAssumptions() []Literal {
	defer p.ready(false)()
	if C.picosat_inconsistent(p.p) != 0 {
		return []Literal{}
	}
	//const int * picosat_maximal_satisfiable_subset_of_assumptions (PicoSAT *);
	litPtr := C.picosat_maximal_satisfiable_subset_of_assumptions(p.p)
	return litArrayToSlice(litPtr, int(C.picosat_variables(p.p)))
}

// This function assumes that you have set up all assumptions with
// 'Assume()'.  Then it calls 'Solve()' internally unless the
// formula is already inconsistent without assumptions, i.e.  it contains
// the empty clause.  After that it extracts a maximal satisfiable subset of
// assumptions.
//
// The result is a zero terminated maximal subset of consistent assumptions
// or an empty list if the formula contains the empty clause and thus no
// more maximal consistent subsets of assumptions can be extracted.  In the
// first case, before returning, a blocking clause is added, that rules out
// the result for the next call.
//
// NOTE: adding the blocking clause changes the CNF.
//
// So the following idiom
//
// var lits []Literal;
// p.Assume(a1)
// p.Assume(a2)
// p.Assume(a3)
// p.Assume(a4)
// for mss := p.NextMaxSatisfiableAssumptions(); len(mss) > 0; mss = p.NextMaxSatisfiableAssumptions() {
//     ProcessResults(mss)
// }
//
// can be used to iterate over all maximal consistent subsets of
// the set of assumptions {a1,a2,a3,a4}.
//
// It could be beneficial to set the default phase of assumptions
// to true (positive). This might speed up the computation.
func (p *Pigosat) NextMaxSatisfiableAssumptions() []Literal {
	defer p.ready(false)()
	if C.picosat_inconsistent(p.p) != 0 {
		return []Literal{}
	}
	// const int *
	// picosat_next_maximal_satisfiable_subset_of_assumptions (PicoSAT *);
	litPtr := C.picosat_next_maximal_satisfiable_subset_of_assumptions(p.p)
	return litArrayToSlice(litPtr, int(C.picosat_variables(p.p)))
}
<commit_msg>Simplify Assume documentation and target it toward Go and PiGoSAT.<commit_after>// Copyright William Schwartz 2014. See the LICENSE file for more information.

package pigosat

// #include "picosat.h"
import "C"
import "unsafe"

// For testing litArrayToSlice in TestlitArrayToSlice since you can't import CGo
// into test files.
var cArray123 = [3]C.int{1, 2, 3}
var cArray1230 = [4]C.int{1, 2, 3, 0}
var cZero C.int = 0

// Assume adds a temporary unit clause, i.e., a clause containing the one
// literal you pass as an argument. An assumption remains valid after the next
// call to Solve returns until a call to AddClauses, Assume, or a second Solve.
// You can add arbitrary many assumptions before the next call to Solve.
// Methods FailedAssumptions, FailedAssumptions, MaxSatisfiableAssumptions, and
// NextMaxSatisfiableAssumptions operate on the current, valid assumptions.
func (p *Pigosat) Assume(lit Literal) {
	defer p.ready(false)()
	// void picosat_assume (PicoSAT *, int lit);
	C.picosat_assume(p.p, C.int(lit))
}

// FailedAssumption returns whether a given literal is a valid, "failed"
// assumption (see Assume), meaning the last call to Solve used the assumption
// to derive unsatisfiability. FailedAssumption always returns false when the
// last call to Solve had status other than Unsatisfiable. This may be an
// over-approximation of the set of assumptions necessary to derive
// unsatisfiability.
func (p *Pigosat) FailedAssumption(lit Literal) bool {
	defer p.ready(true)()
	// picoast_failed_assumption SIGABRTs if the following conditional is true
	if p.res() != Unsatisfiable || lit == 0 {
		return false
	}
	// int picosat_failed_assumption (PicoSAT *, int lit);
	return C.picosat_failed_assumption(p.p, C.int(lit)) != 0
}

// litArrayToSlice converts a 0-terminated C array of ints (of length at most
// maxLen) to a Go slice of Literals.
func litArrayToSlice(litPtr *C.int, maxLen int) []Literal {
	if litPtr == nil {
		panic("NULL pointer")
	}
	lits := []Literal{}
	for *litPtr != 0 {
		if len(lits) > maxLen {
			panic("Array not zero-terminated")
		}
		lits = append(lits, Literal(*litPtr))
		litPtr = (*C.int)(unsafe.Pointer(
			uintptr(unsafe.Pointer(litPtr)) + uintptr(C.sizeof_int)))
	}
	return lits
}

// FailedAssumptions returns a list of failed assumptions, i.e., all the
// for which FailedAssumption returns true. See FailedAssumption.
func (p *Pigosat) FailedAssumptions() []Literal {
	defer p.ready(false)() // Overwrites what becomes litPtr below.
	if p.res() != Unsatisfiable {
		return []Literal{}
	}

	// const int * picosat_failed_assumptions (PicoSAT *);
	litPtr := C.picosat_failed_assumptions(p.p)
	return litArrayToSlice(litPtr, int(C.picosat_variables(p.p)))
}

// MaxSatisfiableAssumptions computes a maximal subset of satisfiable
// assumptions. See Assume's documentation. You need to set the assumptions
// and call Solve() before calling this method. The result is a list of
// assumptions that consistently can be asserted at the same time.
func (p *Pigosat) MaxSatisfiableAssumptions() []Literal {
	defer p.ready(false)()
	if C.picosat_inconsistent(p.p) != 0 {
		return []Literal{}
	}
	// const int * picosat_maximal_satisfiable_subset_of_assumptions (PicoSAT *);
	litPtr := C.picosat_maximal_satisfiable_subset_of_assumptions(p.p)
	return litArrayToSlice(litPtr, int(C.picosat_variables(p.p)))
}

// NextMaxSatisfiableAssumptions is like MaxSatisfiableAssumptions, but operates
// like an iterator to give a different maximal subset of satisfiable
// assumptions upon each call. To do this, it internally calls Solve followed by
// BlockSolution, which modifies the underlying formula (thus changing the
// result of AddedOriginalClauses), and then reassuming the solutions that were
// valid when you called NextMaxSatisfiableAssumptions. Use it as follows.
// First, set your formula and assumptions using AddClauses and Assume. Then
// iterate over the different maximal satisfiable subsets of assumptions with:
//    var mss []Literal
//    for mss := p.NextMaxSatisfiableAssumptions(); len(mss) > 0; mss = p.NextMaxSatisfiableAssumptions() {
//        // Do stuff with mss
//    }
func (p *Pigosat) NextMaxSatisfiableAssumptions() []Literal {
	defer p.ready(false)()
	if C.picosat_inconsistent(p.p) != 0 {
		return []Literal{}
	}
	// const int *
	// picosat_next_maximal_satisfiable_subset_of_assumptions (PicoSAT *);
	litPtr := C.picosat_next_maximal_satisfiable_subset_of_assumptions(p.p)
	return litArrayToSlice(litPtr, int(C.picosat_variables(p.p)))
}
<|endoftext|>


<commit_before>package kawasaki_test

import (
	"errors"

	"github.com/cloudfoundry-incubator/guardian/gardener"
	"github.com/cloudfoundry-incubator/guardian/kawasaki"
	"github.com/cloudfoundry-incubator/guardian/kawasaki/fakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("CompositeNetworker", func() {
	var (
		compositeNetworker kawasaki.Networker
		fakeNetworker      *fakes.FakeNetworker
		extraHooks         []kawasaki.NetworkHooker
	)

	BeforeEach(func() {
		fakeNetworker = new(fakes.FakeNetworker)
		extraHooks = []kawasaki.NetworkHooker{}
	})

	JustBeforeEach(func() {
		compositeNetworker = &kawasaki.CompositeNetworker{fakeNetworker, extraHooks}
	})

	It("delegates to the netwoker", func() {
		_, err := compositeNetworker.Hooks(nil, "handle", "spec", "external-spec")
		Expect(err).NotTo(HaveOccurred())
		Expect(fakeNetworker.HooksCallCount()).To(Equal(1))

		_, handle, spec, externalSpec := fakeNetworker.HooksArgsForCall(0)
		Expect(handle).To(Equal("handle"))
		Expect(spec).To(Equal("spec"))
		Expect(externalSpec).To(Equal("external-spec"))
	})

	It("returns the hooks from the networker", func() {
		fakeNetworker.HooksReturns([]gardener.Hooks{
			gardener.Hooks{},
		}, nil)

		hooks, err := compositeNetworker.Hooks(nil, "handle", "spec", "external-spec")
		Expect(err).NotTo(HaveOccurred())
		Expect(hooks).To(Equal([]gardener.Hooks{
			gardener.Hooks{},
		}))
	})

	Context("when the networker fails", func() {
		It("returns the error", func() {
			fakeNetworker.HooksReturns(nil, errors.New("kaput-hook"))

			_, err := compositeNetworker.Hooks(nil, "handle", "spec", "external-spec")
			Expect(err).To(MatchError("kaput-hook"))
		})
	})

	Context("when extra hooks are supplied", func() {
		var hooker *fakes.FakeNetworkHooker

		BeforeEach(func() {
			hooker = new(fakes.FakeNetworkHooker)
			extraHooks = append(extraHooks, hooker)

			fakeNetworker.HooksReturns(
				[]gardener.Hooks{
					gardener.Hooks{
						Prestart: gardener.Hook{
							Path: "/some/prestaaaarthooook",
							Args: []string{"spiderman"},
						},
					},
				}, nil)

			hooker.HooksReturns(
				gardener.Hooks{
					Prestart: gardener.Hook{
						Path: "/some/prestarthook",
						Args: []string{"yo"},
					},
				}, nil)

		})

		It("delegates the correct arguments to the extra hooks", func() {
			_, err := compositeNetworker.Hooks(nil, "some-handle", "some-spec", "some-extra-spec")
			Expect(err).NotTo(HaveOccurred())

			Expect(hooker.HooksCallCount()).To(Equal(1))
			_, handle, spec, extraSpec := hooker.HooksArgsForCall(0)
			Expect(handle).To(Equal("some-handle"))
			Expect(spec).To(Equal("some-spec"))
			Expect(extraSpec).To(Equal("some-extra-spec"))
		})

		It("returns the netwokrs hooks followed by any extra hooks", func() {
			hooks, err := compositeNetworker.Hooks(nil, "some-handle", "some-spec", "some-extra-spec")
			Expect(err).NotTo(HaveOccurred())

			Expect(hooks).To(ContainElement(gardener.Hooks{
				Prestart: gardener.Hook{
					Path: "/some/prestaaaarthooook",
					Args: []string{"spiderman"},
				},
			}))

			Expect(hooks).To(ContainElement(gardener.Hooks{
				Prestart: gardener.Hook{
					Path: "/some/prestarthook",
					Args: []string{"yo"},
				},
			}))
		})

		Context("when the hook returns an error", func() {
			It("returns the error", func() {
				hooker.HooksReturns(gardener.Hooks{}, errors.New("kaput-hook-2"))
				_, err := compositeNetworker.Hooks(nil, "some-handle", "some-spec", "some-extra-spec")

				Expect(err).To(MatchError("kaput-hook-2"))
			})

		})
	})
})
<commit_msg>Address go-vet errors<commit_after>package kawasaki_test

import (
	"errors"

	"github.com/cloudfoundry-incubator/guardian/gardener"
	"github.com/cloudfoundry-incubator/guardian/kawasaki"
	"github.com/cloudfoundry-incubator/guardian/kawasaki/fakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("CompositeNetworker", func() {
	var (
		compositeNetworker kawasaki.Networker
		fakeNetworker      *fakes.FakeNetworker
		extraHooks         []kawasaki.NetworkHooker
	)

	BeforeEach(func() {
		fakeNetworker = new(fakes.FakeNetworker)
		extraHooks = []kawasaki.NetworkHooker{}
	})

	JustBeforeEach(func() {
		compositeNetworker = &kawasaki.CompositeNetworker{
			Networker:  fakeNetworker,
			ExtraHooks: extraHooks,
		}
	})

	It("delegates to the netwoker", func() {
		_, err := compositeNetworker.Hooks(nil, "handle", "spec", "external-spec")
		Expect(err).NotTo(HaveOccurred())
		Expect(fakeNetworker.HooksCallCount()).To(Equal(1))

		_, handle, spec, externalSpec := fakeNetworker.HooksArgsForCall(0)
		Expect(handle).To(Equal("handle"))
		Expect(spec).To(Equal("spec"))
		Expect(externalSpec).To(Equal("external-spec"))
	})

	It("returns the hooks from the networker", func() {
		fakeNetworker.HooksReturns([]gardener.Hooks{
			gardener.Hooks{},
		}, nil)

		hooks, err := compositeNetworker.Hooks(nil, "handle", "spec", "external-spec")
		Expect(err).NotTo(HaveOccurred())
		Expect(hooks).To(Equal([]gardener.Hooks{
			gardener.Hooks{},
		}))
	})

	Context("when the networker fails", func() {
		It("returns the error", func() {
			fakeNetworker.HooksReturns(nil, errors.New("kaput-hook"))

			_, err := compositeNetworker.Hooks(nil, "handle", "spec", "external-spec")
			Expect(err).To(MatchError("kaput-hook"))
		})
	})

	Context("when extra hooks are supplied", func() {
		var hooker *fakes.FakeNetworkHooker

		BeforeEach(func() {
			hooker = new(fakes.FakeNetworkHooker)
			extraHooks = append(extraHooks, hooker)

			fakeNetworker.HooksReturns(
				[]gardener.Hooks{
					gardener.Hooks{
						Prestart: gardener.Hook{
							Path: "/some/prestaaaarthooook",
							Args: []string{"spiderman"},
						},
					},
				}, nil)

			hooker.HooksReturns(
				gardener.Hooks{
					Prestart: gardener.Hook{
						Path: "/some/prestarthook",
						Args: []string{"yo"},
					},
				}, nil)

		})

		It("delegates the correct arguments to the extra hooks", func() {
			_, err := compositeNetworker.Hooks(nil, "some-handle", "some-spec", "some-extra-spec")
			Expect(err).NotTo(HaveOccurred())

			Expect(hooker.HooksCallCount()).To(Equal(1))
			_, handle, spec, extraSpec := hooker.HooksArgsForCall(0)
			Expect(handle).To(Equal("some-handle"))
			Expect(spec).To(Equal("some-spec"))
			Expect(extraSpec).To(Equal("some-extra-spec"))
		})

		It("returns the netwokrs hooks followed by any extra hooks", func() {
			hooks, err := compositeNetworker.Hooks(nil, "some-handle", "some-spec", "some-extra-spec")
			Expect(err).NotTo(HaveOccurred())

			Expect(hooks).To(ContainElement(gardener.Hooks{
				Prestart: gardener.Hook{
					Path: "/some/prestaaaarthooook",
					Args: []string{"spiderman"},
				},
			}))

			Expect(hooks).To(ContainElement(gardener.Hooks{
				Prestart: gardener.Hook{
					Path: "/some/prestarthook",
					Args: []string{"yo"},
				},
			}))
		})

		Context("when the hook returns an error", func() {
			It("returns the error", func() {
				hooker.HooksReturns(gardener.Hooks{}, errors.New("kaput-hook-2"))
				_, err := compositeNetworker.Hooks(nil, "some-handle", "some-spec", "some-extra-spec")

				Expect(err).To(MatchError("kaput-hook-2"))
			})

		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"time"
	"os"
	"bufio"

	log "github.com/Sirupsen/logrus"
	"github.com/coreos/go-etcd/etcd"
)

// loadEtcdNode reads `r` containing JSON objects representing etcd nodes and
// loads them into server.
func loadEtcdNode(etcdClient *etcd.Client, r io.Reader) error {
	jsonReader := json.NewDecoder(r)
	for {
		node := etcd.Node{}
		if err := jsonReader.Decode(&node); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if node.Key == "" && node.Dir {
			continue // skip root
		}

		// compute a new TTL
		ttl := 0
		if node.Expiration != nil {
			ttl := node.Expiration.Sub(time.Now()).Seconds()
			if ttl < 0 {
				// expired, skip
				continue
			}
		}

		if node.Dir {
			_, err := etcdClient.SetDir(node.Key, uint64(ttl))
			if err != nil {
				return fmt.Errorf("%s: %s", node.Key, err)
			}
		} else {
			_, err := etcdClient.Set(node.Key, node.Value, uint64(ttl))
			if err != nil {
				return fmt.Errorf("%s: %s", node.Key, err)
			}
		}
	}
	return nil
}

func restoreBackup(etcdLocalURL string, backupPath string, backupFile string) error {

	etcdClient := etcd.NewClient([]string{fmt.Sprintf("http://%s:2379", etcdLocalURL)})
	file, err := os.Open(fmt.Sprintf("%s%s", backupPath, backupFile))

	if err != nil {
		return fmt.Errorf("Error opening file: %s", err)
	}

	gzipReader, err := gzip.NewReader(bufio.NewReader(file))
	if err != nil {
		return err
	}
	defer gzipReader.Close()

	if err := loadEtcdNode(etcdClient, gzipReader); err != nil {
		return err
	}
	log.Printf("restore: complete")

	file.Close()
	return nil
}
<commit_msg>now checks if etcd dir tree is already populated<commit_after>package main

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"time"
	"os"
	"bufio"
	"errors"

	log "github.com/Sirupsen/logrus"
	"github.com/coreos/go-etcd/etcd"

)

// loadEtcdNode reads `r` containing JSON objects representing etcd nodes and
// loads them into server.
func loadEtcdNode(etcdClient *etcd.Client, r io.Reader) error {
	jsonReader := json.NewDecoder(r)
	for {
		node := etcd.Node{}
		if err := jsonReader.Decode(&node); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if node.Key == "" && node.Dir {
			continue // skip root
		}

		// compute a new TTL
		ttl := 0
		if node.Expiration != nil {
			ttl := node.Expiration.Sub(time.Now()).Seconds()
			if ttl < 0 {
				// expired, skip
				continue
			}
		}

		if node.Dir {
			_, err := etcdClient.SetDir(node.Key, uint64(ttl))
			if err != nil {
				return fmt.Errorf("%s: %s", node.Key, err)
			}
		} else {
			_, err := etcdClient.Set(node.Key, node.Value, uint64(ttl))
			if err != nil {
				return fmt.Errorf("%s: %s", node.Key, err)
			}
		}
	}
	return nil
}

func checkNodeState(key string, etcdClient *etcd.Client) (int, error) {
	response, err := etcdClient.Get(key, false, false)
	if err != nil {
		return 0, err
	}

	childNodes := response.Node.Nodes
	count := 0

	// enumerate all the child nodes.
	for range childNodes {
		count += 1
	}

	return count, nil
}

func restoreBackup(etcdLocalURL string, backupPath string, backupFile string) error {

	etcdClient := etcd.NewClient([]string{fmt.Sprintf("http://%s:2379", etcdLocalURL)})

	var valueCount int
	var err error

	valueCount, err = checkNodeState("registry", etcdClient)

	if valueCount > 0 {
		return errors.New("etcd dir tree already populated")
	}

	file, err := os.Open(fmt.Sprintf("%s%s", backupPath, backupFile))

	if err != nil {
		return err
	}

	gzipReader, err := gzip.NewReader(bufio.NewReader(file))
	if err != nil {
		return err
	}
	defer gzipReader.Close()

	if err := loadEtcdNode(etcdClient, gzipReader); err != nil {
		return err
	}

	log.Printf("restore: complete")
	file.Close()

	return nil
}
<|endoftext|>


<commit_before><commit_msg>Cleanup<commit_after><|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"time"
)

const (
	ACK        = 0x79
	NAK        = 0x1F
	GET_CMD    = 0x00
	GETID_CMD  = 0x02
	WRITE_CMD  = 0x31
	ERASE_CMD  = 0x43
	EXTERA_CMD = 0x44
	RDUNP_CMD  = 0x92
)

var (
	checkSum uint8
	pending  []byte
	extended bool
)

func uploadSTM32(data []byte) {
	fmt.Print(" Synchronise: ")
	connectToTarget()
	fmt.Println(" OK")

	fmt.Printf(" Boot loader: %02x hex\n", getBootVersion())
	fmt.Printf("   Chip type: %04x hex\n", getChipType())

	fmt.Print("   Unprotect: ")
	sendCmd(RDUNP_CMD)
	wantAck()
	fmt.Println("OK")

	fmt.Print("      Resume: ")
	connectToTarget()
	fmt.Println(" OK")

	fmt.Print("  Mass erase: ")
	massErase()
	fmt.Println("OK")

	fmt.Print("   Uploading: ")
	writeFlash(data)
	fmt.Println(" OK")
}

func getReply() uint8 {
	if len(pending) == 0 {
		pending = readWithTimeout()
	}
	if len(pending) == 0 {
		return 0
	}
	b := pending[0]
	//fmt.Printf("<%2x", b)
	pending = pending[1:]
	return b
}

func connectToTarget() {
	for {
		conn.Flush()
		fmt.Print(".") // auto-baud greeting
		conn.Write([]byte{0x7F})
		r := getReply()
		if r == ACK || r == NAK {
			break
		}
		time.Sleep(time.Second)
	}
	// got a valid reply
}

func wantAck() {
	r := getReply()
	if r != ACK {
		fmt.Println()
		fmt.Println("failed:", r)
		os.Exit(1)
	}
	checkSum = 0
}

func sendByte(b uint8) {
	//fmt.Printf(">%2x", b)
	conn.Write([]byte{b})
	checkSum ^= b
}

func sendCmd(cmd uint8) {
	sendByte(cmd)
	sendByte(^cmd)
	wantAck()
}

func getBootVersion() uint8 {
	sendCmd(GET_CMD)
	n := getReply()
	rev := getReply()
	for i := 0; i < int(n); i++ {
		if getReply() == EXTERA_CMD {
			extended = true
		}
	}
	wantAck()
	return rev
}

func getChipType() uint16 {
	sendCmd(GETID_CMD)
	getReply() // should be 1
	chipType := uint16(getReply()) << 8
	chipType |= uint16(getReply())
	wantAck()
	return chipType
}

func massErase() {
	if extended {
		sendCmd(EXTERA_CMD)
		sendByte(0xFF)
		sendByte(0xFF)
		sendByte(0xFF)
	} else {
		sendCmd(ERASE_CMD)
		sendByte(0xFF)
		sendByte(0x00)
	}
	wantAck()
}

func writeFlash(data []byte) {
	for len(data)%256 != 0 {
		data = append(data, 0xFF)
	}
	for offset := 0; offset < len(data); offset += 256 {
		fmt.Print("+")
		sendCmd(WRITE_CMD)
		addr := 0x08000000 + offset
		sendByte(uint8(addr >> 24))
		sendByte(uint8(addr >> 16))
		sendByte(uint8(addr >> 8))
		sendByte(uint8(addr))
		sendByte(checkSum)
		wantAck()
		sendByte(256 - 1)
		for i := 0; i < 256; i++ {
			sendByte(data[offset+i])
		}
		sendByte(checkSum)
		wantAck()
	}
}
<commit_msg>insert delay before talking to the serial port<commit_after>package main

import (
	"fmt"
	"os"
	"time"
)

const (
	ACK        = 0x79
	NAK        = 0x1F
	GET_CMD    = 0x00
	GETID_CMD  = 0x02
	WRITE_CMD  = 0x31
	ERASE_CMD  = 0x43
	EXTERA_CMD = 0x44
	RDUNP_CMD  = 0x92
)

var (
	checkSum uint8
	pending  []byte
	extended bool
)

func uploadSTM32(data []byte) {
	fmt.Print(" Synchronise: ")
	connectToTarget()
	fmt.Println(" OK")

	fmt.Printf(" Boot loader: %02x hex\n", getBootVersion())
	fmt.Printf("   Chip type: %04x hex\n", getChipType())

	fmt.Print("   Unprotect: ")
	sendCmd(RDUNP_CMD)
	wantAck()
	fmt.Println("OK")

	fmt.Print("      Resume: ")
	connectToTarget()
	fmt.Println(" OK")

	fmt.Print("  Mass erase: ")
	massErase()
	fmt.Println("OK")

	fmt.Print("   Uploading: ")
	writeFlash(data)
	fmt.Println(" OK")
}

func getReply() uint8 {
	if len(pending) == 0 {
		pending = readWithTimeout()
	}
	if len(pending) == 0 {
		return 0
	}
	b := pending[0]
	//fmt.Printf("<%2x", b)
	pending = pending[1:]
	return b
}

func connectToTarget() {
	for {
		time.Sleep(time.Second)
		conn.Flush()
		fmt.Print(".") // auto-baud greeting
		conn.Write([]byte{0x7F})
		r := getReply()
		if r == ACK || r == NAK {
			break
		}
	}
	// got a valid reply
}

func wantAck() {
	r := getReply()
	if r != ACK {
		fmt.Println()
		fmt.Println("failed:", r)
		os.Exit(1)
	}
	checkSum = 0
}

func sendByte(b uint8) {
	//fmt.Printf(">%2x", b)
	conn.Write([]byte{b})
	checkSum ^= b
}

func sendCmd(cmd uint8) {
	sendByte(cmd)
	sendByte(^cmd)
	wantAck()
}

func getBootVersion() uint8 {
	sendCmd(GET_CMD)
	n := getReply()
	rev := getReply()
	for i := 0; i < int(n); i++ {
		if getReply() == EXTERA_CMD {
			extended = true
		}
	}
	wantAck()
	return rev
}

func getChipType() uint16 {
	sendCmd(GETID_CMD)
	getReply() // should be 1
	chipType := uint16(getReply()) << 8
	chipType |= uint16(getReply())
	wantAck()
	return chipType
}

func massErase() {
	if extended {
		sendCmd(EXTERA_CMD)
		sendByte(0xFF)
		sendByte(0xFF)
		sendByte(0xFF)
	} else {
		sendCmd(ERASE_CMD)
		sendByte(0xFF)
		sendByte(0x00)
	}
	wantAck()
}

func writeFlash(data []byte) {
	for len(data)%256 != 0 {
		data = append(data, 0xFF)
	}
	for offset := 0; offset < len(data); offset += 256 {
		fmt.Print("+")
		sendCmd(WRITE_CMD)
		addr := 0x08000000 + offset
		sendByte(uint8(addr >> 24))
		sendByte(uint8(addr >> 16))
		sendByte(uint8(addr >> 8))
		sendByte(uint8(addr))
		sendByte(checkSum)
		wantAck()
		sendByte(256 - 1)
		for i := 0; i < 256; i++ {
			sendByte(data[offset+i])
		}
		sendByte(checkSum)
		wantAck()
	}
}
<|endoftext|>


<commit_before>package system

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

type FileSystem struct{}

func (self *FileSystem) Walk(root string, step filepath.WalkFunc) {
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if self.isMetaDirectory(info) {
			return filepath.SkipDir
		}

		return step(path, info, err)
	})

	if err != nil {
		log.Println("Error while walking file system:", err)
		panic(err)
	}
}

func (self *FileSystem) isMetaDirectory(info os.FileInfo) bool {
	return info.IsDir() && strings.HasPrefix(info.Name(), ".")
}

func (self *FileSystem) Exists(directory string) bool {
	info, err := os.Stat(directory)
	return err == nil && info.IsDir()
}

func NewFileSystem() *FileSystem {
	return &FileSystem{}
}
<commit_msg>Expand ignored directories to match behaviour of the go tool<commit_after>package system

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

type FileSystem struct{}

func (self *FileSystem) Walk(root string, step filepath.WalkFunc) {
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if self.isMetaDirectory(info) {
			return filepath.SkipDir
		}

		return step(path, info, err)
	})

	if err != nil {
		log.Println("Error while walking file system:", err)
		panic(err)
	}
}

func (self *FileSystem) isMetaDirectory(info os.FileInfo) bool {
	name := info.Name()
	return info.IsDir() && (strings.HasPrefix(name, ".") || strings.HasPrefix(name, "_") || name == "testdata")
}

func (self *FileSystem) Exists(directory string) bool {
	info, err := os.Stat(directory)
	return err == nil && info.IsDir()
}

func NewFileSystem() *FileSystem {
	return &FileSystem{}
}
<|endoftext|>


<commit_before>package poloniex

import (
	"encoding/json"
	"strconv"
)

type Volume map[string]float64

type VolumeCollection struct {
	TotalBTC  float64 `json:"totalBTC,string"`
	TotalUSDT float64 `json:"totalUSDT,string"`
	TotalXMR  float64 `json:"totalXMR,string"`
	TotalXUSD float64 `json:"totalXUSD,string"`
	Volumes   map[string]Volume
}

func (tc *VolumeCollection) UnmarshalJSON(b []byte) error {
	m := make(map[string]json.RawMessage)
	if err := json.Unmarshal(b, &m); err != nil {
		return err
	}
	tc.Volumes = make(map[string]Volume)
	for k, v := range m {
		switch k {
		case "totalBTC":
			f, err := parseJSONFloatString(v)
			if err != nil {
				return err
			}
			tc.TotalBTC = f
		case "totalUSDT":
			f, err := parseJSONFloatString(v)
			if err != nil {
				return err
			}
			tc.TotalUSDT = f
		case "totalXMR":
			f, err := parseJSONFloatString(v)
			if err != nil {
				return err
			}
			tc.TotalXMR = f
		case "totalXUSD":
			f, err := parseJSONFloatString(v)
			if err != nil {
				return err
			}
			tc.TotalXUSD = f
		default:
			t := make(Volume)
			if err := json.Unmarshal(v, &t); err != nil {
				return err
			}
			tc.Volumes[k] = t
		}
	}
	return nil
}

func (t *Volume) UnmarshalJSON(b []byte) error {
	m := make(map[string]json.RawMessage)
	if err := json.Unmarshal(b, &m); err != nil {
		return err
	}
	for k, v := range m {
		f, err := parseJSONFloatString(v)
		if err != nil {
			return err
		}
		(*t)[k] = f
	}
	return nil
}

func parseJSONFloatString(b json.RawMessage) (float64, error) {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return 0, err
	}
	return strconv.ParseFloat(s, 64)
}
<commit_msg>fix(volumes): missing case for totalETH<commit_after>package poloniex

import (
	"encoding/json"
	"strconv"
)

type Volume map[string]float64

type VolumeCollection struct {
	TotalBTC  float64 `json:"totalBTC,string"`
	TotalUSDT float64 `json:"totalUSDT,string"`
	TotalXMR  float64 `json:"totalXMR,string"`
	TotalXUSD float64 `json:"totalXUSD,string"`
	Volumes   map[string]Volume
}

func (tc *VolumeCollection) UnmarshalJSON(b []byte) error {
	m := make(map[string]json.RawMessage)
	if err := json.Unmarshal(b, &m); err != nil {
		return err
	}
	tc.Volumes = make(map[string]Volume)
	for k, v := range m {
		switch k {
		case "totalBTC":
			f, err := parseJSONFloatString(v)
			if err != nil {
				return err
			}
			tc.TotalBTC = f
		case "totalETH":
			f, err := parseJSONFloatString(v)
			if err != nil {
				return err
			}
			tc.TotalUSDT = f
		case "totalUSDT":
			f, err := parseJSONFloatString(v)
			if err != nil {
				return err
			}
			tc.TotalUSDT = f
		case "totalXMR":
			f, err := parseJSONFloatString(v)
			if err != nil {
				return err
			}
			tc.TotalXMR = f
		case "totalXUSD":
			f, err := parseJSONFloatString(v)
			if err != nil {
				return err
			}
			tc.TotalXUSD = f
		default:
			t := make(Volume)
			if err := json.Unmarshal(v, &t); err != nil {
				return err
			}
			tc.Volumes[k] = t
		}
	}
	return nil
}

func (t *Volume) UnmarshalJSON(b []byte) error {
	m := make(map[string]json.RawMessage)
	if err := json.Unmarshal(b, &m); err != nil {
		return err
	}
	for k, v := range m {
		f, err := parseJSONFloatString(v)
		if err != nil {
			return err
		}
		(*t)[k] = f
	}
	return nil
}

func parseJSONFloatString(b json.RawMessage) (float64, error) {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return 0, err
	}
	return strconv.ParseFloat(s, 64)
}
<|endoftext|>


<commit_before>package setup

import (
	"database/sql"
	"fmt"
	"strings"
)

/*
CreateUser : Creates a new user in the database
*/
func CreateUser(connDetail ConnectionDetails, userName string, options []string) error {
	var db *sql.DB
	var err error

	if db, err = connect(connDetail); err != nil {
		return err
	}

	var exists bool

	if exists, err = checkIfUserExists(connDetail, userName); err != nil {
		return err
	} else if exists {
		// returns if the user already exists
		return nil
	}

	_, err = db.Exec("CREATE USER dbview " + strings.Join(options, " ") + ";")
	return err
}

func checkIfUserExists(connDetail ConnectionDetails, userName string) (bool, error) {

	var db *sql.DB
	var err error

	found := false

	if db, err = connect(connDetail); err != nil {
		return found, err
	}

	totalRows := 0
	err = db.QueryRow("SELECT count(1) FROM pg_roles WHERE rolname = $1", userName).Scan(&totalRows)

	return (totalRows > 0), err
}

/*
GrantRolesToUser : Grant some roles privilieges for a user
*/
func GrantRolesToUser(connDetail ConnectionDetails, userName string, roles []string) error {
	var db *sql.DB
	var err error

	if db, err = connect(connDetail); err != nil {
		return err
	}

	var exists bool

	if exists, err = checkIfUserExists(connDetail, userName); err != nil {
		return err
	} else if !exists {
		// returns if the user not exists
		return nil
	}

	for _, role := range roles {
		_, err = db.Exec(fmt.Sprintf("GRANT %s to %s;", role, userName))
		if err != nil {
			return err
		}
	}

	return nil
}

/*
SetSearchPathForUser : Set the 'search_path' variable for an user
*/
func SetSearchPathForUser(connDetail ConnectionDetails, userName string, schemas []string) error {

	var db *sql.DB
	var err error

	if db, err = connect(connDetail); err != nil {
		return err
	}

	var exists bool

	if exists, err = checkIfUserExists(connDetail, userName); err != nil {
		return err
	} else if !exists {
		// returns if the user not exists
		return nil
	}

	_, err = db.Exec(fmt.Sprintf("ALTER ROLE %s SET search_path TO %s;", userName, strings.Join(schemas, ",")))
	return err
}
<commit_msg>fix wrong username<commit_after>package setup

import (
	"database/sql"
	"fmt"
	"strings"
)

/*
CreateUser : Creates a new user in the database
*/
func CreateUser(connDetail ConnectionDetails, userName string, options []string) error {
	var db *sql.DB
	var err error

	if db, err = connect(connDetail); err != nil {
		return err
	}

	var exists bool

	if exists, err = checkIfUserExists(connDetail, userName); err != nil {
		return err
	} else if exists {
		// returns if the user already exists
		return nil
	}
	_, err = db.Exec(fmt.Sprintf("CREATE USER %s %s;", userName, strings.Join(options, " ")))
	return err
}

func checkIfUserExists(connDetail ConnectionDetails, userName string) (bool, error) {

	var db *sql.DB
	var err error

	found := false

	if db, err = connect(connDetail); err != nil {
		return found, err
	}

	totalRows := 0
	err = db.QueryRow("SELECT count(1) FROM pg_roles WHERE rolname = $1", userName).Scan(&totalRows)

	return (totalRows > 0), err
}

/*
GrantRolesToUser : Grant some roles privilieges for a user
*/
func GrantRolesToUser(connDetail ConnectionDetails, userName string, roles []string) error {
	var db *sql.DB
	var err error

	if db, err = connect(connDetail); err != nil {
		return err
	}

	var exists bool

	if exists, err = checkIfUserExists(connDetail, userName); err != nil {
		return err
	} else if !exists {
		// returns if the user not exists
		return nil
	}

	for _, role := range roles {
		_, err = db.Exec(fmt.Sprintf("GRANT %s to %s;", role, userName))
		if err != nil {
			return err
		}
	}

	return nil
}

/*
SetSearchPathForUser : Set the 'search_path' variable for an user
*/
func SetSearchPathForUser(connDetail ConnectionDetails, userName string, schemas []string) error {

	var db *sql.DB
	var err error

	if db, err = connect(connDetail); err != nil {
		return err
	}

	var exists bool

	if exists, err = checkIfUserExists(connDetail, userName); err != nil {
		return err
	} else if !exists {
		// returns if the user not exists
		return nil
	}

	_, err = db.Exec(fmt.Sprintf("ALTER ROLE %s SET search_path TO %s;", userName, strings.Join(schemas, ",")))
	return err
}
<|endoftext|>