<commit_before>package generate

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"math"
	"math/rand"
	"time"

	"github.com/netsec-ethz/rains/internal/pkg/algorithmTypes"
	"github.com/netsec-ethz/rains/internal/pkg/keys"
	"github.com/netsec-ethz/rains/internal/pkg/object"
	"github.com/netsec-ethz/rains/internal/pkg/publisher"
	"github.com/netsec-ethz/rains/internal/pkg/section"
	"github.com/netsec-ethz/rains/internal/pkg/zonefile"
	"golang.org/x/crypto/ed25519"
)

var privateKeys map[string]ed25519.PrivateKey

func init() {
	privateKeys = make(map[string]ed25519.PrivateKey)
}

func LeafZone(fileName, zoneName, context string, zoneSize int) {
	zone, err := Zone(zoneName, context, zoneSize, 0, ZoneAsNegProof, publisher.ShardingConfig{}, publisher.PShardingConfig{})
	if err != nil {
		panic(err.Error())
	}
	err = zonefile.Parser{}.EncodeAndStore(fileName, zone)
	if err != nil {
		panic(err.Error())
	}
}

func DelegationZone(fileName, zoneName, context string, zoneSize, maxShardSize int, namesPerPshard, probBound float64) {
	sconf, pconf := shardingConf(zoneSize, maxShardSize, namesPerPshard, probBound)
	zone, err := Zone(zoneName, context, 0, zoneSize, ShardAndPshard, sconf, pconf)
	if err != nil {
		panic(err.Error())
	}
	err = zonefile.Parser{}.EncodeAndStore(fileName, zone)
	if err != nil {
		panic(err.Error())
	}
}

func HybridZone(fileName, zoneName, context string, leafSize, delegSize, maxShardSize int, namesPerPshard,
	probBound float64) {
	sconf, pconf := shardingConf(leafSize+delegSize, maxShardSize, namesPerPshard, probBound)
	zone, err := Zone(zoneName, context, leafSize, delegSize, ShardAndPshard, sconf, pconf)
	if err != nil {
		panic(err.Error())
	}
	err = zonefile.Parser{}.EncodeAndStore(fileName, zone)
	if err != nil {
		panic(err.Error())
	}
}

func Zone(zoneName, context string, leafSize, delegSize int, negProofs NonExistProofs,
	shardingConf publisher.ShardingConfig, pshardingConf publisher.PShardingConfig) (
	*section.Zone, error) {
	zone := &section.Zone{
		Context:     context,
		SubjectZone: zoneName,
	}
	nextName := nameSeq("../../data/names.txt")
	assertions := make([]*section.Assertion, leafSize+delegSize)
	for i := 0; i < leafSize; i++ {
		names, objs := nextObject(Leaf, nextName(), zoneName)
		assertions[i] = &section.Assertion{
			SubjectName: names[0],
			Content:     objs,
		}
	}
	for i := leafSize; i < leafSize+delegSize; i++ {
		names, objs := nextObject(Delegation, nextName(), zoneName)
		for j, obj := range objs {
			assertions[i] = &section.Assertion{
				SubjectName: names[j],
				Content:     []object.Object{obj},
			}
			i++
		}
		i--
	}
	switch negProofs {
	case ShardAndPshard:
		shards, err := publisher.DoSharding(context, zoneName, assertions, nil, shardingConf, true)
		if err != nil {
			return nil, err
		}
		pshards, err := publisher.DoPsharding(context, zoneName, assertions, nil, pshardingConf, true)
		if err != nil {
			return nil, err
		}
		publisher.CreateZone(zone, nil, shards, pshards)
	case ZoneAsNegProof:
		publisher.CreateZone(zone, assertions, nil, nil)
	default:
		return nil, errors.New("unsupported nonExistProofs identifier")
	}
	return zone, nil
}

func nameSeq(path string) func() string {
	i := -1
	names := LoadNames(path)
	j := -1
	return func() string {
		if len(names)-1 == i {
			j++
			i = -1
		}
		i++
		if j == -1 {
			return names[i]
		}
		return fmt.Sprintf("%s-%s", names[j], names[i])
	}
}

//LoadNames returns a slice of unique names loaded from path
func LoadNames(path string) []string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(fmt.Sprintf("Was not able to read file : %s", err.Error()))
	}
	scanner := bufio.NewScanner(bytes.NewReader(data))
	scanner.Split(bufio.ScanWords)
	var names []string
	for scanner.Scan() {
		names = append(names, scanner.Text())
	}
	return names
}

func nextObject(objDistr ObjTypeDistr, name, zoneName string) ([]string, []object.Object) {
	switch objDistr {
	case Delegation:
		return delegationObject(name, zoneName)
	case Leaf:
		return []string{name}, leafObject()
	default:
		panic("unsupported objTypedistribution identifier")
	}
}

func leafObject() []object.Object {
	if rand.Intn(2) == 0 {
		return []object.Object{object.Object{
			Type:  object.OTIP4Addr,
			Value: fmt.Sprintf("%d.%d.%d.%d", rand.Intn(256), rand.Intn(256), rand.Intn(256), rand.Intn(256)),
		}}
	}
	return []object.Object{object.Object{
		Type: object.OTIP6Addr,
		Value: fmt.Sprintf("2001:db8::%d%d%d%d:%d%d%d%d", rand.Intn(10), rand.Intn(10), rand.Intn(10),
			rand.Intn(10), rand.Intn(10), rand.Intn(10), rand.Intn(10), rand.Intn(10)),
	}}
}

func delegationObject(name, zoneName string) ([]string, []object.Object) {
	pubKey, privKey, _ := ed25519.GenerateKey(nil)
	privateKeys[fmt.Sprintf("%s.%s", name, zoneName)] = privKey
	names := []string{name, name, "ns." + name, "ns1." + name}
	objs := make([]object.Object, 4)
	objs[0] = object.Object{
		Type: object.OTDelegation,
		Value: keys.PublicKey{
			PublicKeyID: keys.PublicKeyID{
				Algorithm: algorithmTypes.Ed25519,
				KeyPhase:  1,
				KeySpace:  keys.RainsKeySpace,
			},
			ValidSince: time.Now().Unix(),
			ValidUntil: time.Now().Add(365 * 24 * time.Hour).Unix(),
			Key:        pubKey,
		},
	}
	objs[1] = object.Object{
		Type:  object.OTRedirection,
		Value: "ns." + name,
	}
	objs[2] = object.Object{
		Type: object.OTServiceInfo,
		Value: object.ServiceInfo{
			Name:     "ns1." + name,
			Port:     5022,
			Priority: 0,
		},
	}
	objs[3] = object.Object{
		Type:  object.OTIP4Addr,
		Value: fmt.Sprintf("%d.%d.%d.%d", rand.Intn(256), rand.Intn(256), rand.Intn(256), rand.Intn(256)),
	}
	return names, objs
}

func shardingConf(zoneSize, maxShardSize int, namesPerPshard, probBound float64) (
	publisher.ShardingConfig, publisher.PShardingConfig) {
	//FIXME CFE namesPerPshard is used to calculate the filter parameters, but assertions per pshard
	//should be used to have the correct probability bound!!!
	filterSize := math.Ceil(namesPerPshard * math.Log(1/probBound) / math.Log(1/math.Pow(2, math.Log(2))))
	nofHashFunction := math.Round((filterSize / namesPerPshard) * math.Log(2))
	if zoneSize%4 != 0 {
		panic("size must be a multiple of 4")
	}
	sconf := publisher.ShardingConfig{
		DoSharding:   true,
		MaxShardSize: maxShardSize,
	}
	pconf := publisher.PShardingConfig{
		DoPsharding:            true,
		NofAssertionsPerPshard: int(namesPerPshard),
		BloomFilterConf: publisher.BloomFilterConfig{
			BFOpMode:         section.KirschMitzenmacher1,
			Hashfamily:       []algorithmTypes.Hash{algorithmTypes.Murmur364},
			NofHashFunctions: int(nofHashFunction),
			BloomFilterSize:  int(filterSize),
		},
	}
	return sconf, pconf
}

//ObjTypeDistr is an enumeration of object type distributions
type ObjTypeDistr int

const (
	//Delegation determines the object type distribution of a delegation zone
	Delegation ObjTypeDistr = iota
	//Leaf determines the object type distribution of a leaf zone
	Leaf
)

//NonExistProofs is an enumeration of a zone's possibilities to proof nonexistence
type NonExistProofs int

const (
	ZoneAsNegProof NonExistProofs = iota
	ShardAndPshard
)
<commit_msg>A simple generator for zonefiles<commit_after>package generate

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"math"
	"math/rand"
	"time"

	"github.com/netsec-ethz/rains/internal/pkg/algorithmTypes"
	"github.com/netsec-ethz/rains/internal/pkg/keys"
	"github.com/netsec-ethz/rains/internal/pkg/object"
	"github.com/netsec-ethz/rains/internal/pkg/publisher"
	"github.com/netsec-ethz/rains/internal/pkg/section"
	"github.com/netsec-ethz/rains/internal/pkg/zonefile"
	"golang.org/x/crypto/ed25519"
)

var privateKeys map[string]ed25519.PrivateKey

func init() {
	privateKeys = make(map[string]ed25519.PrivateKey)
}

func Zones() {
	parser := zonefile.Parser{}
	//create root zone
	path := "zonefiles/zones/"
	DelegationZone(path+"root.txt", ".", ".", 8, 500, 10, 10000000) //2 zones
	//create TLDs
	zone, err := parser.LoadZone(path + "root.txt")
	if err != nil {
		panic(err.Error())
	}
	names := createDelegatedZones(path, zone.SubjectZone, zone.Content)
	//create leaf zones
	for _, name := range names {
		LeafZone(path+name+".txt", name, ".", 2)
	}
	//TODO create TLD that delegates all of its commercial names to co.TLDName
}

func createDelegatedZones(path, subjectZone string, content []section.WithSigForward) []string {
	delegNames := []string{}
	for _, sec := range content {
		switch s := sec.(type) {
		case *section.Assertion:
			if s.Content[0].Type == object.OTDelegation {
				names := DelegationZone(path+s.SubjectName+".txt", s.SubjectName, ".", 20, 500, 10, 10000000) //5 zones
				delegNames = append(delegNames, names...)
			}
		case *section.Shard:
			for _, a := range s.Content {
				if a.Content[0].Type == object.OTDelegation {
					names := DelegationZone(path+a.SubjectName+".txt", a.SubjectName, ".", 20, 500, 10, 10000000) //5 zones
					delegNames = append(delegNames, names...)
				}
			}
		}
	}
	return delegNames
}

func LeafZone(fileName, zoneName, context string, zoneSize int) {
	zone, _, err := Zone(zoneName, context, zoneSize, 0, ZoneAsNegProof, publisher.ShardingConfig{}, publisher.PShardingConfig{})
	if err != nil {
		panic(err.Error())
	}
	err = zonefile.Parser{}.EncodeAndStore(fileName, zone)
	if err != nil {
		panic(err.Error())
	}
}

func DelegationZone(fileName, zoneName, context string, zoneSize, maxShardSize int, namesPerPshard,
	probBound float64) []string {
	sconf, pconf := shardingConf(zoneSize, maxShardSize, namesPerPshard, probBound)
	zone, names, err := Zone(zoneName, context, 0, zoneSize, ShardAndPshard, sconf, pconf)
	if err != nil {
		panic(err.Error())
	}
	err = zonefile.Parser{}.EncodeAndStore(fileName, zone)
	if err != nil {
		panic(err.Error())
	}
	return names
}

func HybridZone(fileName, zoneName, context string, leafSize, delegSize, maxShardSize int, namesPerPshard,
	probBound float64) []string {
	sconf, pconf := shardingConf(leafSize+delegSize, maxShardSize, namesPerPshard, probBound)
	zone, names, err := Zone(zoneName, context, leafSize, delegSize, ShardAndPshard, sconf, pconf)
	if err != nil {
		panic(err.Error())
	}
	err = zonefile.Parser{}.EncodeAndStore(fileName, zone)
	if err != nil {
		panic(err.Error())
	}
	return names
}

