<commit_before>package clicommand

import (
	"io/ioutil"
	"os"
	"strings"
	"time"

	"github.com/buildkite/agent/agent"
	"github.com/buildkite/agent/api"
	"github.com/buildkite/agent/cliconfig"
	"github.com/buildkite/agent/logger"
	"github.com/buildkite/agent/retry"
	"github.com/buildkite/agent/stdin"
	"github.com/urfave/cli"
)

var AnnotateHelpDescription = `Usage:

   buildkite-agent annotate <file> [arguments...]

Description:

   Build annotations allow you to customize the Buildkite build interface to
   show information that may surface from your builds. Some examples include:

   - Links to artifacts generated by your jobs
   - Test result summaries
   - Graphs that include analysis about your codebase
   - Helpful information for team members about what happened during a build

   Annotations can be written in either Markdown or HTML.

   You can update an existing annotation's body by running the annotate command
   again and provide the same context as the one you want to update. Or if you
   leave context blank, it will use the default context.

   You can also just update the style of an existing annotation by omitting the
   body and just providing a new style value.

Example:

   $ buildkite-agent annotate "All tests passed! :rocket:"
   $ cat annotation.md | buildkite-agent annotate --style "warning"
   $ buildkite-agent annotate --style "success" --context "junit"
   $ ./script/dynamic_annotation_generator | buildkite-agent annotate --style "success"`

type AnnotateConfig struct {
	Body             string `cli:"arg:0" label:"annotation body"`
	Style            string `cli:"style"`
	Context          string `cli:"context"`
	Job              string `cli:"job" validate:"required"`
	AgentAccessToken string `cli:"agent-access-token" validate:"required"`
	Endpoint         string `cli:"endpoint" validate:"required"`
	NoColor          bool   `cli:"no-color"`
	Debug            bool   `cli:"debug"`
	DebugHTTP        bool   `cli:"debug-http"`
}

var AnnotateCommand = cli.Command{
	Name:        "annotate",
	Usage:       "Annotate the build page within the Buildkite UI with text from within a Buildkite job",
	Description: AnnotateHelpDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:   "context",
			Usage:  "The context of the annotation used to differentiate this annotation from others",
			EnvVar: "BUILDKITE_ANNOTATION_CONTEXT",
		},
		cli.StringFlag{
			Name:   "style",
			Usage:  "The style of the annotation (`success`, `info`, `warning` or `error`)",
			EnvVar: "BUILDKITE_ANNOTATION_STYLE",
		},
		cli.StringFlag{
			Name:   "job",
			Value:  "",
			Usage:  "Which job should the annotation come from",
			EnvVar: "BUILDKITE_JOB_ID",
		},
		AgentAccessTokenFlag,
		EndpointFlag,
		NoColorFlag,
		DebugFlag,
		DebugHTTPFlag,
	},
	Action: func(c *cli.Context) {
		// The configuration will be loaded into this struct
		cfg := AnnotateConfig{}

		// Load the configuration
		loader := cliconfig.Loader{CLI: c, Config: &cfg}
		if err := loader.Load(); err != nil {
			logger.Fatal("%s", err)
		}

		// Setup the any global configuration options
		HandleGlobalFlags(cfg)

		var body string
		var err error

		if cfg.Body != "" {
			body = cfg.Body
		} else if stdin.IsPipe() {
			logger.Info("Reading annotation body from STDIN")

			// Actually read the file from STDIN
			stdin, err := ioutil.ReadAll(os.Stdin)
			if err != nil {
				logger.Fatal("Failed to read from STDIN: %s", err)
			}

			body = string(stdin[:])
		}

		// Trim any whitespace edges on the annotation body
		body = strings.TrimSpace(body)

		// Create the API client
		client := agent.APIClient{
			Endpoint: cfg.Endpoint,
			Token:    cfg.AgentAccessToken,
		}.Create()

		// Create the annotation we'll send to the Buildkite API
		annotation := &api.Annotation{
			Body:    body,
			Style:   cfg.Style,
			Context: cfg.Context,
		}

		// Retry the annotation a few times before giving up
		err = retry.Do(func(s *retry.Stats) error {
			// Attempt ot create the annotation
			resp, err := client.Annotations.Create(cfg.Job, annotation)

			// Don't bother retrying if the response was one of these statuses
			if resp != nil && (resp.StatusCode == 401 || resp.StatusCode == 404 || resp.StatusCode == 400) {
				s.Break()
				return err
			}

			// Show the unexpected error
			if err != nil {
				logger.Warn("%s (%s)", err, s)
			}

			return err
		}, &retry.Config{Maximum: 5, Interval: 1 * time.Second})

		// Show a fatal error if we gave up trying to create the annotation
		if err != nil {
			logger.Fatal("Failed to annotate build: %s", err)
		}

		logger.Info("Successfully annotated build")
	},
}
<commit_msg>Add interval jitter to annotations if they fail<commit_after>package clicommand

import (
	"io/ioutil"
	"os"
	"strings"
	"time"

	"github.com/buildkite/agent/agent"
	"github.com/buildkite/agent/api"
	"github.com/buildkite/agent/cliconfig"
	"github.com/buildkite/agent/logger"
	"github.com/buildkite/agent/retry"
	"github.com/buildkite/agent/stdin"
	"github.com/urfave/cli"
)

var AnnotateHelpDescription = `Usage:

   buildkite-agent annotate <file> [arguments...]

Description:

   Build annotations allow you to customize the Buildkite build interface to
   show information that may surface from your builds. Some examples include:

   - Links to artifacts generated by your jobs
   - Test result summaries
   - Graphs that include analysis about your codebase
   - Helpful information for team members about what happened during a build

   Annotations can be written in either Markdown or HTML.

   You can update an existing annotation's body by running the annotate command
   again and provide the same context as the one you want to update. Or if you
   leave context blank, it will use the default context.

   You can also just update the style of an existing annotation by omitting the
   body and just providing a new style value.

Example:

   $ buildkite-agent annotate "All tests passed! :rocket:"
   $ cat annotation.md | buildkite-agent annotate --style "warning"
   $ buildkite-agent annotate --style "success" --context "junit"
   $ ./script/dynamic_annotation_generator | buildkite-agent annotate --style "success"`

type AnnotateConfig struct {
	Body             string `cli:"arg:0" label:"annotation body"`
	Style            string `cli:"style"`
	Context          string `cli:"context"`
	Job              string `cli:"job" validate:"required"`
	AgentAccessToken string `cli:"agent-access-token" validate:"required"`
	Endpoint         string `cli:"endpoint" validate:"required"`
	NoColor          bool   `cli:"no-color"`
	Debug            bool   `cli:"debug"`
	DebugHTTP        bool   `cli:"debug-http"`
}

var AnnotateCommand = cli.Command{
	Name:        "annotate",
	Usage:       "Annotate the build page within the Buildkite UI with text from within a Buildkite job",
	Description: AnnotateHelpDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:   "context",
			Usage:  "The context of the annotation used to differentiate this annotation from others",
			EnvVar: "BUILDKITE_ANNOTATION_CONTEXT",
		},
		cli.StringFlag{
			Name:   "style",
			Usage:  "The style of the annotation (`success`, `info`, `warning` or `error`)",
			EnvVar: "BUILDKITE_ANNOTATION_STYLE",
		},
		cli.StringFlag{
			Name:   "job",
			Value:  "",
			Usage:  "Which job should the annotation come from",
			EnvVar: "BUILDKITE_JOB_ID",
		},
		AgentAccessTokenFlag,
		EndpointFlag,
		NoColorFlag,
		DebugFlag,
		DebugHTTPFlag,
	},
	Action: func(c *cli.Context) {
		// The configuration will be loaded into this struct
		cfg := AnnotateConfig{}

		// Load the configuration
		loader := cliconfig.Loader{CLI: c, Config: &cfg}
		if err := loader.Load(); err != nil {
			logger.Fatal("%s", err)
		}

		// Setup the any global configuration options
		HandleGlobalFlags(cfg)

		var body string
		var err error

		if cfg.Body != "" {
			body = cfg.Body
		} else if stdin.IsPipe() {
			logger.Info("Reading annotation body from STDIN")

			// Actually read the file from STDIN
			stdin, err := ioutil.ReadAll(os.Stdin)
			if err != nil {
				logger.Fatal("Failed to read from STDIN: %s", err)
			}

			body = string(stdin[:])
		}

		// Trim any whitespace edges on the annotation body
		body = strings.TrimSpace(body)

		// Create the API client
		client := agent.APIClient{
			Endpoint: cfg.Endpoint,
			Token:    cfg.AgentAccessToken,
		}.Create()

		// Create the annotation we'll send to the Buildkite API
		annotation := &api.Annotation{
			Body:    body,
			Style:   cfg.Style,
			Context: cfg.Context,
		}

		// Retry the annotation a few times before giving up
		err = retry.Do(func(s *retry.Stats) error {
			// Attempt ot create the annotation
			resp, err := client.Annotations.Create(cfg.Job, annotation)

			// Don't bother retrying if the response was one of these statuses
			if resp != nil && (resp.StatusCode == 401 || resp.StatusCode == 404 || resp.StatusCode == 400) {
				s.Break()
				return err
			}

			// Show the unexpected error
			if err != nil {
				logger.Warn("%s (%s)", err, s)
			}

			return err
		}, &retry.Config{Maximum: 5, Interval: 1 * time.Second, Jitter: true})

		// Show a fatal error if we gave up trying to create the annotation
		if err != nil {
			logger.Fatal("Failed to annotate build: %s", err)
		}

		logger.Info("Successfully annotated build")
	},
}
<|endoftext|>


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

package binprot

import (
	"bufio"
	"encoding/binary"
	"io"

	"github.com/netflix/rend/client/common"
)

type BinProt struct{}

func consumeResponse(r *bufio.Reader) ([]byte, error) {
	res, err := readRes(r)
	if err != nil {
		return nil, err
	}

	apperr := statusToError(res.Status)

	// read body in regardless of the error in the header
	buf := make([]byte, res.BodyLen)
	io.ReadFull(r, buf)

	// ignore extras for now
	buf = buf[res.ExtraLen:]

	resPool.Put(res)

	if apperr != nil && srsErr(apperr) {
		return buf, apperr
	}

	return buf, err
}

func consumeBatchResponse(r *bufio.Reader) ([][]byte, error) {
	opcode := uint8(Get)
	var apperr error
	ret := make([][]byte, 0)

	for opcode != Noop {
		res, err := readRes(r)
		if err != nil {
			return nil, err
		}

		opcode = res.Opcode
		apperr = statusToError(res.Status)

		// read body in regardless of the error in the header
		buf := make([]byte, res.BodyLen)
		io.ReadFull(r, buf)
		ret = append(ret, buf)
	}

	return ret, apperr
}

func (b BinProt) Set(rw *bufio.ReadWriter, key, value []byte) error {
	// set packet contains the req header, flags, and expiration
	// flags are irrelevant, and are thus zero.
	// expiration could be important, so hammer with random values from 1 sec up to 1 hour

	// Header
	bodylen := 8 + len(key) + len(value)
	writeReq(rw, Set, len(key), 8, bodylen)
	// Extras
	binary.Write(rw, binary.BigEndian, uint32(0))
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body / data
	rw.Write(key)
	rw.Write(value)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}

func (b BinProt) Add(rw *bufio.ReadWriter, key, value []byte) error {
	// add packet contains the req header, flags, and expiration
	// flags are irrelevant, and are thus zero.
	// expiration could be important, so hammer with random values from 1 sec up to 1 hour

	// Header
	bodylen := 8 + len(key) + len(value)
	writeReq(rw, Add, len(key), 8, bodylen)
	// Extras
	binary.Write(rw, binary.BigEndian, uint32(0))
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body / data
	rw.Write(key)
	rw.Write(value)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}

func (b BinProt) Replace(rw *bufio.ReadWriter, key, value []byte) error {
	// replace packet contains the req header, flags, and expiration
	// flags are irrelevant, and are thus zero.
	// expiration could be important, so hammer with random values from 1 sec up to 1 hour

	// Header
	bodylen := 8 + len(key) + len(value)
	writeReq(rw, Replace, len(key), 8, bodylen)
	// Extras
	binary.Write(rw, binary.BigEndian, uint32(0))
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body / data
	rw.Write(key)
	rw.Write(value)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}

func (b BinProt) Get(rw *bufio.ReadWriter, key []byte) ([]byte, error) {
	// Header
	writeReq(rw, Get, len(key), 0, len(key))
	// Body
	rw.Write(key)

	rw.Flush()

	// consume all of the response and return
	return consumeResponse(rw.Reader)
}

