<commit_before>package kube

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"github.com/Sirupsen/logrus"
	"github.com/ghodss/yaml"
	"io/ioutil"
	"net/http"
)

// KubeClient is the interface to access the kubernetes job API
type KubeClient interface {
	CreateJob(job *Job) error
	GetSecret(namespace string, secretName string) (map[string]string, error)
	DeployResourceFile(resourceFile []byte) error
}

// concrete implementation of a job client
type realKubeClient struct {
	*http.Client
	address string
	token   string
}

// NewClient returns a new KubeClient connecting to the address. This uses the service
// account credentials
func NewClient(address string) (KubeClient, error) {
	// create tls client
	cacertFile := "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
	capem, err := ioutil.ReadFile(cacertFile)
	if err != nil {
		return nil, err
	}
	cacert := x509.NewCertPool()
	if !cacert.AppendCertsFromPEM(capem) {
		return nil, fmt.Errorf("unable to load certificate authority")
	}
	config := &tls.Config{RootCAs: cacert}
	transport := &http.Transport{TLSClientConfig: config}

	// read token
	client := &http.Client{Transport: transport}
	tokenFile := "/var/run/secrets/kubernetes.io/serviceaccount/token"
	token, err := ioutil.ReadFile(tokenFile)
	if err != nil {
		return nil, err
	}
	return &realKubeClient{client, address, string(token)}, nil
}

// Create a new kubernetes Job with the given job
func (r *realKubeClient) CreateJob(job *Job) error {
	url := "/apis/extensions/v1beta1/namespaces/" + job.Metadata["namespace"].(string) + "/jobs"
	data, err := json.Marshal(job)
	if err != nil {
		return err
	}
	byteData := bytes.NewReader(data)
	return r.doPost(url, byteData)

}

// DeployResourceFile deploys a Kubernbetes YAML spec file as Kubernetes resources
func (r *realKubeClient) DeployResourceFile(resourceFile []byte) error {

	resources := strings.Split(string(resourceFile), "---")

	for _, resource := range resources {

		// if it's empty, skip
		if strings.TrimSpace(resource) == "" {
			continue
		}

		logrus.Info("deploying to kubernetes: ", resource)

		data, err := yaml.YAMLToJSON([]byte(resource))
		if err != nil {
			logrus.WithError(err).Error("unable to convert yaml to json")
			return err
		}

		var out map[string]interface{}
		err = json.Unmarshal(data, &out)
		if err != nil {
			logrus.WithError(err).Error("unable to unmarshal json to map")
			return err
		}

		logrus.Info("parsed: ", out)

		// if unmarshalled data is nil, skip
		if out == nil {
			continue
		}

		kind := strings.ToLower(out["kind"].(string)) + "s"
		metadata := out["metadata"]
		namespace := "default"
		if metadata != nil {
			namespace = metadata.(map[interface{}]interface{})["namespace"].(string)
		}

		// endpoint is /api/v1/namespaces/{namespace}/{resourceType}
		uri := fmt.Sprintf("/api/v1/namespaces/%s/%s", namespace, kind)
		err = r.doPost(uri, bytes.NewReader(data))
		if err != nil {
			logrus.WithError(err).Error("unable to POST data")
			return err
		}
	}

	return nil
}

//Get secret with a given namespace and secret name
func (r *realKubeClient) GetSecret(namespace string, secretName string) (map[string]string, error) {
	secret := &Secret{}
	secrets := make(map[string]string)

	uri := "/api/v1/namespaces/" + namespace + "/secrets/" + secretName

	err := r.doGet(uri, secret)
	if err != nil {
		return nil, err
	}

	for key, value := range secret.Data {
		decodedValue, err := base64.StdEncoding.DecodeString(value)

		if err != nil {
			logrus.WithError(err).Println("Unable to decode secret", key, value)
			continue
		}
		secrets[key] = string(decodedValue)
	}
	return secrets, nil
}

func (r *realKubeClient) doGet(uri string, response interface{}) error {
	req, err := r.createRequest("GET", uri, nil)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return fmt.Errorf("%d: %s", res.StatusCode, string(body))
	}
	err = json.Unmarshal(body, response)
	if err != nil {
		return err
	}
	return nil
}

func (r *realKubeClient) doPost(uri string, data io.Reader) error {
	req, err := r.createRequest("POST", uri, data)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode == http.StatusCreated || res.StatusCode == http.StatusOK {
		return nil
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	return fmt.Errorf("%d: %s", res.StatusCode, string(body))
}

func (r *realKubeClient) createRequest(method, uri string, data io.Reader) (*http.Request, error) {
	req, err := http.NewRequest(method, r.address+uri, data)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Authorization", "Bearer "+r.token)
	return req, nil
}
<commit_msg>fixed wrong type<commit_after>package kube

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"github.com/Sirupsen/logrus"
	"github.com/ghodss/yaml"
	"io/ioutil"
	"net/http"
)

// KubeClient is the interface to access the kubernetes job API
type KubeClient interface {
	CreateJob(job *Job) error
	GetSecret(namespace string, secretName string) (map[string]string, error)
	DeployResourceFile(resourceFile []byte) error
}

// concrete implementation of a job client
type realKubeClient struct {
	*http.Client
	address string
	token   string
}

// NewClient returns a new KubeClient connecting to the address. This uses the service
// account credentials
func NewClient(address string) (KubeClient, error) {
	// create tls client
	cacertFile := "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
	capem, err := ioutil.ReadFile(cacertFile)
	if err != nil {
		return nil, err
	}
	cacert := x509.NewCertPool()
	if !cacert.AppendCertsFromPEM(capem) {
		return nil, fmt.Errorf("unable to load certificate authority")
	}
	config := &tls.Config{RootCAs: cacert}
	transport := &http.Transport{TLSClientConfig: config}

	// read token
	client := &http.Client{Transport: transport}
	tokenFile := "/var/run/secrets/kubernetes.io/serviceaccount/token"
	token, err := ioutil.ReadFile(tokenFile)
	if err != nil {
		return nil, err
	}
	return &realKubeClient{client, address, string(token)}, nil
}

// Create a new kubernetes Job with the given job
func (r *realKubeClient) CreateJob(job *Job) error {
	url := "/apis/extensions/v1beta1/namespaces/" + job.Metadata["namespace"].(string) + "/jobs"
	data, err := json.Marshal(job)
	if err != nil {
		return err
	}
	byteData := bytes.NewReader(data)
	return r.doPost(url, byteData)

}

// DeployResourceFile deploys a Kubernbetes YAML spec file as Kubernetes resources
func (r *realKubeClient) DeployResourceFile(resourceFile []byte) error {

	resources := strings.Split(string(resourceFile), "---")

	for _, resource := range resources {

		// if it's empty, skip
		if strings.TrimSpace(resource) == "" {
			continue
		}

		logrus.Info("deploying to kubernetes: ", resource)

		data, err := yaml.YAMLToJSON([]byte(resource))
		if err != nil {
			logrus.WithError(err).Error("unable to convert yaml to json")
			return err
		}

		var out map[string]interface{}
		err = json.Unmarshal(data, &out)
		if err != nil {
			logrus.WithError(err).Error("unable to unmarshal json to map")
			return err
		}

		logrus.Info("parsed: ", out)

		// if unmarshalled data is nil, skip
		if out == nil {
			continue
		}

		kind := strings.ToLower(out["kind"].(string)) + "s"
		metadata := out["metadata"]
		namespace := "default"
		if metadata != nil {
			namespace = metadata.(map[string]interface{})["namespace"].(string)
		}

		// endpoint is /api/v1/namespaces/{namespace}/{resourceType}
		uri := fmt.Sprintf("/api/v1/namespaces/%s/%s", namespace, kind)
		err = r.doPost(uri, bytes.NewReader(data))
		if err != nil {
			logrus.WithError(err).Error("unable to POST data")
			return err
		}
	}

	return nil
}

//Get secret with a given namespace and secret name
func (r *realKubeClient) GetSecret(namespace string, secretName string) (map[string]string, error) {
	secret := &Secret{}
	secrets := make(map[string]string)

	uri := "/api/v1/namespaces/" + namespace + "/secrets/" + secretName

	err := r.doGet(uri, secret)
	if err != nil {
		return nil, err
	}

	for key, value := range secret.Data {
		decodedValue, err := base64.StdEncoding.DecodeString(value)

		if err != nil {
			logrus.WithError(err).Println("Unable to decode secret", key, value)
			continue
		}
		secrets[key] = string(decodedValue)
	}
	return secrets, nil
}

func (r *realKubeClient) doGet(uri string, response interface{}) error {
	req, err := r.createRequest("GET", uri, nil)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return fmt.Errorf("%d: %s", res.StatusCode, string(body))
	}
	err = json.Unmarshal(body, response)
	if err != nil {
		return err
	}
	return nil
}