//Zone creates and stores a zonefile according to the given configuration. It returns the names of
//all delegation assertions.
func Zone(zoneName, context string, leafSize, delegSize int, negProofs NonExistProofs,
	shardingConf publisher.ShardingConfig, pshardingConf publisher.PShardingConfig) (
	*section.Zone, []string, error) {
	delegNames := []string{}
	zone := &section.Zone{
		Context:     context,
		SubjectZone: zoneName,
	}
	nextName := nameSeq("../../data/names.txt")
	assertions := make([]*section.Assertion, leafSize+delegSize)
	for i := 0; i < leafSize; i++ {
		names, objs := nextObject(Leaf, nextName(), zoneName)
		assertions[i] = &section.Assertion{
			SubjectName: names[0],
			Content:     objs,
		}
	}
	for i := leafSize; i < leafSize+delegSize; i++ {
		names, objs := nextObject(Delegation, nextName(), zoneName)
		delegNames = append(delegNames, names[0]+"."+zoneName)
		for j, obj := range objs {
			assertions[i] = &section.Assertion{
				SubjectName: names[j],
				Content:     []object.Object{obj},
			}
			i++
		}
		i--
	}
	switch negProofs {
	case ShardAndPshard:
		shards, err := publisher.DoSharding(context, zoneName, assertions, nil, shardingConf, true)
		if err != nil {
			return nil, nil, err
		}
		pshards, err := publisher.DoPsharding(context, zoneName, assertions, nil, pshardingConf, true)
		if err != nil {
			return nil, nil, err
		}
		publisher.CreateZone(zone, nil, shards, pshards)
	case ZoneAsNegProof:
		publisher.CreateZone(zone, assertions, nil, nil)
	default:
		return nil, nil, errors.New("unsupported nonExistProofs identifier")
	}
	return zone, delegNames, nil
}

func nameSeq(path string) func() string {
	i := -1
	names := LoadNames(path)
	j := -1
	return func() string {
		if len(names)-1 == i {
			j++
			i = -1
		}
		i++
		if j == -1 {
			return names[i]
		}
		return fmt.Sprintf("%s-%s", names[j], names[i])
	}
}

//LoadNames returns a slice of unique names loaded from path
func LoadNames(path string) []string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(fmt.Sprintf("Was not able to read file : %s", err.Error()))
	}
	scanner := bufio.NewScanner(bytes.NewReader(data))
	scanner.Split(bufio.ScanWords)
	var names []string
	for scanner.Scan() {
		names = append(names, scanner.Text())
	}
	return names
}

func nextObject(objDistr ObjTypeDistr, name, zoneName string) ([]string, []object.Object) {
	switch objDistr {
	case Delegation:
		return delegationObject(name, zoneName)
	case Leaf:
		return []string{name}, leafObject()
	default:
		panic("unsupported objTypedistribution identifier")
	}
}

func leafObject() []object.Object {
	if rand.Intn(2) == 0 {
		return []object.Object{object.Object{
			Type:  object.OTIP4Addr,
			Value: fmt.Sprintf("%d.%d.%d.%d", rand.Intn(256), rand.Intn(256), rand.Intn(256), rand.Intn(256)),
		}}
	}
	return []object.Object{object.Object{
		Type: object.OTIP6Addr,
		Value: fmt.Sprintf("2001:db8::%d%d%d%d:%d%d%d%d", rand.Intn(10), rand.Intn(10), rand.Intn(10),
			rand.Intn(10), rand.Intn(10), rand.Intn(10), rand.Intn(10), rand.Intn(10)),
	}}
}

func delegationObject(name, zoneName string) ([]string, []object.Object) {
	pubKey, privKey, _ := ed25519.GenerateKey(nil)
	privateKeys[fmt.Sprintf("%s.%s", name, zoneName)] = privKey
	names := []string{name, name, "ns." + name, "ns1." + name}
	objs := make([]object.Object, 4)
	objs[0] = object.Object{
		Type: object.OTDelegation,
		Value: keys.PublicKey{
			PublicKeyID: keys.PublicKeyID{
				Algorithm: algorithmTypes.Ed25519,
				KeyPhase:  1,
				KeySpace:  keys.RainsKeySpace,
			},
			ValidSince: time.Now().Unix(),
			ValidUntil: time.Now().Add(365 * 24 * time.Hour).Unix(),
			Key:        pubKey,
		},
	}
	objs[1] = object.Object{
		Type:  object.OTRedirection,
		Value: "ns." + name,
	}
	objs[2] = object.Object{
		Type: object.OTServiceInfo,
		Value: object.ServiceInfo{
			Name:     "ns1." + name,
			Port:     5022,
			Priority: 0,
		},
	}
	objs[3] = object.Object{
		Type:  object.OTIP4Addr,
		Value: fmt.Sprintf("%d.%d.%d.%d", rand.Intn(256), rand.Intn(256), rand.Intn(256), rand.Intn(256)),
	}
	return names, objs
}

func shardingConf(zoneSize, maxShardSize int, namesPerPshard, probBound float64) (
	publisher.ShardingConfig, publisher.PShardingConfig) {
	//FIXME CFE namesPerPshard is used to calculate the filter parameters, but assertions per pshard
	//should be used to have the correct probability bound!!!
	filterSize := math.Ceil(namesPerPshard * math.Log(1/probBound) / math.Log(1/math.Pow(2, math.Log(2))))
	nofHashFunction := math.Round((filterSize / namesPerPshard) * math.Log(2))
	if zoneSize%4 != 0 {
		panic("size must be a multiple of 4")
	}
	sconf := publisher.ShardingConfig{
		DoSharding:   true,
		MaxShardSize: maxShardSize,
	}
	pconf := publisher.PShardingConfig{
		DoPsharding:            true,
		NofAssertionsPerPshard: int(namesPerPshard),
		BloomFilterConf: publisher.BloomFilterConfig{
			BFOpMode:         section.KirschMitzenmacher1,
			Hashfamily:       []algorithmTypes.Hash{algorithmTypes.Murmur364},
			NofHashFunctions: int(nofHashFunction),
			BloomFilterSize:  int(filterSize),
		},
	}
	return sconf, pconf
}

//ObjTypeDistr is an enumeration of object type distributions
type ObjTypeDistr int

const (
	//Delegation determines the object type distribution of a delegation zone
	Delegation ObjTypeDistr = iota
	//Leaf determines the object type distribution of a leaf zone
	Leaf
)

//NonExistProofs is an enumeration of a zone's possibilities to proof nonexistence
type NonExistProofs int

const (
	ZoneAsNegProof NonExistProofs = iota
	ShardAndPshard
)
<|endoftext|>


<commit_before>// Copyright © 2022 Banzai Cloud
//
// 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 vault

import (
	"context"
	"fmt"
	"net/http"
	"strings"

	"emperror.dev/errors"
	"github.com/hashicorp/vault/api"
	"github.com/mitchellh/mapstructure"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cast"

	vaultpkg "github.com/banzaicloud/bank-vaults/pkg/sdk/vault"
)

func isOverwriteProhibitedError(err error) bool {
	return strings.Contains(err.Error(), "delete them before reconfiguring")
}

// secretEnginesWihtoutNameConfig holds the secret engine types where
// the name shouldn't be part of the config path
var secretEnginesWihtoutNameConfig = map[string]bool{
	"ad":       true,
	"alicloud": true,
	"azure":    true,
	"gcp":      true,
	"gcpkms":   true,
	"kv":       true,
}

type secretEngine struct {
	Path          string                 `mapstructure:"path"`
	Type          string                 `mapstructure:"type"`
	Description   string                 `mapstructure:"description"`
	Configuration map[string]interface{} `mapstructure:"configuration"`
	Config        map[string]interface{} `mapstructure:"config"`
	Options       map[string]string      `mapstructure:"options"`
	PluginName    string                 `mapstructure:"plugin_name"`
	Local         bool                   `mapstructure:"local"`
	SealWrap      bool                   `mapstructure:"seal_wrap"`
}

func initSecretsEnginesConfig(configs []secretEngine) []secretEngine {
	for index, config := range configs {
		if config.Path == "" {
			configs[index].Path = config.Type
		}

		configs[index].Path = strings.Trim(configs[index].Path, "/")
	}

	return configs
}

func (se *secretEngine) getMountConfigInput() (api.MountConfigInput, error) {
	var mountConfigInput api.MountConfigInput

	if err := mapstructure.Decode(se.Config, &mountConfigInput); err != nil {
		return mountConfigInput, errors.Wrap(err, "error parsing config for secret engine")
	}

	// Bank-Vaults supported options outside config to be used options in the mount request
	// so for now, to preserve backward compatibility we overwrite the options inside config
	// with the options outside.
	mountConfigInput.Options = se.Options

	return mountConfigInput, nil
}

func (v *vault) mountExists(path string) (bool, error) {
	mounts, err := v.cl.Sys().ListMounts()
	if err != nil {
		return false, errors.Wrap(err, "error reading mounts from vault")
	}
	logrus.Debugf("already existing mounts: %+v", mounts)

	return mounts[path+"/"] != nil, nil
}

func (v *vault) rotateSecretEngineCredentials(secretEngineType, path, name, configPath string) error {
	var rotatePath string
	switch secretEngineType {
	case "aws":
		rotatePath = fmt.Sprintf("%s/config/rotate-root", path)
	case "database":
		rotatePath = fmt.Sprintf("%s/rotate-root/%s", path, name)
	case "gcp":
		rotatePath = fmt.Sprintf("%s/%s/rotate", path, name)
	default:
		return errors.Errorf("secret engine type '%s' doesn't support credential rotation", secretEngineType)
	}

	if _, ok := v.rotateCache[rotatePath]; !ok {
		logrus.Infoln("doing credential rotation at", rotatePath)

		_, err := v.writeWithWarningCheck(rotatePath, nil)
		if err != nil {
			return errors.Wrapf(err, "error rotating credentials for '%s' config in vault", configPath)
		}

		logrus.Infoln("credential got rotated at", rotatePath)

		v.rotateCache[rotatePath] = true
	} else {
		logrus.Infoln("credentials were rotated previously for", rotatePath)
	}

	return nil
}

// NOTE: Maybe we could convert "getExisting*" and "getUnmanaged*" methods to generic functions
// since probably they will be the same for all config types.

// getExistingSecretsEngines gets all secrets engines that are already in Vault.
func (v *vault) getExistingSecretsEngines() (map[string]bool, error) {
	existingSecretsEngines := make(map[string]bool)

	existingSecretsEnginesList, err := v.cl.Sys().ListMounts()
	if err != nil {
		return nil, errors.Wrapf(err, "unable to list existing secrets engines")
	}

	for existingSecretEnginePath := range existingSecretsEnginesList {
		existingSecretsEngines[strings.Trim(existingSecretEnginePath, "/")] = true
	}

	return existingSecretsEngines, nil
}

// getUnmanagedSecretsEngines gets unmanaged secrets engines by comparing what's already in Vault
// and what's in the externalConfig.
func (v *vault) getUnmanagedSecretsEngines(managedSecretsEngines []secretEngine) map[string]bool {
	unmanagedSecretsEngines, _ := v.getExistingSecretsEngines()

	// Ignore system mounts.
	delete(unmanagedSecretsEngines, "sys")
	delete(unmanagedSecretsEngines, "identity")
	delete(unmanagedSecretsEngines, "cubbyhole")

	// Remove managed secret engine form the items since the reset will be removed.
	for _, managedSecretEngine := range managedSecretsEngines {
		delete(unmanagedSecretsEngines, managedSecretEngine.Path)
	}

	return unmanagedSecretsEngines
}