func (b BinProt) BatchGet(rw *bufio.ReadWriter, keys [][]byte) ([][]byte, error) {
	for _, key := range keys {
		// Header
		writeReq(rw, GetQ, len(key), 0, len(key))
		// Body
		rw.Write(key)
	}

	writeReq(rw, Noop, 0, 0, 0)

	rw.Flush()

	// consume all of the responses and return
	return consumeBatchResponse(rw.Reader)
}

func (b BinProt) GAT(rw *bufio.ReadWriter, key []byte) ([]byte, error) {
	// Header
	writeReq(rw, GAT, len(key), 4, len(key))
	// Extras
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body
	rw.Write(key)

	rw.Flush()

	// consume all of the response and return
	return consumeResponse(rw.Reader)
}

func (b BinProt) Delete(rw *bufio.ReadWriter, key []byte) error {
	// Header
	writeReq(rw, Delete, len(key), 0, len(key))
	// Body
	rw.Write(key)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}

func (b BinProt) Touch(rw *bufio.ReadWriter, key []byte) error {
	// Header
	writeReq(rw, Touch, len(key), 4, len(key)+4)
	// Extras
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body
	rw.Write(key)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}
<commit_msg>filter extras in batch response<commit_after>// Copyright 2015 Netflix, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package binprot

import (
	"bufio"
	"encoding/binary"
	"io"

	"github.com/netflix/rend/client/common"
)

type BinProt struct{}

func consumeResponse(r *bufio.Reader) ([]byte, error) {
	res, err := readRes(r)
	if err != nil {
		return nil, err
	}

	apperr := statusToError(res.Status)

	// read body in regardless of the error in the header
	buf := make([]byte, res.BodyLen)
	io.ReadFull(r, buf)

	// ignore extras for now
	buf = buf[res.ExtraLen:]

	resPool.Put(res)

	if apperr != nil && srsErr(apperr) {
		return buf, apperr
	}

	return buf, err
}

func consumeBatchResponse(r *bufio.Reader) ([][]byte, error) {
	opcode := uint8(Get)
	var apperr error
	ret := make([][]byte, 0)

	for opcode != Noop {
		res, err := readRes(r)
		if err != nil {
			return nil, err
		}

		opcode = res.Opcode
		apperr = statusToError(res.Status)

		// read body in regardless of the error in the header
		buf := make([]byte, res.BodyLen)
		io.ReadFull(r, buf)

		// ignore extras for now
		buf = buf[res.ExtraLen:]

		ret = append(ret, buf)
	}

	return ret, apperr
}

func (b BinProt) Set(rw *bufio.ReadWriter, key, value []byte) error {
	// set packet contains the req header, flags, and expiration
	// flags are irrelevant, and are thus zero.
	// expiration could be important, so hammer with random values from 1 sec up to 1 hour

	// Header
	bodylen := 8 + len(key) + len(value)
	writeReq(rw, Set, len(key), 8, bodylen)
	// Extras
	binary.Write(rw, binary.BigEndian, uint32(0))
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body / data
	rw.Write(key)
	rw.Write(value)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}

func (b BinProt) Add(rw *bufio.ReadWriter, key, value []byte) error {
	// add packet contains the req header, flags, and expiration
	// flags are irrelevant, and are thus zero.
	// expiration could be important, so hammer with random values from 1 sec up to 1 hour

	// Header
	bodylen := 8 + len(key) + len(value)
	writeReq(rw, Add, len(key), 8, bodylen)
	// Extras
	binary.Write(rw, binary.BigEndian, uint32(0))
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body / data
	rw.Write(key)
	rw.Write(value)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}

func (b BinProt) Replace(rw *bufio.ReadWriter, key, value []byte) error {
	// replace packet contains the req header, flags, and expiration
	// flags are irrelevant, and are thus zero.
	// expiration could be important, so hammer with random values from 1 sec up to 1 hour

	// Header
	bodylen := 8 + len(key) + len(value)
	writeReq(rw, Replace, len(key), 8, bodylen)
	// Extras
	binary.Write(rw, binary.BigEndian, uint32(0))
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body / data
	rw.Write(key)
	rw.Write(value)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}

func (b BinProt) Get(rw *bufio.ReadWriter, key []byte) ([]byte, error) {
	// Header
	writeReq(rw, Get, len(key), 0, len(key))
	// Body
	rw.Write(key)

	rw.Flush()

	// consume all of the response and return
	return consumeResponse(rw.Reader)
}

func (b BinProt) BatchGet(rw *bufio.ReadWriter, keys [][]byte) ([][]byte, error) {
	for _, key := range keys {
		// Header
		writeReq(rw, GetQ, len(key), 0, len(key))
		// Body
		rw.Write(key)
	}

	writeReq(rw, Noop, 0, 0, 0)

	rw.Flush()

	// consume all of the responses and return
	return consumeBatchResponse(rw.Reader)
}

func (b BinProt) GAT(rw *bufio.ReadWriter, key []byte) ([]byte, error) {
	// Header
	writeReq(rw, GAT, len(key), 4, len(key))
	// Extras
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body
	rw.Write(key)

	rw.Flush()

	// consume all of the response and return
	return consumeResponse(rw.Reader)
}

func (b BinProt) Delete(rw *bufio.ReadWriter, key []byte) error {
	// Header
	writeReq(rw, Delete, len(key), 0, len(key))
	// Body
	rw.Write(key)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}

func (b BinProt) Touch(rw *bufio.ReadWriter, key []byte) error {
	// Header
	writeReq(rw, Touch, len(key), 4, len(key)+4)
	// Extras
	binary.Write(rw, binary.BigEndian, common.Exp())
	// Body
	rw.Write(key)

	rw.Flush()

	// consume all of the response and discard
	_, err := consumeResponse(rw.Reader)
	return err
}
<|endoftext|>


<commit_before>// Copyright 2017 The etcd 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 naming provides an etcd-backed gRPC resolver for discovering gRPC services.
//
// To use, first import the packages:
//
//	import (
//		"go.etcd.io/etcd/clientv3"
//		etcdnaming "go.etcd.io/etcd/clientv3/naming"
//
//		"google.golang.org/grpc"
//		"google.golang.org/grpc/naming"
//	)
//
// First, register new endpoint addresses for a service:
//
//	func etcdAdd(c *clientv3.Client, service, addr string) error {
//		r := &etcdnaming.GRPCResolver{Client: c}
//		return r.Update(c.Ctx(), service, naming.Update{Op: naming.Add, Addr: addr})
//	}
//
// Dial an RPC service using the etcd gRPC resolver and a gRPC Balancer:
//
//	func etcdDial(c *clientv3.Client, service string) (*grpc.ClientConn, error) {
//		r := &etcdnaming.GRPCResolver{Client: c}
//		b := grpc.RoundRobin(r)
//		return grpc.Dial(service, grpc.WithBalancer(b))
//	}
//
// Optionally, force delete an endpoint:
//
//	func etcdDelete(c *clientv3, service, addr string) error {
//		r := &etcdnaming.GRPCResolver{Client: c}
//		return r.Update(c.Ctx(), "my-service", naming.Update{Op: naming.Delete, Addr: "1.2.3.4"})
//	}
//
// Or register an expiring endpoint with a lease:
//
//	func etcdLeaseAdd(c *clientv3.Client, lid clientv3.LeaseID, service, addr string) error {
//		r := &etcdnaming.GRPCResolver{Client: c}
//		return r.Update(c.Ctx(), service, naming.Update{Op: naming.Add, Addr: addr}, clientv3.WithLease(lid))
//	}
//
package naming
<commit_msg>doc: fix document example error<commit_after>// Copyright 2017 The etcd 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 naming provides an etcd-backed gRPC resolver for discovering gRPC services.
//
// To use, first import the packages:
//
//	import (
//		"go.etcd.io/etcd/clientv3"
//		etcdnaming "go.etcd.io/etcd/clientv3/naming"
//
//		"google.golang.org/grpc"
//		"google.golang.org/grpc/naming"
//	)
//
// First, register new endpoint addresses for a service:
//
//	func etcdAdd(c *clientv3.Client, service, addr string) error {
//		r := &etcdnaming.GRPCResolver{Client: c}
//		return r.Update(c.Ctx(), service, naming.Update{Op: naming.Add, Addr: addr})
//	}
//
// Dial an RPC service using the etcd gRPC resolver and a gRPC Balancer:
//
//	func etcdDial(c *clientv3.Client, service string) (*grpc.ClientConn, error) {
//		r := &etcdnaming.GRPCResolver{Client: c}
//		b := grpc.RoundRobin(r)
//		return grpc.Dial(service, grpc.WithBalancer(b))
//	}
//
// Optionally, force delete an endpoint:
//
//	func etcdDelete(c *clientv3, service, addr string) error {
//		r := &etcdnaming.GRPCResolver{Client: c}
//		return r.Update(c.Ctx(), service, naming.Update{Op: naming.Delete, Addr: "1.2.3.4"})
//	}
//
// Or register an expiring endpoint with a lease:
//
//	func etcdLeaseAdd(c *clientv3.Client, lid clientv3.LeaseID, service, addr string) error {
//		r := &etcdnaming.GRPCResolver{Client: c}
//		return r.Update(c.Ctx(), service, naming.Update{Op: naming.Add, Addr: addr}, clientv3.WithLease(lid))
//	}
//
package naming
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package app

import (
	"flag"
	"fmt"
	"strings"

	"github.com/spf13/viper"
)

const (
	suffixWorkers             = "workers"
	suffixServerQueueSize     = "server-queue-size"
	suffixServerMaxPacketSize = "server-max-packet-size"
	suffixServerHostPort      = "server-host-port"
	collectorHostPort         = "collector.host-port"
	httpServerHostPort        = "http-server.host-port"
	discoveryMinPeers         = "discovery.min-peers"
)

var defaultProcessors = []struct {
	model    model
	protocol protocol
	port     string
}{
	{model: "zipkin", protocol: "compact", port: ":5775"},
	{model: "jaeger", protocol: "compact", port: ":6831"},
	{model: "jaeger", protocol: "binary", port: ":6832"},
}

// AddFlags adds flags for Builder.
func AddFlags(flags *flag.FlagSet) {
	for _, processor := range defaultProcessors {
		prefix := fmt.Sprintf("processor.%s-%s.", processor.model, processor.protocol)
		flags.Int(prefix+suffixWorkers, defaultServerWorkers, "how many workers the processor should run")
		flags.Int(prefix+suffixServerQueueSize, defaultQueueSize, "length of the queue for the UDP server")
		flags.Int(prefix+suffixServerMaxPacketSize, defaultMaxPacketSize, "max packet size for the UDP server")
		flags.String(prefix+suffixServerHostPort, processor.port, "host:port for the UDP server")
	}
	flags.String(
		collectorHostPort,
		"",
		"comma-separated string representing host:ports of a static list of collectors to connect to directly (e.g. when not using service discovery)")
	flags.String(
		httpServerHostPort,
		defaultHTTPServerHostPort,
		"host:port of the http server (e.g. for /sampling point and /baggage endpoint)")
	flags.Int(
		discoveryMinPeers,
		defaultMinPeers,
		"if using service discovery, the min number of connections to maintain to the backend")
}