func (r *realKubeClient) doPost(uri string, data io.Reader) error {
	req, err := r.createRequest("POST", uri, data)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode == http.StatusCreated || res.StatusCode == http.StatusOK {
		return nil
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	return fmt.Errorf("%d: %s", res.StatusCode, string(body))
}

func (r *realKubeClient) createRequest(method, uri string, data io.Reader) (*http.Request, error) {
	req, err := http.NewRequest(method, r.address+uri, data)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Authorization", "Bearer "+r.token)
	return req, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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.
*/

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main // import "github.com/jmccarty3/kube2consul"

import (
	"flag"
	"fmt"
	"strings"
	"net/url"
	"os"
	"time"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	kcache "k8s.io/kubernetes/pkg/client/unversioned/cache"
	kclientcmd "k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	kframework "k8s.io/kubernetes/pkg/controller/framework"
	kcontrollerFramework "k8s.io/kubernetes/pkg/controller/framework"
	kSelector "k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/util"
	"github.com/golang/glog"
	consulapi "github.com/hashicorp/consul/api"

  "k8s.io/kubernetes/pkg/api"
  "k8s.io/kubernetes/pkg/fields"
  "k8s.io/kubernetes/pkg/labels"
)

var (
	argConsulAgent         = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile         = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl       = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

type kube2consul struct {
	// Consul client.
	consulClient *consulapi.Client
	// DNS domain name.
	domain string

	//Nodes Name / valid
	nodes map[string]bool

	//DNS IDS
	ids map[string][]string
}

func Newkube2consul() *kube2consul {
	var k kube2consul
	k.nodes = make(map[string]bool)
	k.ids = make(map[string][]string)

	return &k
}

func (ks *kube2consul) removeDNS(record string) error {
	glog.V(2).Infof("Removing %s from DNS", record)
	return ks.consulClient.Agent().ServiceDeregister(record)
}

func (ks *kube2consul) addDNS(record string, service *kapi.Service) error {
	if strings.Contains(record, ".") {
		glog.V(1).Infof("Service names containing '.' are not supported: %s\n", service.Name)
		return nil
	}

	// if ClusterIP is not set, do not create a DNS records
	if !kapi.IsServiceIPSet(service) {
		glog.V(1).Infof("Skipping dns record for headless service: %s\n", service.Name)
		return nil
	}

	for i := range service.Spec.Ports {
		for n,s := range ks.nodes {
			if s {
				newId := n+record + service.Spec.Ports[i].Name

				asr := &consulapi.AgentServiceRegistration{
					ID:			 newId,
					Name: 	 record + "-" + service.Spec.Ports[i].Name,
					Address: n,
					Port:    service.Spec.Ports[i].NodePort,
				}

				glog.V(2).Infof("Setting DNS record: %v -> %v:%d\n", record, service.Spec.ClusterIP, service.Spec.Ports[i].Port)
				if err := ks.consulClient.Agent().ServiceRegister(asr); err != nil {
					return err
				}
				ks.ids[record] = append(ks.ids[record], newId)
			}
		}
	}
	return nil
}

func newConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
			glog.Infof("Error parsing Consul url")
			return nil, err
	}

	if consulAgentUrl.Host != "" {
	  consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
			glog.Infof("Error creating Consul client")
			return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}
	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}
	return parsedUrl.String(), nil
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	var config *kclient.Config
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}
	if *argKubecfgFile == "" {
		config = &kclient.Config{
			Host:    masterUrl,
			Version: "v1",
		}
	} else {
		var err error
		overrides := &kclientcmd.ConfigOverrides{}
		overrides.ClusterInfo.Server = masterUrl                                     // might be "", but that is OK
		rules := &kclientcmd.ClientConfigLoadingRules{ExplicitPath: *argKubecfgFile} // might be "", but that is OK
		if config, err = kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides).ClientConfig(); err != nil {
			return nil, err
		}
	}
	glog.Infof("Using %s for kubernetes master", config.Host)
	glog.Infof("Using kubernetes API %s", config.Version)
	return kclient.New(config)
}

func buildNameString(service, namespace string) string {
	//glog.Infof("Name String: %s  %s", service, namespace)
	return fmt.Sprintf("%s.%s", service, namespace)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

// Returns a cache.ListWatch that gets all changes to services.
func createNodeLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "nodes", kapi.NamespaceAll, kSelector.Everything())
}

func (ks *kube2consul) newService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.addDNS(s.Name, s); err != nil {
			glog.V(1).Infof("Failed to add service: %v due to: %v", name, err)
		}
	}
}

func (ks *kube2consul) removeService(obj interface{}) {
	glog.Info("Service remove")
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.removeDNS(s.Name); err != nil {
			glog.V(1).Infof("Failed to remove service: %v due to: %v", name, err)
		}
	}
}

func (ks *kube2consul) updateNode(oldObj, newObj interface{}) {
	if n, ok := oldObj.(*kapi.Node); ok {
		name := n.Name
		ready := n.Status.Conditions[0].Status == kapi.ConditionTrue

		ks.nodes[name] = ready
	}

	if n, ok := newObj.(*kapi.Node); ok {
		name := n.Name
		ready := n.Status.Conditions[0].Status == kapi.ConditionTrue

		ks.nodes[name] = ready
	}
}

func watchForServices(kubeClient *kclient.Client, ks *kube2consul) {
	var serviceController *kcontrollerFramework.Controller
	_, serviceController = kframework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    ks.newService,
			DeleteFunc: ks.removeService,
			UpdateFunc: func(oldObj, newObj interface{}) {
				ks.newService(newObj)
			},
		},
	)
	glog.Info("About to call run!")
	go serviceController.Run(util.NeverStop)
}

func watchForNodes(kubeClient *kclient.Client, ks *kube2consul) kcache.Store {
	store, serviceController := kframework.NewInformer(
		createNodeLW(kubeClient),
		&kapi.Node{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    func(newObj interface{}) {
				glog.Info("Adding node!")
			},
			DeleteFunc: func(oldObj interface{}) {
				glog.Info("Node Removed!!")
			},
			UpdateFunc: ks.updateNode,
		},
	)
	glog.Info("About to call run!")
	go serviceController.Run(util.NeverStop)
	return store
}

func main() {
	flag.Parse()
	var err error
	// TODO: Validate input flags.
	ks := Newkube2consul()

	if ks.consulClient, err = newConsulClient(*argConsulAgent); err != nil {
		glog.Fatalf("Failed to create Consul client - %v", err)
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		glog.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	glog.Info(kubeClient.ServerVersion())
	glog.Info(kubeClient.Services(kapi.NamespaceAll).Get("sensu-core"))

	pods, err := kubeClient.Pods(api.NamespaceDefault).List(labels.Everything(), fields.Everything())
	if err != nil {
	  for pod := range pods.Items {
			glog.Info(pod)
		}
	}

	watchForServices(kubeClient, ks)
	watchForNodes(kubeClient, ks)
	glog.Info("Watchers running")
	select {}
}
<commit_msg>Updated service removal<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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.
*/

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main // import "github.com/jmccarty3/kube2consul"

import (
	"flag"
	"fmt"
	"strings"
	"net/url"
	"os"
	"time"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	kcache "k8s.io/kubernetes/pkg/client/unversioned/cache"
	kclientcmd "k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	kframework "k8s.io/kubernetes/pkg/controller/framework"
	kcontrollerFramework "k8s.io/kubernetes/pkg/controller/framework"
	kSelector "k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/util"
	"github.com/golang/glog"
	consulapi "github.com/hashicorp/consul/api"

  "k8s.io/kubernetes/pkg/api"
  "k8s.io/kubernetes/pkg/fields"
  "k8s.io/kubernetes/pkg/labels"
)

var (
	argConsulAgent         = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile         = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl       = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

type kube2consul struct {
	// Consul client.
	consulClient *consulapi.Client

	// DNS domain name.
	domain string

	//Nodes Name / valid
	nodes map[string]bool

	//DNS IDS
	ids map[string][]string
}

func Newkube2consul() *kube2consul {
	var k kube2consul
	k.nodes = make(map[string]bool)
	k.ids = make(map[string][]string)

	return &k
}

func (ks *kube2consul) removeDNS(record string) error {
	glog.V(2).Infof("Removing %s from DNS", record)
	for _,id := range ks.ids[record] {
		ks.consulClient.Agent().ServiceDeregister(id)
	}
	ks.ids[record] = []string{}
	return nil
}

func (ks *kube2consul) addDNS(record string, service *kapi.Service) error {
	if strings.Contains(record, ".") {
		glog.V(1).Infof("Service names containing '.' are not supported: %s\n", service.Name)
		return nil
	}

	// if ClusterIP is not set, do not create a DNS records
	if !kapi.IsServiceIPSet(service) {
		glog.V(1).Infof("Skipping dns record for headless service: %s\n", service.Name)
		return nil
	}

	for i := range service.Spec.Ports {
		for n,s := range ks.nodes {
			if s {
				newId := n+record + service.Spec.Ports[i].Name

				asr := &consulapi.AgentServiceRegistration{
					ID:			 newId,
					Name: 	 record + "-" + service.Spec.Ports[i].Name,
					Address: n,
					//Port:    service.Spec.Ports[i].NodePort,
				}

				glog.V(2).Infof("Setting DNS record: %v -> %v:%d\n", record, service.Spec.ClusterIP, service.Spec.Ports[i].Port)
				if err := ks.consulClient.Agent().ServiceRegister(asr); err != nil {
					return err
				}
				ks.ids[record] = append(ks.ids[record], newId)
			}
		}
	}
	return nil
}

func newConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
			glog.Infof("Error parsing Consul url")
			return nil, err
	}

	if consulAgentUrl.Host != "" {
	  consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
			glog.Infof("Error creating Consul client")
			return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}
	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}
	return parsedUrl.String(), nil
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	var config *kclient.Config
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}
	if *argKubecfgFile == "" {
		config = &kclient.Config{
			Host:    masterUrl,
			Version: "v1",
		}
	} else {
		var err error
		overrides := &kclientcmd.ConfigOverrides{}
		overrides.ClusterInfo.Server = masterUrl                                     // might be "", but that is OK
		rules := &kclientcmd.ClientConfigLoadingRules{ExplicitPath: *argKubecfgFile} // might be "", but that is OK
		if config, err = kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides).ClientConfig(); err != nil {
			return nil, err
		}
	}
	glog.Infof("Using %s for kubernetes master", config.Host)
	glog.Infof("Using kubernetes API %s", config.Version)
	return kclient.New(config)
}

func buildNameString(service, namespace string) string {
	//glog.Infof("Name String: %s  %s", service, namespace)
	return fmt.Sprintf("%s.%s", service, namespace)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

// Returns a cache.ListWatch that gets all changes to services.
func createNodeLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "nodes", kapi.NamespaceAll, kSelector.Everything())
}

func (ks *kube2consul) newService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.addDNS(s.Name, s); err != nil {
			glog.V(1).Infof("Failed to add service: %v due to: %v", name, err)
		}
	}
}

func (ks *kube2consul) removeService(obj interface{}) {
	glog.Info("Service remove")
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.removeDNS(s.Name); err != nil {
			glog.V(1).Infof("Failed to remove service: %v due to: %v", name, err)
		}
	}
}

func (ks *kube2consul) updateNode(oldObj, newObj interface{}) {
	if n, ok := oldObj.(*kapi.Node); ok {
		name := n.Name
		ready := n.Status.Conditions[0].Status == kapi.ConditionTrue

		ks.nodes[name] = ready
	}

	if n, ok := newObj.(*kapi.Node); ok {
		name := n.Name
		ready := n.Status.Conditions[0].Status == kapi.ConditionTrue

		ks.nodes[name] = ready
	}
}

func watchForServices(kubeClient *kclient.Client, ks *kube2consul) {
	var serviceController *kcontrollerFramework.Controller
	_, serviceController = kframework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    ks.newService,
			DeleteFunc: ks.removeService,
			UpdateFunc: func(oldObj, newObj interface{}) {
				ks.newService(newObj)
			},
		},
	)
	glog.Info("About to call run!")
	go serviceController.Run(util.NeverStop)
}

func watchForNodes(kubeClient *kclient.Client, ks *kube2consul) kcache.Store {
	store, serviceController := kframework.NewInformer(
		createNodeLW(kubeClient),
		&kapi.Node{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    func(newObj interface{}) {
				glog.Info("Adding node!")
			},
			DeleteFunc: func(oldObj interface{}) {
				glog.Info("Node Removed!!")
			},
			UpdateFunc: ks.updateNode,
		},
	)
	glog.Info("About to call run!")
	go serviceController.Run(util.NeverStop)
	return store
}