func configNeedsNoName(secretEngineType string, configOption string) bool {
	if configOption == "config" {
		_, ok := secretEnginesWihtoutNameConfig[secretEngineType]
		return ok
	}

	if secretEngineType == "aws" && configOption == "config/root" {
		return true
	}

	return false
}

func (v *vault) addManagedSecretsEngines(managedSecretsEngines []secretEngine) error {
	for _, secretEngine := range managedSecretsEngines {
		mountExists, err := v.mountExists(secretEngine.Path)
		if err != nil {
			return err
		}

		mountConfigInput, err := secretEngine.getMountConfigInput()
		if err != nil {
			return err
		}

		if !mountExists {
			// Mount the secret engine if it's not already there.
			mountInput := api.MountInput{
				Type:        secretEngine.Type,
				Description: secretEngine.Description,
				PluginName:  secretEngine.PluginName,
				Config:      mountConfigInput,
				Options:     mountConfigInput.Options, // options needs to be sent here first time
				Local:       secretEngine.Local,
				SealWrap:    secretEngine.SealWrap,
			}

			logrus.Infof("adding secret engine %s (%s)", secretEngine.Path, secretEngine.Type)
			logrus.Debugf("secret engine input %#v", mountInput)
			err = v.cl.Sys().Mount(secretEngine.Path, &mountInput)
			if err != nil {
				return errors.Wrapf(err, "error mounting %s into vault", secretEngine.Path)
			}
		} else {
			// If the secret engine is already mounted, only update its config in place.
			logrus.Infof("tuning already existing secret engine %s/", secretEngine.Path)
			err = v.cl.Sys().TuneMount(secretEngine.Path, mountConfigInput)
			if err != nil {
				return errors.Wrapf(err, "error tuning %s in vault", secretEngine.Path)
			}
		}

		// Configuration of the Secret Engine in a very generic manner, YAML config file should have the proper format
		for configOption, configData := range secretEngine.Configuration {
			configData, err := cast.ToSliceE(configData)
			if err != nil {
				return errors.Wrap(err, "error converting config data for secret engine")
			}
			for _, subConfigData := range configData {
				subConfigData, err := cast.ToStringMapE(subConfigData)
				if err != nil {
					return errors.Wrap(err, "error converting sub config data for secret engine")
				}

				name, ok := subConfigData["name"]
				if !ok && !configNeedsNoName(secretEngine.Type, configOption) {
					return errors.Errorf("error finding sub config data name for secret engine: %s/%s", secretEngine.Path, configOption)
				}

				// config data can have a child dict. But it will cause:
				// `json: unsupported type: map[interface {}]interface {}`
				// So check and replace by `map[string]interface{}` before using it.
				for k, v := range subConfigData {
					switch val := v.(type) {
					case map[interface{}]interface{}:
						subConfigData[k] = cast.ToStringMap(val)
					}
				}

				var configPath string
				if name != nil {
					configPath = fmt.Sprintf("%s/%s/%s", secretEngine.Path, configOption, name)
				} else {
					configPath = fmt.Sprintf("%s/%s", secretEngine.Path, configOption)
				}

				// Control if the configs should be updated or just Created once and skipped later on
				// This is a workaround to secrets backend like GCP that will destroy and recreate secrets at every iteration
				createOnly := cast.ToBool(subConfigData["create_only"])
				// Delete the create_only key from the map, so we don't push it to vault
				delete(subConfigData, "create_only")

				rotate := cast.ToBool(subConfigData["rotate"])
				// Delete the rotate key from the map, so we don't push it to vault
				delete(subConfigData, "rotate")

				saveTo := cast.ToString(subConfigData["save_to"])
				// Delete the rotate key from the map, so we don't push it to vault
				delete(subConfigData, "save_to")

				shouldUpdate := true
				if (createOnly || rotate) && mountExists {
					secretExists := false
					if configOption == "root/generate" { // the pki generate call is a different beast
						req := v.cl.NewRequest("GET", fmt.Sprintf("/v1/%s/ca", secretEngine.Path))
						resp, err := v.cl.RawRequestWithContext(context.Background(), req)
						if resp != nil {
							defer resp.Body.Close()
						}
						if err != nil {
							return errors.Wrapf(err, "failed to check pki CA")
						}
						if resp.StatusCode == http.StatusOK {
							secretExists = true
						}
					} else {
						secret, err := v.cl.Logical().Read(configPath)
						if err != nil {
							return errors.Wrapf(err, "error reading configPath %s", configPath)
						}
						if secret != nil && secret.Data != nil {
							secretExists = true
						}
					}

					if secretExists {
						reason := "rotate"
						if createOnly {
							reason = "create_only"
						}
						logrus.Infof("Secret at configpath %s already exists, %s was set so this will not be updated", configPath, reason)
						shouldUpdate = false
					}
				}

				if shouldUpdate {
					sec, err := v.writeWithWarningCheck(configPath, subConfigData)
					if err != nil {
						if isOverwriteProhibitedError(err) {
							logrus.Infoln("can't reconfigure", configPath, "please delete it manually")

							continue
						}
						return errors.Wrapf(err, "error configuring %s config in vault", configPath)
					}

					if saveTo != "" {
						_, err = v.writeWithWarningCheck(saveTo, vaultpkg.NewData(0, sec.Data))
						if err != nil {
							return errors.Wrapf(err, "error saving secret in vault to %s", saveTo)
						}
					}
				}

				// For secret engines where the root credentials are rotatable we don't wan't to reconfigure again
				// with the old credentials, because that would cause access denied issues. Currently these are:
				// - AWS
				// - Database
				if rotate && mountExists &&
					((secretEngine.Type == "database" && configOption == "config") ||
						(secretEngine.Type == "aws" && configOption == "config/root")) {
					// TODO we need to find out if it was rotated or not
					err = v.rotateSecretEngineCredentials(secretEngine.Type, secretEngine.Path, name.(string), configPath)
					if err != nil {
						return errors.Wrapf(err, "error rotating credentials for '%s' config in vault", configPath)
					}
				}
			}
		}
	}

	return nil
}

func (v *vault) removeUnmanagedSecretsEngines(unmanagedSecretsEngines map[string]bool) error {
	if len(unmanagedSecretsEngines) == 0 || !extConfig.PurgeUnmanagedConfig.Enabled ||
		extConfig.PurgeUnmanagedConfig.Exclude.Secrets {
		return nil
	}

	for secretEnginePath := range unmanagedSecretsEngines {
		logrus.Infof("removing secret engine path %s ", secretEnginePath)
		if err := v.cl.Sys().Unmount(secretEnginePath); err != nil {
			return errors.Wrapf(err, "error unmounting %s secret engine from vault", secretEnginePath)
		}
	}

	return nil
}

func (v *vault) configureSecretsEngines() error {
	managedSecretsEngines := initSecretsEnginesConfig(extConfig.Secrets)
	unmanagedSecretsEngines := v.getUnmanagedSecretsEngines(managedSecretsEngines)

	if err := v.addManagedSecretsEngines(managedSecretsEngines); err != nil {
		return errors.Wrap(err, "error adding secrets engines")
	}

	if err := v.removeUnmanagedSecretsEngines(unmanagedSecretsEngines); err != nil {
		return errors.Wrap(err, "error removing secrets engines")
	}

	return nil
}
<commit_msg>configurer: Don't crash trying to rotate aws config/root. (#1702)<commit_after>// Copyright © 2022 Banzai Cloud
//
// 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 vault

import (
	"context"
	"fmt"
	"net/http"
	"strings"

	"emperror.dev/errors"
	"github.com/hashicorp/vault/api"
	"github.com/mitchellh/mapstructure"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cast"

	vaultpkg "github.com/banzaicloud/bank-vaults/pkg/sdk/vault"
)

func isOverwriteProhibitedError(err error) bool {
	return strings.Contains(err.Error(), "delete them before reconfiguring")
}

// secretEnginesWihtoutNameConfig holds the secret engine types where
// the name shouldn't be part of the config path
var secretEnginesWihtoutNameConfig = map[string]bool{
	"ad":       true,
	"alicloud": true,
	"azure":    true,
	"gcp":      true,
	"gcpkms":   true,
	"kv":       true,
}

type secretEngine struct {
	Path          string                 `mapstructure:"path"`
	Type          string                 `mapstructure:"type"`
	Description   string                 `mapstructure:"description"`
	Configuration map[string]interface{} `mapstructure:"configuration"`
	Config        map[string]interface{} `mapstructure:"config"`
	Options       map[string]string      `mapstructure:"options"`
	PluginName    string                 `mapstructure:"plugin_name"`
	Local         bool                   `mapstructure:"local"`
	SealWrap      bool                   `mapstructure:"seal_wrap"`
}

func initSecretsEnginesConfig(configs []secretEngine) []secretEngine {
	for index, config := range configs {
		if config.Path == "" {
			configs[index].Path = config.Type
		}

		configs[index].Path = strings.Trim(configs[index].Path, "/")
	}

	return configs
}

func (se *secretEngine) getMountConfigInput() (api.MountConfigInput, error) {
	var mountConfigInput api.MountConfigInput

	if err := mapstructure.Decode(se.Config, &mountConfigInput); err != nil {
		return mountConfigInput, errors.Wrap(err, "error parsing config for secret engine")
	}

	// Bank-Vaults supported options outside config to be used options in the mount request
	// so for now, to preserve backward compatibility we overwrite the options inside config
	// with the options outside.
	mountConfigInput.Options = se.Options

	return mountConfigInput, nil
}

func (v *vault) mountExists(path string) (bool, error) {
	mounts, err := v.cl.Sys().ListMounts()
	if err != nil {
		return false, errors.Wrap(err, "error reading mounts from vault")
	}
	logrus.Debugf("already existing mounts: %+v", mounts)

	return mounts[path+"/"] != nil, nil
}

func (v *vault) rotateSecretEngineCredentials(secretEngineType, path, name, configPath string) error {
	var rotatePath string
	switch secretEngineType {
	case "aws":
		rotatePath = fmt.Sprintf("%s/config/rotate-root", path)
	case "database":
		rotatePath = fmt.Sprintf("%s/rotate-root/%s", path, name)
	case "gcp":
		rotatePath = fmt.Sprintf("%s/%s/rotate", path, name)
	default:
		return errors.Errorf("secret engine type '%s' doesn't support credential rotation", secretEngineType)
	}

	if _, ok := v.rotateCache[rotatePath]; !ok {
		logrus.Infoln("doing credential rotation at", rotatePath)

		_, err := v.writeWithWarningCheck(rotatePath, nil)
		if err != nil {
			return errors.Wrapf(err, "error rotating credentials for '%s' config in vault", configPath)
		}

		logrus.Infoln("credential got rotated at", rotatePath)

		v.rotateCache[rotatePath] = true
	} else {
		logrus.Infoln("credentials were rotated previously for", rotatePath)
	}

	return nil
}

// NOTE: Maybe we could convert "getExisting*" and "getUnmanaged*" methods to generic functions
// since probably they will be the same for all config types.

// getExistingSecretsEngines gets all secrets engines that are already in Vault.
func (v *vault) getExistingSecretsEngines() (map[string]bool, error) {
	existingSecretsEngines := make(map[string]bool)

	existingSecretsEnginesList, err := v.cl.Sys().ListMounts()
	if err != nil {
		return nil, errors.Wrapf(err, "unable to list existing secrets engines")
	}

	for existingSecretEnginePath := range existingSecretsEnginesList {
		existingSecretsEngines[strings.Trim(existingSecretEnginePath, "/")] = true
	}

	return existingSecretsEngines, nil
}