// InitFromViper initializes Builder with properties retrieved from Viper.
func (b *Builder) InitFromViper(v *viper.Viper) {
	b.Metrics.InitFromViper(v)

	for _, processor := range defaultProcessors {
		prefix := fmt.Sprintf("processor.%s-%s.", processor.model, processor.protocol)
		p := &ProcessorConfiguration{Model: processor.model, Protocol: processor.protocol}
		p.Workers = v.GetInt(prefix + suffixWorkers)
		p.Server.QueueSize = v.GetInt(prefix + suffixServerQueueSize)
		p.Server.MaxPacketSize = v.GetInt(prefix + suffixServerMaxPacketSize)
		p.Server.HostPort = v.GetString(prefix + suffixServerHostPort)
		b.Processors = append(b.Processors, *p)
	}

	if len(v.GetString(collectorHostPort)) > 0 {
		b.CollectorHostPorts = strings.Split(v.GetString(collectorHostPort), ",")
	}
	b.HTTPServer.HostPort = v.GetString(httpServerHostPort)
	b.DiscoveryMinPeers = v.GetInt(discoveryMinPeers)
}
<commit_msg>Fix hostPort variable name (#291)<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package app

import (
	"flag"
	"fmt"
	"strings"

	"github.com/spf13/viper"
)

const (
	suffixWorkers             = "workers"
	suffixServerQueueSize     = "server-queue-size"
	suffixServerMaxPacketSize = "server-max-packet-size"
	suffixServerHostPort      = "server-host-port"
	collectorHostPort         = "collector.host-port"
	httpServerHostPort        = "http-server.host-port"
	discoveryMinPeers         = "discovery.min-peers"
)

var defaultProcessors = []struct {
	model    model
	protocol protocol
	hostPort string
}{
	{model: "zipkin", protocol: "compact", hostPort: ":5775"},
	{model: "jaeger", protocol: "compact", hostPort: ":6831"},
	{model: "jaeger", protocol: "binary", hostPort: ":6832"},
}

// AddFlags adds flags for Builder.
func AddFlags(flags *flag.FlagSet) {
	for _, processor := range defaultProcessors {
		prefix := fmt.Sprintf("processor.%s-%s.", processor.model, processor.protocol)
		flags.Int(prefix+suffixWorkers, defaultServerWorkers, "how many workers the processor should run")
		flags.Int(prefix+suffixServerQueueSize, defaultQueueSize, "length of the queue for the UDP server")
		flags.Int(prefix+suffixServerMaxPacketSize, defaultMaxPacketSize, "max packet size for the UDP server")
		flags.String(prefix+suffixServerHostPort, processor.hostPort, "host:port for the UDP server")
	}
	flags.String(
		collectorHostPort,
		"",
		"comma-separated string representing host:ports of a static list of collectors to connect to directly (e.g. when not using service discovery)")
	flags.String(
		httpServerHostPort,
		defaultHTTPServerHostPort,
		"host:port of the http server (e.g. for /sampling point and /baggage endpoint)")
	flags.Int(
		discoveryMinPeers,
		defaultMinPeers,
		"if using service discovery, the min number of connections to maintain to the backend")
}

// InitFromViper initializes Builder with properties retrieved from Viper.
func (b *Builder) InitFromViper(v *viper.Viper) {
	b.Metrics.InitFromViper(v)

	for _, processor := range defaultProcessors {
		prefix := fmt.Sprintf("processor.%s-%s.", processor.model, processor.protocol)
		p := &ProcessorConfiguration{Model: processor.model, Protocol: processor.protocol}
		p.Workers = v.GetInt(prefix + suffixWorkers)
		p.Server.QueueSize = v.GetInt(prefix + suffixServerQueueSize)
		p.Server.MaxPacketSize = v.GetInt(prefix + suffixServerMaxPacketSize)
		p.Server.HostPort = v.GetString(prefix + suffixServerHostPort)
		b.Processors = append(b.Processors, *p)
	}

	if len(v.GetString(collectorHostPort)) > 0 {
		b.CollectorHostPorts = strings.Split(v.GetString(collectorHostPort), ",")
	}
	b.HTTPServer.HostPort = v.GetString(httpServerHostPort)
	b.DiscoveryMinPeers = v.GetInt(discoveryMinPeers)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"sync"
	"time"

	"gopkg.in/adsi.v0"
	"gopkg.in/dfsr.v0/config"
	"gopkg.in/dfsr.v0/core"
	"gopkg.in/dfsr.v0/helper"
)

var domainFlag string
var groupFlag regexSlice
var fromFlag regexSlice
var toFlag regexSlice
var memberFlag regexSlice
var loopFlag uintOrInf
var delaySecondsFlag uint
var cacheSecondsFlag uint
var skipFlag regexSlice
var minFlag uint
var verboseFlag bool

const defaultLoopValue = 1

func init() {
	flag.StringVar(&domainFlag, "d", "", "domain to query")
	flag.Var(&groupFlag, "g", "group to query")
	flag.Var(&fromFlag, "f", "regex of source hostname")
	flag.Var(&toFlag, "t", "regex of dest hostname")
	flag.Var(&memberFlag, "m", "regex of member hostname (matches either dest or source)")
	flag.Var(&loopFlag, "loop", "number of iterations or \"infinite\"")
	flag.UintVar(&delaySecondsFlag, "delay", 5, "number of seconds to deley between loops")
	flag.UintVar(&cacheSecondsFlag, "cache", 5, "number of seconds to cache vectors")
	flag.Var(&skipFlag, "skip", "regex of hostname to skip")
	flag.UintVar(&minFlag, "min", 0, "minimum backlog to display")
	flag.BoolVar(&verboseFlag, "v", false, "verbose")
}

func main() {
	flag.Parse()

	if !loopFlag.Present {
		loopFlag.Value = defaultLoopValue
	}

	if !loopFlag.Inf && loopFlag.Value == 0 {
		return
	}

	domain, connections, err := setup(domainFlag, groupFlag, fromFlag, toFlag, memberFlag, skipFlag)
	if err != nil {
		log.Fatal(err)
	}

	client, err := helper.NewClient()
	if err != nil {
		log.Fatal(err)
	}

	if cacheSecondsFlag > 0 {
		client.Cache(time.Duration(cacheSecondsFlag) * time.Second)
	}

	if loopFlag.Inf {
		for loop := uint(0); ; loop++ {
			run(domain, loop, minFlag, client, connections)
			time.Sleep(time.Duration(delaySecondsFlag) * time.Second)
		}
	} else {
		for loop := uint(0); loop < loopFlag.Value; loop++ {
			run(domain, loop, minFlag, client, connections)
			if loop+1 < loopFlag.Value {
				fmt.Println("")
				time.Sleep(time.Duration(delaySecondsFlag) * time.Second)
			}
		}
	}
}

func run(domain string, iteration uint, min uint, client *helper.Client, connections []core.Backlog) {
	var wg sync.WaitGroup
	wg.Add(len(connections))

	//fmt.Printf("[query %v] %s\n", iteration, domain)
	fmt.Printf("%-50s %-50s %-50s %-15s %s\n", "Group", "Source", "Destination", "Backlog", "Time")
	fmt.Printf("%-50s %-50s %-50s %-15s %s\n", "-----", "------", "-----------", "-------", "----")

	start := time.Now()

	for i := 0; i < len(connections); i++ {
		go computeBacklog(client, &connections[i], &wg)
	}

	wg.Wait()

	finish := time.Now()

	for i := 0; i < len(connections); i++ {
		c := &connections[i]

		if c.Sum() < min {
			continue
		}

		fmt.Printf("%-50s %-50s %-50s ", c.Group.Name, c.From, c.To)
		if c.Err != nil {
			fmt.Printf("%-15v ", c.Err)
			//fmt.Printf("%-15v ", c.Call)
		} else {
			fmt.Printf("%-15s ", fmt.Sprint(c.Sum()))
		}
		fmt.Printf("%v\n", c.Call.Duration())
		if verboseFlag {
			fmt.Printf("Call: %v\n", c.Call)
		}
	}

	fmt.Printf("Total Time: %v\n", finish.Sub(start))
}

func setup(domain string, groupRegex, fromRegex, toRegex, memberRegex, skipRegex regexSlice) (dom string, connections []core.Backlog, err error) {
	client, err := adsi.NewClient()
	if err != nil {
		return "", nil, err
	}
	defer client.Close()

	if domain == "" {
		domain, err = dnc(client)
		if err != nil {
			return "", nil, err
		}
	}
	dom = domain

	d, err := config.Domain(client, domain)
	if err != nil {
		return domain, nil, err
	}

	for g := 0; g < len(d.Groups); g++ {
		group := &d.Groups[g]
		if !isMatch(group.Name, groupRegex, true) {
			continue
		}

		for m := 0; m < len(group.Members); m++ {
			member := &group.Members[m]
			to := member.Computer.Host
			if to == "" {
				continue
			}
			if isMatch(to, skipRegex, false) {
				continue
			}
			if !isMatch(to, toRegex, true) {
				continue
			}

			for c := 0; c < len(member.Connections); c++ {
				conn := &member.Connections[c]
				from := conn.Computer.Host
				if from == "" {
					continue
				}
				if !conn.Enabled {
					continue
				}
				if isMatch(from, skipRegex, false) {
					continue
				}
				if !isMatch(from, fromRegex, true) {
					continue
				}
				if !isMatch(from, memberRegex, true) && !isMatch(to, memberRegex, true) {
					continue
				}

				connections = append(connections, core.Backlog{
					Group: group,
					From:  from,
					To:    to,
				})
			}
		}
	}
	return
}

func computeBacklog(client *helper.Client, backlog *core.Backlog, wg *sync.WaitGroup) {
	var values []int
	values, backlog.Call, backlog.Err = client.Backlog(backlog.From, backlog.To, *backlog.Group.ID)
	if n := len(values); n == len(backlog.Group.Folders) {
		backlog.Folders = make([]core.FolderBacklog, n)
		for v := 0; v < n; v++ {
			backlog.Folders[v].Folder = &backlog.Group.Folders[v]
			backlog.Folders[v].Backlog = values[v]
		}
	}
	wg.Done()
}
<commit_msg>Updated  backlog command to use context<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"sync"
	"time"

	"gopkg.in/adsi.v0"
	"gopkg.in/dfsr.v0/config"
	"gopkg.in/dfsr.v0/core"
	"gopkg.in/dfsr.v0/helper"
)

var (
	domainFlag         string
	groupFlag          regexSlice
	fromFlag           regexSlice
	toFlag             regexSlice
	memberFlag         regexSlice
	loopFlag           uintOrInf
	timeoutSecondsFlag uintOrInf
	delaySecondsFlag   uint
	cacheSecondsFlag   uint
	skipFlag           regexSlice
	minFlag            uint
	verboseFlag        bool
)

const (
	defaultLoopValue           = 1
	defaultTimeoutSecondsValue = 30
)

func init() {
	flag.StringVar(&domainFlag, "d", "", "domain to query")
	flag.Var(&groupFlag, "g", "group to query")
	flag.Var(&fromFlag, "f", "regex of source hostname")
	flag.Var(&toFlag, "t", "regex of dest hostname")
	flag.Var(&memberFlag, "m", "regex of member hostname (matches either dest or source)")
	flag.Var(&loopFlag, "loop", "number of iterations or \"infinite\"")
	flag.Var(&timeoutSecondsFlag, "timeout", "number of seconds before timeout occurs or \"infininte\"")
	flag.UintVar(&delaySecondsFlag, "delay", 5, "number of seconds to delay between loops")
	flag.UintVar(&cacheSecondsFlag, "cache", 5, "number of seconds to cache vectors")
	flag.Var(&skipFlag, "skip", "regex of hostname to skip")
	flag.UintVar(&minFlag, "min", 0, "minimum backlog to display")
	flag.BoolVar(&verboseFlag, "v", false, "verbose")
}

func main() {
	flag.Parse()

	if !loopFlag.Present {
		loopFlag.Value = defaultLoopValue
	}

	if !loopFlag.Inf && loopFlag.Value == 0 {
		return
	}

	if !timeoutSecondsFlag.Present {
		timeoutSecondsFlag.Value = defaultTimeoutSecondsValue
	}

	domain, connections, err := setup(domainFlag, groupFlag, fromFlag, toFlag, memberFlag, skipFlag)
	if err != nil {
		log.Fatal(err)
	}

	client, err := helper.NewClient()
	if err != nil {
		log.Fatal(err)
	}

	if cacheSecondsFlag > 0 {
		client.Cache(time.Duration(cacheSecondsFlag) * time.Second)
	}

	var (
		ctx    context.Context
		cancel context.CancelFunc
	)
	if !timeoutSecondsFlag.Inf && timeoutSecondsFlag.Value != 0 {
		ctx, cancel = context.WithTimeout(context.Background(), time.Duration(timeoutSecondsFlag.Value)*time.Second)
	} else {
		ctx, cancel = context.WithCancel(context.Background())
	}

	defer cancel()

	if loopFlag.Inf {
		for loop := uint(0); ; loop++ {
			run(ctx, domain, loop, minFlag, client, connections)
			time.Sleep(time.Duration(delaySecondsFlag) * time.Second)
		}
	} else {
		for loop := uint(0); loop < loopFlag.Value; loop++ {
			run(ctx, domain, loop, minFlag, client, connections)
			if loop+1 < loopFlag.Value {
				fmt.Println("")
				time.Sleep(time.Duration(delaySecondsFlag) * time.Second)
			}
		}
	}
}

func run(ctx context.Context, domain string, iteration uint, min uint, client *helper.Client, connections []core.Backlog) {
	var wg sync.WaitGroup
	wg.Add(len(connections))

	//fmt.Printf("[query %v] %s\n", iteration, domain)
	fmt.Printf("%-50s %-50s %-50s %-15s %s\n", "Group", "Source", "Destination", "Backlog", "Time")
	fmt.Printf("%-50s %-50s %-50s %-15s %s\n", "-----", "------", "-----------", "-------", "----")

	start := time.Now()

	for i := 0; i < len(connections); i++ {
		go computeBacklog(ctx, client, &connections[i], &wg)
	}

	wg.Wait()

	finish := time.Now()

	for i := 0; i < len(connections); i++ {
		c := &connections[i]

		if c.Sum() < min {
			continue
		}

		fmt.Printf("%-50s %-50s %-50s ", c.Group.Name, c.From, c.To)
		if c.Err != nil {
			fmt.Printf("%-15v ", c.Err)
			//fmt.Printf("%-15v ", c.Call)
		} else {
			fmt.Printf("%-15s ", fmt.Sprint(c.Sum()))
		}
		fmt.Printf("%v\n", c.Call.Duration())
		if verboseFlag {
			fmt.Printf("Call: %v\n", c.Call)
		}
	}

	fmt.Printf("Total Time: %v\n", finish.Sub(start))
}