func main() {
	flag.Parse()
	var err error
	// TODO: Validate input flags.
	ks := Newkube2consul()

	if ks.consulClient, err = newConsulClient(*argConsulAgent); err != nil {
		glog.Fatalf("Failed to create Consul client - %v", err)
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		glog.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	glog.Info(kubeClient.ServerVersion())
	glog.Info(kubeClient.Services(kapi.NamespaceAll).Get("sensu-core"))

	pods, err := kubeClient.Pods(api.NamespaceDefault).List(labels.Everything(), fields.Everything())
	if err != nil {
	  for pod := range pods.Items {
			glog.Info(pod)
		}
	}

	watchForServices(kubeClient, ks)
	watchForNodes(kubeClient, ks)
	glog.Info("Watchers running")
	select {}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 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 main

import (
	"fmt"
	"os"
	"strings"

	"github.com/spf13/cobra"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/registry"
	"k8s.io/kops/util/pkg/tables"
	k8sapi "k8s.io/kubernetes/pkg/api"
)

type GetClustersCmd struct {
	FullSpec bool
}

var getClustersCmd GetClustersCmd

func init() {
	cmd := &cobra.Command{
		Use:     "clusters",
		Aliases: []string{"cluster"},
		Short:   "get clusters",
		Long:    `List or get clusters.`,
		Run: func(cmd *cobra.Command, args []string) {
			err := getClustersCmd.Run(args)
			if err != nil {
				exitWithError(err)
			}
		},
	}

	getCmd.cobraCommand.AddCommand(cmd)

	cmd.Flags().BoolVar(&getClustersCmd.FullSpec, "full", false, "Show fully populated configuration")
}

func (c *GetClustersCmd) Run(args []string) error {
	client, err := rootCommand.Clientset()
	if err != nil {
		return err
	}

	clusterList, err := client.Clusters().List(k8sapi.ListOptions{})
	if err != nil {
		return err
	}

	var clusters []*api.Cluster
	if len(args) != 0 {
		m := make(map[string]*api.Cluster)
		for i := range clusterList.Items {
			c := &clusterList.Items[i]
			m[c.ObjectMeta.Name] = c
		}
		for _, arg := range args {
			ig := m[arg]
			if ig == nil {
				return fmt.Errorf("cluster not found %q", arg)
			}

			clusters = append(clusters, ig)
		}
	} else {
		for i := range clusterList.Items {
			c := &clusterList.Items[i]
			clusters = append(clusters, c)
		}
	}

	if len(clusters) == 0 {
		fmt.Fprintf(os.Stderr, "No clusters found\n")
		return nil
	}

	if c.FullSpec {
		var err error
		clusters, err = fullClusterSpecs(clusters)
		if err != nil {
			return err
		}
	}

	switch getCmd.output {
	case OutputTable:

		t := &tables.Table{}
		t.AddColumn("NAME", func(c *api.Cluster) string {
			return c.ObjectMeta.Name
		})
		t.AddColumn("CLOUD", func(c *api.Cluster) string {
			return c.Spec.CloudProvider
		})
		t.AddColumn("SUBNETS", func(c *api.Cluster) string {
			var subnetNames []string
			for _, s := range c.Spec.Subnets {
				subnetNames = append(subnetNames, s.Name)
			}
			return strings.Join(subnetNames, ",")
		})
		return t.Render(clusters, os.Stdout, "NAME", "CLOUD", "SUBNETS")

	case OutputYaml:
		for _, cluster := range clusters {
			if err := marshalToWriter(cluster, marshalYaml, os.Stdout); err != nil {
				return err
			}
		}
		return nil
	case OutputJSON:
		for _, cluster := range clusters {
			if err := marshalToWriter(cluster, marshalJSON, os.Stdout); err != nil {
				return err
			}
		}
		return nil

	default:
		return fmt.Errorf("Unknown output format: %q", getCmd.output)
	}
}

func fullClusterSpecs(clusters []*api.Cluster) ([]*api.Cluster, error) {
	var fullSpecs []*api.Cluster
	for _, cluster := range clusters {
		configBase, err := registry.ConfigBase(cluster)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		fullSpec := &api.Cluster{}
		err = registry.ReadConfigDeprecated(configBase.Join(registry.PathClusterCompleted), fullSpec)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		fullSpecs = append(fullSpecs, fullSpec)
	}
	return fullSpecs, nil
}
<commit_msg>Respect the name flag on cluster get<commit_after>/*
Copyright 2016 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 main

import (
	"fmt"
	"os"
	"strings"

	"github.com/spf13/cobra"
	"io"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/registry"
	"k8s.io/kops/util/pkg/tables"
	k8sapi "k8s.io/kubernetes/pkg/api"
)

type GetClusterOptions struct {
	// FullSpec determines if we should output the completed (fully populated) spec
	FullSpec bool

	// ClusterNames is a list of cluster names to show; if not specified all clusters will be shown
	ClusterNames []string
}

func init() {
	var options GetClusterOptions

	cmd := &cobra.Command{
		Use:     "clusters",
		Aliases: []string{"cluster"},
		Short:   "get clusters",
		Long:    `List or get clusters.`,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 0 {
				options.ClusterNames = append(options.ClusterNames, args...)
			}

			if rootCommand.clusterName != "" {
				if len(args) != 0 {
					exitWithError(fmt.Errorf("cannot mix --name for cluster with positional arguments"))
				}

				options.ClusterNames = append(options.ClusterNames, rootCommand.clusterName)
			}

			err := RunGetClusters(&rootCommand, os.Stdout, &options)
			if err != nil {
				exitWithError(err)
			}
		},
	}

	cmd.Flags().BoolVar(&options.FullSpec, "full", options.FullSpec, "Show fully populated configuration")

	getCmd.cobraCommand.AddCommand(cmd)
}

func RunGetClusters(context Factory, out io.Writer, options *GetClusterOptions) error {
	client, err := context.Clientset()
	if err != nil {
		return err
	}

	clusterList, err := client.Clusters().List(k8sapi.ListOptions{})
	if err != nil {
		return err
	}

	var clusters []*api.Cluster
	if len(options.ClusterNames) != 0 {
		m := make(map[string]*api.Cluster)
		for i := range clusterList.Items {
			c := &clusterList.Items[i]
			m[c.ObjectMeta.Name] = c
		}
		for _, clusterName := range options.ClusterNames {
			c := m[clusterName]
			if c == nil {
				return fmt.Errorf("cluster not found %q", clusterName)
			}

			clusters = append(clusters, c)
		}
	} else {
		for i := range clusterList.Items {
			c := &clusterList.Items[i]
			clusters = append(clusters, c)
		}
	}

	if len(clusters) == 0 {
		fmt.Fprintf(os.Stderr, "No clusters found\n")
		return nil
	}

	if options.FullSpec {
		var err error
		clusters, err = fullClusterSpecs(clusters)
		if err != nil {
			return err
		}
	}

	switch getCmd.output {
	case OutputTable:

		t := &tables.Table{}
		t.AddColumn("NAME", func(c *api.Cluster) string {
			return c.ObjectMeta.Name
		})
		t.AddColumn("CLOUD", func(c *api.Cluster) string {
			return c.Spec.CloudProvider
		})
		t.AddColumn("SUBNETS", func(c *api.Cluster) string {
			var subnetNames []string
			for _, s := range c.Spec.Subnets {
				subnetNames = append(subnetNames, s.Name)
			}
			return strings.Join(subnetNames, ",")
		})
		return t.Render(clusters, out, "NAME", "CLOUD", "SUBNETS")

	case OutputYaml:
		for _, cluster := range clusters {
			if err := marshalToWriter(cluster, marshalYaml, out); err != nil {
				return err
			}
		}
		return nil
	case OutputJSON:
		for _, cluster := range clusters {
			if err := marshalToWriter(cluster, marshalJSON, out); err != nil {
				return err
			}
		}
		return nil

	default:
		return fmt.Errorf("Unknown output format: %q", getCmd.output)
	}
}

func fullClusterSpecs(clusters []*api.Cluster) ([]*api.Cluster, error) {
	var fullSpecs []*api.Cluster
	for _, cluster := range clusters {
		configBase, err := registry.ConfigBase(cluster)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		fullSpec := &api.Cluster{}
		err = registry.ReadConfigDeprecated(configBase.Join(registry.PathClusterCompleted), fullSpec)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		fullSpecs = append(fullSpecs, fullSpec)
	}
	return fullSpecs, nil
}
<|endoftext|>


<commit_before>// +build debug

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"

	"github.com/spf13/cobra"

	"github.com/restic/restic/internal/errors"
	"github.com/restic/restic/internal/pack"
	"github.com/restic/restic/internal/repository"
	"github.com/restic/restic/internal/restic"
)

var cmdDebug = &cobra.Command{
	Use:   "debug",
	Short: "Debug commands",
}

var cmdDebugDump = &cobra.Command{
	Use:   "dump [indexes|snapshots|all|packs]",
	Short: "Dump data structures",
	Long: `
The "dump" command dumps data structures from the repository as JSON objects. It
is used for debugging purposes only.

EXIT STATUS
===========

Exit status is 0 if the command was successful, and non-zero if there was any error.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runDebugDump(globalOptions, args)
	},
}

func init() {
	cmdRoot.AddCommand(cmdDebug)
	cmdDebug.AddCommand(cmdDebugDump)
}

func prettyPrintJSON(wr io.Writer, item interface{}) error {
	buf, err := json.MarshalIndent(item, "", "  ")
	if err != nil {
		return err
	}

	_, err = wr.Write(append(buf, '\n'))
	return err
}

func debugPrintSnapshots(ctx context.Context, repo *repository.Repository, wr io.Writer) error {
	return restic.ForAllSnapshots(ctx, repo, nil, func(id restic.ID, snapshot *restic.Snapshot, err error) error {
		if err != nil {
			return err
		}

		fmt.Fprintf(wr, "snapshot_id: %v\n", id)

		return prettyPrintJSON(wr, snapshot)
	})
}

// Pack is the struct used in printPacks.
type Pack struct {
	Name string `json:"name"`

	Blobs []Blob `json:"blobs"`
}

// Blob is the struct used in printPacks.
type Blob struct {
	Type   restic.BlobType `json:"type"`
	Length uint            `json:"length"`
	ID     restic.ID       `json:"id"`
	Offset uint            `json:"offset"`
}

func printPacks(ctx context.Context, repo *repository.Repository, wr io.Writer) error {

	return repo.List(ctx, restic.PackFile, func(id restic.ID, size int64) error {
		h := restic.Handle{Type: restic.PackFile, Name: id.String()}

		blobs, _, err := pack.List(repo.Key(), restic.ReaderAt(ctx, repo.Backend(), h), size)
		if err != nil {
			Warnf("error for pack %v: %v\n", id.Str(), err)
			return nil
		}

		p := Pack{
			Name:  id.String(),
			Blobs: make([]Blob, len(blobs)),
		}
		for i, blob := range blobs {
			p.Blobs[i] = Blob{
				Type:   blob.Type,
				Length: blob.Length,
				ID:     blob.ID,
				Offset: blob.Offset,
			}
		}

		return prettyPrintJSON(wr, p)
	})
}

func dumpIndexes(ctx context.Context, repo restic.Repository, wr io.Writer) error {
	return repository.ForAllIndexes(ctx, repo, func(id restic.ID, idx *repository.Index, oldFormat bool, err error) error {
		Printf("index_id: %v\n", id)
		if err != nil {
			return err
		}

		return idx.Dump(wr)
	})
}

func runDebugDump(gopts GlobalOptions, args []string) error {
	if len(args) != 1 {
		return errors.Fatal("type not specified")
	}

	repo, err := OpenRepository(gopts)
	if err != nil {
		return err
	}

	if !gopts.NoLock {
		lock, err := lockRepo(gopts.ctx, repo)
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	}

	tpe := args[0]

	switch tpe {
	case "indexes":
		return dumpIndexes(gopts.ctx, repo, gopts.stdout)
	case "snapshots":
		return debugPrintSnapshots(gopts.ctx, repo, gopts.stdout)
	case "packs":
		return printPacks(gopts.ctx, repo, gopts.stdout)
	case "all":
		Printf("snapshots:\n")
		err := debugPrintSnapshots(gopts.ctx, repo, gopts.stdout)
		if err != nil {
			return err
		}

		Printf("\nindexes:\n")
		err = dumpIndexes(gopts.ctx, repo, gopts.stdout)
		if err != nil {
			return err
		}

		return nil
	default:
		return errors.Fatalf("no such type %q", tpe)
	}
}
<commit_msg>Add 'debug examine' command to debug #1999<commit_after>// +build debug

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sort"

	"github.com/spf13/cobra"

	"github.com/restic/restic/internal/errors"
	"github.com/restic/restic/internal/pack"
	"github.com/restic/restic/internal/repository"
	"github.com/restic/restic/internal/restic"
)

var cmdDebug = &cobra.Command{
	Use:   "debug",
	Short: "Debug commands",
}

var cmdDebugDump = &cobra.Command{
	Use:   "dump [indexes|snapshots|all|packs]",
	Short: "Dump data structures",
	Long: `
The "dump" command dumps data structures from the repository as JSON objects. It
is used for debugging purposes only.

EXIT STATUS
===========

Exit status is 0 if the command was successful, and non-zero if there was any error.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runDebugDump(globalOptions, args)
	},
}

func init() {
	cmdRoot.AddCommand(cmdDebug)
	cmdDebug.AddCommand(cmdDebugDump)
	cmdDebug.AddCommand(cmdDebugExamine)
}

func prettyPrintJSON(wr io.Writer, item interface{}) error {
	buf, err := json.MarshalIndent(item, "", "  ")
	if err != nil {
		return err
	}

	_, err = wr.Write(append(buf, '\n'))
	return err
}

func debugPrintSnapshots(ctx context.Context, repo *repository.Repository, wr io.Writer) error {
	return restic.ForAllSnapshots(ctx, repo, nil, func(id restic.ID, snapshot *restic.Snapshot, err error) error {
		if err != nil {
			return err
		}

		fmt.Fprintf(wr, "snapshot_id: %v\n", id)

		return prettyPrintJSON(wr, snapshot)
	})
}

// Pack is the struct used in printPacks.
type Pack struct {
	Name string `json:"name"`

	Blobs []Blob `json:"blobs"`
}

// Blob is the struct used in printPacks.
type Blob struct {
	Type   restic.BlobType `json:"type"`
	Length uint            `json:"length"`
	ID     restic.ID       `json:"id"`
	Offset uint            `json:"offset"`
}

func printPacks(ctx context.Context, repo *repository.Repository, wr io.Writer) error {

	return repo.List(ctx, restic.PackFile, func(id restic.ID, size int64) error {
		h := restic.Handle{Type: restic.PackFile, Name: id.String()}

		blobs, _, err := pack.List(repo.Key(), restic.ReaderAt(ctx, repo.Backend(), h), size)
		if err != nil {
			Warnf("error for pack %v: %v\n", id.Str(), err)
			return nil
		}

		p := Pack{
			Name:  id.String(),
			Blobs: make([]Blob, len(blobs)),
		}
		for i, blob := range blobs {
			p.Blobs[i] = Blob{
				Type:   blob.Type,
				Length: blob.Length,
				ID:     blob.ID,
				Offset: blob.Offset,
			}
		}

		return prettyPrintJSON(wr, p)
	})
}

func dumpIndexes(ctx context.Context, repo restic.Repository, wr io.Writer) error {
	return repository.ForAllIndexes(ctx, repo, func(id restic.ID, idx *repository.Index, oldFormat bool, err error) error {
		Printf("index_id: %v\n", id)
		if err != nil {
			return err
		}

		return idx.Dump(wr)
	})
}

func runDebugDump(gopts GlobalOptions, args []string) error {
	if len(args) != 1 {
		return errors.Fatal("type not specified")
	}

	repo, err := OpenRepository(gopts)
	if err != nil {
		return err
	}

	if !gopts.NoLock {
		lock, err := lockRepo(gopts.ctx, repo)
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	}

	tpe := args[0]

	switch tpe {
	case "indexes":
		return dumpIndexes(gopts.ctx, repo, gopts.stdout)
	case "snapshots":
		return debugPrintSnapshots(gopts.ctx, repo, gopts.stdout)
	case "packs":
		return printPacks(gopts.ctx, repo, gopts.stdout)
	case "all":
		Printf("snapshots:\n")
		err := debugPrintSnapshots(gopts.ctx, repo, gopts.stdout)
		if err != nil {
			return err
		}

		Printf("\nindexes:\n")
		err = dumpIndexes(gopts.ctx, repo, gopts.stdout)
		if err != nil {
			return err
		}

		return nil
	default:
		return errors.Fatalf("no such type %q", tpe)
	}
}

var cmdDebugExamine = &cobra.Command{
	Use:               "examine",
	Short:             "Examine a pack file",
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runDebugExamine(globalOptions, args)
	},
}

func loadBlobs(ctx context.Context, repo restic.Repository, pack string, list []restic.PackedBlob) error {
	be := repo.Backend()
	for _, blob := range list {
		fmt.Printf("      loading blob %v at %v (length %v)\n", blob.ID, blob.Offset, blob.Length)
		buf := make([]byte, blob.Length)
		h := restic.Handle{
			Name: pack,
			Type: restic.PackFile,
		}
		err := be.Load(ctx, h, int(blob.Length), int64(blob.Offset), func(rd io.Reader) error {
			n, err := io.ReadFull(rd, buf)
			if err != nil {
				fmt.Fprintf(os.Stderr, "read error after %d bytes: %v\n", n, err)
				return err
			}
			return nil
		})
		if err != nil {
			fmt.Fprintf(os.Stderr, "error read: %v\n", err)
			continue
		}

		key := repo.Key()

		nonce, buf := buf[:key.NonceSize()], buf[key.NonceSize():]
		buf, err = key.Open(buf[:0], nonce, buf, nil)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error decrypting blob: %v\n", err)
			continue
		}

		id := restic.Hash(buf)
		fmt.Printf("         successfully decrypted blob (length %v), hash is %v\n", len(buf), id)
		if !id.Equal(blob.ID) {
			fmt.Printf("         IDs do not match, want %v, got %v\n", blob.ID, id)
		} else {
			fmt.Printf("         IDs match\n")
		}
	}

	return nil
}

func runDebugExamine(gopts GlobalOptions, args []string) error {
	repo, err := OpenRepository(gopts)
	if err != nil {
		return err
	}

	if !gopts.NoLock {
		lock, err := lockRepo(gopts.ctx, repo)
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	}

	err = repo.LoadIndex(gopts.ctx)
	if err != nil {
		return err
	}

	for _, name := range args {
		fmt.Printf("examine %v\n", name)
		id, err := restic.ParseID(name)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
		}

		h := restic.Handle{
			Type: restic.PackFile,
			Name: name,
		}
		fi, err := repo.Backend().Stat(gopts.ctx, h)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
		}

		fmt.Printf("  file size is %v\n", fi.Size)

		// examine all data the indexes have for the pack file
		for _, idx := range repo.Index().(*repository.MasterIndex).All() {
			idxIDs, err := idx.IDs()
			if err != nil {
				idxIDs = restic.IDs{}
			}

			blobs := idx.ListPack(id)
			if len(blobs) == 0 {
				fmt.Printf("    index %v does not contain the file\n", idxIDs)
				continue
			}

			fmt.Printf("    index %v:\n", idxIDs)

			// track current size and offset
			var size, offset uint64

			sort.Slice(blobs, func(i, j int) bool {
				return blobs[i].Offset < blobs[j].Offset
			})

			for _, pb := range blobs {
				fmt.Printf("      %v blob %v, offset %-6d, raw length %-6d\n", pb.Type, pb.ID, pb.Offset, pb.Length)
				if offset != uint64(pb.Offset) {
					fmt.Printf("      hole in file, want offset %v, got %v\n", offset, pb.Offset)
				}
				offset += uint64(pb.Length)
				size += uint64(pb.Length)
			}

			// compute header size, per blob: 1 byte type, 4 byte length, 32 byte id
			size += uint64(restic.CiphertextLength(len(blobs) * (1 + 4 + 32)))
			// length in uint32 little endian
			size += 4

			if uint64(fi.Size) != size {
				fmt.Printf("      file sizes do not match: computed %v from index, file size is %v\n", size, fi.Size)
			} else {
				fmt.Printf("      file sizes match\n")
			}

			err = loadBlobs(gopts.ctx, repo, name, blobs)
			if err != nil {
				fmt.Fprintf(os.Stderr, "error: %v\n", err)
			}
		}

		// inspect the pack file itself
		blobs, _, err := pack.List(repo.Key(), restic.ReaderAt(gopts.ctx, repo.Backend(), h), fi.Size)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error for pack %v: %v\n", id.Str(), err)
			return nil
		}

		// track current size and offset
		var size, offset uint64

		sort.Slice(blobs, func(i, j int) bool {
			return blobs[i].Offset < blobs[j].Offset
		})

		for _, pb := range blobs {
			fmt.Printf("      %v blob %v, offset %-6d, raw length %-6d\n", pb.Type, pb.ID, pb.Offset, pb.Length)
			if offset != uint64(pb.Offset) {
				fmt.Printf("      hole in file, want offset %v, got %v\n", offset, pb.Offset)
			}
			offset += uint64(pb.Length)
			size += uint64(pb.Length)
		}

		// compute header size, per blob: 1 byte type, 4 byte length, 32 byte id
		size += uint64(restic.CiphertextLength(len(blobs) * (1 + 4 + 32)))
		// length in uint32 little endian
		size += 4

		if uint64(fi.Size) != size {
			fmt.Printf("      file sizes do not match: computed %v from index, file size is %v\n", size, fi.Size)
		} else {
			fmt.Printf("      file sizes match\n")
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"github.com/restic/restic/internal/restic"
	"github.com/spf13/cobra"
)

var cmdStats = &cobra.Command{
	Use:   "stats",
	Short: "Scan the repository and show basic statistics",
	Long: `
The "stats" command walks one or all snapshots in a repository and
accumulates statistics about the data stored therein. It reports on
the number of unique files and their sizes, according to one of
the counting modes as given by a flag.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runStats(globalOptions, args)
	},
}

var countModeFlag []string

func init() {
	cmdRoot.AddCommand(cmdStats)

	f := cmdStats.Flags()
	f.BoolVar(&countModeRestoreSize, "count-restore-size", false, "count the size of files that would be restored (default)")
	f.BoolVar(&countModeUniqueFilesByContent, "count-files-by-contents", false, "count files as unique by their contents")
	f.BoolVar(&countModeBlobsPerFile, "count-blobs-per-file", false, "count sizes of blobs by filename")
	f.BoolVar(&countModeRawData, "count-raw-data", false, "count unique blob sizes irrespective of files referencing them")
	f.StringVar(&snapshotByHost, "host", "", "filter latest snapshot by this hostname")
}

func runStats(gopts GlobalOptions, args []string) error {
	err := verifyStatsInput(gopts, args)
	if err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(gopts.ctx)
	defer cancel()

	repo, err := OpenRepository(gopts)
	if err != nil {
		return err
	}

	if err = repo.LoadIndex(ctx); err != nil {
		return err
	}

	if !gopts.NoLock {
		lock, err := lockRepo(repo)
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	}

	// create a container for the stats (and other needed state)
	stats := &statsContainer{
		uniqueFiles: make(map[fileID]struct{}),
		idSet:       make(restic.IDSet),
		fileBlobs:   make(map[string]restic.IDSet),
		blobs:       restic.NewBlobSet(),
		blobsSeen:   restic.NewBlobSet(),
	}

	if snapshotIDString != "" {
		// scan just a single snapshot

		var sID restic.ID
		if snapshotIDString == "latest" {
			sID, err = restic.FindLatestSnapshot(ctx, repo, []string{}, []restic.TagList{}, snapshotByHost)
			if err != nil {
				Exitf(1, "latest snapshot for criteria not found: %v", err)
			}
		} else {
			sID, err = restic.FindSnapshot(repo, snapshotIDString)
			if err != nil {
				return err
			}
		}

		snapshot, err := restic.LoadSnapshot(ctx, repo, sID)
		if err != nil {
			return err
		}

		err = statsWalkSnapshot(ctx, snapshot, repo, stats)
	} else {
		// iterate every snapshot in the repo
		err = repo.List(ctx, restic.SnapshotFile, func(snapshotID restic.ID, size int64) error {
			snapshot, err := restic.LoadSnapshot(ctx, repo, snapshotID)
			if err != nil {
				return fmt.Errorf("Error loading snapshot %s: %v", snapshotID.Str(), err)
			}
			return statsWalkSnapshot(ctx, snapshot, repo, stats)
		})
	}
	if err != nil {
		return err
	}

	if countModeRawData {
		// the blob handles have been collected, but not yet counted
		for blobHandle := range stats.blobs {
			blobSize, found := repo.LookupBlobSize(blobHandle.ID, blobHandle.Type)
			if !found {
				return fmt.Errorf("blob %v not found", blobHandle)
			}
			stats.TotalSize += uint64(blobSize)
			stats.TotalBlobCount++
		}
	}

	if gopts.JSON {
		err = json.NewEncoder(os.Stdout).Encode(stats)
		if err != nil {
			return fmt.Errorf("encoding output: %v", err)
		}
		return nil
	}

	if stats.TotalBlobCount > 0 {
		Printf("  Total Blob Count:   %d\n", stats.TotalBlobCount)
	}
	if stats.TotalFileCount > 0 {
		Printf("  Total File Count:   %d\n", stats.TotalFileCount)
	}
	Printf("        Total Size:   %-5s\n", formatBytes(stats.TotalSize))

	return nil
}

func statsWalkSnapshot(ctx context.Context, snapshot *restic.Snapshot, repo restic.Repository, stats *statsContainer) error {
	if snapshot.Tree == nil {
		return fmt.Errorf("snapshot %s has nil tree", snapshot.ID().Str())
	}

	if countModeRawData {
		// count just the sizes of unique blobs; we don't need to walk the tree
		// ourselves in this case, since a nifty function does it for us
		return restic.FindUsedBlobs(ctx, repo, *snapshot.Tree, stats.blobs, stats.blobsSeen)
	}

	err := statsWalkTree(ctx, repo, *snapshot.Tree, stats, string(filepath.Separator))
	if err != nil {
		return fmt.Errorf("walking tree %s: %v", *snapshot.Tree, err)
	}
	return nil
}

func statsWalkTree(ctx context.Context, repo restic.Repository, treeID restic.ID, stats *statsContainer, fpath string) error {
	// don't visit a tree we've already walked
	if stats.idSet.Has(treeID) {
		return nil
	}
	stats.idSet.Insert(treeID)

	tree, err := repo.LoadTree(ctx, treeID)
	if err != nil {
		return fmt.Errorf("loading tree: %v", err)
	}

	for _, node := range tree.Nodes {
		if countModeUniqueFilesByContent || countModeBlobsPerFile {
			// only count this file if we haven't visited it before
			fid := makeFileIDByContents(node)
			if _, ok := stats.uniqueFiles[fid]; !ok {
				// mark the file as visited
				stats.uniqueFiles[fid] = struct{}{}

				if countModeUniqueFilesByContent {
					// simply count the size of each unique file (unique by contents only)
					stats.TotalSize += node.Size
					stats.TotalFileCount++
				}
				if countModeBlobsPerFile {
					// count the size of each unique blob reference, which is
					// by unique file (unique by contents and file path)
					for _, blobID := range node.Content {
						// ensure we have this file (by path) in our map; in this
						// mode, a file is unique by both contents and path
						nodePath := filepath.Join(fpath, node.Name)
						if _, ok := stats.fileBlobs[nodePath]; !ok {
							stats.fileBlobs[nodePath] = restic.NewIDSet()
							stats.TotalFileCount++
						}
						if _, ok := stats.fileBlobs[nodePath][blobID]; !ok {
							// TODO: Is the blob type always 'data' in this case?
							blobSize, found := repo.LookupBlobSize(blobID, restic.DataBlob)
							if !found {
								return fmt.Errorf("blob %s not found for tree %s", blobID, treeID)
							}

							// count the blob's size, then add this blob by this
							// file (path) so we don't double-count it
							stats.TotalSize += uint64(blobSize)
							stats.fileBlobs[nodePath].Insert(blobID)

							// this mode also counts total unique blob _references_ per file
							stats.TotalBlobCount++
						}
					}
				}
			}
		}

		if countModeRestoreSize {
			// as this is a file in the snapshot, we can simply count its
			// size without worrying about uniqueness, since duplicate files
			// will still be restored
			stats.TotalSize += node.Size
			stats.TotalFileCount++
		}

		// visit subtrees (i.e. directory contents)
		if node.Subtree != nil {
			err = statsWalkTree(ctx, repo, *node.Subtree, stats, filepath.Join(fpath, node.Name))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// makeFileIDByContents returns a hash of the blob IDs of the
// node's Content in sequence.
func makeFileIDByContents(node *restic.Node) fileID {
	var bb []byte
	for _, c := range node.Content {
		bb = append(bb, []byte(c[:])...)
	}
	return sha256.Sum256(bb)
}

func verifyStatsInput(gopts GlobalOptions, args []string) error {
	// ensure only one counting mode was specified, for clarity
	var countModes int
	if countModeRestoreSize {
		countModes++
	}
	if countModeUniqueFilesByContent {
		countModes++
	}
	if countModeBlobsPerFile {
		countModes++
	}
	if countModeRawData {
		countModes++
	}
	if countModes > 1 {
		return fmt.Errorf("only one counting mode may be used")
	}
	// set a default count mode if none were specified
	if countModes == 0 {
		countModeRestoreSize = true
	}
	// ensure one or none snapshots were specified
	if len(args) > 1 {
		return fmt.Errorf("only one snapshot may be specified")
	}
	// set the snapshot to scan, if one was specified
	if len(args) == 1 {
		snapshotIDString = args[0]
	}
	return nil
}

// statsContainer holds information during a walk of a repository
// to collect information about it, as well as state needed
// for a successful and efficient walk.
type statsContainer struct {
	TotalSize      uint64 `json:"total_size"`
	TotalFileCount uint64 `json:"total_file_count"`
	TotalBlobCount uint64 `json:"total_blob_count,omitempty"`

	// idSet marks visited trees, to avoid repeated walks
	idSet restic.IDSet

	// uniqueFiles marks visited files according to their
	// contents (hashed sequence of content blob IDs)
	uniqueFiles map[fileID]struct{}

	// fileBlobs maps a file name (path) to the set of
	// blobs that have been seen as a part of the file
	fileBlobs map[string]restic.IDSet

	// blobs and blobsSeen are used to count indiviudal
	// unique blobs, independent of references to files
	blobs, blobsSeen restic.BlobSet
}

// fileID is a 256-bit hash that distinguishes unique files.
type fileID [32]byte

var (
	countModeRestoreSize          bool
	countModeUniqueFilesByContent bool
	countModeBlobsPerFile         bool
	countModeRawData              bool

	// the snapshot to scan, as given by the user
	snapshotIDString string

	// snapshotByHost is the host to filter latest
	// snapshot by, if given by user
	snapshotByHost string
)
<commit_msg>Update comment now that question was answered<commit_after>package main

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"github.com/restic/restic/internal/restic"
	"github.com/spf13/cobra"
)

var cmdStats = &cobra.Command{
	Use:   "stats",
	Short: "Scan the repository and show basic statistics",
	Long: `
The "stats" command walks one or all snapshots in a repository and
accumulates statistics about the data stored therein. It reports on
the number of unique files and their sizes, according to one of
the counting modes as given by a flag.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runStats(globalOptions, args)
	},
}

var countModeFlag []string

func init() {
	cmdRoot.AddCommand(cmdStats)

	f := cmdStats.Flags()
	f.BoolVar(&countModeRestoreSize, "count-restore-size", false, "count the size of files that would be restored (default)")
	f.BoolVar(&countModeUniqueFilesByContent, "count-files-by-contents", false, "count files as unique by their contents")
	f.BoolVar(&countModeBlobsPerFile, "count-blobs-per-file", false, "count sizes of blobs by filename")
	f.BoolVar(&countModeRawData, "count-raw-data", false, "count unique blob sizes irrespective of files referencing them")
	f.StringVar(&snapshotByHost, "host", "", "filter latest snapshot by this hostname")
}

func runStats(gopts GlobalOptions, args []string) error {
	err := verifyStatsInput(gopts, args)
	if err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(gopts.ctx)
	defer cancel()

	repo, err := OpenRepository(gopts)
	if err != nil {
		return err
	}

	if err = repo.LoadIndex(ctx); err != nil {
		return err
	}

	if !gopts.NoLock {
		lock, err := lockRepo(repo)
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	}

	// create a container for the stats (and other needed state)
	stats := &statsContainer{
		uniqueFiles: make(map[fileID]struct{}),
		idSet:       make(restic.IDSet),
		fileBlobs:   make(map[string]restic.IDSet),
		blobs:       restic.NewBlobSet(),
		blobsSeen:   restic.NewBlobSet(),
	}

	if snapshotIDString != "" {
		// scan just a single snapshot

		var sID restic.ID
		if snapshotIDString == "latest" {
			sID, err = restic.FindLatestSnapshot(ctx, repo, []string{}, []restic.TagList{}, snapshotByHost)
			if err != nil {
				Exitf(1, "latest snapshot for criteria not found: %v", err)
			}
		} else {
			sID, err = restic.FindSnapshot(repo, snapshotIDString)
			if err != nil {
				return err
			}
		}

		snapshot, err := restic.LoadSnapshot(ctx, repo, sID)
		if err != nil {
			return err
		}

		err = statsWalkSnapshot(ctx, snapshot, repo, stats)
	} else {
		// iterate every snapshot in the repo
		err = repo.List(ctx, restic.SnapshotFile, func(snapshotID restic.ID, size int64) error {
			snapshot, err := restic.LoadSnapshot(ctx, repo, snapshotID)
			if err != nil {
				return fmt.Errorf("Error loading snapshot %s: %v", snapshotID.Str(), err)
			}
			return statsWalkSnapshot(ctx, snapshot, repo, stats)
		})
	}
	if err != nil {
		return err
	}

	if countModeRawData {
		// the blob handles have been collected, but not yet counted
		for blobHandle := range stats.blobs {
			blobSize, found := repo.LookupBlobSize(blobHandle.ID, blobHandle.Type)
			if !found {
				return fmt.Errorf("blob %v not found", blobHandle)
			}
			stats.TotalSize += uint64(blobSize)
			stats.TotalBlobCount++
		}
	}

	if gopts.JSON {
		err = json.NewEncoder(os.Stdout).Encode(stats)
		if err != nil {
			return fmt.Errorf("encoding output: %v", err)
		}
		return nil
	}

	if stats.TotalBlobCount > 0 {
		Printf("  Total Blob Count:   %d\n", stats.TotalBlobCount)
	}
	if stats.TotalFileCount > 0 {
		Printf("  Total File Count:   %d\n", stats.TotalFileCount)
	}
	Printf("        Total Size:   %-5s\n", formatBytes(stats.TotalSize))

	return nil
}

func statsWalkSnapshot(ctx context.Context, snapshot *restic.Snapshot, repo restic.Repository, stats *statsContainer) error {
	if snapshot.Tree == nil {
		return fmt.Errorf("snapshot %s has nil tree", snapshot.ID().Str())
	}

	if countModeRawData {
		// count just the sizes of unique blobs; we don't need to walk the tree
		// ourselves in this case, since a nifty function does it for us
		return restic.FindUsedBlobs(ctx, repo, *snapshot.Tree, stats.blobs, stats.blobsSeen)
	}

	err := statsWalkTree(ctx, repo, *snapshot.Tree, stats, string(filepath.Separator))
	if err != nil {
		return fmt.Errorf("walking tree %s: %v", *snapshot.Tree, err)
	}
	return nil
}

func statsWalkTree(ctx context.Context, repo restic.Repository, treeID restic.ID, stats *statsContainer, fpath string) error {
	// don't visit a tree we've already walked
	if stats.idSet.Has(treeID) {
		return nil
	}
	stats.idSet.Insert(treeID)

	tree, err := repo.LoadTree(ctx, treeID)
	if err != nil {
		return fmt.Errorf("loading tree: %v", err)
	}

	for _, node := range tree.Nodes {
		if countModeUniqueFilesByContent || countModeBlobsPerFile {
			// only count this file if we haven't visited it before
			fid := makeFileIDByContents(node)
			if _, ok := stats.uniqueFiles[fid]; !ok {
				// mark the file as visited
				stats.uniqueFiles[fid] = struct{}{}

				if countModeUniqueFilesByContent {
					// simply count the size of each unique file (unique by contents only)
					stats.TotalSize += node.Size
					stats.TotalFileCount++
				}
				if countModeBlobsPerFile {
					// count the size of each unique blob reference, which is
					// by unique file (unique by contents and file path)
					for _, blobID := range node.Content {
						// ensure we have this file (by path) in our map; in this
						// mode, a file is unique by both contents and path
						nodePath := filepath.Join(fpath, node.Name)
						if _, ok := stats.fileBlobs[nodePath]; !ok {
							stats.fileBlobs[nodePath] = restic.NewIDSet()
							stats.TotalFileCount++
						}
						if _, ok := stats.fileBlobs[nodePath][blobID]; !ok {
							// is always a data blob since we're accessing it via a file's Content array
							blobSize, found := repo.LookupBlobSize(blobID, restic.DataBlob)
							if !found {
								return fmt.Errorf("blob %s not found for tree %s", blobID, treeID)
							}

							// count the blob's size, then add this blob by this
							// file (path) so we don't double-count it
							stats.TotalSize += uint64(blobSize)
							stats.fileBlobs[nodePath].Insert(blobID)

							// this mode also counts total unique blob _references_ per file
							stats.TotalBlobCount++
						}
					}
				}
			}
		}

		if countModeRestoreSize {
			// as this is a file in the snapshot, we can simply count its
			// size without worrying about uniqueness, since duplicate files
			// will still be restored
			stats.TotalSize += node.Size
			stats.TotalFileCount++
		}

		// visit subtrees (i.e. directory contents)
		if node.Subtree != nil {
			err = statsWalkTree(ctx, repo, *node.Subtree, stats, filepath.Join(fpath, node.Name))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// makeFileIDByContents returns a hash of the blob IDs of the
// node's Content in sequence.
func makeFileIDByContents(node *restic.Node) fileID {
	var bb []byte
	for _, c := range node.Content {
		bb = append(bb, []byte(c[:])...)
	}
	return sha256.Sum256(bb)
}

func verifyStatsInput(gopts GlobalOptions, args []string) error {
	// ensure only one counting mode was specified, for clarity
	var countModes int
	if countModeRestoreSize {
		countModes++
	}
	if countModeUniqueFilesByContent {
		countModes++
	}
	if countModeBlobsPerFile {
		countModes++
	}
	if countModeRawData {
		countModes++
	}
	if countModes > 1 {
		return fmt.Errorf("only one counting mode may be used")
	}
	// set a default count mode if none were specified
	if countModes == 0 {
		countModeRestoreSize = true
	}
	// ensure one or none snapshots were specified
	if len(args) > 1 {
		return fmt.Errorf("only one snapshot may be specified")
	}
	// set the snapshot to scan, if one was specified
	if len(args) == 1 {
		snapshotIDString = args[0]
	}
	return nil
}

// statsContainer holds information during a walk of a repository
// to collect information about it, as well as state needed
// for a successful and efficient walk.
type statsContainer struct {
	TotalSize      uint64 `json:"total_size"`
	TotalFileCount uint64 `json:"total_file_count"`
	TotalBlobCount uint64 `json:"total_blob_count,omitempty"`

	// idSet marks visited trees, to avoid repeated walks
	idSet restic.IDSet

	// uniqueFiles marks visited files according to their
	// contents (hashed sequence of content blob IDs)
	uniqueFiles map[fileID]struct{}

	// fileBlobs maps a file name (path) to the set of
	// blobs that have been seen as a part of the file
	fileBlobs map[string]restic.IDSet

	// blobs and blobsSeen are used to count indiviudal
	// unique blobs, independent of references to files
	blobs, blobsSeen restic.BlobSet
}

// fileID is a 256-bit hash that distinguishes unique files.
type fileID [32]byte

var (
	countModeRestoreSize          bool
	countModeUniqueFilesByContent bool
	countModeBlobsPerFile         bool
	countModeRawData              bool

	// the snapshot to scan, as given by the user
	snapshotIDString string

	// snapshotByHost is the host to filter latest
	// snapshot by, if given by user
	snapshotByHost string
)
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/rpc"
	"time"

	wrpc "github.com/disq/werify/rpc"
)

type client struct {
	env     string
	server  string
	timeout time.Duration

	conn *rpc.Client
}

func (c *client) connect() error {
	connection, err := net.DialTimeout("tcp", c.server, c.timeout)
	if err != nil {
		return err
	}

	//fmt.Printf("Connected to %s\n", c.server)

	c.conn = rpc.NewClient(connection)
	return nil
}

func (c *client) parseCommand(command string, args []string) error {
	cmdCfg, ok := wrpc.Commands[command]
	if !ok {
		return fmt.Errorf("Unknown command %s", command)
	}
	if cmdCfg.NumArgs != len(args) {
		return fmt.Errorf("Invalid number of arguments for %s: Expected %d but got %d", command, cmdCfg.NumArgs, len(args))
	}

	rpcCmd := wrpc.BuildMethod(cmdCfg.RpcMethod)
	ci := c.newCommonInput()

	switch command {
	case "add":
		out := wrpc.AddHostOutput{}
		err := c.conn.Call(rpcCmd, wrpc.AddHostInput{CommonInput: ci, Endpoint: wrpc.Endpoint(args[0])}, &out)
		if err != nil {
			return err
		}
		if out.Ok {
			fmt.Printf("Added host %s\n", args[0])
		} else {
			fmt.Printf("Could not add host %s\n", args[0])
		}

	case "del":
		out := wrpc.RemoveHostOutput{}
		err := c.conn.Call(rpcCmd, wrpc.RemoveHostInput{CommonInput: ci, Endpoint: wrpc.Endpoint(args[0])}, &out)
		if err != nil {
			return err
		}
		if out.Ok {
			fmt.Printf("Removed host %s\n", args[0])
		} else {
			fmt.Printf("Could not remove host %s\n", args[0])
		}

	case "refresh":
		out := wrpc.RefreshOutput{}
		err := c.conn.Call(rpcCmd, wrpc.RefreshInput{CommonInput: ci}, &out)
		if err != nil {
			return err
		}
		if out.Ok {
			fmt.Println("Initiated refresh")
		} else {
			fmt.Println("Could not initiate refresh")
		}

	// FIXME: repeating ugly code below
	case "list":
		fallthrough
	case "listactive":
		fallthrough
	case "listinactive":
		out := wrpc.ListHostsOutput{}
		in := wrpc.ListHostsInput{CommonInput: ci, ListActive: true, ListInactive: true}
		if command == "listactive" {
			in.ListInactive = false
		} else if command == "listinactive" {
			in.ListActive = false
		}
		err := c.conn.Call(rpcCmd, in, &out)
		if err != nil {
			return err
		}

		if command == "list" || command == "listactive" {
			fmt.Printf("Active hosts (%d)\n", len(out.ActiveHosts))
			for _, e := range out.ActiveHosts {
				fmt.Println(e)
			}
		}
		if command == "list" || command == "listinactive" {
			fmt.Printf("Inactive hosts (%d)\n", len(out.InactiveHosts))
			for _, e := range out.InactiveHosts {
				fmt.Println(e)
			}
		}
		fmt.Println("End of list")

	case "operation":
		b, err := ioutil.ReadFile(args[0])
		if err != nil {
			return fmt.Errorf("Reading %s: %s", args[0], err.Error())
		}

		in := wrpc.OperationInput{
			CommonInput: ci,
			Forward:     true,
		}

		err = json.Unmarshal(b, &in.Ops)
		if err != nil {
			return fmt.Errorf("Parsing %s: %s", args[0], err.Error())
		}

		out := wrpc.OperationOutput{}

		// TODO: make it async? Get an identifier, run another command to read so-far-collected results and status
		err = c.conn.Call(rpcCmd, in, &out)
		if err != nil {
			return err
		}

		if out.Handle != "" {
			fmt.Printf("Operation submitted. To check progress, run: ./werifyctl get %s\n", out.Handle)
		} else {
			c.displayOperation(out)
		}

	case "get":
		out := wrpc.OperationStatusCheckOutput{}
		err := c.conn.Call(rpcCmd, wrpc.OperationStatusCheckInput{CommonInput: ci, Handle: args[0]}, &out)
		if err != nil {
			return err
		}
		c.displayOperation(wrpc.OperationOutput(out))

	default:
		return fmt.Errorf("Unhandled command %s", command)
	}

	return nil
}

// newCommonInput initializes and returns a CommonInput struct using client's information
func (c *client) newCommonInput() wrpc.CommonInput {
	return wrpc.CommonInput{
		EnvTag: c.env,
	}
}

func (c *client) displayOperation(o wrpc.OperationOutput) {
	for id, res := range o.Results {
		for name, result := range res {
			if result.Err != "" {
				fmt.Printf("Host:%s Operation:%s Error:%s\n", id, name, result.Err)
			} else {
				fmt.Printf("Host:%s Operation:%s Success:%t\n", id, name, result.Success)
			}
		}
	}

	if o.EndedAt != nil {
		fmt.Printf("Operation ended, took %v\n", o.EndedAt.Sub(o.StartedAt))
	} else {
		fmt.Printf("Operation still running...\n")
	}
}
<commit_msg>werifyctl: Remove done TODO<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/rpc"
	"time"

	wrpc "github.com/disq/werify/rpc"
)

type client struct {
	env     string
	server  string
	timeout time.Duration

	conn *rpc.Client
}

func (c *client) connect() error {
	connection, err := net.DialTimeout("tcp", c.server, c.timeout)
	if err != nil {
		return err
	}

	//fmt.Printf("Connected to %s\n", c.server)

	c.conn = rpc.NewClient(connection)
	return nil
}

func (c *client) parseCommand(command string, args []string) error {
	cmdCfg, ok := wrpc.Commands[command]
	if !ok {
		return fmt.Errorf("Unknown command %s", command)
	}
	if cmdCfg.NumArgs != len(args) {
		return fmt.Errorf("Invalid number of arguments for %s: Expected %d but got %d", command, cmdCfg.NumArgs, len(args))
	}

	rpcCmd := wrpc.BuildMethod(cmdCfg.RpcMethod)
	ci := c.newCommonInput()

	switch command {
	case "add":
		out := wrpc.AddHostOutput{}
		err := c.conn.Call(rpcCmd, wrpc.AddHostInput{CommonInput: ci, Endpoint: wrpc.Endpoint(args[0])}, &out)
		if err != nil {
			return err
		}
		if out.Ok {
			fmt.Printf("Added host %s\n", args[0])
		} else {
			fmt.Printf("Could not add host %s\n", args[0])
		}

	case "del":
		out := wrpc.RemoveHostOutput{}
		err := c.conn.Call(rpcCmd, wrpc.RemoveHostInput{CommonInput: ci, Endpoint: wrpc.Endpoint(args[0])}, &out)
		if err != nil {
			return err
		}
		if out.Ok {
			fmt.Printf("Removed host %s\n", args[0])
		} else {
			fmt.Printf("Could not remove host %s\n", args[0])
		}

	case "refresh":
		out := wrpc.RefreshOutput{}
		err := c.conn.Call(rpcCmd, wrpc.RefreshInput{CommonInput: ci}, &out)
		if err != nil {
			return err
		}
		if out.Ok {
			fmt.Println("Initiated refresh")
		} else {
			fmt.Println("Could not initiate refresh")
		}

	// FIXME: repeating ugly code below
	case "list":
		fallthrough
	case "listactive":
		fallthrough
	case "listinactive":
		out := wrpc.ListHostsOutput{}
		in := wrpc.ListHostsInput{CommonInput: ci, ListActive: true, ListInactive: true}
		if command == "listactive" {
			in.ListInactive = false
		} else if command == "listinactive" {
			in.ListActive = false
		}
		err := c.conn.Call(rpcCmd, in, &out)
		if err != nil {
			return err
		}

		if command == "list" || command == "listactive" {
			fmt.Printf("Active hosts (%d)\n", len(out.ActiveHosts))
			for _, e := range out.ActiveHosts {
				fmt.Println(e)
			}
		}
		if command == "list" || command == "listinactive" {
			fmt.Printf("Inactive hosts (%d)\n", len(out.InactiveHosts))
			for _, e := range out.InactiveHosts {
				fmt.Println(e)
			}
		}
		fmt.Println("End of list")

	case "operation":
		b, err := ioutil.ReadFile(args[0])
		if err != nil {
			return fmt.Errorf("Reading %s: %s", args[0], err.Error())
		}

		in := wrpc.OperationInput{
			CommonInput: ci,
			Forward:     true,
		}

		err = json.Unmarshal(b, &in.Ops)
		if err != nil {
			return fmt.Errorf("Parsing %s: %s", args[0], err.Error())
		}

		out := wrpc.OperationOutput{}

		err = c.conn.Call(rpcCmd, in, &out)
		if err != nil {
			return err
		}

		if out.Handle != "" {
			fmt.Printf("Operation submitted. To check progress, run: ./werifyctl get %s\n", out.Handle)
		} else {
			c.displayOperation(out)
		}

	case "get":
		out := wrpc.OperationStatusCheckOutput{}
		err := c.conn.Call(rpcCmd, wrpc.OperationStatusCheckInput{CommonInput: ci, Handle: args[0]}, &out)
		if err != nil {
			return err
		}
		c.displayOperation(wrpc.OperationOutput(out))

	default:
		return fmt.Errorf("Unhandled command %s", command)
	}

	return nil
}

// newCommonInput initializes and returns a CommonInput struct using client's information
func (c *client) newCommonInput() wrpc.CommonInput {
	return wrpc.CommonInput{
		EnvTag: c.env,
	}
}

func (c *client) displayOperation(o wrpc.OperationOutput) {
	for id, res := range o.Results {
		for name, result := range res {
			if result.Err != "" {
				fmt.Printf("Host:%s Operation:%s Error:%s\n", id, name, result.Err)
			} else {
				fmt.Printf("Host:%s Operation:%s Success:%t\n", id, name, result.Success)
			}
		}
	}

	if o.EndedAt != nil {
		fmt.Printf("Operation ended, took %v\n", o.EndedAt.Sub(o.StartedAt))
	} else {
		fmt.Printf("Operation still running...\n")
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"regexp"
	"testing"

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

func TestAccAWSLambdaAlias_basic(t *testing.T) {
	var conf lambda.AliasConfiguration

	rString := acctest.RandString(8)
	roleName := fmt.Sprintf("tf_acc_role_lambda_alias_basic_%s", rString)
	policyName := fmt.Sprintf("tf_acc_policy_lambda_alias_basic_%s", rString)
	attachmentName := fmt.Sprintf("tf_acc_attachment_%s", rString)
	funcName := fmt.Sprintf("tf_acc_lambda_func_alias_basic_%s", rString)
	aliasName := fmt.Sprintf("tf_acc_lambda_alias_basic_%s", rString)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAwsLambdaAliasDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAwsLambdaAliasConfig(roleName, policyName, attachmentName, funcName, aliasName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaAliasExists("aws_lambda_alias.lambda_alias_test", &conf),
					testAccCheckAwsLambdaAttributes(&conf),
					resource.TestMatchResourceAttr("aws_lambda_alias.lambda_alias_test", "arn",
						regexp.MustCompile(`^arn:aws:lambda:[a-z]+-[a-z]+-[0-9]+:\d{12}:function:`+funcName+`:`+aliasName+`$`)),
				),
			},
			resource.TestStep{
				Config: testAccAwsLambdaAliasConfigWithRoutingConfig(roleName, policyName, attachmentName, funcName, aliasName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaAliasExists("aws_lambda_alias.lambda_alias_test", &conf),
					testAccCheckAwsLambdaAttributes(&conf),
					testAccCheckAwsLambdaAliasRoutingConfig(&conf),
					resource.TestMatchResourceAttr("aws_lambda_alias.lambda_alias_test", "arn",
						regexp.MustCompile(`^arn:aws:lambda:[a-z]+-[a-z]+-[0-9]+:\d{12}:function:`+funcName+`:`+aliasName+`$`)),
				),
			},
		},
	})
}

func testAccCheckAwsLambdaAliasDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).lambdaconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_lambda_alias" {
			continue
		}

		_, err := conn.GetAlias(&lambda.GetAliasInput{
			FunctionName: aws.String(rs.Primary.ID),
		})

		if err == nil {
			return fmt.Errorf("Lambda alias was not deleted")
		}

	}

	return nil
}

func testAccCheckAwsLambdaAliasExists(n string, mapping *lambda.AliasConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Lambda alias not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("Lambda alias not set")
		}

		conn := testAccProvider.Meta().(*AWSClient).lambdaconn

		params := &lambda.GetAliasInput{
			FunctionName: aws.String(rs.Primary.ID),
			Name:         aws.String(rs.Primary.Attributes["name"]),
		}

		getAliasConfiguration, err := conn.GetAlias(params)
		if err != nil {
			return err
		}

		*mapping = *getAliasConfiguration

		return nil
	}
}

func testAccCheckAwsLambdaAttributes(mapping *lambda.AliasConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		name := *mapping.Name
		arn := *mapping.AliasArn
		if arn == "" {
			return fmt.Errorf("Could not read Lambda alias ARN")
		}
		if name == "" {
			return fmt.Errorf("Could not read Lambda alias name")
		}
		return nil
	}
}

func testAccCheckAwsLambdaAliasRoutingConfig(mapping *lambda.AliasConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		routingConfig := mapping.RoutingConfig

		if routingConfig == nil {
			return fmt.Errorf("Could not read Lambda alias routing config")
		}
		if len(routingConfig.AdditionalVersionWeights) != 1 {
			return fmt.Errorf("Could not read Lambda alias additional version weights")
		}
		return nil
	}
}

func testAccAwsLambdaAliasConfig(roleName, policyName, attachmentName, funcName, aliasName string) string {
	return fmt.Sprintf(`
resource "aws_iam_role" "iam_for_lambda" {
  name = "%s"

  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "lambda.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_policy" "policy_for_role" {
  name        = "%s"
  path        = "/"
  description = "IAM policy for for Lamda alias testing"

  policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
      {
          "Effect": "Allow",
          "Action": [
            "lambda:*"
          ],
          "Resource": "*"
      }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "policy_attachment_for_role" {
  name       = "%s"
  roles      = ["${aws_iam_role.iam_for_lambda.name}"]
  policy_arn = "${aws_iam_policy.policy_for_role.arn}"
}

resource "aws_lambda_function" "lambda_function_test_create" {
  filename      	 = "test-fixtures/lambdatest.zip"
  function_name 	 = "%s"
  role          	 = "${aws_iam_role.iam_for_lambda.arn}"
  handler       	 = "exports.example"
  runtime       	 = "nodejs4.3"
  source_code_hash = "${base64sha256(file("test-fixtures/lambdatest.zip"))}"
  publish          = "true"
}

resource "aws_lambda_alias" "lambda_alias_test" {
  name             = "%s"
  description      = "a sample description"
  function_name    = "${aws_lambda_function.lambda_function_test_create.arn}"
  function_version = "1"
}`, roleName, policyName, attachmentName, funcName, aliasName)
}

func testAccAwsLambdaAliasConfigWithRoutingConfig(roleName, policyName, attachmentName, funcName, aliasName string) string {
	return fmt.Sprintf(`
resource "aws_iam_role" "iam_for_lambda" {
  name = "%s"

  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "lambda.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_policy" "policy_for_role" {
  name        = "%s"
  path        = "/"
  description = "IAM policy for for Lamda alias testing"

  policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
      {
          "Effect": "Allow",
          "Action": [
            "lambda:*"
          ],
          "Resource": "*"
      }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "policy_attachment_for_role" {
  name       = "%s"
  roles      = ["${aws_iam_role.iam_for_lambda.name}"]
  policy_arn = "${aws_iam_policy.policy_for_role.arn}"
}

resource "aws_lambda_function" "lambda_function_test_create" {
  filename      	 = "test-fixtures/lambdatest_modified.zip"
  function_name 	 = "%s"
  role          	 = "${aws_iam_role.iam_for_lambda.arn}"
  handler       	 = "exports.example"
  runtime       	 = "nodejs4.3"
  source_code_hash = "${base64sha256(file("test-fixtures/lambdatest_modified.zip"))}"
  publish          = "true"
}

resource "aws_lambda_alias" "lambda_alias_test" {
  name             = "%s"
  description      = "a sample description"
  function_name    = "${aws_lambda_function.lambda_function_test_create.arn}"
  function_version = "1"
	routing_config   = {
		additional_version_weights = {
			"2" = 0.5
		}
	}
}`, roleName, policyName, attachmentName, funcName, aliasName)
}
<commit_msg>fmt'd<commit_after>package aws

import (
	"fmt"
	"regexp"
	"testing"

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

func TestAccAWSLambdaAlias_basic(t *testing.T) {
	var conf lambda.AliasConfiguration

	rString := acctest.RandString(8)
	roleName := fmt.Sprintf("tf_acc_role_lambda_alias_basic_%s", rString)
	policyName := fmt.Sprintf("tf_acc_policy_lambda_alias_basic_%s", rString)
	attachmentName := fmt.Sprintf("tf_acc_attachment_%s", rString)
	funcName := fmt.Sprintf("tf_acc_lambda_func_alias_basic_%s", rString)
	aliasName := fmt.Sprintf("tf_acc_lambda_alias_basic_%s", rString)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAwsLambdaAliasDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAwsLambdaAliasConfig(roleName, policyName, attachmentName, funcName, aliasName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaAliasExists("aws_lambda_alias.lambda_alias_test", &conf),
					testAccCheckAwsLambdaAttributes(&conf),
					resource.TestMatchResourceAttr("aws_lambda_alias.lambda_alias_test", "arn",
						regexp.MustCompile(`^arn:aws:lambda:[a-z]+-[a-z]+-[0-9]+:\d{12}:function:`+funcName+`:`+aliasName+`$`)),
				),
			},
			resource.TestStep{
				Config: testAccAwsLambdaAliasConfigWithRoutingConfig(roleName, policyName, attachmentName, funcName, aliasName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaAliasExists("aws_lambda_alias.lambda_alias_test", &conf),
					testAccCheckAwsLambdaAttributes(&conf),
					testAccCheckAwsLambdaAliasRoutingConfig(&conf),
					resource.TestMatchResourceAttr("aws_lambda_alias.lambda_alias_test", "arn",
						regexp.MustCompile(`^arn:aws:lambda:[a-z]+-[a-z]+-[0-9]+:\d{12}:function:`+funcName+`:`+aliasName+`$`)),
				),
			},
		},
	})
}

func testAccCheckAwsLambdaAliasDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).lambdaconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_lambda_alias" {
			continue
		}

		_, err := conn.GetAlias(&lambda.GetAliasInput{
			FunctionName: aws.String(rs.Primary.ID),
		})

		if err == nil {
			return fmt.Errorf("Lambda alias was not deleted")
		}

	}

	return nil
}

func testAccCheckAwsLambdaAliasExists(n string, mapping *lambda.AliasConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Lambda alias not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("Lambda alias not set")
		}

		conn := testAccProvider.Meta().(*AWSClient).lambdaconn

		params := &lambda.GetAliasInput{
			FunctionName: aws.String(rs.Primary.ID),
			Name:         aws.String(rs.Primary.Attributes["name"]),
		}

		getAliasConfiguration, err := conn.GetAlias(params)
		if err != nil {
			return err
		}

		*mapping = *getAliasConfiguration

		return nil
	}
}

func testAccCheckAwsLambdaAttributes(mapping *lambda.AliasConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		name := *mapping.Name
		arn := *mapping.AliasArn
		if arn == "" {
			return fmt.Errorf("Could not read Lambda alias ARN")
		}
		if name == "" {
			return fmt.Errorf("Could not read Lambda alias name")
		}
		return nil
	}
}

func testAccCheckAwsLambdaAliasRoutingConfig(mapping *lambda.AliasConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		routingConfig := mapping.RoutingConfig

		if routingConfig == nil {
			return fmt.Errorf("Could not read Lambda alias routing config")
		}
		if len(routingConfig.AdditionalVersionWeights) != 1 {
			return fmt.Errorf("Could not read Lambda alias additional version weights")
		}
		return nil
	}
}

func testAccAwsLambdaAliasConfig(roleName, policyName, attachmentName, funcName, aliasName string) string {
	return fmt.Sprintf(`
resource "aws_iam_role" "iam_for_lambda" {
  name = "%s"

  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "lambda.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_policy" "policy_for_role" {
  name        = "%s"
  path        = "/"
  description = "IAM policy for for Lamda alias testing"

  policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
      {
          "Effect": "Allow",
          "Action": [
            "lambda:*"
          ],
          "Resource": "*"
      }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "policy_attachment_for_role" {
  name       = "%s"
  roles      = ["${aws_iam_role.iam_for_lambda.name}"]
  policy_arn = "${aws_iam_policy.policy_for_role.arn}"
}

resource "aws_lambda_function" "lambda_function_test_create" {
  filename         = "test-fixtures/lambdatest.zip"
  function_name    = "%s"
  role             = "${aws_iam_role.iam_for_lambda.arn}"
  handler          = "exports.example"
  runtime          = "nodejs4.3"
  source_code_hash = "${base64sha256(file("test-fixtures/lambdatest.zip"))}"
  publish          = "true"
}

resource "aws_lambda_alias" "lambda_alias_test" {
  name             = "%s"
  description      = "a sample description"
  function_name    = "${aws_lambda_function.lambda_function_test_create.arn}"
  function_version = "1"
}`, roleName, policyName, attachmentName, funcName, aliasName)
}

func testAccAwsLambdaAliasConfigWithRoutingConfig(roleName, policyName, attachmentName, funcName, aliasName string) string {
	return fmt.Sprintf(`
resource "aws_iam_role" "iam_for_lambda" {
  name = "%s"

  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "lambda.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_policy" "policy_for_role" {
  name        = "%s"
  path        = "/"
  description = "IAM policy for for Lamda alias testing"

  policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
      {
          "Effect": "Allow",
          "Action": [
            "lambda:*"
          ],
          "Resource": "*"
      }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "policy_attachment_for_role" {
  name       = "%s"
  roles      = ["${aws_iam_role.iam_for_lambda.name}"]
  policy_arn = "${aws_iam_policy.policy_for_role.arn}"
}

resource "aws_lambda_function" "lambda_function_test_create" {
  filename         = "test-fixtures/lambdatest_modified.zip"
  function_name    = "%s"
  role             = "${aws_iam_role.iam_for_lambda.arn}"
  handler          = "exports.example"
  runtime          = "nodejs4.3"
  source_code_hash = "${base64sha256(file("test-fixtures/lambdatest_modified.zip"))}"
  publish          = "true"
}

resource "aws_lambda_alias" "lambda_alias_test" {
  name             = "%s"
  description      = "a sample description"
  function_name    = "${aws_lambda_function.lambda_function_test_create.arn}"
  function_version = "1"
  routing_config   = {
    additional_version_weights = {
      "2" = 0.5
    }
  }
}`, roleName, policyName, attachmentName, funcName, aliasName)
}
<|endoftext|>