// getUnmanagedSecretsEngines gets unmanaged secrets engines by comparing what's already in Vault
// and what's in the externalConfig.
func (v *vault) getUnmanagedSecretsEngines(managedSecretsEngines []secretEngine) map[string]bool {
	unmanagedSecretsEngines, _ := v.getExistingSecretsEngines()

	// Ignore system mounts.
	delete(unmanagedSecretsEngines, "sys")
	delete(unmanagedSecretsEngines, "identity")
	delete(unmanagedSecretsEngines, "cubbyhole")

	// Remove managed secret engine form the items since the reset will be removed.
	for _, managedSecretEngine := range managedSecretsEngines {
		delete(unmanagedSecretsEngines, managedSecretEngine.Path)
	}

	return unmanagedSecretsEngines
}

func configNeedsNoName(secretEngineType string, configOption string) bool {
	if configOption == "config" {
		_, ok := secretEnginesWihtoutNameConfig[secretEngineType]
		return ok
	}

	if secretEngineType == "aws" && configOption == "config/root" {
		return true
	}

	return false
}

func (v *vault) addManagedSecretsEngines(managedSecretsEngines []secretEngine) error {
	for _, secretEngine := range managedSecretsEngines {
		mountExists, err := v.mountExists(secretEngine.Path)
		if err != nil {
			return err
		}

		mountConfigInput, err := secretEngine.getMountConfigInput()
		if err != nil {
			return err
		}

		if !mountExists {
			// Mount the secret engine if it's not already there.
			mountInput := api.MountInput{
				Type:        secretEngine.Type,
				Description: secretEngine.Description,
				PluginName:  secretEngine.PluginName,
				Config:      mountConfigInput,
				Options:     mountConfigInput.Options, // options needs to be sent here first time
				Local:       secretEngine.Local,
				SealWrap:    secretEngine.SealWrap,
			}

			logrus.Infof("adding secret engine %s (%s)", secretEngine.Path, secretEngine.Type)
			logrus.Debugf("secret engine input %#v", mountInput)
			err = v.cl.Sys().Mount(secretEngine.Path, &mountInput)
			if err != nil {
				return errors.Wrapf(err, "error mounting %s into vault", secretEngine.Path)
			}
		} else {
			// If the secret engine is already mounted, only update its config in place.
			logrus.Infof("tuning already existing secret engine %s/", secretEngine.Path)
			err = v.cl.Sys().TuneMount(secretEngine.Path, mountConfigInput)
			if err != nil {
				return errors.Wrapf(err, "error tuning %s in vault", secretEngine.Path)
			}
		}

		// Configuration of the Secret Engine in a very generic manner, YAML config file should have the proper format
		for configOption, configData := range secretEngine.Configuration {
			configData, err := cast.ToSliceE(configData)
			if err != nil {
				return errors.Wrap(err, "error converting config data for secret engine")
			}
			for _, subConfigData := range configData {
				subConfigData, err := cast.ToStringMapE(subConfigData)
				if err != nil {
					return errors.Wrap(err, "error converting sub config data for secret engine")
				}

				name, ok := subConfigData["name"]
				if !ok && !configNeedsNoName(secretEngine.Type, configOption) {
					return errors.Errorf("error finding sub config data name for secret engine: %s/%s", secretEngine.Path, configOption)
				}

				// config data can have a child dict. But it will cause:
				// `json: unsupported type: map[interface {}]interface {}`
				// So check and replace by `map[string]interface{}` before using it.
				for k, v := range subConfigData {
					switch val := v.(type) {
					case map[interface{}]interface{}:
						subConfigData[k] = cast.ToStringMap(val)
					}
				}

				var configPath string
				if name != nil {
					configPath = fmt.Sprintf("%s/%s/%s", secretEngine.Path, configOption, name)
				} else {
					configPath = fmt.Sprintf("%s/%s", secretEngine.Path, configOption)
				}

				// Control if the configs should be updated or just Created once and skipped later on
				// This is a workaround to secrets backend like GCP that will destroy and recreate secrets at every iteration
				createOnly := cast.ToBool(subConfigData["create_only"])
				// Delete the create_only key from the map, so we don't push it to vault
				delete(subConfigData, "create_only")

				rotate := cast.ToBool(subConfigData["rotate"])
				// Delete the rotate key from the map, so we don't push it to vault
				delete(subConfigData, "rotate")

				saveTo := cast.ToString(subConfigData["save_to"])
				// Delete the rotate key from the map, so we don't push it to vault
				delete(subConfigData, "save_to")

				shouldUpdate := true
				if (createOnly || rotate) && mountExists {
					secretExists := false
					if configOption == "root/generate" { // the pki generate call is a different beast
						req := v.cl.NewRequest("GET", fmt.Sprintf("/v1/%s/ca", secretEngine.Path))
						resp, err := v.cl.RawRequestWithContext(context.Background(), req)
						if resp != nil {
							defer resp.Body.Close()
						}
						if err != nil {
							return errors.Wrapf(err, "failed to check pki CA")
						}
						if resp.StatusCode == http.StatusOK {
							secretExists = true
						}
					} else {
						secret, err := v.cl.Logical().Read(configPath)
						if err != nil {
							return errors.Wrapf(err, "error reading configPath %s", configPath)
						}
						if secret != nil && secret.Data != nil {
							secretExists = true
						}
					}

					if secretExists {
						reason := "rotate"
						if createOnly {
							reason = "create_only"
						}
						logrus.Infof("Secret at configpath %s already exists, %s was set so this will not be updated", configPath, reason)
						shouldUpdate = false
					}
				}

				if shouldUpdate {
					sec, err := v.writeWithWarningCheck(configPath, subConfigData)
					if err != nil {
						if isOverwriteProhibitedError(err) {
							logrus.Infoln("can't reconfigure", configPath, "please delete it manually")

							continue
						}
						return errors.Wrapf(err, "error configuring %s config in vault", configPath)
					}

					if saveTo != "" {
						_, err = v.writeWithWarningCheck(saveTo, vaultpkg.NewData(0, sec.Data))
						if err != nil {
							return errors.Wrapf(err, "error saving secret in vault to %s", saveTo)
						}
					}
				}

				// For secret engines where the root credentials are rotatable we don't want to reconfigure again
				// with the old credentials, because that would cause access denied issues. Currently these are:
				// - AWS
				// - Database
				if rotate && mountExists &&
					((secretEngine.Type == "database" && configOption == "config") ||
						(secretEngine.Type == "aws" && configOption == "config/root")) {
					// TODO we need to find out if it was rotated or not
					nameStr := ""
					if name != nil {
						nameStr = name.(string)
					}
					err = v.rotateSecretEngineCredentials(secretEngine.Type, secretEngine.Path, nameStr, configPath)
					if err != nil {
						return errors.Wrapf(err, "error rotating credentials for '%s' config in vault", configPath)
					}
				}
			}
		}
	}

	return nil
}

func (v *vault) removeUnmanagedSecretsEngines(unmanagedSecretsEngines map[string]bool) error {
	if len(unmanagedSecretsEngines) == 0 || !extConfig.PurgeUnmanagedConfig.Enabled ||
		extConfig.PurgeUnmanagedConfig.Exclude.Secrets {
		return nil
	}

	for secretEnginePath := range unmanagedSecretsEngines {
		logrus.Infof("removing secret engine path %s ", secretEnginePath)
		if err := v.cl.Sys().Unmount(secretEnginePath); err != nil {
			return errors.Wrapf(err, "error unmounting %s secret engine from vault", secretEnginePath)
		}
	}

	return nil
}

func (v *vault) configureSecretsEngines() error {
	managedSecretsEngines := initSecretsEnginesConfig(extConfig.Secrets)
	unmanagedSecretsEngines := v.getUnmanagedSecretsEngines(managedSecretsEngines)

	if err := v.addManagedSecretsEngines(managedSecretsEngines); err != nil {
		return errors.Wrap(err, "error adding secrets engines")
	}

	if err := v.removeUnmanagedSecretsEngines(unmanagedSecretsEngines); err != nil {
		return errors.Wrap(err, "error removing secrets engines")
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"time"

	"github.com/codegangsta/cli"
	"github.com/kardianos/osext"
)

var InitialPath string
var InitialGoPath string

var Verbose bool

var SpinnerCharSet = 14
var SpinnerInterval = 50 * time.Millisecond

func main() {
	currentExecutable, _ := osext.Executable()
	vendoredBunchPath := path.Join(".vendor", "bin", "bunch")

	fi1, errStat1 := os.Stat(currentExecutable)
	fi2, errStat2 := os.Stat(vendoredBunchPath)

	if exists, _ := pathExists(vendoredBunchPath); errStat1 == nil && errStat2 == nil && exists && os.SameFile(fi1, fi2) {
		cmd := exec.Command(vendoredBunchPath, os.Args[1:]...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr

		err := cmd.Run()
		if err != nil {
			fmt.Println("running vendored bin/bunch was unsuccessful")
			os.Exit(1)
		} else {
			if cmd.ProcessState.Success() {
				os.Exit(0)
			} else {
				os.Exit(1)
			}
		}
	}

	InitialPath = os.Getenv("PATH")
	InitialGoPath = os.Getenv("GOPATH")

	app := cli.NewApp()
	app.Name = "bunch"
	app.Usage = "npm-like tool for managing Go dependencies"
	app.Version = "0.0.1"
	app.Authors = []cli.Author{cli.Author{Name: "Daniil Kulchenko", Email: "daniil@kulchenko.com"}}

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose",
			Usage: "output more information",
		},
	}

	app.Before = func(context *cli.Context) error {
		Verbose = context.GlobalBool("verbose")

		return nil
	}

	app.Commands = []cli.Command{
		{
			Name:    "install",
			Aliases: []string{"i"},
			Usage:   "install package(s)",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "save",
					Usage: "save installed package to Bunchfile",
				},
				cli.BoolFlag{
					Name:  "g",
					Usage: "install package to global $GOPATH instead of vendored directory",
				},
			},
			Action: func(c *cli.Context) {
				installCommand(c, false, true)
			},
		},
		{
			Name:    "update",
			Aliases: []string{"u"},
			Usage:   "update package(s)",
			Action: func(c *cli.Context) {
				installCommand(c, true, true)
			},
		},
		{
			Name:    "uninstall",
			Aliases: []string{"r"},
			Usage:   "uninstall package(s)",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "save",
					Usage: "save uninstalled package to Bunchfile",
				},
				cli.BoolFlag{
					Name:  "g",
					Usage: "uninstall package from global $GOPATH instead of vendored directory",
				},
			},
			Action: func(c *cli.Context) {
				uninstallCommand(c)
			},
		},
		{
			Name:  "prune",
			Usage: "remove packages not referenced in Bunchfile",
			Action: func(c *cli.Context) {
				pruneCommand(c)
			},
		},
		{
			Name:  "outdated",
			Usage: "list outdated packages",
			Action: func(c *cli.Context) {
				outdatedCommand(c)
			},
		},
		{
			Name:  "lock",
			Usage: "generate a file locking down current versions of dependencies",
			Action: func(c *cli.Context) {
				lockCommand(c)
			},
		},
		{
			Name:  "rebuild",
			Usage: "rebuild all dependencies",
			Action: func(c *cli.Context) {
				installCommand(c, true, false)
			},
		},
		{
			Name:  "generate",
			Usage: "generate a Bunchfile based on package imports in current directory",
			Action: func(c *cli.Context) {
				generateCommand(c)
			},
		},
		{
			Name:            "go",
			Usage:           "run a Go command within the vendor environment (e.g. bunch go fmt)",
			SkipFlagParsing: true,
			Action: func(c *cli.Context) {
				goCommand(c)
			},
		},
		{
			Name:            "exec",
			Usage:           "run any command within the vendor environment (e.g. bunch exec make)",
			SkipFlagParsing: true,
			Action: func(c *cli.Context) {
				execCommand(c)
			},
		},
		{
			Name:            "shell",
			Usage:           "start a shell within the vendor environment",
			SkipFlagParsing: true,
			Action: func(c *cli.Context) {
				shellCommand(c)
			},
		},
		{
			Name:  "shim",
			Usage: "sourced in .bash_profile to alias the 'go' tool",
			Action: func(c *cli.Context) {
				shimCommand(c)
			},
		},
	}

	app.Run(os.Args)
}
<commit_msg>only spawn if they're not the same file<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"time"

	"github.com/codegangsta/cli"
	"github.com/kardianos/osext"
)