func setup(domain string, groupRegex, fromRegex, toRegex, memberRegex, skipRegex regexSlice) (dom string, connections []core.Backlog, err error) {
	client, err := adsi.NewClient()
	if err != nil {
		return "", nil, err
	}
	defer client.Close()

	if domain == "" {
		domain, err = dnc(client)
		if err != nil {
			return "", nil, err
		}
	}
	dom = domain

	d, err := config.Domain(client, domain)
	if err != nil {
		return domain, nil, err
	}

	for g := 0; g < len(d.Groups); g++ {
		group := &d.Groups[g]
		if !isMatch(group.Name, groupRegex, true) {
			continue
		}

		for m := 0; m < len(group.Members); m++ {
			member := &group.Members[m]
			to := member.Computer.Host
			if to == "" {
				continue
			}
			if isMatch(to, skipRegex, false) {
				continue
			}
			if !isMatch(to, toRegex, true) {
				continue
			}

			for c := 0; c < len(member.Connections); c++ {
				conn := &member.Connections[c]
				from := conn.Computer.Host
				if from == "" {
					continue
				}
				if !conn.Enabled {
					continue
				}
				if isMatch(from, skipRegex, false) {
					continue
				}
				if !isMatch(from, fromRegex, true) {
					continue
				}
				if !isMatch(from, memberRegex, true) && !isMatch(to, memberRegex, true) {
					continue
				}

				connections = append(connections, core.Backlog{
					Group: group,
					From:  from,
					To:    to,
				})
			}
		}
	}
	return
}

func computeBacklog(ctx context.Context, client *helper.Client, backlog *core.Backlog, wg *sync.WaitGroup) {
	var values []int
	values, backlog.Call, backlog.Err = client.Backlog(ctx, backlog.From, backlog.To, *backlog.Group.ID)
	if n := len(values); n == len(backlog.Group.Folders) {
		backlog.Folders = make([]core.FolderBacklog, n)
		for v := 0; v < n; v++ {
			backlog.Folders[v].Folder = &backlog.Group.Folders[v]
			backlog.Folders[v].Backlog = values[v]
		}
	}
	wg.Done()
}
<|endoftext|>


<commit_before>package main

//-----------------------------------------------------------------------------
// Package factored import statement:
//-----------------------------------------------------------------------------

import (

	// Stdlib:
	"os"
	"path"
	"regexp"
	"runtime"
	"strings"

	// Local:
	"github.com/katosys/kato/providers/cloud/ec2"
	"github.com/katosys/kato/providers/cloud/pkt"
	"github.com/katosys/kato/providers/dns/ns1"
	"github.com/katosys/kato/udata"

	// Community:
	log "github.com/Sirupsen/logrus"
	"gopkg.in/alecthomas/kingpin.v2"
)

//-----------------------------------------------------------------------------
// katoctl root level command:
//-----------------------------------------------------------------------------

var app = kingpin.New("katoctl", "Katoctl defines and deploys Kato's infrastructure.")

//----------------------------------------------------------------------------
// func init() is called after all the variable declarations in the package
// have evaluated their initializers, and those are evaluated only after all
// the imported packages have been initialized:
//----------------------------------------------------------------------------

func init() {

	// Customize kingpin:
	app.Version("v0.1.0-beta").Author("Marc Villacorta Morera")
	app.UsageTemplate(usageTemplate)
	app.HelpFlag.Short('h')

	// Customize the default logger:
	log.SetFormatter(&log.TextFormatter{ForceColors: true})
	log.SetOutput(os.Stderr)
	log.SetLevel(log.InfoLevel)
	log.AddHook(contextHook{})
}

//----------------------------------------------------------------------------
// Entry point:
//----------------------------------------------------------------------------

func main() {

	// Sub-command selector:
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {

	//---------------
	// katoctl udata
	//---------------

	case cmdUdata.FullCommand():

		udata := udata.Data{
			QuorumCount:         *flUdataQuorumCount,
			ClusterID:           *flUdataClusterID,
			HostName:            *flUdataHostName,
			HostID:              *flUdataHostID,
			Domain:              *flUdataDomain,
			Roles:               strings.Split(*flUdataRoles, ","),
			Ns1ApiKey:           *flUdataNs1Apikey,
			CaCert:              *flUdataCaCert,
			EtcdToken:           *flUdataEtcdToken,
			GzipUdata:           *flUdataGzipUdata,
			FlannelNetwork:      *flUdataFlannelNetwork,
			FlannelSubnetLen:    *flUdataFlannelSubnetLen,
			FlannelSubnetMin:    *flUdataFlannelSubnetMin,
			FlannelSubnetMax:    *flUdataFlannelSubnetMax,
			FlannelBackend:      *flUdataFlannelBackend,
			RexrayStorageDriver: *flUdataRexrayStorageDriver,
			RexrayEndpointIP:    *flUdataRexrayEndpointIP,
			Ec2Region:           *flUdataEc2Region,
			IaasProvider:        *flUdataIaasProvider,
			SysdigAccessKey:     *flUdataSysdigAccessKey,
			DatadogAPIKey:       *flUdataDatadogAPIKey,
			StubZones:           *flUdataStubZones,
			Prometheus:          *flUdataPrometheus,
		}

		udata.Render()

	//--------------------
	// katoctl pkt deploy
	//--------------------

	case cmdPktDeploy.FullCommand():

		pkt := pkt.Data{}
		pkt.Deploy()

	//-------------------
	// katoctl pkt setup
	//-------------------

	case cmdPktSetup.FullCommand():

		pkt := pkt.Data{}
		pkt.Setup()

	//-----------------
	// katoctl pkt run
	//-----------------

	case cmdPktRun.FullCommand():

		pkt := pkt.Data{
			APIKey:    *flPktRunAPIKey,
			HostName:  *flPktRunHostName,
			ProjectID: *flPktRunProjectID,
			Plan:      *flPktRunPlan,
			OS:        *flPktRunOS,
			Facility:  *flPktRunFacility,
			Billing:   *flPktRunBilling,
		}

		pkt.Run()

	//--------------------
	// katoctl ec2 deploy
	//--------------------

	case cmdEc2Deploy.FullCommand():

		ec2 := ec2.Data{
			State: ec2.State{
				ClusterID:        *flEc2DeployClusterID,
				Channel:          *flEc2DeployChannel,
				EtcdToken:        *flEc2DeployEtcdToken,
				Ns1ApiKey:        *flEc2DeployNs1ApiKey,
				SysdigAccessKey:  *flEc2DeploySysdigAccessKey,
				DatadogAPIKey:    *flEc2DeployDatadogAPIKey,
				CaCert:           *flEc2DeployCaCert,
				Domain:           *flEc2DeployDomain,
				Region:           *flEc2DeployRegion,
				Zone:             *flEc2DeployZone,
				KeyPair:          *flEc2DeployKeyPair,
				VpcCidrBlock:     *flEc2DeployVpcCidrBlock,
				IntSubnetCidr:    *flEc2DeployIntSubnetCidr,
				ExtSubnetCidr:    *flEc2DeployExtSubnetCidr,
				FlannelNetwork:   *flEc2DeployFlannelNetwork,
				FlannelSubnetLen: *flEc2DeployFlannelSubnetLen,
				FlannelSubnetMin: *flEc2DeployFlannelSubnetMin,
				FlannelSubnetMax: *flEc2DeployFlannelSubnetMax,
				FlannelBackend:   *flEc2DeployFlannelBackend,
				StubZones:        *flEc2DeployStubZones,
				Quadruplets:      *arEc2DeployQuadruplet,
			},
		}

		ec2.Deploy()

	//-------------------
	// katoctl ec2 setup
	//-------------------

	case cmdEc2Setup.FullCommand():

		ec2 := ec2.Data{
			State: ec2.State{
				ClusterID:     *flEc2SetupClusterID,
				Domain:        *flEc2SetupDomain,
				Region:        *flEc2SetupRegion,
				Zone:          *flEc2SetupZone,
				VpcCidrBlock:  *flEc2SetupVpcCidrBlock,
				IntSubnetCidr: *flEc2SetupIntSubnetCidr,
				ExtSubnetCidr: *flEc2SetupExtSubnetCidr,
			},
		}

		ec2.Setup()

	//-----------------
	// katoctl ec2 add
	//-----------------

	case cmdEc2Add.FullCommand():

		ec2 := ec2.Data{
			State: ec2.State{
				ClusterID: *flEc2AddCluserID,
			},
			Instance: ec2.Instance{
				Roles:        *flEc2AddRoles,
				HostName:     *flEc2AddHostName,
				HostID:       *flEc2AddHostID,
				AmiID:        *flEc2AddAmiID,
				InstanceType: *flEc2AddInsanceType,
			},
		}

		ec2.Add()

	//-----------------
	// katoctl ec2 run
	//-----------------

	case cmdEc2Run.FullCommand():

		ec2 := ec2.Data{
			State: ec2.State{
				Region:  *flEc2RunRegion,
				Zone:    *flEc2RunZone,
				KeyPair: *flEc2RunKeyPair,
			},
			Instance: ec2.Instance{
				SubnetID:     *flEc2RunSubnetID,
				SecGrpIDs:    *flEc2RunSecGrpIDs,
				InstanceType: *flEc2RunInstanceType,
				TagName:      *flEc2RunTagName,
				PublicIP:     *flEc2RunPublicIP,
				IAMRole:      *flEc2RunIAMRole,
				SrcDstCheck:  *flEc2RunSrcDstCheck,
				AmiID:        *flEc2RunAmiID,
				ELBName:      *flEc2RunELBName,
				PrivateIP:    *flEc2RunPrivateIP,
			},
		}

		ec2.Run()

	//----------------------
	// katoctl ns1 zone add
	//----------------------

	case cmdNs1ZoneAdd.FullCommand():

		ns1 := ns1.Data{
			APIKey: *flNs1APIKey,
			Link:   *flNs1ZoneAddLink,
			Zones:  *arNs1ZoneAddName,
		}

		ns1.AddZones()

	//------------------------
	// katoctl ns1 record add
	//------------------------

	case cmdNs1RecordAdd.FullCommand():

		ns1 := ns1.Data{
			APIKey:  *flNs1APIKey,
			Zone:    *flNs1RecordAddZone,
			Records: *arNs1RecordAddName,
		}

		ns1.AddRecords()
	}
}

//-----------------------------------------------------------------------------
// Log filename and line number:
//-----------------------------------------------------------------------------

type contextHook struct{}

func (hook contextHook) Levels() []log.Level {
	levels := []log.Level{log.ErrorLevel, log.FatalLevel}
	return levels
}

func (hook contextHook) Fire(entry *log.Entry) error {
	pc := make([]uintptr, 3, 3)
	cnt := runtime.Callers(6, pc)

	for i := 0; i < cnt; i++ {
		fu := runtime.FuncForPC(pc[i] - 1)
		name := fu.Name()
		if !strings.Contains(name, "github.com/Sirupsen/logrus") {
			file, line := fu.FileLine(pc[i] - 1)
			entry.Data["file"] = path.Base(file)
			entry.Data["func"] = path.Base(name)
			entry.Data["line"] = line
			break
		}
	}
	return nil
}

//-----------------------------------------------------------------------------
// Regular expression custom parser:
//-----------------------------------------------------------------------------

type regexpMatchValue struct {
	value  string
	regexp string
}

func (id *regexpMatchValue) Set(value string) error {

	if match, _ := regexp.MatchString(id.regexp, value); !match {
		log.WithField("value", value).Fatal("Value must match: " + id.regexp)
	}

	id.value = value
	return nil
}

func (id *regexpMatchValue) String() string {
	return id.value
}

func regexpMatch(s kingpin.Settings, regexp string) *string {
	target := &regexpMatchValue{}
	target.regexp = regexp
	s.SetValue(target)
	return &target.value
}

//-----------------------------------------------------------------------------
// Quadruplets custom parser:
//-----------------------------------------------------------------------------

type quadList []string

func (q *quadList) Set(value string) error {
	*q = append(*q, value)
	return nil
}

func (q *quadList) String() string {
	return ""
}

func (q *quadList) IsCumulative() bool {
	return true
}

func quadruplets(s kingpin.Settings) (target *[]string) {
	target = new([]string)
	s.SetValue((*quadList)(target))
	return
}
<commit_msg>Quadruplet parser implementation (wip)<commit_after>package main

//-----------------------------------------------------------------------------
// Package factored import statement:
//-----------------------------------------------------------------------------

import (

	// Stdlib:
	"os"
	"path"
	"regexp"
	"runtime"
	"strconv"
	"strings"

	// Local:
	"github.com/katosys/kato/providers/cloud/ec2"
	"github.com/katosys/kato/providers/cloud/pkt"
	"github.com/katosys/kato/providers/dns/ns1"
	"github.com/katosys/kato/udata"

	// Community:
	log "github.com/Sirupsen/logrus"
	"gopkg.in/alecthomas/kingpin.v2"
)

//-----------------------------------------------------------------------------
// katoctl root level command:
//-----------------------------------------------------------------------------

var app = kingpin.New("katoctl", "Katoctl defines and deploys Kato's infrastructure.")

//----------------------------------------------------------------------------
// func init() is called after all the variable declarations in the package
// have evaluated their initializers, and those are evaluated only after all
// the imported packages have been initialized:
//----------------------------------------------------------------------------

func init() {

	// Customize kingpin:
	app.Version("v0.1.0-beta").Author("Marc Villacorta Morera")
	app.UsageTemplate(usageTemplate)
	app.HelpFlag.Short('h')

	// Customize the default logger:
	log.SetFormatter(&log.TextFormatter{ForceColors: true})
	log.SetOutput(os.Stderr)
	log.SetLevel(log.InfoLevel)
	log.AddHook(contextHook{})
}

//----------------------------------------------------------------------------
// Entry point:
//----------------------------------------------------------------------------

func main() {

	// Sub-command selector:
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {

	//---------------
	// katoctl udata
	//---------------

	case cmdUdata.FullCommand():

		udata := udata.Data{
			QuorumCount:         *flUdataQuorumCount,
			ClusterID:           *flUdataClusterID,
			HostName:            *flUdataHostName,
			HostID:              *flUdataHostID,
			Domain:              *flUdataDomain,
			Roles:               strings.Split(*flUdataRoles, ","),
			Ns1ApiKey:           *flUdataNs1Apikey,
			CaCert:              *flUdataCaCert,
			EtcdToken:           *flUdataEtcdToken,
			GzipUdata:           *flUdataGzipUdata,
			FlannelNetwork:      *flUdataFlannelNetwork,
			FlannelSubnetLen:    *flUdataFlannelSubnetLen,
			FlannelSubnetMin:    *flUdataFlannelSubnetMin,
			FlannelSubnetMax:    *flUdataFlannelSubnetMax,
			FlannelBackend:      *flUdataFlannelBackend,
			RexrayStorageDriver: *flUdataRexrayStorageDriver,
			RexrayEndpointIP:    *flUdataRexrayEndpointIP,
			Ec2Region:           *flUdataEc2Region,
			IaasProvider:        *flUdataIaasProvider,
			SysdigAccessKey:     *flUdataSysdigAccessKey,
			DatadogAPIKey:       *flUdataDatadogAPIKey,
			StubZones:           *flUdataStubZones,
			Prometheus:          *flUdataPrometheus,
		}

		udata.Render()

	//--------------------
	// katoctl pkt deploy
	//--------------------

	case cmdPktDeploy.FullCommand():

		pkt := pkt.Data{}
		pkt.Deploy()

	//-------------------
	// katoctl pkt setup
	//-------------------

	case cmdPktSetup.FullCommand():

		pkt := pkt.Data{}
		pkt.Setup()

	//-----------------
	// katoctl pkt run
	//-----------------

	case cmdPktRun.FullCommand():

		pkt := pkt.Data{
			APIKey:    *flPktRunAPIKey,
			HostName:  *flPktRunHostName,
			ProjectID: *flPktRunProjectID,
			Plan:      *flPktRunPlan,
			OS:        *flPktRunOS,
			Facility:  *flPktRunFacility,
			Billing:   *flPktRunBilling,
		}

		pkt.Run()

	//--------------------
	// katoctl ec2 deploy
	//--------------------

	case cmdEc2Deploy.FullCommand():

		ec2 := ec2.Data{
			State: ec2.State{
				ClusterID:        *flEc2DeployClusterID,
				Channel:          *flEc2DeployChannel,
				EtcdToken:        *flEc2DeployEtcdToken,
				Ns1ApiKey:        *flEc2DeployNs1ApiKey,
				SysdigAccessKey:  *flEc2DeploySysdigAccessKey,
				DatadogAPIKey:    *flEc2DeployDatadogAPIKey,
				CaCert:           *flEc2DeployCaCert,
				Domain:           *flEc2DeployDomain,
				Region:           *flEc2DeployRegion,
				Zone:             *flEc2DeployZone,
				KeyPair:          *flEc2DeployKeyPair,
				VpcCidrBlock:     *flEc2DeployVpcCidrBlock,
				IntSubnetCidr:    *flEc2DeployIntSubnetCidr,
				ExtSubnetCidr:    *flEc2DeployExtSubnetCidr,
				FlannelNetwork:   *flEc2DeployFlannelNetwork,
				FlannelSubnetLen: *flEc2DeployFlannelSubnetLen,
				FlannelSubnetMin: *flEc2DeployFlannelSubnetMin,
				FlannelSubnetMax: *flEc2DeployFlannelSubnetMax,
				FlannelBackend:   *flEc2DeployFlannelBackend,
				StubZones:        *flEc2DeployStubZones,
				Quadruplets:      *arEc2DeployQuadruplet,
			},
		}

		ec2.Deploy()

	//-------------------
	// katoctl ec2 setup
	//-------------------

	case cmdEc2Setup.FullCommand():

		ec2 := ec2.Data{
			State: ec2.State{
				ClusterID:     *flEc2SetupClusterID,
				Domain:        *flEc2SetupDomain,
				Region:        *flEc2SetupRegion,
				Zone:          *flEc2SetupZone,
				VpcCidrBlock:  *flEc2SetupVpcCidrBlock,
				IntSubnetCidr: *flEc2SetupIntSubnetCidr,
				ExtSubnetCidr: *flEc2SetupExtSubnetCidr,
			},
		}

		ec2.Setup()

	//-----------------
	// katoctl ec2 add
	//-----------------

	case cmdEc2Add.FullCommand():

		ec2 := ec2.Data{
			State: ec2.State{
				ClusterID: *flEc2AddCluserID,
			},
			Instance: ec2.Instance{
				Roles:        *flEc2AddRoles,
				HostName:     *flEc2AddHostName,
				HostID:       *flEc2AddHostID,
				AmiID:        *flEc2AddAmiID,
				InstanceType: *flEc2AddInsanceType,
			},
		}

		ec2.Add()

	//-----------------
	// katoctl ec2 run
	//-----------------

	case cmdEc2Run.FullCommand():

		ec2 := ec2.Data{
			State: ec2.State{
				Region:  *flEc2RunRegion,
				Zone:    *flEc2RunZone,
				KeyPair: *flEc2RunKeyPair,
			},
			Instance: ec2.Instance{
				SubnetID:     *flEc2RunSubnetID,
				SecGrpIDs:    *flEc2RunSecGrpIDs,
				InstanceType: *flEc2RunInstanceType,
				TagName:      *flEc2RunTagName,
				PublicIP:     *flEc2RunPublicIP,
				IAMRole:      *flEc2RunIAMRole,
				SrcDstCheck:  *flEc2RunSrcDstCheck,
				AmiID:        *flEc2RunAmiID,
				ELBName:      *flEc2RunELBName,
				PrivateIP:    *flEc2RunPrivateIP,
			},
		}

		ec2.Run()

	//----------------------
	// katoctl ns1 zone add
	//----------------------

	case cmdNs1ZoneAdd.FullCommand():

		ns1 := ns1.Data{
			APIKey: *flNs1APIKey,
			Link:   *flNs1ZoneAddLink,
			Zones:  *arNs1ZoneAddName,
		}

		ns1.AddZones()

	//------------------------
	// katoctl ns1 record add
	//------------------------

	case cmdNs1RecordAdd.FullCommand():

		ns1 := ns1.Data{
			APIKey:  *flNs1APIKey,
			Zone:    *flNs1RecordAddZone,
			Records: *arNs1RecordAddName,
		}

		ns1.AddRecords()
	}
}

//-----------------------------------------------------------------------------
// Log filename and line number:
//-----------------------------------------------------------------------------

type contextHook struct{}

func (hook contextHook) Levels() []log.Level {
	levels := []log.Level{log.ErrorLevel, log.FatalLevel}
	return levels
}

func (hook contextHook) Fire(entry *log.Entry) error {
	pc := make([]uintptr, 3, 3)
	cnt := runtime.Callers(6, pc)

	for i := 0; i < cnt; i++ {
		fu := runtime.FuncForPC(pc[i] - 1)
		name := fu.Name()
		if !strings.Contains(name, "github.com/Sirupsen/logrus") {
			file, line := fu.FileLine(pc[i] - 1)
			entry.Data["file"] = path.Base(file)
			entry.Data["func"] = path.Base(name)
			entry.Data["line"] = line
			break
		}
	}
	return nil
}

//-----------------------------------------------------------------------------
// Regular expression custom parser:
//-----------------------------------------------------------------------------

type regexpMatchValue struct {
	value  string
	regexp string
}

func (id *regexpMatchValue) Set(value string) error {

	if match, _ := regexp.MatchString(id.regexp, value); !match {
		log.WithField("value", value).Fatal("Value must match: " + id.regexp)
	}

	id.value = value
	return nil
}

func (id *regexpMatchValue) String() string {
	return id.value
}

func regexpMatch(s kingpin.Settings, regexp string) *string {
	target := &regexpMatchValue{}
	target.regexp = regexp
	s.SetValue(target)
	return &target.value
}

//-----------------------------------------------------------------------------
// Quadruplets custom parser:
//-----------------------------------------------------------------------------

/*type quadrupletsValue struct {
	quadList []string
	types    *[]string
	roles    *[]string
}*/

type quadList []string

func (q *quadList) Set(value string) error {

	// 1. Four elements:
	if quad := strings.Split(value, ":"); len(quad) != 4 {
		log.WithField("value", value).
			Fatal("Expected 4 elements, but got " + strconv.Itoa(len(quad)))

		// 2. Positive integer:
	} else if i, err := strconv.Atoi(quad[0]); err != nil || i < 0 {
		log.WithField("value", value).
			Fatal("First quadruplet element must be a positive integer, but got: " + quad[0])

		// 3. Valid instance type:
	} else if 1 == 1 {
		log.WithField("value", value).
			Fatal("Second quadruplet element must be a valid instance type, but got: " + quad[1])

		// 4. Valid DNS name:
	} else if match, err := regexp.MatchString("[a-z0-9-]*", value); err != nil || !match {
		log.WithField("value", value).
			Fatal("Third quadruplet element must matmatch [a-z0-9-]*, but got: " + quad[2])

		// 5. Valid Káto roles:
	} else if 1 == 1 {
		log.WithField("value", value).
			Fatal("Fourth quadruplet element must be a valid list of Káto roles, but got: " + quad[3])
	}

	//roles := strings.Split(quad[3], ",")
	log.WithField("value", value).Fatal("STOP")

	// All tests ok:
	*q = append(*q, value)
	return nil
}

func (q *quadList) String() string {
	return ""
}

func (q *quadList) IsCumulative() bool {
	return true
}

func quadruplets(s kingpin.Settings) (target *[]string) {
	target = new([]string)
	s.SetValue((*quadList)(target))
	return
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/user"
	"path"
	"strings"

	"github.com/vpn-kill-switch/killswitch"
)

func PadRight(str, pad string, lenght int) string {
	for {
		str += pad
		if len(str) > lenght {
			return str[0:lenght]
		}
	}
}

func exit1(err error) {
	fmt.Println(err)
	os.Exit(1)
}

var version string