var InitialPath string
var InitialGoPath string

var Verbose bool

var SpinnerCharSet = 14
var SpinnerInterval = 50 * time.Millisecond

func main() {
	currentExecutable, _ := osext.Executable()
	vendoredBunchPath := path.Join(".vendor", "bin", "bunch")

	fi1, errStat1 := os.Stat(currentExecutable)
	fi2, errStat2 := os.Stat(vendoredBunchPath)

	if exists, _ := pathExists(vendoredBunchPath); errStat1 == nil && errStat2 == nil && exists && !os.SameFile(fi1, fi2) {
		cmd := exec.Command(vendoredBunchPath, os.Args[1:]...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr

		err := cmd.Run()
		if err != nil {
			fmt.Println("running vendored bin/bunch was unsuccessful")
			os.Exit(1)
		} else {
			if cmd.ProcessState.Success() {
				os.Exit(0)
			} else {
				os.Exit(1)
			}
		}
	}

	InitialPath = os.Getenv("PATH")
	InitialGoPath = os.Getenv("GOPATH")

	app := cli.NewApp()
	app.Name = "bunch"
	app.Usage = "npm-like tool for managing Go dependencies"
	app.Version = "0.0.1"
	app.Authors = []cli.Author{cli.Author{Name: "Daniil Kulchenko", Email: "daniil@kulchenko.com"}}

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose",
			Usage: "output more information",
		},
	}

	app.Before = func(context *cli.Context) error {
		Verbose = context.GlobalBool("verbose")

		return nil
	}

	app.Commands = []cli.Command{
		{
			Name:    "install",
			Aliases: []string{"i"},
			Usage:   "install package(s)",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "save",
					Usage: "save installed package to Bunchfile",
				},
				cli.BoolFlag{
					Name:  "g",
					Usage: "install package to global $GOPATH instead of vendored directory",
				},
			},
			Action: func(c *cli.Context) {
				installCommand(c, false, true)
			},
		},
		{
			Name:    "update",
			Aliases: []string{"u"},
			Usage:   "update package(s)",
			Action: func(c *cli.Context) {
				installCommand(c, true, true)
			},
		},
		{
			Name:    "uninstall",
			Aliases: []string{"r"},
			Usage:   "uninstall package(s)",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "save",
					Usage: "save uninstalled package to Bunchfile",
				},
				cli.BoolFlag{
					Name:  "g",
					Usage: "uninstall package from global $GOPATH instead of vendored directory",
				},
			},
			Action: func(c *cli.Context) {
				uninstallCommand(c)
			},
		},
		{
			Name:  "prune",
			Usage: "remove packages not referenced in Bunchfile",
			Action: func(c *cli.Context) {
				pruneCommand(c)
			},
		},
		{
			Name:  "outdated",
			Usage: "list outdated packages",
			Action: func(c *cli.Context) {
				outdatedCommand(c)
			},
		},
		{
			Name:  "lock",
			Usage: "generate a file locking down current versions of dependencies",
			Action: func(c *cli.Context) {
				lockCommand(c)
			},
		},
		{
			Name:  "rebuild",
			Usage: "rebuild all dependencies",
			Action: func(c *cli.Context) {
				installCommand(c, true, false)
			},
		},
		{
			Name:  "generate",
			Usage: "generate a Bunchfile based on package imports in current directory",
			Action: func(c *cli.Context) {
				generateCommand(c)
			},
		},
		{
			Name:            "go",
			Usage:           "run a Go command within the vendor environment (e.g. bunch go fmt)",
			SkipFlagParsing: true,
			Action: func(c *cli.Context) {
				goCommand(c)
			},
		},
		{
			Name:            "exec",
			Usage:           "run any command within the vendor environment (e.g. bunch exec make)",
			SkipFlagParsing: true,
			Action: func(c *cli.Context) {
				execCommand(c)
			},
		},
		{
			Name:            "shell",
			Usage:           "start a shell within the vendor environment",
			SkipFlagParsing: true,
			Action: func(c *cli.Context) {
				shellCommand(c)
			},
		},
		{
			Name:  "shim",
			Usage: "sourced in .bash_profile to alias the 'go' tool",
			Action: func(c *cli.Context) {
				shimCommand(c)
			},
		},
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package packet

import (
	"testing"

	"github.com/yosssi/gmq/mqtt"
)

func Test_publish_setFixedHeader(t *testing.T) {
	p := &publish{
		dup:    true,
		retain: true,
	}

	p.variableHeader = []byte{0x00}

	p.payload = []byte{0x00, 0x00}

	p.setFixedHeader()

	want := []byte{0x39, 0x03}

	if len(p.fixedHeader) != len(want) || p.fixedHeader[0] != want[0] || p.fixedHeader[1] != want[1] {
		t.Errorf("p.fixedHeader => %v, want => %v", p.fixedHeader, want)
	}
}

func Test_setVariableHeader(t *testing.T) {
	p := &publish{
		qos:       mqtt.QoS1,
		topicName: []byte("topicName"),
		packetID:  1,
	}

	p.setVariableHeader()

	want := []byte{0x00, 0x09, 0x74, 0x6F, 0x70, 0x69, 0x63, 0x4E, 0x061, 0x6D, 0x65, 0x00, 0x01}

	if len(p.variableHeader) != len(want) {
		t.Errorf("p.variableHeader => %v, want => %v", p.variableHeader, want)
		return
	}

	for i, b := range p.variableHeader {
		if b != want[i] {
			t.Errorf("p.variableHeader => %v, want => %v", p.variableHeader, want)
			return
		}
	}
}

func Test_setPayload(t *testing.T) {
	p := &publish{
		message: []byte{0x00, 0x01},
	}

	p.setPayload()

	if len(p.payload) != len(p.message) || p.payload[0] != p.message[0] || p.payload[1] != p.message[1] {
		t.Errorf("p.payload => %v, want => %v", p.payload, p.message)
	}
}
<commit_msg>Update mqtt/packet/publish_test.go<commit_after>package packet

import (
	"testing"

	"github.com/yosssi/gmq/mqtt"
)

func Test_publish_setFixedHeader(t *testing.T) {
	p := &publish{
		dup:    true,
		retain: true,
	}

	p.variableHeader = []byte{0x00}

	p.payload = []byte{0x00, 0x00}

	p.setFixedHeader()

	want := []byte{0x39, 0x03}

	if len(p.fixedHeader) != len(want) || p.fixedHeader[0] != want[0] || p.fixedHeader[1] != want[1] {
		t.Errorf("p.fixedHeader => %v, want => %v", p.fixedHeader, want)
	}
}

func Test_publish_setVariableHeader(t *testing.T) {
	p := &publish{
		qos:       mqtt.QoS1,
		topicName: []byte("topicName"),
		packetID:  1,
	}

	p.setVariableHeader()

	want := []byte{0x00, 0x09, 0x74, 0x6F, 0x70, 0x69, 0x63, 0x4E, 0x061, 0x6D, 0x65, 0x00, 0x01}

	if len(p.variableHeader) != len(want) {
		t.Errorf("p.variableHeader => %v, want => %v", p.variableHeader, want)
		return
	}

	for i, b := range p.variableHeader {
		if b != want[i] {
			t.Errorf("p.variableHeader => %v, want => %v", p.variableHeader, want)
			return
		}
	}
}

func Test_publish_setPayload(t *testing.T) {
	p := &publish{
		message: []byte{0x00, 0x01},
	}

	p.setPayload()

	if len(p.payload) != len(p.message) || p.payload[0] != p.message[0] || p.payload[1] != p.message[1] {
		t.Errorf("p.payload => %v, want => %v", p.payload, p.message)
	}
}

func TestNewPUBLISH_optsNil(t *testing.T) {
	if _, err := NewPUBLISH(nil); err != nil {
		nilErrorExpected(t, err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Mender Software AS
//
//    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 s3

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/pkg/errors"

	"github.com/mendersoftware/deployments/resources/images"
	"github.com/mendersoftware/deployments/resources/images/model"
)

const (
	ExpireMaxLimit                 = 7 * 24 * time.Hour
	ExpireMinLimit                 = 1 * time.Minute
	ErrCodeBucketAlreadyOwnedByYou = "BucketAlreadyOwnedByYou"
)

// SimpleStorageService - AWS S3 client.
// Data layer for file storage.
// Implements model.FileStorage interface
type SimpleStorageService struct {
	client *s3.S3
	bucket string
}

// NewSimpleStorageServiceStatic create new S3 client model.
// AWS authentication keys are automatically reloaded from env variables.
func NewSimpleStorageServiceStatic(bucket, key, secret, region, token, uri string) (*SimpleStorageService, error) {

	credentials := credentials.NewStaticCredentials(key, secret, token)
	config := aws.NewConfig().WithCredentials(credentials).WithRegion(region)

	if len(uri) > 0 {
		sslDisabled := !strings.HasPrefix(uri, "https://")
		config = config.WithDisableSSL(sslDisabled).WithEndpoint(uri)
	}

	config.S3ForcePathStyle = aws.Bool(true)
	sess := session.New(config)

	client := s3.New(sess)

	// minio requires explicit bucket creation
	cparams := &s3.CreateBucketInput{
		Bucket: aws.String(bucket), // Required
	}

	_, err := client.CreateBucket(cparams)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() != ErrCodeBucketAlreadyOwnedByYou {
				return nil, err
			}
		}
	}

	return &SimpleStorageService{
		client: client,
		bucket: bucket,
	}, nil
}

// NewSimpleStorageServiceDefaults create new S3 client model.
// Use default authentication provides which looks at env variables,
// Aws profile file and ec2 iam role
func NewSimpleStorageServiceDefaults(bucket, region string) (*SimpleStorageService, error) {

	sess := session.New(aws.NewConfig().WithRegion(region))
	client := s3.New(sess)

	// minio requires explicit bucket creation
	cparams := &s3.CreateBucketInput{
		Bucket: aws.String(bucket), // Required
	}

	_, err := client.CreateBucket(cparams)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() != ErrCodeBucketAlreadyOwnedByYou {
				return nil, err
			}
		}
	}

	return &SimpleStorageService{
		client: client,
		bucket: bucket,
	}, nil
}

// Delete removes delected file from storage.
// Noop if ID does not exist.
func (s *SimpleStorageService) Delete(ctx context.Context, objectID string) error {

	params := &s3.DeleteObjectInput{
		// Required
		Bucket: aws.String(s.bucket),
		Key:    aws.String(objectID),

		// Optional
		RequestPayer: aws.String(s3.RequestPayerRequester),
	}

	// ignore return response which contains charing info
	// and file versioning data which are not in interest
	_, err := s.client.DeleteObject(params)
	if err != nil {
		return errors.Wrap(err, "Removing file")
	}

	return nil
}

// Exists check if selected object exists in the storage
func (s *SimpleStorageService) Exists(ctx context.Context, objectID string) (bool, error) {

	params := &s3.ListObjectsInput{
		// Required
		Bucket: aws.String(s.bucket),

		// Optional
		MaxKeys: aws.Int64(1),
		Prefix:  aws.String(objectID),
	}

	resp, err := s.client.ListObjects(params)
	if err != nil {
		return false, errors.Wrap(err, "Searching for file")
	}

	if len(resp.Contents) == 0 {
		return false, nil
	}

	// Note: Response should contain max 1 object (MaxKetys=1)
	// Double check if it's exact match as object search matches prefix.
	if *resp.Contents[0].Key == objectID {
		return true, nil
	}

	return false, nil
}