func main() {

	var (
		ip = flag.String("ip", "", "VPN peer `IPv4`, killswitch tries to find this automatically")
		d  = flag.Bool("d", false, "`Disable` load /etc/pf.conf rules")
		e  = flag.Bool("e", false, "`Enable` load the pf rules")
		v  = flag.Bool("v", false, fmt.Sprintf("Print version: %s", version))
	)

	flag.Parse()

	if *v {
		fmt.Printf("%s\n", version)
		os.Exit(0)
	}

	if *d {
		exec.Command("pfctl", "-e").CombinedOutput()
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		fmt.Println("# Loading /etc/pf.conf rules")
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		out, _ := exec.Command("pfctl",
			"-Fa",
			"-f",
			"/etc/pf.conf").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl", "-sr").CombinedOutput()
		fmt.Printf("%s\n", out)
		return
	}

	ks, err := killswitch.New(*ip)
	if err != nil {
		exit1(err)
	}

	err = ks.GetActive()
	if err != nil {
		exit1(err)
	}

	if len(ks.UpInterfaces) == 0 {
		exit1(fmt.Errorf("No active interfaces found, verify network settings, use (\"%s -h\") for help.\n", os.Args[0]))
	}

	fmt.Println("Interface  MAC address         IP")
	for k, v := range ks.UpInterfaces {
		fmt.Printf("%s %s   %s\n", PadRight(k, " ", 10), v[0], v[1])
	}
	for k, v := range ks.P2PInterfaces {
		fmt.Printf("%s %s   %s\n", PadRight(k, " ", 10), PadRight(v[0], " ", 17), v[1])
	}
	// check for DNS leaks
	if ipDNS, err := killswitch.WhoamiDNS(); err == nil {
		if ipWWW, err := killswitch.WhoamiWWW(); err == nil {
			if ipDNS != ipWWW {
				fmt.Printf("\n%s:\n", killswitch.Red("DNS leaking"))
				fmt.Printf("Public IP address (DNS): %s\n", killswitch.Red(ipDNS))
				fmt.Printf("Public IP address (WWW): %s\n", killswitch.Red(ipWWW))
			} else {
				fmt.Printf("\nPublic IP address: %s\n", ipDNS)
			}
		}
	}

	// add some space
	println()

	if len(ks.P2PInterfaces) == 0 {
		exit1(fmt.Errorf("No VPN interface found, verify VPN is connected"))
	}

	fmt.Printf("PEER IP address:   %s\n", killswitch.Yellow(ks.PeerIP))

	if *ip != "" {
		if ipv4 := net.ParseIP(*ip); ipv4.To4() == nil {
			exit1(fmt.Errorf("%s is not a valid IPv4 address, use (\"%s -h\") for help.\n", *ip, os.Args[0]))
		}
	}

	ks.CreatePF()

	fmt.Printf("\n%s: %s\n", "To enable the kill switch run", killswitch.Green("sudo killswitch -e"))
	fmt.Printf("\n%s: %s\n\n", "To disable run", killswitch.Yellow("sudo pfctl -Fa -f /etc/pf.conf"))
	fmt.Printf("PF rules to be loaded:\n")
	fmt.Println(ks.PFRules.String())

	usr, err := user.Current()
	if err != nil {
		exit1(err)
	}
	if err = ioutil.WriteFile(path.Join(usr.HomeDir, ".killswitch.pf.conf"),
		ks.PFRules.Bytes(),
		0644,
	); err != nil {
		exit1(err)
	}

	if *e {
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		fmt.Println("# Loading rules")
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		out, _ := exec.Command("pfctl", "-e").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl",
			"-Fa",
			"-f",
			path.Join(usr.HomeDir, ".killswitch.pf.conf")).CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl", "-sr").CombinedOutput()
		fmt.Printf("%s\n", out)
	}
}
<commit_msg>print how to disable the kill switch<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/user"
	"path"
	"strings"

	"github.com/vpn-kill-switch/killswitch"
)

func PadRight(str, pad string, lenght int) string {
	for {
		str += pad
		if len(str) > lenght {
			return str[0:lenght]
		}
	}
}

func exit1(err error) {
	fmt.Println(err)
	os.Exit(1)
}

var version string

func main() {

	var (
		ip = flag.String("ip", "", "VPN peer `IPv4`, killswitch tries to find this automatically")
		d  = flag.Bool("d", false, "`Disable` load /etc/pf.conf rules")
		e  = flag.Bool("e", false, "`Enable` load the pf rules")
		v  = flag.Bool("v", false, fmt.Sprintf("Print version: %s", version))
	)

	flag.Parse()

	if *v {
		fmt.Printf("%s\n", version)
		os.Exit(0)
	}

	if *d {
		exec.Command("pfctl", "-e").CombinedOutput()
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		fmt.Println("# Loading /etc/pf.conf rules")
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		out, _ := exec.Command("pfctl",
			"-Fa",
			"-f",
			"/etc/pf.conf").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl", "-sr").CombinedOutput()
		fmt.Printf("%s\n", out)
		return
	}

	ks, err := killswitch.New(*ip)
	if err != nil {
		exit1(err)
	}

	err = ks.GetActive()
	if err != nil {
		exit1(err)
	}

	if len(ks.UpInterfaces) == 0 {
		exit1(fmt.Errorf("No active interfaces found, verify network settings, use (\"%s -h\") for help.\n", os.Args[0]))
	}

	fmt.Println("Interface  MAC address         IP")
	for k, v := range ks.UpInterfaces {
		fmt.Printf("%s %s   %s\n", PadRight(k, " ", 10), v[0], v[1])
	}
	for k, v := range ks.P2PInterfaces {
		fmt.Printf("%s %s   %s\n", PadRight(k, " ", 10), PadRight(v[0], " ", 17), v[1])
	}
	// check for DNS leaks
	if ipDNS, err := killswitch.WhoamiDNS(); err == nil {
		if ipWWW, err := killswitch.WhoamiWWW(); err == nil {
			if ipDNS != ipWWW {
				fmt.Printf("\n%s:\n", killswitch.Red("DNS leaking"))
				fmt.Printf("Public IP address (DNS): %s\n", killswitch.Red(ipDNS))
				fmt.Printf("Public IP address (WWW): %s\n", killswitch.Red(ipWWW))
			} else {
				fmt.Printf("\nPublic IP address: %s\n", ipDNS)
			}
		}
	}

	// add some space
	println()

	if len(ks.P2PInterfaces) == 0 {
		exit1(fmt.Errorf("No VPN interface found, verify VPN is connected"))
	}

	fmt.Printf("PEER IP address:   %s\n", killswitch.Yellow(ks.PeerIP))

	if *ip != "" {
		if ipv4 := net.ParseIP(*ip); ipv4.To4() == nil {
			exit1(fmt.Errorf("%s is not a valid IPv4 address, use (\"%s -h\") for help.\n", *ip, os.Args[0]))
		}
	}

	ks.CreatePF()

	fmt.Printf("\n%s: %s\n", "To enable the kill switch run", killswitch.Green("sudo killswitch -e"))
	fmt.Printf("%s: %s\n\n", "To disable", killswitch.Yellow("sudo killswitch -d"))
	fmt.Printf("PF rules to be loaded:\n")
	fmt.Println(ks.PFRules.String())

	usr, err := user.Current()
	if err != nil {
		exit1(err)
	}
	if err = ioutil.WriteFile(path.Join(usr.HomeDir, ".killswitch.pf.conf"),
		ks.PFRules.Bytes(),
		0644,
	); err != nil {
		exit1(err)
	}

	if *e {
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		fmt.Println("# Loading rules")
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		out, _ := exec.Command("pfctl", "-e").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl",
			"-Fa",
			"-f",
			path.Join(usr.HomeDir, ".killswitch.pf.conf")).CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl", "-sr").CombinedOutput()
		fmt.Printf("%s\n", out)
	}
}
<|endoftext|>


<commit_before>package operate

import (
	"encoding/json"
	"fmt"

	"github.com/go-errors/errors"
	"github.com/itchio/butler/buse"
	"github.com/itchio/butler/comm"
	"github.com/itchio/butler/mansion"
	"github.com/sourcegraph/jsonrpc2"
)

func Start(ctx *mansion.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) (*buse.OperationResult, error) {
	params := &buse.OperationStartParams{}
	err := json.Unmarshal(*req.Params, params)
	if err != nil {
		return nil, errors.Wrap(err, 0)
	}

	if params.StagingFolder == "" {
		return nil, errors.New("No staging folder specified")
	}

	oc := LoadContext(conn, ctx, comm.NewStateConsumer(), params.StagingFolder)

	meta := &MetaSubcontext{
		data: params,
	}

	if meta.data.Operation == "" {
		return nil, errors.New("No operation specified")
	}

	oc.Save(meta)

	switch params.Operation {
	case "install":
		ires, err := install(oc, meta)
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}

		return &buse.OperationResult{
			Success:       true,
			InstallResult: ires,
		}, nil
	}

	return nil, fmt.Errorf("Unknown operation '%s'", params.Operation)
}

type MetaSubcontext struct {
	data *buse.OperationStartParams
}

var _ Subcontext = (*MetaSubcontext)(nil)

func (mt *MetaSubcontext) Key() string {
	return "meta"
}

func (mt *MetaSubcontext) Data() interface{} {
	return &mt.data
}
<commit_msg>retire staging folder on operation success<commit_after>package operate

import (
	"encoding/json"
	"fmt"

	"github.com/go-errors/errors"
	"github.com/itchio/butler/buse"
	"github.com/itchio/butler/comm"
	"github.com/itchio/butler/mansion"
	"github.com/sourcegraph/jsonrpc2"
)

func Start(ctx *mansion.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) (*buse.OperationResult, error) {
	params := &buse.OperationStartParams{}
	err := json.Unmarshal(*req.Params, params)
	if err != nil {
		return nil, errors.Wrap(err, 0)
	}

	if params.StagingFolder == "" {
		return nil, errors.New("No staging folder specified")
	}

	oc := LoadContext(conn, ctx, comm.NewStateConsumer(), params.StagingFolder)

	meta := &MetaSubcontext{
		data: params,
	}

	if meta.data.Operation == "" {
		return nil, errors.New("No operation specified")
	}

	oc.Save(meta)

	switch params.Operation {
	case "install":
		ires, err := install(oc, meta)
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}

		err = oc.Retire()
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}

		return &buse.OperationResult{
			Success:       true,
			InstallResult: ires,
		}, nil
	}

	return nil, fmt.Errorf("Unknown operation '%s'", params.Operation)
}

type MetaSubcontext struct {
	data *buse.OperationStartParams
}

var _ Subcontext = (*MetaSubcontext)(nil)

func (mt *MetaSubcontext) Key() string {
	return "meta"
}

func (mt *MetaSubcontext) Data() interface{} {
	return &mt.data
}
<|endoftext|>


<commit_before>/*
	This file is part of go-ethereum

	go-ethereum is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	go-ethereum is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with go-ethereum.  If not, see <http://www.gnu.org/licenses/>.
*/
package main

import (
	"crypto/elliptic"
	"log"
	"os"

	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/p2p"
)

func main() {
	logger.AddLogSystem(logger.NewStdLogSystem(os.Stdout, log.LstdFlags, logger.InfoLevel))
	key, _ := crypto.GenerateKey()
	marshaled := elliptic.Marshal(crypto.S256(), key.PublicKey.X, key.PublicKey.Y)

	srv := p2p.Server{
		MaxPeers:   100,
		Identity:   p2p.NewSimpleClientIdentity("Ethereum(G)", "0.1", "Peer Server Two", marshaled),
		ListenAddr: ":30301",
		NAT:        p2p.UPNP(),
		NoDial:     true,
	}
	if err := srv.Start(); err != nil {
		log.Fatal("could not start server:", err)
	}
	select {}
}
<commit_msg>cmd/peerserver: add some command line switches<commit_after>/*
	This file is part of go-ethereum

	go-ethereum is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	go-ethereum is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with go-ethereum.  If not, see <http://www.gnu.org/licenses/>.
*/
package main

import (
	"crypto/elliptic"
	"flag"
	"log"
	"os"

	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/p2p"
)

var (
	natType    = flag.String("nat", "", "NAT traversal implementation")
	pmpGateway = flag.String("gateway", "", "gateway address for NAT-PMP")
	listenAddr = flag.String("addr", ":30301", "listen address")
)

func main() {
	flag.Parse()
	nat, err := p2p.ParseNAT(*natType, *pmpGateway)
	if err != nil {
		log.Fatal("invalid nat:", err)
	}

	logger.AddLogSystem(logger.NewStdLogSystem(os.Stdout, log.LstdFlags, logger.InfoLevel))
	key, _ := crypto.GenerateKey()
	marshaled := elliptic.Marshal(crypto.S256(), key.PublicKey.X, key.PublicKey.Y)

	srv := p2p.Server{
		MaxPeers:   100,
		Identity:   p2p.NewSimpleClientIdentity("Ethereum(G)", "0.1", "Peer Server Two", marshaled),
		ListenAddr: *listenAddr,
		NAT:        nat,
		NoDial:     true,
	}
	if err := srv.Start(); err != nil {
		log.Fatal("could not start server:", err)
	}
	select {}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/jessevdk/go-flags"
	"github.com/moul/shikaku"
)

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