// UploadArtifact uploads given artifact into the file server (AWS S3 or minio)
// using objectID as a key
func (s *SimpleStorageService) UploadArtifact(ctx context.Context,
	objectID string, size int64, artifact io.Reader, contentType string) error {

	params := &s3.PutObjectInput{
		// Required
		Bucket: aws.String(s.bucket),
		Key:    aws.String(objectID),
	}

	// Ignore out object
	r, _ := s.client.PutObjectRequest(params)

	// Presign request
	uri, err := r.Presign(10 * time.Second)
	if err != nil {
		return err
	}

	client := &http.Client{}
	request, err := http.NewRequest(http.MethodPut, uri, artifact)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", contentType)
	request.ContentLength = size
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	if resp.StatusCode != http.StatusOK {
		return errors.New("Artifact upload failed: " + resp.Status)
	}

	return nil
}

// PutRequest duration is limited to 7 days (AWS limitation)
func (s *SimpleStorageService) PutRequest(ctx context.Context, objectID string,
	duration time.Duration) (*images.Link, error) {

	if err := s.validateDurationLimits(duration); err != nil {
		return nil, err
	}

	params := &s3.PutObjectInput{
		// Required
		Bucket: aws.String(s.bucket),
		Key:    aws.String(objectID),
	}

	// Ignore out object
	req, _ := s.client.PutObjectRequest(params)

	uri, err := req.Presign(duration)
	if err != nil {
		return nil, errors.Wrap(err, "Signing PUT request")
	}

	return images.NewLink(uri, req.Time.Add(req.ExpireTime)), nil
}

// GetRequest duration is limited to 7 days (AWS limitation)
func (s *SimpleStorageService) GetRequest(ctx context.Context, objectID string,
	duration time.Duration, responseContentType string) (*images.Link, error) {

	if err := s.validateDurationLimits(duration); err != nil {
		return nil, err
	}

	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(objectID),
	}

	if responseContentType != "" {
		params.ResponseContentType = &responseContentType
	}

	// Ignore out object
	req, _ := s.client.GetObjectRequest(params)

	uri, err := req.Presign(duration)
	if err != nil {
		return nil, errors.Wrap(err, "Signing GET request")
	}

	return images.NewLink(uri, req.Time.Add(req.ExpireTime)), nil
}

func (s *SimpleStorageService) validateDurationLimits(duration time.Duration) error {
	if duration > ExpireMaxLimit || duration < ExpireMinLimit {
		return fmt.Errorf("Expire duration out of range: allowed %d-%d[ns]",
			ExpireMinLimit, ExpireMaxLimit)
	}

	return nil
}

// LastModified returns last file modification time.
// If object not found return ErrFileStorageFileNotFound
func (s *SimpleStorageService) LastModified(ctx context.Context, objectID string) (time.Time, error) {

	params := &s3.ListObjectsInput{
		// Required
		Bucket: aws.String(s.bucket),

		// Optional
		MaxKeys: aws.Int64(1),
		Prefix:  aws.String(objectID),
	}

	resp, err := s.client.ListObjects(params)
	if err != nil {
		return time.Time{}, errors.Wrap(err, "Searching for file")
	}

	if len(resp.Contents) == 0 {
		return time.Time{}, model.ErrFileStorageFileNotFound
	}

	// Note: Response should contain max 1 object (MaxKetys=1)
	// Double check if it's exact match as object search matches prefix.
	if *resp.Contents[0].Key != objectID {
		return time.Time{}, model.ErrFileStorageFileNotFound
	}

	return *resp.Contents[0].LastModified, nil
}
<commit_msg>images/s3: close S3 response body when uploading image<commit_after>// Copyright 2016 Mender Software AS
//
//    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 s3

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/pkg/errors"

	"github.com/mendersoftware/deployments/resources/images"
	"github.com/mendersoftware/deployments/resources/images/model"
)

const (
	ExpireMaxLimit                 = 7 * 24 * time.Hour
	ExpireMinLimit                 = 1 * time.Minute
	ErrCodeBucketAlreadyOwnedByYou = "BucketAlreadyOwnedByYou"
)

// SimpleStorageService - AWS S3 client.
// Data layer for file storage.
// Implements model.FileStorage interface
type SimpleStorageService struct {
	client *s3.S3
	bucket string
}

// NewSimpleStorageServiceStatic create new S3 client model.
// AWS authentication keys are automatically reloaded from env variables.
func NewSimpleStorageServiceStatic(bucket, key, secret, region, token, uri string) (*SimpleStorageService, error) {

	credentials := credentials.NewStaticCredentials(key, secret, token)
	config := aws.NewConfig().WithCredentials(credentials).WithRegion(region)

	if len(uri) > 0 {
		sslDisabled := !strings.HasPrefix(uri, "https://")
		config = config.WithDisableSSL(sslDisabled).WithEndpoint(uri)
	}

	config.S3ForcePathStyle = aws.Bool(true)
	sess := session.New(config)

	client := s3.New(sess)

	// minio requires explicit bucket creation
	cparams := &s3.CreateBucketInput{
		Bucket: aws.String(bucket), // Required
	}

	_, err := client.CreateBucket(cparams)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() != ErrCodeBucketAlreadyOwnedByYou {
				return nil, err
			}
		}
	}

	return &SimpleStorageService{
		client: client,
		bucket: bucket,
	}, nil
}

// NewSimpleStorageServiceDefaults create new S3 client model.
// Use default authentication provides which looks at env variables,
// Aws profile file and ec2 iam role
func NewSimpleStorageServiceDefaults(bucket, region string) (*SimpleStorageService, error) {

	sess := session.New(aws.NewConfig().WithRegion(region))
	client := s3.New(sess)

	// minio requires explicit bucket creation
	cparams := &s3.CreateBucketInput{
		Bucket: aws.String(bucket), // Required
	}

	_, err := client.CreateBucket(cparams)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() != ErrCodeBucketAlreadyOwnedByYou {
				return nil, err
			}
		}
	}

	return &SimpleStorageService{
		client: client,
		bucket: bucket,
	}, nil
}

// Delete removes delected file from storage.
// Noop if ID does not exist.
func (s *SimpleStorageService) Delete(ctx context.Context, objectID string) error {

	params := &s3.DeleteObjectInput{
		// Required
		Bucket: aws.String(s.bucket),
		Key:    aws.String(objectID),

		// Optional
		RequestPayer: aws.String(s3.RequestPayerRequester),
	}

	// ignore return response which contains charing info
	// and file versioning data which are not in interest
	_, err := s.client.DeleteObject(params)
	if err != nil {
		return errors.Wrap(err, "Removing file")
	}

	return nil
}

// Exists check if selected object exists in the storage
func (s *SimpleStorageService) Exists(ctx context.Context, objectID string) (bool, error) {

	params := &s3.ListObjectsInput{
		// Required
		Bucket: aws.String(s.bucket),

		// Optional
		MaxKeys: aws.Int64(1),
		Prefix:  aws.String(objectID),
	}

	resp, err := s.client.ListObjects(params)
	if err != nil {
		return false, errors.Wrap(err, "Searching for file")
	}

	if len(resp.Contents) == 0 {
		return false, nil
	}

	// Note: Response should contain max 1 object (MaxKetys=1)
	// Double check if it's exact match as object search matches prefix.
	if *resp.Contents[0].Key == objectID {
		return true, nil
	}

	return false, nil
}

// UploadArtifact uploads given artifact into the file server (AWS S3 or minio)
// using objectID as a key
func (s *SimpleStorageService) UploadArtifact(ctx context.Context,
	objectID string, size int64, artifact io.Reader, contentType string) error {

	params := &s3.PutObjectInput{
		// Required
		Bucket: aws.String(s.bucket),
		Key:    aws.String(objectID),
	}

	// Ignore out object
	r, _ := s.client.PutObjectRequest(params)

	// Presign request
	uri, err := r.Presign(10 * time.Second)
	if err != nil {
		return err
	}

	client := &http.Client{}
	request, err := http.NewRequest(http.MethodPut, uri, artifact)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", contentType)
	request.ContentLength = size
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return errors.New("Artifact upload failed: " + resp.Status)
	}

	return nil
}

// PutRequest duration is limited to 7 days (AWS limitation)
func (s *SimpleStorageService) PutRequest(ctx context.Context, objectID string,
	duration time.Duration) (*images.Link, error) {

	if err := s.validateDurationLimits(duration); err != nil {
		return nil, err
	}

	params := &s3.PutObjectInput{
		// Required
		Bucket: aws.String(s.bucket),
		Key:    aws.String(objectID),
	}

	// Ignore out object
	req, _ := s.client.PutObjectRequest(params)

	uri, err := req.Presign(duration)
	if err != nil {
		return nil, errors.Wrap(err, "Signing PUT request")
	}

	return images.NewLink(uri, req.Time.Add(req.ExpireTime)), nil
}

// GetRequest duration is limited to 7 days (AWS limitation)
func (s *SimpleStorageService) GetRequest(ctx context.Context, objectID string,
	duration time.Duration, responseContentType string) (*images.Link, error) {

	if err := s.validateDurationLimits(duration); err != nil {
		return nil, err
	}

	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(objectID),
	}

	if responseContentType != "" {
		params.ResponseContentType = &responseContentType
	}

	// Ignore out object
	req, _ := s.client.GetObjectRequest(params)

	uri, err := req.Presign(duration)
	if err != nil {
		return nil, errors.Wrap(err, "Signing GET request")
	}

	return images.NewLink(uri, req.Time.Add(req.ExpireTime)), nil
}

func (s *SimpleStorageService) validateDurationLimits(duration time.Duration) error {
	if duration > ExpireMaxLimit || duration < ExpireMinLimit {
		return fmt.Errorf("Expire duration out of range: allowed %d-%d[ns]",
			ExpireMinLimit, ExpireMaxLimit)
	}

	return nil
}