func main() {
	var opts struct {
		Width  int `short:"W" long:"width" description:"Width of the grid" required:"true"`
		Height int `short:"H" long:"height" description:"Height of the grid" required:"true"`
		Blocks int `short:"B" long:"blocks" description:"Blocks in the grid" required:"true"`

		DrawMap      bool `short:"m" long:"draw-map" description:"Draw the map in ascii-art"`
		DrawSolution bool `short:"s" long:"draw-solution" description:"Draw the solution in ascii-art"`
	}

	if _, err := flags.Parse(&opts); err != nil {
		log.Fatalf("Parsing error: %v", err)
	}

	shikakuMap := shikaku.NewShikakuMap(opts.Width, opts.Height, 0, 0)
	if err := shikakuMap.GenerateBlocks(opts.Blocks); err != nil {
		log.Fatalf("Failed to generate %d blocks: %v", opts.Blocks, err)
	}
	fmt.Println(shikakuMap.String())
	if opts.DrawMap {
		fmt.Println()
		fmt.Println(shikakuMap.DrawMap())
	}
	if opts.DrawSolution {
		fmt.Println()
		fmt.Println(shikakuMap.DrawSolution())
	}
}
<commit_msg>Better output + add --no-machine-output<commit_after>package main

import (
	"fmt"
	"log"
	"math/rand"
	"strings"
	"time"

	"github.com/jessevdk/go-flags"
	"github.com/moul/shikaku"
)

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

func main() {
	var opts struct {
		Width  int `short:"W" long:"width" description:"Width of the grid" required:"true"`
		Height int `short:"H" long:"height" description:"Height of the grid" required:"true"`
		Blocks int `short:"B" long:"blocks" description:"Blocks in the grid" required:"true"`

		DrawMap         bool `short:"m" long:"draw-map" description:"Draw the map in ascii-art"`
		DrawSolution    bool `short:"s" long:"draw-solution" description:"Draw the solution in ascii-art"`
		NoMachineOutput bool `short:"q" long:"no-machine-output" description:"No machine output"`
	}

	if _, err := flags.Parse(&opts); err != nil {
		log.Fatalf("Parsing error: %v", err)
	}

	shikakuMap := shikaku.NewShikakuMap(opts.Width, opts.Height, 0, 0)
	if err := shikakuMap.GenerateBlocks(opts.Blocks); err != nil {
		log.Fatalf("Failed to generate %d blocks: %v", opts.Blocks, err)
	}

	outputs := []string{}

	if !opts.NoMachineOutput {
		outputs = append(outputs, shikakuMap.String())
	}
	if opts.DrawMap {
		outputs = append(outputs, shikakuMap.DrawMap())
	}
	if opts.DrawSolution {
		outputs = append(outputs, shikakuMap.DrawSolution())
	}
	fmt.Println(strings.Join(outputs, "\n\n"))
}
<|endoftext|>


<commit_before>package system

// On Windows user is implemented via syscalls and does not require a C compiler
import "os/user"

import (
	"os"
	"path/filepath"
	"strings"
	"syscall"

	bosherr "github.com/cloudfoundry/bosh-utils/errors"
)

func (fs *osFileSystem) currentHomeDir() (string, error) {
	t, err := syscall.OpenCurrentProcessToken()
	if err != nil {
		return "", err
	}
	defer t.Close()
	return t.GetUserProfileDirectory()
}

func (fs *osFileSystem) homeDir(username string) (string, error) {
	u, err := user.Current()
	if err != nil {
		return "", err
	}
	// On Windows, looking up the home directory
	// is only supported for the current user.
	if username != "" && !strings.EqualFold(username, u.Name) {
		return "", bosherr.Errorf("Failed to get user '%s' home directory", username)
	}
	return u.HomeDir, nil
}

func (fs *osFileSystem) chown(path, username string) error {
	return bosherr.WrapError(error(syscall.EWINDOWS), "Chown not supported on Windows")
}

func isSlash(c uint8) bool { return c == '\\' || c == '/' }

func absPath(path string) (string, error) {
	if filepath.IsAbs(path) {
		return filepath.Clean(path), nil
	}
	wd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	if len(path) > 0 && isSlash(path[0]) {
		return filepath.Join(filepath.VolumeName(wd), path), nil
	}
	return filepath.Join(wd, path), nil
}

func (fs *osFileSystem) symlinkPaths(oldPath, newPath string) (old, new string, err error) {
	// note: the type of the returned error is not *os.LinkError
	old, err = absPath(oldPath)
	if err != nil {
		return
	}
	new, err = absPath(newPath)
	if err != nil {
		return
	}
	return
}
<commit_msg>Stop windows filesystem from returning a 'not implemented' error when chowning<commit_after>package system

// On Windows user is implemented via syscalls and does not require a C compiler
import "os/user"

import (
	"os"
	"path/filepath"
	"strings"
	"syscall"

	bosherr "github.com/cloudfoundry/bosh-utils/errors"
)

func (fs *osFileSystem) currentHomeDir() (string, error) {
	t, err := syscall.OpenCurrentProcessToken()
	if err != nil {
		return "", err
	}
	defer t.Close()
	return t.GetUserProfileDirectory()
}

func (fs *osFileSystem) homeDir(username string) (string, error) {
	u, err := user.Current()
	if err != nil {
		return "", err
	}
	// On Windows, looking up the home directory
	// is only supported for the current user.
	if username != "" && !strings.EqualFold(username, u.Name) {
		return "", bosherr.Errorf("Failed to get user '%s' home directory", username)
	}
	return u.HomeDir, nil
}

func (fs *osFileSystem) chown(path, username string) error {
	return nil
}

func isSlash(c uint8) bool { return c == '\\' || c == '/' }

func absPath(path string) (string, error) {
	if filepath.IsAbs(path) {
		return filepath.Clean(path), nil
	}
	wd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	if len(path) > 0 && isSlash(path[0]) {
		return filepath.Join(filepath.VolumeName(wd), path), nil
	}
	return filepath.Join(wd, path), nil
}

func (fs *osFileSystem) symlinkPaths(oldPath, newPath string) (old, new string, err error) {
	// note: the type of the returned error is not *os.LinkError
	old, err = absPath(oldPath)
	if err != nil {
		return
	}
	new, err = absPath(newPath)
	if err != nil {
		return
	}
	return
}
<|endoftext|>


<commit_before>// Utilities to Generate SQL Scanners for custom Go primatives
//
// Usage:
// go:generate sqltype -type=<type> -primative=<primative>
//
// type is the name of the custom defined type
//
// primative is one of:
//    int64
//    float64
//    bool
//    []byte
//    string
//    time.Time
//
// Primative should make the primative of the custom defined type
//
// nil - for NULL values
package main

import (
	"bytes"
	"flag"
	"fmt"
	//	pickle "github.com/hydrogen18/stalecucumber"
	"github.com/pkg/errors"
	"go/build"
	"go/format"
	"io"
	"log"
	"os"
	"strings"
	"text/template"
)

const StringTemplate = ``

const ScannerTemplate = `
// DO NOT EDIT - Auto generated by sqltype for {{ .Type }}

package {{ .Package }}

import (
	"fmt"
	"reflect"
	"database/sql/driver"
	{{if or (eq .Primative "int") (eq .Primative "bool") }}
	"strconv"
	{{end}}
	{{if or (eq .Primative "pythondict") (eq .Primative "pythonstringdict") }}
	"bytes"
	pickle "github.com/hydrogen18/stalecucumber"
	{{else if eq .Primative "pythonlist"}}
	"bytes"
	pickle "github.com/hydrogen18/stalecucumber"
	{{end}}
)

func (t *{{ .Type }}) Scan(value interface{}) error {

	if value == nil {
		return nil
	}

	switch v := value.(type) {
	{{if eq .Primative "string"}}
		case []byte:
			*t = {{ .Type }}(string(v))

		case string:
			*t = {{ .Type }}(v)

		case *string:
			*t = {{ .Type }}(*v)

	{{else if eq .Primative "int"}}
		case []byte:
			i, err := strconv.Atoi(string(v))
			if err != nil {
				return fmt.Errorf("%s Can't convert '%v' to int %v", reflect.TypeOf(t), value, err)
			}
			*t = {{ .Type }}(i)

		case int:
			*t = {{ .Type }}(v)

		case *int:
			*t = {{ .Type }}(*v)

	{{else if eq .Primative "bool"}}
		case []byte:
			b, err := strconv.ParseBool(string(v))
			if err != nil {
				return fmt.Errorf("%s Can't convert '%v' to bool %v", reflect.TypeOf(t), value, err)
			}
			*t = {{ .Type }}(b)

		case bool:
			*t = {{ .Type }}(v)

		case *bool:
			*t = {{ .Type }}(*v)

	{{else if eq .Primative  "pythondict"}}
		case []byte:
			dict, err := pickle.Dict(pickle.Unpickle(bytes.NewReader(v)))
			if err != nil {
				return fmt.Errorf("%s Can't convert '%v' to dict %v", reflect.TypeOf(t), value, err)
			}
			*t = {{ .Type }}(dict)

	{{else if eq .Primative  "pythonstringdict"}}
		case []byte:
			dict, err := pickle.DictString(pickle.Unpickle(bytes.NewReader(v)))
			if err != nil {
				return fmt.Errorf("%s Can't convert '%v' to dict %v", reflect.TypeOf(t), value, err)
			}
			*t = {{ .Type }}(dict)

	{{else if eq .Primative  "pythonlist"}}
		case []byte:
			list := make({{ .Type }},0)
			err := pickle.UnpackInto(&list).From(pickle.Unpickle(bytes.NewReader(v)))
			if err != nil {
				return fmt.Errorf("%s Can't convert '%v' to list %v", reflect.TypeOf(t), value, err)
			}
			*t = list

	{{end}}
		default:
			return fmt.Errorf("%s Can't convert '%v' to {{ .Primative }}", reflect.TypeOf(t), value)
	}

	return nil
}

func (t {{ .Type }}) Value() (driver.Value, error) {

	{{if or (eq .Primative "pythondict") (eq .Primative "pythonlist") (eq .Primative "pythonstringdict") }}
		buf := new(bytes.Buffer)
		if _, err := pickle.NewPickler(buf).Pickle(t); err != nil {
			return nil, err
		}
		return buf.Bytes(), nil
	{{else}}
		return {{ .Primative }}(t), nil
	{{end}}
}
`

/*
	if value != nil {
		 return nil
	}

	switch v := value.(type) {
	case string:
		*t = {{ .Type }}(v)
		return nil
	case []byte:
		*t = {{ .Type }}(v)
		return nil
	case nil:
		return nil
	default:
		return fmt.Errorf("{{ .Type }}: Underlying value %v is not convertable to a string", reflect.TypeOf(value))
	}
}
*/

type Scanner struct {
	Package   string
	Primative string
	Type      string
}

var (
	primativeType string
	typeName      string
)

func init() {
	flag.StringVar(&primativeType, "primative", "", "Corresponding primative type")
	flag.StringVar(&typeName, "type", "", "Name of type")
}

func main() {
	log.SetFlags(0)
	flag.Parse()

	if typeName == "" {
		log.Fatal("Need to provide --type")
	}

	if primativeType == "" {
		log.Fatal("Need to provide --primative")
	}

	packageName, err := getPackageName()
	if err != nil {
		log.Fatal(err)
	}

	if err := validatePrimative(primativeType); err != nil {
		log.Fatal(err)
	}

	scanner := Scanner{
		Package:   packageName,
		Primative: primativeType,
		Type:      typeName,
	}

	filename := fmt.Sprintf("sql_%s.go", strings.ToLower(typeName))
	file, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	if err := SqlType(file, scanner); err != nil {
		log.Fatal(err)
	}
}

func SqlType(writer io.Writer, scanner Scanner) error {
	t := template.Must(template.New("scanner").Parse(ScannerTemplate))

	buf := &bytes.Buffer{}
	if err := t.Execute(buf, scanner); err != nil {
		return errors.Wrap(err, "error excuting template")
	}

	b, err := format.Source(buf.Bytes())
	if err != nil {
		return errors.Wrap(err, "error sourcing go struct")
	}

	if _, err := writer.Write(b); err != nil {
		return errors.Wrap(err, "error writing to file")
	}

	return nil
}

func getPackageName() (string, error) {
	directory := "."

	pkg, err := build.Default.ImportDir(directory, 0)
	if err != nil {
		return "", err
	}

	return pkg.Name, err
}

func validatePrimative(primativeType string) error {
	validTypes := map[string]bool{
		"int":              true,
		"int64":            true,
		"string":           true,
		"bool":             true,
		"pythonstringdict": true,
		"pythondict":       true,
		"pythonlist":       true,
	}

	if _, found := validTypes[primativeType]; !found {
		return errors.Errorf("Invalid primative '%s'", primativeType)
	}

	return nil
}

/*
func (t *{{ .Type }}) Scan(value interface{}) error {
	fmt.Println("VALUE", value, reflect.TypeOf(value))

	s, err := ConvertString(value)
	if err != nil {
		return errors.Wrapf(err, "%s Can't convert '%v' to string %s", reflect.TypeOf(t), value, s)
	}
	*t = RewardTargetId(s)
	return nil
}

func (t RewardTargetId) Value() (driver.Value, error) {
	return string(t), nil
}

func ConvertString(value interface{}) (string, error) {
	switch t := value.(type) {
	case string:
		return t, nil
	case []byte:
		return string(t), nil
	case nil:
		return "", nil
	default:
		return "", errors.Errorf("Underlying value %v is not a string")
	}
}

func ConvertInt(value interface{}) (string, error) {
	switch t := value.(type) {
	case string:
		return t, nil
	case []byte:
		return string(t), nil
	case nil:
		return "", nil
	default:
		return "", errors.Errorf("Underlying value %v is not a string")
	}
}

func ConvertFloat(value interface{}) (string, error) {
	switch t := value.(type) {
	case string:
		return t, nil
	case []byte:
		return string(t), nil
	case nil:
		return "", nil
	default:
		return "", errors.Errorf("Underlying value %v is not a string")
	}
}

func ConvertBool(value interface{}) (bool, error) {
	b, err := driver.Bool.ConvertValue(value)
	if err != nil {
		return err
	}
	return b.(bool)
}

func (t *StringList) Scan(value interface{}) error {

	if value == nil {
		return nil
	}

	buf, ok := value.([]byte)
	if !ok {
		return errors.Errorf("%s Can't cast '%v' to []byte", reflect.TypeOf(t), value)
	}

	reader := bytes.NewReader(buf)

	list, err := pickle.ListOrTuple(pickle.Unpickle(reader))
	if err != nil {
		return errors.Wrapf(err, "%s Can't convert '%v' to list", reflect.TypeOf(t), value)
	}

	fmt.Println("LIST", list)

	slist := make([]string, 0, len(list))
	for _, l := range list {
		slist = append(slist, l.(string))
	}
	*t = slist
	return nil
}

func (t StringList) Value() (driver.Value, error) {
	buf := new(bytes.Buffer)
	if _, err := pickle.NewPickler(buf).Pickle(t); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func (t StringListList) Value() (driver.Value, error) {
	buf := new(bytes.Buffer)
	if _, err := pickle.NewPickler(buf).Pickle(t); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}
*/
<commit_msg>fixing json generation<commit_after>// Utilities to Generate SQL Scanners for custom Go primatives
//
// Usage:
// go:generate sqltype -type=<type> -primative=<primative>
//
// type is the name of the custom defined type
//
// primative is one of:
//    int64
//    float64
//    bool
//    []byte
//    string
//    time.Time
//
// Primative should make the primative of the custom defined type
//
// nil - for NULL values
package main

import (
	"bytes"
	"flag"
	"fmt"
	//	pickle "github.com/hydrogen18/stalecucumber"
	"github.com/pkg/errors"
	"go/build"
	"go/format"
	"io"
	"log"
	"os"
	"strings"
	"text/template"
)

const StringTemplate = ``

const ScannerTemplate = `
// DO NOT EDIT - Auto generated by sqltype for {{ .Type }}

package {{ .Package }}

import (
	"fmt"
	"reflect"
	"database/sql/driver"
	{{if or (eq .Primative "int") (eq .Primative "bool") }}
	"strconv"
	{{end}}
	{{if or (eq .Primative "pythondict") (eq .Primative "pythonstringdict") }}
	"bytes"
	pickle "github.com/hydrogen18/stalecucumber"
	{{else if eq .Primative "pythonlist"}}
	"bytes"
	pickle "github.com/hydrogen18/stalecucumber"
	{{else if eq .Primative "pythonlist"}}
	"encoding/json"
	{{else if eq .Primative "json"}}
	"encoding/json"
	{{end}}
)

func (t *{{ .Type }}) Scan(value interface{}) error {

	if value == nil {
		return nil
	}

	switch v := value.(type) {
	{{if eq .Primative "string"}}
		case []byte:
			*t = {{ .Type }}(string(v))

		case string:
			*t = {{ .Type }}(v)

		case *string:
			*t = {{ .Type }}(*v)

	{{else if eq .Primative "int"}}
		case []byte:
			i, err := strconv.Atoi(string(v))
			if err != nil {
				return fmt.Errorf("sqltype: %s Can't convert '%v' to int %v", reflect.TypeOf(t), value, err)
			}
			*t = {{ .Type }}(i)

		case int:
			*t = {{ .Type }}(v)

		case *int:
			*t = {{ .Type }}(*v)

	{{else if eq .Primative "bool"}}
		case []byte:
			b, err := strconv.ParseBool(string(v))
			if err != nil {
				return fmt.Errorf("sqltype: %s Can't convert '%v' to bool %v", reflect.TypeOf(t), value, err)
			}
			*t = {{ .Type }}(b)

		case bool:
			*t = {{ .Type }}(v)

		case *bool:
			*t = {{ .Type }}(*v)

	{{else if eq .Primative  "pythondict"}}
		case []byte:
			dict, err := pickle.Dict(pickle.Unpickle(bytes.NewReader(v)))
			if err != nil {
				return fmt.Errorf("sqltype: %s Can't convert '%v' to dict %v", reflect.TypeOf(t), value, err)
			}
			*t = {{ .Type }}(dict)

	{{else if eq .Primative  "pythonstringdict"}}
		case []byte:
			dict, err := pickle.DictString(pickle.Unpickle(bytes.NewReader(v)))
			if err != nil {
				return fmt.Errorf("sqltype: %s Can't convert '%v' to dict %v", reflect.TypeOf(t), value, err)
			}
			*t = {{ .Type }}(dict)

	{{else if eq .Primative  "pythonlist"}}
		case []byte:
			list := make({{ .Type }},0)
			err := pickle.UnpackInto(&list).From(pickle.Unpickle(bytes.NewReader(v)))
			if err != nil {
				return fmt.Errorf("sqltype: %s Can't convert '%v' to list %v", reflect.TypeOf(t), value, err)
			}
			*t = list

	{{else if eq .Primative "json"}}
		case []byte:
			var obj {{ .Type }}
			if err := json.Unmarshal(v, &obj); err != nil {
				return fmt.Errorf("sqltype: can't convert json '%v' to %s %v", value, reflect.TypeOf(t), err)
			}
			*t = obj

	{{end}}
		default:
			return fmt.Errorf("sqltype: %s Can't convert '%v' to {{ .Primative }}", reflect.TypeOf(t), value)
	}

	return nil
}

func (t {{ .Type }}) Value() (driver.Value, error) {

	{{if or (eq .Primative "pythondict") (eq .Primative "pythonlist") (eq .Primative "pythonstringdict") }}
		buf := new(bytes.Buffer)
		if _, err := pickle.NewPickler(buf).Pickle(t); err != nil {
			return nil, err
		}
		return buf.Bytes(), nil
	{{else if eq .Primative "json"}}
		b, err := json.Marshal(t)
		if err != nil {
			return nil, err
		}
		return b, nil
	{{else}}
		return {{ .Primative }}(t), nil
	{{end}}
}
`

/*
	if value != nil {
		 return nil
	}

	switch v := value.(type) {
	case string:
		*t = {{ .Type }}(v)
		return nil
	case []byte:
		*t = {{ .Type }}(v)
		return nil
	case nil:
		return nil
	default:
		return fmt.Errorf("{{ .Type }}: Underlying value %v is not convertable to a string", reflect.TypeOf(value))
	}
}
*/

type Scanner struct {
	Package   string
	Primative string
	Type      string
}

var (
	primativeType string
	typeName      string
	version       bool

	Sha    string
	Branch string
)

func init() {
	flag.StringVar(&primativeType, "primative", "", "Corresponding primative type")
	flag.StringVar(&typeName, "type", "", "Name of type")
	flag.BoolVar(&version, "version", false, "Version of sqltype")
}

func main() {
	log.SetFlags(0)
	flag.Parse()

	if version {
		fmt.Printf("Sha: %s, Branch: %s\n", Sha, Branch)
		os.Exit(-1)
	}

	if typeName == "" {
		log.Fatal("Need to provide --type")
	}

	if primativeType == "" {
		log.Fatal("Need to provide --primative")
	}

	packageName, err := getPackageName()
	if err != nil {
		log.Fatal(err)
	}

	if err := validatePrimative(primativeType); err != nil {
		log.Fatal(err)
	}

	scanner := Scanner{
		Package:   packageName,
		Primative: primativeType,
		Type:      typeName,
	}

	filename := fmt.Sprintf("sql_%s.go", strings.ToLower(typeName))
	file, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	if err := SqlType(file, scanner); err != nil {
		log.Fatal(err)
	}
}

func SqlType(writer io.Writer, scanner Scanner) error {
	t := template.Must(template.New("scanner").Parse(ScannerTemplate))

	buf := &bytes.Buffer{}
	if err := t.Execute(buf, scanner); err != nil {
		return errors.Wrap(err, "error excuting template")
	}

	b, err := format.Source(buf.Bytes())
	if err != nil {
		return errors.Wrap(err, "error sourcing go struct")
	}

	if _, err := writer.Write(b); err != nil {
		return errors.Wrap(err, "error writing to file")
	}

	return nil
}

func getPackageName() (string, error) {
	directory := "."

	pkg, err := build.Default.ImportDir(directory, 0)
	if err != nil {
		return "", err
	}

	return pkg.Name, err
}

func validatePrimative(primativeType string) error {
	validTypes := map[string]bool{
		"int":              true,
		"int64":            true,
		"string":           true,
		"bool":             true,
		"pythonstringdict": true,
		"pythondict":       true,
		"pythonlist":       true,
		"json":             true,
	}

	if _, found := validTypes[primativeType]; !found {
		return errors.Errorf("Invalid primative '%s'", primativeType)
	}

	return nil
}

/*
func (t *{{ .Type }}) Scan(value interface{}) error {
	fmt.Println("VALUE", value, reflect.TypeOf(value))

	s, err := ConvertString(value)
	if err != nil {
		return errors.Wrapf(err, "%s Can't convert '%v' to string %s", reflect.TypeOf(t), value, s)
	}
	*t = RewardTargetId(s)
	return nil
}

func (t RewardTargetId) Value() (driver.Value, error) {
	return string(t), nil
}

func ConvertString(value interface{}) (string, error) {
	switch t := value.(type) {
	case string:
		return t, nil
	case []byte:
		return string(t), nil
	case nil:
		return "", nil
	default:
		return "", errors.Errorf("Underlying value %v is not a string")
	}
}

func ConvertInt(value interface{}) (string, error) {
	switch t := value.(type) {
	case string:
		return t, nil
	case []byte:
		return string(t), nil
	case nil:
		return "", nil
	default:
		return "", errors.Errorf("Underlying value %v is not a string")
	}
}

func ConvertFloat(value interface{}) (string, error) {
	switch t := value.(type) {
	case string:
		return t, nil
	case []byte:
		return string(t), nil
	case nil:
		return "", nil
	default:
		return "", errors.Errorf("Underlying value %v is not a string")
	}
}

func ConvertBool(value interface{}) (bool, error) {
	b, err := driver.Bool.ConvertValue(value)
	if err != nil {
		return err
	}
	return b.(bool)
}

func (t *StringList) Scan(value interface{}) error {

	if value == nil {
		return nil
	}

	buf, ok := value.([]byte)
	if !ok {
		return errors.Errorf("%s Can't cast '%v' to []byte", reflect.TypeOf(t), value)
	}

	reader := bytes.NewReader(buf)

	list, err := pickle.ListOrTuple(pickle.Unpickle(reader))
	if err != nil {
		return errors.Wrapf(err, "%s Can't convert '%v' to list", reflect.TypeOf(t), value)
	}

	fmt.Println("LIST", list)

	slist := make([]string, 0, len(list))
	for _, l := range list {
		slist = append(slist, l.(string))
	}
	*t = slist
	return nil
}

func (t StringList) Value() (driver.Value, error) {
	buf := new(bytes.Buffer)
	if _, err := pickle.NewPickler(buf).Pickle(t); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func (t StringListList) Value() (driver.Value, error) {
	buf := new(bytes.Buffer)
	if _, err := pickle.NewPickler(buf).Pickle(t); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}
*/
<|endoftext|>