// LastModified returns last file modification time.
// If object not found return ErrFileStorageFileNotFound
func (s *SimpleStorageService) LastModified(ctx context.Context, objectID string) (time.Time, error) {

	params := &s3.ListObjectsInput{
		// Required
		Bucket: aws.String(s.bucket),

		// Optional
		MaxKeys: aws.Int64(1),
		Prefix:  aws.String(objectID),
	}

	resp, err := s.client.ListObjects(params)
	if err != nil {
		return time.Time{}, errors.Wrap(err, "Searching for file")
	}

	if len(resp.Contents) == 0 {
		return time.Time{}, model.ErrFileStorageFileNotFound
	}

	// Note: Response should contain max 1 object (MaxKetys=1)
	// Double check if it's exact match as object search matches prefix.
	if *resp.Contents[0].Key != objectID {
		return time.Time{}, model.ErrFileStorageFileNotFound
	}

	return *resp.Contents[0].LastModified, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package mungers

import (
	"fmt"
	"path"
	"regexp"
	"strings"

	"k8s.io/kubernetes/pkg/util/sets"
	"k8s.io/test-infra/mungegithub/features"
	"k8s.io/test-infra/mungegithub/github"
	"k8s.io/test-infra/mungegithub/options"

	"github.com/golang/glog"
	githubapi "github.com/google/go-github/github"
)

const (
	labelSizePrefix = "size/"
)

var (
	sizeRE = regexp.MustCompile("Labelling this PR as " + labelSizePrefix + "(XS|S|M|L|XL|XXL)")
)

// SizeMunger will update a label on a PR based on how many lines are changed.
// It will exclude certain files in it's calculations based on the config
// file provided in --generated-files-config
type SizeMunger struct {
	generatedFilesFile string
	genFilePaths       sets.String
	genFilePrefixes    sets.String
	genFileNames       sets.String
	genPathPrefixes    sets.String
}

func init() {
	s := &SizeMunger{}
	RegisterMungerOrDie(s)
	RegisterStaleIssueComments(s)
}

// Name is the name usable in --pr-mungers
func (*SizeMunger) Name() string { return "size" }

// RequiredFeatures is a slice of 'features' that must be provided
func (*SizeMunger) RequiredFeatures() []string { return []string{} }

// Initialize will initialize the munger
func (s *SizeMunger) Initialize(config *github.Config, features *features.Features) error {
	return nil
}

// EachLoop is called at the start of every munge loop
func (*SizeMunger) EachLoop() error { return nil }

// RegisterOptions registers options for this munger; returns any that require a restart when changed.
func (s *SizeMunger) RegisterOptions(opts *options.Options) sets.String {
	opts.RegisterString(&s.generatedFilesFile, "generated-files-config", "", "file in the repo containing the generated file rules")
	opts.RegisterUpdateCallback(func(changed sets.String) error {
		if changed.Has("generated-files-config") {
			// Force 'getGeneratedFiles' to re-read config next time it is called.
			s.genFilePaths = nil
			s.genFilePrefixes = nil
			s.genFileNames = nil
			s.genPathPrefixes = nil
		}
		return nil
	})

	return nil
}

// getGeneratedFiles returns a list of all automatically generated files in the repo. These include
// docs, deep_copy, and conversions
//
// It would be 'better' to call this for every commit but that takes
// a whole lot of time for almost always the same information, and if
// our results are slightly wrong, who cares? Instead look for the
// generated files once and if someone changed what files are generated
// we'll size slightly wrong. No biggie.
func (s *SizeMunger) getGeneratedFiles(obj *github.MungeObject) {
	if s.genFilePaths != nil {
		return
	}
	if s.genFilePrefixes != nil {
		return
	}
	if s.genFileNames != nil {
		return
	}
	if s.genPathPrefixes != nil {
		return
	}

	// Don't write into s yet, since this might fail, and these are indicators
	// to not repeat this function.
	paths := sets.NewString()
	filePrefixes := sets.NewString()
	fileNames := sets.NewString()
	pathPrefixes := sets.NewString()

	file := s.generatedFilesFile
	if len(file) == 0 {
		return
	}

	contents, err := obj.GetFileContents(file, "")
	if err != nil {
		glog.Errorf("Failed to read generated-files-config %q: %v", file, err)
	}

	lines := strings.Split(contents, "\n")
	for i, line := range lines {
		line := strings.TrimSpace(line)
		if strings.HasPrefix(line, "#") || line == "" {
			continue
		}
		fields := strings.Fields(line)
		if len(fields) != 2 {
			glog.Errorf("Invalid line %d in generated docs config %s: %q", i, file, line)
			continue
		}
		key := fields[0]
		val := fields[1]
		if key == "prefix" || key == "path-prefix" {
			pathPrefixes.Insert(val)
		} else if key == "file-prefix" {
			filePrefixes.Insert(val)
		} else if key == "file-name" {
			fileNames.Insert(val)
		} else if key == "path" {
			paths.Insert(val)
		} else if key == "paths-from-repo" {
			docs, err := obj.GetFileContents(val, "")
			if err != nil {
				continue
			}
			docSlice := strings.Split(docs, "\n")
			paths.Insert(docSlice...)
		} else {
			glog.Errorf("Invalid line %d in generated docs config, unknown key: %s, %q", i, key, line)
			continue
		}
	}

	// Save the results so we don't repeat this function.
	s.genFilePaths = paths
	s.genFilePrefixes = filePrefixes
	s.genFileNames = fileNames
	s.genPathPrefixes = pathPrefixes

	return
}

// Munge is the workhorse the will actually make updates to the PR
func (s *SizeMunger) Munge(obj *github.MungeObject) {
	if !obj.IsPR() {
		return
	}

	issue := obj.Issue

	s.getGeneratedFiles(obj)

	files, ok := obj.ListFiles()
	if !ok {
		return
	}

	adds := 0
	dels := 0
	for _, f := range files {
		skip := false
		for p := range s.genPathPrefixes {
			if strings.HasPrefix(*f.Filename, p) {
				skip = true
				break
			}
		}
		if skip {
			continue
		}
		if s.genFilePaths.Has(*f.Filename) {
			continue
		}
		_, filename := path.Split(*f.Filename)
		if hasPrefix(filename, s.genFilePrefixes) {
			continue
		}
		if s.genFileNames.Has(filename) {
			continue
		}
		if f.Additions != nil {
			adds += *f.Additions
		}
		if f.Deletions != nil {
			dels += *f.Deletions
		}
	}

	newSize := calculateSize(adds, dels)
	newLabel := labelSizePrefix + newSize

	existing := github.GetLabelsWithPrefix(issue.Labels, labelSizePrefix)
	needsUpdate := true
	for _, l := range existing {
		if l == newLabel {
			needsUpdate = false
			continue
		}
		obj.RemoveLabel(l)
	}
	if needsUpdate {
		obj.AddLabels([]string{newLabel})

		body := fmt.Sprintf("Labelling this PR as %s", newLabel)
		obj.WriteComment(body)
	}
}

func hasPrefix(filename string, filePrefixes sets.String) bool {
	for pfx := range filePrefixes {
		if strings.HasPrefix(filename, pfx) {
			return true
		}
	}
	return false
}

const (
	sizeXS  = "XS"
	sizeS   = "S"
	sizeM   = "M"
	sizeL   = "L"
	sizeXL  = "XL"
	sizeXXL = "XXL"
)

func calculateSize(adds, dels int) string {
	lines := adds + dels

	// This is a totally arbitrary heuristic and is open for tweaking.
	if lines < 10 {
		return sizeXS
	}
	if lines < 30 {
		return sizeS
	}
	if lines < 100 {
		return sizeM
	}
	if lines < 500 {
		return sizeL
	}
	if lines < 1000 {
		return sizeXL
	}
	return sizeXXL
}

func (s *SizeMunger) isStaleIssueComment(obj *github.MungeObject, comment *githubapi.IssueComment) bool {
	if !obj.IsRobot(comment.User) {
		return false
	}
	stale := sizeRE.MatchString(*comment.Body)
	if stale {
		glog.V(6).Infof("Found stale SizeMunger comment")
	}
	return stale
}

// StaleIssueComments returns a slice of stale issue comments.
func (s *SizeMunger) StaleIssueComments(obj *github.MungeObject, comments []*githubapi.IssueComment) []*githubapi.IssueComment {
	return forEachCommentTest(obj, comments, s.isStaleIssueComment)
}
<commit_msg>Don't comment when labeling PR size.<commit_after>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package mungers

import (
	"path"
	"strings"

	"k8s.io/kubernetes/pkg/util/sets"
	"k8s.io/test-infra/mungegithub/features"
	"k8s.io/test-infra/mungegithub/github"
	"k8s.io/test-infra/mungegithub/options"

	"github.com/golang/glog"
)

const (
	labelSizePrefix = "size/"
)

// SizeMunger will update a label on a PR based on how many lines are changed.
// It will exclude certain files in it's calculations based on the config
// file provided in --generated-files-config
type SizeMunger struct {
	generatedFilesFile string
	genFilePaths       sets.String
	genFilePrefixes    sets.String
	genFileNames       sets.String
	genPathPrefixes    sets.String
}

func init() {
	s := &SizeMunger{}
	RegisterMungerOrDie(s)
}

// Name is the name usable in --pr-mungers
func (*SizeMunger) Name() string { return "size" }

// RequiredFeatures is a slice of 'features' that must be provided
func (*SizeMunger) RequiredFeatures() []string { return []string{} }

// Initialize will initialize the munger
func (s *SizeMunger) Initialize(config *github.Config, features *features.Features) error {
	return nil
}

// EachLoop is called at the start of every munge loop
func (*SizeMunger) EachLoop() error { return nil }

// RegisterOptions registers options for this munger; returns any that require a restart when changed.
func (s *SizeMunger) RegisterOptions(opts *options.Options) sets.String {
	opts.RegisterString(&s.generatedFilesFile, "generated-files-config", "", "file in the repo containing the generated file rules")
	opts.RegisterUpdateCallback(func(changed sets.String) error {
		if changed.Has("generated-files-config") {
			// Force 'getGeneratedFiles' to re-read config next time it is called.
			s.genFilePaths = nil
			s.genFilePrefixes = nil
			s.genFileNames = nil
			s.genPathPrefixes = nil
		}
		return nil
	})

	return nil
}

// getGeneratedFiles returns a list of all automatically generated files in the repo. These include
// docs, deep_copy, and conversions
//
// It would be 'better' to call this for every commit but that takes
// a whole lot of time for almost always the same information, and if
// our results are slightly wrong, who cares? Instead look for the
// generated files once and if someone changed what files are generated
// we'll size slightly wrong. No biggie.
func (s *SizeMunger) getGeneratedFiles(obj *github.MungeObject) {
	if s.genFilePaths != nil {
		return
	}
	if s.genFilePrefixes != nil {
		return
	}
	if s.genFileNames != nil {
		return
	}
	if s.genPathPrefixes != nil {
		return
	}

	// Don't write into s yet, since this might fail, and these are indicators
	// to not repeat this function.
	paths := sets.NewString()
	filePrefixes := sets.NewString()
	fileNames := sets.NewString()
	pathPrefixes := sets.NewString()

	file := s.generatedFilesFile
	if len(file) == 0 {
		return
	}

	contents, err := obj.GetFileContents(file, "")
	if err != nil {
		glog.Errorf("Failed to read generated-files-config %q: %v", file, err)
	}

	lines := strings.Split(contents, "\n")
	for i, line := range lines {
		line := strings.TrimSpace(line)
		if strings.HasPrefix(line, "#") || line == "" {
			continue
		}
		fields := strings.Fields(line)
		if len(fields) != 2 {
			glog.Errorf("Invalid line %d in generated docs config %s: %q", i, file, line)
			continue
		}
		key := fields[0]
		val := fields[1]
		if key == "prefix" || key == "path-prefix" {
			pathPrefixes.Insert(val)
		} else if key == "file-prefix" {
			filePrefixes.Insert(val)
		} else if key == "file-name" {
			fileNames.Insert(val)
		} else if key == "path" {
			paths.Insert(val)
		} else if key == "paths-from-repo" {
			docs, err := obj.GetFileContents(val, "")
			if err != nil {
				continue
			}
			docSlice := strings.Split(docs, "\n")
			paths.Insert(docSlice...)
		} else {
			glog.Errorf("Invalid line %d in generated docs config, unknown key: %s, %q", i, key, line)
			continue
		}
	}

	// Save the results so we don't repeat this function.
	s.genFilePaths = paths
	s.genFilePrefixes = filePrefixes
	s.genFileNames = fileNames
	s.genPathPrefixes = pathPrefixes

	return
}

// Munge is the workhorse the will actually make updates to the PR
func (s *SizeMunger) Munge(obj *github.MungeObject) {
	if !obj.IsPR() {
		return
	}

	issue := obj.Issue

	s.getGeneratedFiles(obj)

	files, ok := obj.ListFiles()
	if !ok {
		return
	}

	adds := 0
	dels := 0
	for _, f := range files {
		skip := false
		for p := range s.genPathPrefixes {
			if strings.HasPrefix(*f.Filename, p) {
				skip = true
				break
			}
		}
		if skip {
			continue
		}
		if s.genFilePaths.Has(*f.Filename) {
			continue
		}
		_, filename := path.Split(*f.Filename)
		if hasPrefix(filename, s.genFilePrefixes) {
			continue
		}
		if s.genFileNames.Has(filename) {
			continue
		}
		if f.Additions != nil {
			adds += *f.Additions
		}
		if f.Deletions != nil {
			dels += *f.Deletions
		}
	}

	newSize := calculateSize(adds, dels)
	newLabel := labelSizePrefix + newSize

	existing := github.GetLabelsWithPrefix(issue.Labels, labelSizePrefix)
	needsUpdate := true
	for _, l := range existing {
		if l == newLabel {
			needsUpdate = false
			continue
		}
		obj.RemoveLabel(l)
	}
	if needsUpdate {
		obj.AddLabels([]string{newLabel})
	}
}

func hasPrefix(filename string, filePrefixes sets.String) bool {
	for pfx := range filePrefixes {
		if strings.HasPrefix(filename, pfx) {
			return true
		}
	}
	return false
}

const (
	sizeXS  = "XS"
	sizeS   = "S"
	sizeM   = "M"
	sizeL   = "L"
	sizeXL  = "XL"
	sizeXXL = "XXL"
)

func calculateSize(adds, dels int) string {
	lines := adds + dels

	// This is a totally arbitrary heuristic and is open for tweaking.
	if lines < 10 {
		return sizeXS
	}
	if lines < 30 {
		return sizeS
	}
	if lines < 100 {
		return sizeM
	}
	if lines < 500 {
		return sizeL
	}
	if lines < 1000 {
		return sizeXL
	}
	return sizeXXL
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/opsworks"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

// These tests assume the existence of predefined Opsworks IAM roles named `aws-opsworks-ec2-role`
// and `aws-opsworks-service-role`.

func TestAccAWSOpsworksCustomLayer(t *testing.T) {
	stackName := fmt.Sprintf("tf-opsworks-acc-%d", acctest.RandInt())
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAwsOpsworksCustomLayerDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAwsOpsworksCustomLayerConfigCreate(stackName),
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "name", "tf-ops-acc-custom-layer",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "auto_assign_elastic_ips", "false",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "auto_healing", "true",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "drain_elb_on_shutdown", "true",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "instance_shutdown_timeout", "300",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "custom_security_group_ids.#", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.#", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.1368285564", "git",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.2937857443", "golang",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.#", "1",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.type", "gp2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.number_of_disks", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.mount_point", "/home",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.size", "100",
					),
				),
			},
			resource.TestStep{
				Config: testAccAwsOpsworksCustomLayerConfigUpdate(stackName),
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "name", "tf-ops-acc-custom-layer",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "drain_elb_on_shutdown", "false",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "instance_shutdown_timeout", "120",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "custom_security_group_ids.#", "3",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.#", "3",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.1368285564", "git",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.2937857443", "golang",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.4101929740", "subversion",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.#", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.type", "gp2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.number_of_disks", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.mount_point", "/home",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.size", "100",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.type", "io1",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.number_of_disks", "4",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.mount_point", "/var",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.size", "100",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.raid_level", "1",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.iops", "3000",
					),
				),
			},
		},
	})
}

func testAccCheckAwsOpsworksCustomLayerDestroy(s *terraform.State) error {
	opsworksconn := testAccProvider.Meta().(*AWSClient).opsworksconn
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_opsworks_custom_layer" {
			continue
		}
		req := &opsworks.DescribeLayersInput{
			LayerIds: []*string{
				aws.String(rs.Primary.ID),
			},
		}

		_, err := opsworksconn.DescribeLayers(req)
		if err != nil {
			if awserr, ok := err.(awserr.Error); ok {
				if awserr.Code() == "ResourceNotFoundException" {
					// not found, good to go
					return nil
				}
			}
			return err
		}
	}

	return fmt.Errorf("Fall through error on OpsWorks custom layer test")
}

func testAccAwsOpsworksCustomLayerSecurityGroups(name string) string {
	return fmt.Sprintf(`
resource "aws_security_group" "tf-ops-acc-layer1" {
  name = "%s-layer1"
  ingress {
    from_port = 8
    to_port = -1
    protocol = "icmp"
    cidr_blocks = ["0.0.0.0/0"]
  }
}
resource "aws_security_group" "tf-ops-acc-layer2" {
  name = "%s-layer2"
  ingress {
    from_port = 8
    to_port = -1
    protocol = "icmp"
    cidr_blocks = ["0.0.0.0/0"]
  }
}`, name)
}

func testAccAwsOpsworksCustomLayerConfigCreate(name string) string {
	return fmt.Sprintf(`
provider "aws" {
	region = "us-east-1"
}

resource "aws_opsworks_custom_layer" "tf-acc" {
  stack_id = "${aws_opsworks_stack.tf-acc.id}"
  name = "tf-ops-acc-custom-layer"
  short_name = "tf-ops-acc-custom-layer"
  auto_assign_public_ips = true
  custom_security_group_ids = [
    "${aws_security_group.tf-ops-acc-layer1.id}",
    "${aws_security_group.tf-ops-acc-layer2.id}",
  ]
  drain_elb_on_shutdown = true
  instance_shutdown_timeout = 300
  system_packages = [
    "git",
    "golang",
  ]
  ebs_volume {
    type = "gp2"
    number_of_disks = 2
    mount_point = "/home"
    size = 100
    raid_level = 0
  }
}

%s

%s

`, name, testAccAwsOpsworksStackConfigNoVpcCreate(name), testAccAwsOpsworksCustomLayerSecurityGroups(name))
}

func testAccAwsOpsworksCustomLayerConfigUpdate(name string) string {
	return fmt.Sprintf(`
provider "aws" {
  region = "us-east-1"
}

resource "aws_security_group" "tf-ops-acc-layer3" {
  name = "tf-ops-acc-layer3"
  ingress {
    from_port = 8
    to_port = -1
    protocol = "icmp"
    cidr_blocks = ["0.0.0.0/0"]
  }
}
resource "aws_opsworks_custom_layer" "tf-acc" {
  stack_id = "${aws_opsworks_stack.tf-acc.id}"
  name = "tf-ops-acc-custom-layer"
  short_name = "tf-ops-acc-custom-layer"
  auto_assign_public_ips = true
  custom_security_group_ids = [
    "${aws_security_group.tf-ops-acc-layer1.id}",
    "${aws_security_group.tf-ops-acc-layer2.id}",
    "${aws_security_group.tf-ops-acc-layer3.id}",
  ]
  drain_elb_on_shutdown = false
  instance_shutdown_timeout = 120
  system_packages = [
    "git",
    "golang",
    "subversion",
  ]
  ebs_volume {
    type = "gp2"
    number_of_disks = 2
    mount_point = "/home"
    size = 100
    raid_level = 0
  }
  ebs_volume {
    type = "io1"
    number_of_disks = 4
    mount_point = "/var"
    size = 100
    raid_level = 1
    iops = 3000
  }
}

%s

%s

`, name, testAccAwsOpsworksStackConfigNoVpcCreate(name), testAccAwsOpsworksCustomLayerSecurityGroups(name))
}
<commit_msg>fix some typos<commit_after>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/opsworks"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

// These tests assume the existence of predefined Opsworks IAM roles named `aws-opsworks-ec2-role`
// and `aws-opsworks-service-role`.

func TestAccAWSOpsworksCustomLayer(t *testing.T) {
	stackName := fmt.Sprintf("tf-opsworks-acc-%d", acctest.RandInt())
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAwsOpsworksCustomLayerDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAwsOpsworksCustomLayerConfigCreate(stackName),
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "name", "tf-ops-acc-custom-layer",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "auto_assign_elastic_ips", "false",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "auto_healing", "true",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "drain_elb_on_shutdown", "true",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "instance_shutdown_timeout", "300",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "custom_security_group_ids.#", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.#", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.1368285564", "git",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.2937857443", "golang",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.#", "1",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.type", "gp2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.number_of_disks", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.mount_point", "/home",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.size", "100",
					),
				),
			},
			resource.TestStep{
				Config: testAccAwsOpsworksCustomLayerConfigUpdate(stackName),
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "name", "tf-ops-acc-custom-layer",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "drain_elb_on_shutdown", "false",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "instance_shutdown_timeout", "120",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "custom_security_group_ids.#", "3",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.#", "3",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.1368285564", "git",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.2937857443", "golang",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "system_packages.4101929740", "subversion",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.#", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.type", "gp2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.number_of_disks", "2",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.mount_point", "/home",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.3575749636.size", "100",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.type", "io1",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.number_of_disks", "4",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.mount_point", "/var",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.size", "100",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.raid_level", "1",
					),
					resource.TestCheckResourceAttr(
						"aws_opsworks_custom_layer.tf-acc", "ebs_volume.1266957920.iops", "3000",
					),
				),
			},
		},
	})
}

func testAccCheckAwsOpsworksCustomLayerDestroy(s *terraform.State) error {
	opsworksconn := testAccProvider.Meta().(*AWSClient).opsworksconn
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_opsworks_custom_layer" {
			continue
		}
		req := &opsworks.DescribeLayersInput{
			LayerIds: []*string{
				aws.String(rs.Primary.ID),
			},
		}

		_, err := opsworksconn.DescribeLayers(req)
		if err != nil {
			if awserr, ok := err.(awserr.Error); ok {
				if awserr.Code() == "ResourceNotFoundException" {
					// not found, good to go
					return nil
				}
			}
			return err
		}
	}

	return fmt.Errorf("Fall through error on OpsWorks custom layer test")
}

func testAccAwsOpsworksCustomLayerSecurityGroups(name string) string {
	return fmt.Sprintf(`
resource "aws_security_group" "tf-ops-acc-layer1" {
  name = "%s-layer1"
  ingress {
    from_port = 8
    to_port = -1
    protocol = "icmp"
    cidr_blocks = ["0.0.0.0/0"]
  }
}
resource "aws_security_group" "tf-ops-acc-layer2" {
  name = "%s-layer2"
  ingress {
    from_port = 8
    to_port = -1
    protocol = "icmp"
    cidr_blocks = ["0.0.0.0/0"]
  }
}`, name, name)
}

func testAccAwsOpsworksCustomLayerConfigCreate(name string) string {
	return fmt.Sprintf(`
provider "aws" {
	region = "us-east-1"
}

resource "aws_opsworks_custom_layer" "tf-acc" {
  stack_id = "${aws_opsworks_stack.tf-acc.id}"
  name = "%s"
  short_name = "tf-ops-acc-custom-layer"
  auto_assign_public_ips = true
  custom_security_group_ids = [
    "${aws_security_group.tf-ops-acc-layer1.id}",
    "${aws_security_group.tf-ops-acc-layer2.id}",
  ]
  drain_elb_on_shutdown = true
  instance_shutdown_timeout = 300
  system_packages = [
    "git",
    "golang",
  ]
  ebs_volume {
    type = "gp2"
    number_of_disks = 2
    mount_point = "/home"
    size = 100
    raid_level = 0
  }
}

%s

%s

`, name, testAccAwsOpsworksStackConfigNoVpcCreate(name), testAccAwsOpsworksCustomLayerSecurityGroups(name))
}

func testAccAwsOpsworksCustomLayerConfigUpdate(name string) string {
	return fmt.Sprintf(`
provider "aws" {
  region = "us-east-1"
}

resource "aws_security_group" "tf-ops-acc-layer3" {
  name = "tf-ops-acc-layer3"
  ingress {
    from_port = 8
    to_port = -1
    protocol = "icmp"
    cidr_blocks = ["0.0.0.0/0"]
  }
}
resource "aws_opsworks_custom_layer" "tf-acc" {
  stack_id = "${aws_opsworks_stack.tf-acc.id}"
  name = "%s"
  short_name = "tf-ops-acc-custom-layer"
  auto_assign_public_ips = true
  custom_security_group_ids = [
    "${aws_security_group.tf-ops-acc-layer1.id}",
    "${aws_security_group.tf-ops-acc-layer2.id}",
    "${aws_security_group.tf-ops-acc-layer3.id}",
  ]
  drain_elb_on_shutdown = false
  instance_shutdown_timeout = 120
  system_packages = [
    "git",
    "golang",
    "subversion",
  ]
  ebs_volume {
    type = "gp2"
    number_of_disks = 2
    mount_point = "/home"
    size = 100
    raid_level = 0
  }
  ebs_volume {
    type = "io1"
    number_of_disks = 4
    mount_point = "/var"
    size = 100
    raid_level = 1
    iops = 3000
  }
}

%s

%s

`, name, testAccAwsOpsworksStackConfigNoVpcCreate(name), testAccAwsOpsworksCustomLayerSecurityGroups(name))
}
<|endoftext|>