<commit_before>package main

import (
	"flag"
	"github.com/codegangsta/cli"
	"io/ioutil"
	"os"
	"strings"
	"testing"

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

func capture(args string) (string, error) {
	stdout := os.Stdout

	r, w, err := os.Pipe()
	if err != nil {
		return "", err
	}

	defer func() {
		w.Close()
		os.Stdout = stdout
	}()

	os.Stdout = w
	err = app.Run(strings.Split(args, " "))
	if err != nil {
		return "", err
	}

	w.Close()
	output, err := ioutil.ReadAll(r)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(output)), nil
}

func flagSet(name string, flags []cli.Flag) *flag.FlagSet {
	set := flag.NewFlagSet(name, flag.ContinueOnError)
	for _, f := range flags {
		f.Apply(set)
	}
	return set
}

func TestCLI(t *testing.T) {
	Convey("Command should execute as expected", t, func() {
		set := flagSet(appTest.Name, appTest.Flags)
		ctx := cli.NewContext(appTest, set, set)

		Convey("Invalid flags should cause an error", func() {
			_, err := capture("envetcd -shmaltz delicious")
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldEqual, "flag provided but not defined: -shmaltz")
		})

		Convey("Version should be printed", func() {
			output, err := capture("envetcd --version")
			So(err, ShouldBeNil)
			So(output, ShouldEqual, "envetcd version "+app.Version)
		})
		Convey("Initlogger should be printed", func() {
			initLogger(ctx)
			//os.Args = []string{"./envetcd", "no-upcase=true", "no-sync=true", "-o", "ooooo", "--system", "nsq", "-c", "env"}
			//fmt.Println(appTest.Run(os.Args))
		})

	})
}

var (
	appTest = cli.NewApp()
)

//Set up a new test app with some predetermined values
func init() {
	os.Setenv("ENVETCD_CLEAN_ENV", "true")
	os.Setenv("ENVETCD_NO_SANITIZE", "true")
	os.Setenv("ENVETCD_NO_UPCASE", "true")
	appTest.Name = "testApp"
	appTest.Author = "Karl Dominguez"
	appTest.Email = "kdominguez@zvelo.com"
	appTest.Version = "0.0.4"
	appTest.Usage = "get environment variables from etcd"
	appTest.Flags = []cli.Flag{
		cli.StringSliceFlag{
			Name:   "peers, C",
			EnvVar: "ENVETCD_PEERS",
			Value:  &cli.StringSlice{"127.0.0.1:4001"},
			Usage:  "a comma-delimited list of machine addresses in the cluster (default: \"127.0.0.1:4001\")",
		},
		cli.StringFlag{
			Name:   "ca-file",
			EnvVar: "ENVETCD_CA_FILE",
			Usage:  "certificate authority file",
		},
		cli.StringFlag{
			Name:   "cert-file",
			EnvVar: "ENVETCD_CERT_FILE",
			Usage:  "tls client certificate file",
		},
		cli.StringFlag{
			Name:   "key-file",
			EnvVar: "ENVETCD_KEY_FILE",
			Usage:  "tls client key file",
		},
		cli.StringFlag{
			Name:   "hostname",
			EnvVar: "HOSTNAME",
			Value:  "env",
			Usage:  "computer hostname for host specific configuration",
		},
		cli.StringFlag{
			Name:   "system",
			EnvVar: "ENVETCD_SYSTEM",
			Value:  "systemtest",
			Usage:  "system name for system specific configuration",
		},
		cli.StringFlag{
			Name:   "service",
			EnvVar: "ENVETCD_SERVICE",
			Value:  "servicetest",
			Usage:  "service name for service specific configuration",
		},
		cli.StringFlag{
			Name:   "prefix",
			EnvVar: "ENVETCD_PREFIX",
			Value:  "/config",
			Usage:  "etcd prefix for all keys",
		},
		cli.StringFlag{
			Name:   "log-level, l",
			EnvVar: "ENVETCD_LOG_LEVEL",
			Value:  "DEBUG",
			Usage:  "set log level (DEBUG, INFO, WARN, ERR)",
		},
		cli.StringFlag{
			Name:   "output, o",
			Value:  "testOut.txt",
			EnvVar: "ENVETCD_OUTPUT",
			Usage:  "write stdout from the command to this file",
		},
		cli.BoolFlag{
			Name:   "no-sync",
			EnvVar: "ENVETCD_NO_SYNC",
			Usage:  "don't synchronize cluster information before sending request",
		},
		cli.BoolFlag{
			Name:   "clean-env, c",
			EnvVar: "ENVETCD_CLEAN_ENV",
			Usage:  "don't inherit any environment variables other than those pulled from etcd",
		},
		cli.BoolFlag{
			Name:   "no-sanitize",
			EnvVar: "ENVETCD_NO_SANITIZE",
			Usage:  "don't remove bad characters from environment keys",
		},
		cli.BoolFlag{
			Name:   "no-upcase",
			EnvVar: "ENVETCD_NO_UPCASE",
			Usage:  "don't convert all environment keys to uppercase",
		},
	}
	appTest.Action = run
}
<commit_msg>testing wercker box<commit_after>package main

import (
	"flag"
	"github.com/codegangsta/cli"
	"io/ioutil"
	"os"
	"strings"
	"testing"

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

func capture(args string) (string, error) {
	stdout := os.Stdout

	r, w, err := os.Pipe()
	if err != nil {
		return "", err
	}

	defer func() {
		w.Close()
		os.Stdout = stdout
	}()

	os.Stdout = w
	err = app.Run(strings.Split(args, " "))
	if err != nil {
		return "", err
	}

	w.Close()
	output, err := ioutil.ReadAll(r)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(output)), nil
}

func flagSet(name string, flags []cli.Flag) *flag.FlagSet {
	set := flag.NewFlagSet(name, flag.ContinueOnError)
	for _, f := range flags {
		f.Apply(set)
	}
	return set
}

func TestCLI(t *testing.T) {
	Convey("Command should execute as expected", t, func() {
		set := flagSet(appTest.Name, appTest.Flags)
		ctx := cli.NewContext(appTest, set, set)

		Convey("Invalid flags should cause an error", func() {
			_, err := capture("envetcd -shmaltz delicious")
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldEqual, "flag provided but not defined: -shmaltz")
		})

		Convey("Version should be printed", func() {
			output, err := capture("envetcd --version")
			So(err, ShouldBeNil)
			So(output, ShouldEqual, "envetcd version "+app.Version)
		})
		Convey("Initlogger should be printed", func() {
			initLogger(ctx)
			//os.Args = []string{"./envetcd", "no-upcase=true", "no-sync=true", "-o", "ooooo", "--system", "nsq", "-c", "env"}
			//fmt.Println(appTest.Run(os.Args))
		})

	})
}

var (
	appTest = cli.NewApp()
)

//Set up a new test app with some predetermined values
func init() {
	os.Setenv("ENVETCD_CLEAN_ENV", "true")
	os.Setenv("ENVETCD_NO_SANITIZE", "true")
	os.Setenv("ENVETCD_NO_UPCASE", "true")

	appTest.Name = "testApp"
	appTest.Author = "Karl Dominguez"
	appTest.Email = "kdominguez@zvelo.com"
	appTest.Version = "0.0.4"
	appTest.Usage = "get environment variables from etcd"
	appTest.Flags = []cli.Flag{
		cli.StringSliceFlag{
			Name:   "peers, C",
			EnvVar: "WERCKER_ETCD_HOST",
			Value:  &cli.StringSlice{"127.0.0.1:4001"},
			Usage:  "a comma-delimited list of machine addresses in the cluster (default: \"127.0.0.1:4001\")",
		},
		cli.StringFlag{
			Name:   "ca-file",
			EnvVar: "ENVETCD_CA_FILE",
			Usage:  "certificate authority file",
		},
		cli.StringFlag{
			Name:   "cert-file",
			EnvVar: "ENVETCD_CERT_FILE",
			Usage:  "tls client certificate file",
		},
		cli.StringFlag{
			Name:   "key-file",
			EnvVar: "ENVETCD_KEY_FILE",
			Usage:  "tls client key file",
		},
		cli.StringFlag{
			Name:   "hostname",
			EnvVar: "HOSTNAME",
			Value:  "env",
			Usage:  "computer hostname for host specific configuration",
		},
		cli.StringFlag{
			Name:   "system",
			EnvVar: "ENVETCD_SYSTEM",
			Value:  "systemtest",
			Usage:  "system name for system specific configuration",
		},
		cli.StringFlag{
			Name:   "service",
			EnvVar: "ENVETCD_SERVICE",
			Value:  "servicetest",
			Usage:  "service name for service specific configuration",
		},
		cli.StringFlag{
			Name:   "prefix",
			EnvVar: "ENVETCD_PREFIX",
			Value:  "/config",
			Usage:  "etcd prefix for all keys",
		},
		cli.StringFlag{
			Name:   "log-level, l",
			EnvVar: "ENVETCD_LOG_LEVEL",
			Value:  "DEBUG",
			Usage:  "set log level (DEBUG, INFO, WARN, ERR)",
		},
		cli.StringFlag{
			Name:   "output, o",
			Value:  "testOut.txt",
			EnvVar: "ENVETCD_OUTPUT",
			Usage:  "write stdout from the command to this file",
		},
		cli.BoolFlag{
			Name:   "no-sync",
			EnvVar: "ENVETCD_NO_SYNC",
			Usage:  "don't synchronize cluster information before sending request",
		},
		cli.BoolFlag{
			Name:   "clean-env, c",
			EnvVar: "ENVETCD_CLEAN_ENV",
			Usage:  "don't inherit any environment variables other than those pulled from etcd",
		},
		cli.BoolFlag{
			Name:   "no-sanitize",
			EnvVar: "ENVETCD_NO_SANITIZE",
			Usage:  "don't remove bad characters from environment keys",
		},
		cli.BoolFlag{
			Name:   "no-upcase",
			EnvVar: "ENVETCD_NO_UPCASE",
			Usage:  "don't convert all environment keys to uppercase",
		},
	}
	appTest.Action = run
}
<|endoftext|>


<commit_before>package cmd

import (
	"github.com/sirupsen/logrus"

	chClient "github.com/containerum/chkit/pkg/client"
	"github.com/containerum/chkit/pkg/model"
	cli "gopkg.in/urfave/cli.v2"
)

const (
	FlagAPIaddr    = "apiaddr"
	FlagConfigFile = "config"
)

var Configuration = struct {
	ConfigPath   string
	ConfigFile   string
	TokenFile    string
	ClientConfig model.Config
}{}

var ChkitClient chClient.Client
var log = &logrus.Logger{
	Formatter: &logrus.TextFormatter{},
}

var App = &cli.App{
	Name: "chkit",
	Action: func(ctx *cli.Context) error {
		return nil
	},
}
<commit_msg>add configuration initilaisation<commit_after>package cmd

import (
	"path"

	"github.com/blang/semver"

	"github.com/containerum/chkit/pkg/model"
	"github.com/sirupsen/logrus"
	cli "gopkg.in/urfave/cli.v2"
)

const (
	Version        = "3.0.0-alpha"
	FlagAPIaddr    = "apiaddr"
	FlagConfigFile = "config"
)

var (
	Configuration = model.Config{}
	log           = &logrus.Logger{
		Formatter: &logrus.TextFormatter{},
	}
)

var App = &cli.App{
	Name:    "chkit",
	Version: semver.MustParse(Version).String(),
	Action: func(ctx *cli.Context) error {

		return nil
	},
	Before: func(ctx *cli.Context) error {
		err := initConfig()
		if err != nil {
			log.WithError(err).
				Errorf("error while getting homedir path")
			return err
		}
		return nil
	},
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:    "config",
			Aliases: []string{"c"},
			Value:   path.Join(Configuration.ConfigPath, "config.file"),
		},
	},
}
<|endoftext|>


<commit_before>package main

import (
	rsa "github.com/mrvine/encryptor-rsa"
	"flag"
	"github.com/fatih/color"
	"fmt"
	"github.com/pkg/errors"
	"io/ioutil"
)

var (
	redConsole    = color.New(color.FgRed)
	yellowConsole = color.New(color.FgYellow)
	greenConsole  = color.New(color.FgGreen)
)

func generate(args RsaArgs) (string, error) {

	var e rsa.RsaEncryptor
	var err error

	if args.PrivateKeyPassword == "" {
		e, err = rsa.Init(args.KeyLength)
	} else {
		e, err = rsa.InitWithPassword(args.KeyLength, args.PrivateKeyPassword)
	}

	if err != nil {
		return "", errors.Wrap(err, "can not initialize rsa encryptor")
	}

	if args.PublicKeyPath == "" {
		args.PublicKeyPath = "public_" + getRandomString(16) + ".txt"
	}

	if args.PrivateKeyPath == "" {
		args.PrivateKeyPath = "private_" + getRandomString(16) + ".txt"
	}

	err = e.SavePublicKeyInPem(args.PublicKeyPath)
	if err != nil {
		return "", errors.Wrap(err, "can not create public key file")
	}

	err = e.SavePrivateKeyInPem(args.PrivateKeyPath)
	if err != nil {
		return "", errors.Wrap(err, "can not create private key file")
	}

	return fmt.Sprintf(
			"new public key has been saved to: [%s]\n" +
			"new private key has been saved to: [%s]",
			args.PublicKeyPath,
			args.PrivateKeyPath,),
		nil
}

func encrypt(args RsaArgs) (string, error){

	e := rsa.InitEmpty()

	bytes, err := ioutil.ReadFile(args.PublicKeyPath)
	if err != nil {
		return "", errors.Wrap(err,
			"can not read a content of the file with public key",
		)
	}

	err = e.SetPublicKeyFromPem(string(bytes))
	if err != nil {
		return "", errors.Wrap(err, "can not set public key")
	}

	return e.Encrypt(args.Message)
}

func decrypt(args RsaArgs) (string, error) {

	var e rsa.RsaEncryptor

	if args.PrivateKeyPassword == "" {
		e = rsa.InitEmpty()
	} else {
		e = rsa.InitEmptyWithPassword(args.PrivateKeyPassword)
	}

	bytes, err := ioutil.ReadFile(args.PrivateKeyPath)
	if err != nil {
		return "", errors.Wrap(err,
			"can not read a content of the file with private key",
		)
	}

	err = e.SetPrivateKeyFromPem(string(bytes))
	if err != nil {
		return "", errors.Wrap(err, "can not set private key")
	}

	return e.Decrypt(args.Message)
}

func getHandler(command string) func(args RsaArgs) (string, error) {

	switch command {
	case "generate":
		return generate
	case "encrypt":
		return encrypt
	case "decrypt":
		return decrypt
	default:
		return func(args RsaArgs) (string, error) {
			return "", errors.New(
				fmt.Sprintf("unknown command: %s\n", args.Command),
			)
		}
	}
}

func main() {

	args, err := getArgs()
	if err != nil {
		redConsole.Println("args parsing error:", err)
		flag.PrintDefaults()
		return
	}

	hander := getHandler(args.Command)

	result, err := hander(args)
	if err != nil {
		redConsole.Printf(
			"[%s] command execution is failed. reason: %s\n",
			args.Command,
			err,
		)
		return
	}

	if result != "" {
		greenConsole.Println(result)
	}
}
<commit_msg>update code<commit_after>package main

import (
	rsa "github.com/mrvine/encryptor-rsa"
	"flag"
	"github.com/fatih/color"
	"fmt"
	"github.com/pkg/errors"
	"io/ioutil"
)

var (
	redConsole    = color.New(color.FgRed)
	yellowConsole = color.New(color.FgYellow)
	greenConsole  = color.New(color.FgGreen)
)

func generate(args RsaArgs) (string, error) {

	var e rsa.RsaEncryptor
	var err error

	if args.PrivateKeyPassword == "" {
		e, err = rsa.InitWithKeyLength(args.KeyLength)
	} else {
		e, err = rsa.InitWithPassword(args.KeyLength, args.PrivateKeyPassword)
	}

	if err != nil {
		return "", errors.Wrap(err, "can not initialize rsa encryptor")
	}

	if args.PublicKeyPath == "" {
		args.PublicKeyPath = "public_" + getRandomString(16) + ".txt"
	}

	if args.PrivateKeyPath == "" {
		args.PrivateKeyPath = "private_" + getRandomString(16) + ".txt"
	}

	err = e.SavePublicKeyInPem(args.PublicKeyPath)
	if err != nil {
		return "", errors.Wrap(err, "can not create public key file")
	}

	err = e.SavePrivateKeyInPem(args.PrivateKeyPath)
	if err != nil {
		return "", errors.Wrap(err, "can not create private key file")
	}

	return fmt.Sprintf(
			"new public key has been saved to: [%s]\n" +
			"new private key has been saved to: [%s]",
			args.PublicKeyPath,
			args.PrivateKeyPath,),
		nil
}

func encrypt(args RsaArgs) (string, error){

	e := rsa.InitEmpty()

	bytes, err := ioutil.ReadFile(args.PublicKeyPath)
	if err != nil {
		return "", errors.Wrap(err,
			"can not read a content of the file with public key",
		)
	}

	err = e.SetPublicKeyFromPem(string(bytes))
	if err != nil {
		return "", errors.Wrap(err, "can not set public key")
	}

	return e.Encrypt(args.Message)
}

func decrypt(args RsaArgs) (string, error) {

	var e rsa.RsaEncryptor

	if args.PrivateKeyPassword == "" {
		e = rsa.InitEmpty()
	} else {
		e = rsa.InitEmptyWithPassword(args.PrivateKeyPassword)
	}

	bytes, err := ioutil.ReadFile(args.PrivateKeyPath)
	if err != nil {
		return "", errors.Wrap(err,
			"can not read a content of the file with private key",
		)
	}

	err = e.SetPrivateKeyFromPem(string(bytes))
	if err != nil {
		return "", errors.Wrap(err, "can not set private key")
	}

	return e.Decrypt(args.Message)
}

func getHandler(command string) func(args RsaArgs) (string, error) {

	switch command {
	case "generate":
		return generate
	case "encrypt":
		return encrypt
	case "decrypt":
		return decrypt
	default:
		return func(args RsaArgs) (string, error) {
			return "", errors.New(
				fmt.Sprintf("unknown command: %s\n", args.Command),
			)
		}
	}
}

func main() {

	args, err := getArgs()
	if err != nil {
		redConsole.Println("args parsing error:", err)
		flag.PrintDefaults()
		return
	}

	hander := getHandler(args.Command)

	result, err := hander(args)
	if err != nil {
		redConsole.Printf(
			"[%s] command execution is failed. reason: %s\n",
			args.Command,
			err,
		)
		return
	}

	if result != "" {
		greenConsole.Println(result)
	}
}
<|endoftext|>


<commit_before>// Copyright © 2017 Josh Dvir <josh@dvir.uk>
//
// 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 cmd

import (
	"context"
	"fmt"
	"os"
	"strconv"
	"sync"
	"time"

	"strings"

	cron "github.com/robfig/cron"
	"github.com/spf13/cobra"
	elastic "gopkg.in/olivere/elastic.v5"
)

var (
	olderThanInDays int
	esURL           string
	prefix          string
	wg              sync.WaitGroup
	ctx             context.Context
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "delete-aws-es-incidents",
	Short: "Delete ELK incidents on AWS ES 5.1",
	Long:  "",
	Run: func(cmd *cobra.Command, args []string) {
		if esURL == "" {
			println("No Elasticsearch URL present, can't continue.")
			os.Exit(0)
		}

		client, err := elastic.NewClient(
			elastic.SetURL(esURL),
			elastic.SetSniff(false),
		)
		if err != nil {
			panic(err)
		}

		esversion, err := client.ElasticsearchVersion(esURL)
		if err != nil {
			panic(err)
		}
		fmt.Printf("Elasticsearch version %s\n", esversion)

		runCommand()
		var wgm sync.WaitGroup
		cron := cron.New()
		cron.AddFunc("@hourly", func() { runCommand() })
		cron.Start()
		println("Cron run started...")
		wgm.Add(1)
		wgm.Wait()
	},
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	RootCmd.Flags().IntVarP(&olderThanInDays, "older-than-in-days", "d", 14, "delete incidents older then in days")
	RootCmd.Flags().StringVarP(&esURL, "es-url", "e", "", "Elasticsearch URL, eg. https://path-to-es.aws.com/")
	RootCmd.Flags().StringVarP(&prefix, "prefix", "p", "logstash-", "prefix for indexs. default is 'logstash-'")
}

func runCommand() {
	println("Starting deleting incidents run...")
	ctx = context.Background()
	client, err := elastic.NewClient(
		elastic.SetURL(esURL),
		elastic.SetSniff(false),
	)
	if err != nil {
		panic(err)
	}

	indexNames, err := client.IndexNames()
	if err != nil {
		panic(err)
	}

	for _, indexName := range indexNames {
		date := strings.TrimPrefix(indexName, prefix)
		dateArr := strings.Split(date, ".")
		nowTime := time.Now()
		indexYear, _ := strconv.Atoi(dateArr[0])
		indexMonth, _ := strconv.Atoi(dateArr[1])
		indexDay, _ := strconv.Atoi(dateArr[2])
		incidentTime := time.Date(indexYear, time.Month(indexMonth), indexDay, 0, 0, 0, 0, nowTime.Location())
		if daysDiff(nowTime, incidentTime) > olderThanInDays {
			wg.Add(1)
			go deleteIncident(ctx, client, indexName)
		}
	}

	wg.Wait()
	println("Ending deleting incidents run...")
}

func deleteIncident(ctx context.Context, client *elastic.Client, indexName string) {
	deleteIndex, err := client.DeleteIndex(indexName).Do(ctx)
	if err != nil {
		fmt.Printf("Error deleting index %s\n", indexName)
	}

	if deleteIndex.Acknowledged {
		fmt.Printf("index %s deleted.\n", indexName)
	}

	defer wg.Done()
}

func lastDayOfYear(t time.Time) time.Time {
	return time.Date(t.Year(), 12, 31, 0, 0, 0, 0, t.Location())
}

func firstDayOfNextYear(t time.Time) time.Time {
	return time.Date(t.Year()+1, 1, 1, 0, 0, 0, 0, t.Location())
}

// a - b in days
func daysDiff(a, b time.Time) (days int) {
	cur := b
	for cur.Year() < a.Year() {
		// add 1 to count the last day of the year too.
		days += lastDayOfYear(cur).YearDay() - cur.YearDay() + 1
		cur = firstDayOfNextYear(cur)
	}
	days += a.YearDay() - cur.YearDay()
	if b.AddDate(0, 0, days).After(a) {
		days--
	}
	return days
}
<commit_msg>only indexes that start with prefix<commit_after>// Copyright © 2017 Josh Dvir <josh@dvir.uk>
//
// 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 cmd

import (
	"context"
	"fmt"
	"os"
	"strconv"
	"sync"
	"time"

	"strings"

	cron "github.com/robfig/cron"
	"github.com/spf13/cobra"
	elastic "gopkg.in/olivere/elastic.v5"
)

var (
	olderThanInDays int
	esURL           string
	prefix          string
	wg              sync.WaitGroup
	ctx             context.Context
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "delete-aws-es-incidents",
	Short: "Delete ELK incidents on AWS ES 5.1",
	Long:  "",
	Run: func(cmd *cobra.Command, args []string) {
		if esURL == "" {
			println("No Elasticsearch URL present, can't continue.")
			os.Exit(0)
		}

		client, err := elastic.NewClient(
			elastic.SetURL(esURL),
			elastic.SetSniff(false),
		)
		if err != nil {
			panic(err)
		}

		esversion, err := client.ElasticsearchVersion(esURL)
		if err != nil {
			panic(err)
		}
		fmt.Printf("Elasticsearch version %s\n", esversion)

		runCommand()
		var wgm sync.WaitGroup
		cron := cron.New()
		cron.AddFunc("@hourly", func() { runCommand() })
		cron.Start()
		println("Cron run started...")
		wgm.Add(1)
		wgm.Wait()
	},
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	RootCmd.Flags().IntVarP(&olderThanInDays, "older-than-in-days", "d", 14, "delete incidents older then in days")
	RootCmd.Flags().StringVarP(&esURL, "es-url", "e", "", "Elasticsearch URL, eg. https://path-to-es.aws.com/")
	RootCmd.Flags().StringVarP(&prefix, "prefix", "p", "logstash-", "prefix for indexs. default is 'logstash-'")
}

func runCommand() {
	println("Starting deleting incidents run...")
	ctx = context.Background()
	client, err := elastic.NewClient(
		elastic.SetURL(esURL),
		elastic.SetSniff(false),
	)
	if err != nil {
		panic(err)
	}

	indexNames, err := client.IndexNames()
	if err != nil {
		panic(err)
	}

	for _, indexName := range indexNames {
		if strings.HasPrefix(indexName, prefix) {
			date := strings.TrimPrefix(indexName, prefix)
			dateArr := strings.Split(date, ".")
			nowTime := time.Now()
			indexYear, _ := strconv.Atoi(dateArr[0])
			indexMonth, _ := strconv.Atoi(dateArr[1])
			indexDay, _ := strconv.Atoi(dateArr[2])
			incidentTime := time.Date(indexYear, time.Month(indexMonth), indexDay, 0, 0, 0, 0, nowTime.Location())
			if daysDiff(nowTime, incidentTime) > olderThanInDays {
				wg.Add(1)
				go deleteIncident(ctx, client, indexName)
			}
		}
	}

	wg.Wait()
	println("Ending deleting incidents run...")
}

func deleteIncident(ctx context.Context, client *elastic.Client, indexName string) {
	deleteIndex, err := client.DeleteIndex(indexName).Do(ctx)
	if err != nil {
		fmt.Printf("Error deleting index %s\n", indexName)
	}

	if deleteIndex.Acknowledged {
		fmt.Printf("index %s deleted.\n", indexName)
	}

	defer wg.Done()
}

func lastDayOfYear(t time.Time) time.Time {
	return time.Date(t.Year(), 12, 31, 0, 0, 0, 0, t.Location())
}

func firstDayOfNextYear(t time.Time) time.Time {
	return time.Date(t.Year()+1, 1, 1, 0, 0, 0, 0, t.Location())
}

// a - b in days
func daysDiff(a, b time.Time) (days int) {
	cur := b
	for cur.Year() < a.Year() {
		// add 1 to count the last day of the year too.
		days += lastDayOfYear(cur).YearDay() - cur.YearDay() + 1
		cur = firstDayOfNextYear(cur)
	}
	days += a.YearDay() - cur.YearDay()
	if b.AddDate(0, 0, days).After(a) {
		days--
	}
	return days
}
<|endoftext|>


<commit_before>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/wandoulabs/redis-port/pkg/libs/atomic2"
	"github.com/wandoulabs/redis-port/pkg/libs/io/pipe"
	"github.com/wandoulabs/redis-port/pkg/libs/log"
	"github.com/wandoulabs/redis-port/pkg/libs/stats"
	"github.com/wandoulabs/redis-port/pkg/redis"
)

type cmdSync struct {
	rbytes, wbytes, nentry, ignore atomic2.Int64

	forward, nbypass atomic2.Int64
}

type cmdSyncStat struct {
	rbytes, wbytes, nentry, ignore int64

	forward, nbypass int64
}

func (cmd *cmdSync) Stat() *cmdSyncStat {
	return &cmdSyncStat{
		rbytes: cmd.rbytes.Get(),
		wbytes: cmd.wbytes.Get(),
		nentry: cmd.nentry.Get(),
		ignore: cmd.ignore.Get(),

		forward: cmd.forward.Get(),
		nbypass: cmd.nbypass.Get(),
	}
}

func (cmd *cmdSync) Main() {
	from, target := args.from, args.target
	if len(from) == 0 {
		log.Panic("invalid argument: from")
	}
	if len(target) == 0 {
		log.Panic("invalid argument: target")
	}

	log.Infof("sync from '%s' to '%s'\n", from, target)

	var sockfile *os.File
	if len(args.sockfile) != 0 {
		sockfile = openReadWriteFile(args.sockfile)
		defer sockfile.Close()
	}

	var input io.ReadCloser
	var nsize int64
	if args.psync {
		input, nsize = cmd.SendPSyncCmd(from, args.passwd)
	} else {
		input, nsize = cmd.SendSyncCmd(from, args.passwd)
	}
	defer input.Close()

	log.Infof("rdb file = %d\n", nsize)

	if sockfile != nil {
		r, w := pipe.NewFilePipe(int(args.filesize), sockfile)
		defer r.Close()
		go func(r io.Reader) {
			defer w.Close()
			p := make([]byte, ReaderBufferSize)
			for {
				iocopy(r, w, p, len(p))
			}
		}(input)
		input = r
	}

	reader := bufio.NewReaderSize(input, ReaderBufferSize)

	cmd.SyncRDBFile(reader, target, args.auth, nsize)
	cmd.SyncCommand(reader, target, args.auth)
}

func (cmd *cmdSync) SendSyncCmd(master, passwd string) (net.Conn, int64) {
	c, wait := openSyncConn(master, passwd)
	for {
		select {
		case nsize := <-wait:
			if nsize == 0 {
				log.Info("+")
			} else {
				return c, nsize
			}
		case <-time.After(time.Second):
			log.Info("-")
		}
	}
}

func (cmd *cmdSync) SendPSyncCmd(master, passwd string) (pipe.Reader, int64) {
	c := openNetConn(master, passwd)
	br := bufio.NewReaderSize(c, ReaderBufferSize)
	bw := bufio.NewWriterSize(c, WriterBufferSize)

	runid, offset, wait := sendPSyncFullsync(br, bw)
	log.Infof("psync runid = %s offset = %d, fullsync", runid, offset)

	var nsize int64
	for nsize == 0 {
		select {
		case nsize = <-wait:
			if nsize == 0 {
				log.Info("+")
			}
		case <-time.After(time.Second):
			log.Info("-")
		}
	}

	piper, pipew := pipe.NewSize(ReaderBufferSize)

	go func() {
		defer pipew.Close()
		p := make([]byte, 8192)
		for rdbsize := int(nsize); rdbsize != 0; {
			rdbsize -= iocopy(br, pipew, p, rdbsize)
		}
		for {
			n, err := cmd.PSyncPipeCopy(c, br, bw, offset, pipew)
			if err != nil {
				log.PanicErrorf(err, "psync runid = %s, offset = %d, pipe is broken", runid, offset)
			}
			offset += n
			for {
				time.Sleep(time.Second)
				c = openNetConnSoft(master, passwd)
				if c != nil {
					log.Infof("psync reopen connection, offset = %d", offset)
					break
				} else {
					log.Infof("psync reopen connection, failed")
				}
			}
			authPassword(c, passwd)
			br = bufio.NewReaderSize(c, ReaderBufferSize)
			bw = bufio.NewWriterSize(c, WriterBufferSize)
			sendPSyncContinue(br, bw, runid, offset)
		}
	}()
	return piper, nsize
}

func (cmd *cmdSync) PSyncPipeCopy(c net.Conn, br *bufio.Reader, bw *bufio.Writer, offset int64, copyto io.Writer) (int64, error) {
	defer c.Close()
	var nread atomic2.Int64
	go func() {
		defer c.Close()
		for {
			time.Sleep(time.Second * 5)
			if err := sendPSyncAck(bw, offset+nread.Get()); err != nil {
				return
			}
		}
	}()

	var p = make([]byte, 8192)
	for {
		n, err := br.Read(p)
		if err != nil {
			return nread.Get(), nil
		}
		if _, err := copyto.Write(p[:n]); err != nil {
			return nread.Get(), err
		}
		nread.Add(int64(n))
	}
}

func (cmd *cmdSync) SyncRDBFile(reader *bufio.Reader, target, passwd string, nsize int64) {
	pipe := newRDBLoader(reader, &cmd.rbytes, args.parallel*32)
	wait := make(chan struct{})
	go func() {
		defer close(wait)
		group := make(chan int, args.parallel)
		for i := 0; i < cap(group); i++ {
			go func() {
				defer func() {
					group <- 0
				}()
				c := openRedisConn(target, passwd)
				defer c.Close()
				var lastdb uint32 = 0
				for e := range pipe {
					if !acceptDB(e.DB) {
						cmd.ignore.Incr()
					} else {
						cmd.nentry.Incr()
						if e.DB != lastdb {
							lastdb = e.DB
							selectDB(c, lastdb)
						}
						restoreRdbEntry(c, e)
					}
				}
			}()
		}
		for i := 0; i < cap(group); i++ {
			<-group
		}
	}()

	for done := false; !done; {
		select {
		case <-wait:
			done = true
		case <-time.After(time.Second):
		}
		stat := cmd.Stat()
		var b bytes.Buffer
		fmt.Fprintf(&b, "total=%d - %12d [%3d%%]", nsize, stat.rbytes, 100*stat.rbytes/nsize)
		fmt.Fprintf(&b, "  entry=%-12d", stat.nentry)
		if stat.ignore != 0 {
			fmt.Fprintf(&b, "  ignore=%-12d", stat.ignore)
		}
		log.Info(b.String())
	}
	log.Info("sync rdb done")
}

func (cmd *cmdSync) SyncCommand(reader *bufio.Reader, target, passwd string) {
	c := openNetConn(target, passwd)
	defer c.Close()

	writer := bufio.NewWriterSize(stats.NewCountWriter(c, &cmd.wbytes), WriterBufferSize)
	defer flushWriter(writer)

	go func() {
		p := make([]byte, ReaderBufferSize)
		for {
			iocopy(c, ioutil.Discard, p, len(p))
		}
	}()

	go func() {
		var bypass bool = false
		for {
			resp := redis.MustDecode(reader)
			if scmd, args, err := redis.ParseArgs(resp); err != nil {
				log.PanicError(err, "parse command arguments failed")
			} else if scmd != "ping" {
				if scmd == "select" {
					if len(args) != 1 {
						log.Panicf("select command len(args) = %d", len(args))
					}
					s := string(args[0])
					n, err := parseInt(s, MinDB, MaxDB)
					if err != nil {
						log.PanicErrorf(err, "parse db = %s failed", s)
					}
					bypass = !acceptDB(uint32(n))
				}
				if bypass {
					cmd.nbypass.Incr()
					continue
				}
			}
			cmd.forward.Incr()
			redis.MustEncode(writer, resp)
			flushWriter(writer)
		}
	}()

	for lstat := cmd.Stat(); ; {
		time.Sleep(time.Second)
		nstat := cmd.Stat()
		var b bytes.Buffer
		fmt.Fprintf(&b, "sync: ")
		fmt.Fprintf(&b, " +forward=%-6d", nstat.forward-lstat.forward)
		fmt.Fprintf(&b, " +nbypass=%-6d", nstat.nbypass-lstat.nbypass)
		fmt.Fprintf(&b, " +nbytes=%d", nstat.wbytes-lstat.wbytes)
		log.Info(b.String())
		lstat = nstat
	}
}
<commit_msg>ignore big key+value<commit_after>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/wandoulabs/redis-port/pkg/libs/atomic2"
	"github.com/wandoulabs/redis-port/pkg/libs/io/pipe"
	"github.com/wandoulabs/redis-port/pkg/libs/log"
	"github.com/wandoulabs/redis-port/pkg/libs/stats"
	"github.com/wandoulabs/redis-port/pkg/redis"
)

type cmdSync struct {
	rbytes, wbytes, nentry, ignore atomic2.Int64

	forward, nbypass atomic2.Int64
}

type cmdSyncStat struct {
	rbytes, wbytes, nentry, ignore int64

	forward, nbypass int64
}

func (cmd *cmdSync) Stat() *cmdSyncStat {
	return &cmdSyncStat{
		rbytes: cmd.rbytes.Get(),
		wbytes: cmd.wbytes.Get(),
		nentry: cmd.nentry.Get(),
		ignore: cmd.ignore.Get(),

		forward: cmd.forward.Get(),
		nbypass: cmd.nbypass.Get(),
	}
}

func (cmd *cmdSync) Main() {
	from, target := args.from, args.target
	if len(from) == 0 {
		log.Panic("invalid argument: from")
	}
	if len(target) == 0 {
		log.Panic("invalid argument: target")
	}

	log.Infof("sync from '%s' to '%s'\n", from, target)

	var sockfile *os.File
	if len(args.sockfile) != 0 {
		sockfile = openReadWriteFile(args.sockfile)
		defer sockfile.Close()
	}

	var input io.ReadCloser
	var nsize int64
	if args.psync {
		input, nsize = cmd.SendPSyncCmd(from, args.passwd)
	} else {
		input, nsize = cmd.SendSyncCmd(from, args.passwd)
	}
	defer input.Close()

	log.Infof("rdb file = %d\n", nsize)

	if sockfile != nil {
		r, w := pipe.NewFilePipe(int(args.filesize), sockfile)
		defer r.Close()
		go func(r io.Reader) {
			defer w.Close()
			p := make([]byte, ReaderBufferSize)
			for {
				iocopy(r, w, p, len(p))
			}
		}(input)
		input = r
	}

	reader := bufio.NewReaderSize(input, ReaderBufferSize)

	cmd.SyncRDBFile(reader, target, args.auth, nsize)
	cmd.SyncCommand(reader, target, args.auth)
}

func (cmd *cmdSync) SendSyncCmd(master, passwd string) (net.Conn, int64) {
	c, wait := openSyncConn(master, passwd)
	for {
		select {
		case nsize := <-wait:
			if nsize == 0 {
				log.Info("+")
			} else {
				return c, nsize
			}
		case <-time.After(time.Second):
			log.Info("-")
		}
	}
}

func (cmd *cmdSync) SendPSyncCmd(master, passwd string) (pipe.Reader, int64) {
	c := openNetConn(master, passwd)
	br := bufio.NewReaderSize(c, ReaderBufferSize)
	bw := bufio.NewWriterSize(c, WriterBufferSize)

	runid, offset, wait := sendPSyncFullsync(br, bw)
	log.Infof("psync runid = %s offset = %d, fullsync", runid, offset)

	var nsize int64
	for nsize == 0 {
		select {
		case nsize = <-wait:
			if nsize == 0 {
				log.Info("+")
			}
		case <-time.After(time.Second):
			log.Info("-")
		}
	}

	piper, pipew := pipe.NewSize(ReaderBufferSize)

	go func() {
		defer pipew.Close()
		p := make([]byte, 8192)
		for rdbsize := int(nsize); rdbsize != 0; {
			rdbsize -= iocopy(br, pipew, p, rdbsize)
		}
		for {
			n, err := cmd.PSyncPipeCopy(c, br, bw, offset, pipew)
			if err != nil {
				log.PanicErrorf(err, "psync runid = %s, offset = %d, pipe is broken", runid, offset)
			}
			offset += n
			for {
				time.Sleep(time.Second)
				c = openNetConnSoft(master, passwd)
				if c != nil {
					log.Infof("psync reopen connection, offset = %d", offset)
					break
				} else {
					log.Infof("psync reopen connection, failed")
				}
			}
			authPassword(c, passwd)
			br = bufio.NewReaderSize(c, ReaderBufferSize)
			bw = bufio.NewWriterSize(c, WriterBufferSize)
			sendPSyncContinue(br, bw, runid, offset)
		}
	}()
	return piper, nsize
}

func (cmd *cmdSync) PSyncPipeCopy(c net.Conn, br *bufio.Reader, bw *bufio.Writer, offset int64, copyto io.Writer) (int64, error) {
	defer c.Close()
	var nread atomic2.Int64
	go func() {
		defer c.Close()
		for {
			time.Sleep(time.Second * 5)
			if err := sendPSyncAck(bw, offset+nread.Get()); err != nil {
				return
			}
		}
	}()

	var p = make([]byte, 8192)
	for {
		n, err := br.Read(p)
		if err != nil {
			return nread.Get(), nil
		}
		if _, err := copyto.Write(p[:n]); err != nil {
			return nread.Get(), err
		}
		nread.Add(int64(n))
	}
}

func (cmd *cmdSync) SyncRDBFile(reader *bufio.Reader, target, passwd string, nsize int64) {
	pipe := newRDBLoader(reader, &cmd.rbytes, args.parallel*32)
	wait := make(chan struct{})
	go func() {
		defer close(wait)
		group := make(chan int, args.parallel)
		for i := 0; i < cap(group); i++ {
			go func() {
				defer func() {
					group <- 0
				}()
				c := openRedisConn(target, passwd)
				defer c.Close()
				var lastdb uint32 = 0
				for e := range pipe {
					if !acceptDB(e.DB) {
						cmd.ignore.Incr()
					} else {
						cmd.nentry.Incr()
						if e.DB != lastdb {
							lastdb = e.DB
							selectDB(c, lastdb)
						}
						if len(e.Value) > 512*1024*1024 {
							s := e.Value[:1024]
							fmt.Printf(">>>>>>>>>>>>> key:[%s] %v, has value len=%d, value dump=%v\n", e.Key, e.Key, len(e.Value), s)
						} else {
							restoreRdbEntry(c, e)
						}
					}
				}
			}()
		}
		for i := 0; i < cap(group); i++ {
			<-group
		}
	}()

	for done := false; !done; {
		select {
		case <-wait:
			done = true
		case <-time.After(time.Second):
		}
		stat := cmd.Stat()
		var b bytes.Buffer
		fmt.Fprintf(&b, "total=%d - %12d [%3d%%]", nsize, stat.rbytes, 100*stat.rbytes/nsize)
		fmt.Fprintf(&b, "  entry=%-12d", stat.nentry)
		if stat.ignore != 0 {
			fmt.Fprintf(&b, "  ignore=%-12d", stat.ignore)
		}
		log.Info(b.String())
	}
	log.Info("sync rdb done")
}

func (cmd *cmdSync) SyncCommand(reader *bufio.Reader, target, passwd string) {
	c := openNetConn(target, passwd)
	defer c.Close()

	writer := bufio.NewWriterSize(stats.NewCountWriter(c, &cmd.wbytes), WriterBufferSize)
	defer flushWriter(writer)

	go func() {
		p := make([]byte, ReaderBufferSize)
		for {
			iocopy(c, ioutil.Discard, p, len(p))
		}
	}()

	go func() {
		var bypass bool = false
		for {
			resp := redis.MustDecode(reader)
			if scmd, args, err := redis.ParseArgs(resp); err != nil {
				log.PanicError(err, "parse command arguments failed")
			} else if scmd != "ping" {
				if scmd == "select" {
					if len(args) != 1 {
						log.Panicf("select command len(args) = %d", len(args))
					}
					s := string(args[0])
					n, err := parseInt(s, MinDB, MaxDB)
					if err != nil {
						log.PanicErrorf(err, "parse db = %s failed", s)
					}
					bypass = !acceptDB(uint32(n))
				}
				if bypass {
					cmd.nbypass.Incr()
					continue
				}
			}
			cmd.forward.Incr()
			redis.MustEncode(writer, resp)
			flushWriter(writer)
		}
	}()

	for lstat := cmd.Stat(); ; {
		time.Sleep(time.Second)
		nstat := cmd.Stat()
		var b bytes.Buffer
		fmt.Fprintf(&b, "sync: ")
		fmt.Fprintf(&b, " +forward=%-6d", nstat.forward-lstat.forward)
		fmt.Fprintf(&b, " +nbypass=%-6d", nstat.nbypass-lstat.nbypass)
		fmt.Fprintf(&b, " +nbytes=%d", nstat.wbytes-lstat.wbytes)
		log.Info(b.String())
		lstat = nstat
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	es "github.com/bebanjo/elastigo/lib"
	"github.com/spf13/cobra"
)

// takeCmd represents the snapshot take command
var takeCmd = &cobra.Command{
	Use:   "take",
	Short: "Take a snapshot",
	Long: `You are required to set a destination. It will create a snapshot
on the destination repository. If repository does not exist, you can create
it with the provided flag.`,
	Run: func(cmd *cobra.Command, args []string) {
		var conn = es.NewConn()
		var date = time.Now().Format("20060102150405")
		var state = "STARTING"
		var query interface{}

		// A destination is required
		if *destination == "" {
			fmt.Fprintf(os.Stderr, "take: destination required\n")
			os.Exit(1)
		}

		// Create repository if --create-repository flag is enabled
		if *createRepositoryTake {
			log.Println("creating repository", *destination)
			if err := createRepository(conn, *destination); err != nil {
				fmt.Fprintf(os.Stderr, "create repository: error for %s %v", *destination, err)
				os.Exit(1)
			}
		}

		// Select only destinationTake-related indices if --all flag is not used
		if !*allIndices {
			indicesInfo := conn.GetCatIndexInfo(fmt.Sprintf("%s*", *destination))
			indicesNamesString := strings.Join(indicesNames(indicesInfo), ",")
			query = map[string]interface{}{"indices": indicesNamesString}
		}

		// Take Snapshot
		_, err := conn.TakeSnapshot(*destination, date, nil, query)
		if err != nil {
			fmt.Fprintf(os.Stderr, "take: error %v\n", err)
			os.Exit(1)
		}

		// Poll for Snapshot status until it is done
		log.Println("waiting for snapshot", date, "to be ready...", state)
		for state != "SUCCESS" {
			snapshots, err := conn.GetSnapshotByName(*destination, date, nil)
			if err != nil {
				fmt.Fprintf(os.Stderr, "take: error getting snapshot %s %v\n", *destination, err)
				os.Exit(1)
			}

			if len(snapshots.Snapshots) < 1 {
				break
			}

			state = snapshots.Snapshots[0].State
			log.Println("waiting for snapshot", date, "to be ready...", state)
			time.Sleep(5 * time.Second)
		}

	},
}

func init() {
	RootCmd.AddCommand(takeCmd)

	createRepositoryTake = takeCmd.PersistentFlags().BoolP("create-repository", "r", false, "Create repository")
	allIndices = takeCmd.PersistentFlags().BoolP("all", "a", false,
		"Take snapshot of all indices. Otherwise, only those matching the destination")
}

func indicesNames(catIndexInfo []es.CatIndexInfo) []string {
	var names []string
	for _, cii := range catIndexInfo {
		names = append(names, cii.Name)
	}
	return names
}
<commit_msg>Log indices being snapshotted<commit_after>package cmd

import (
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	es "github.com/bebanjo/elastigo/lib"
	"github.com/spf13/cobra"
)

// takeCmd represents the snapshot take command
var takeCmd = &cobra.Command{
	Use:   "take",
	Short: "Take a snapshot",
	Long: `You are required to set a destination. It will create a snapshot
on the destination repository. If repository does not exist, you can create
it with the provided flag.`,
	Run: func(cmd *cobra.Command, args []string) {
		var conn = es.NewConn()
		var date = time.Now().Format("20060102150405")
		var state = "STARTING"
		var query interface{}
		var indicesNamesString string

		// A destination is required
		if *destination == "" {
			fmt.Fprintf(os.Stderr, "take: destination required\n")
			os.Exit(1)
		}

		// Create repository if --create-repository flag is enabled
		if *createRepositoryTake {
			log.Println("creating repository", *destination)
			if err := createRepository(conn, *destination); err != nil {
				fmt.Fprintf(os.Stderr, "create repository: error for %s %v", *destination, err)
				os.Exit(1)
			}
		}

		// Select only destinationTake-related indices if --all flag is not used
		if !*allIndices {
			indicesInfo := conn.GetCatIndexInfo(fmt.Sprintf("%s*", *destination))
			indicesNamesString = strings.Join(indicesNames(indicesInfo), ",")
			query = map[string]interface{}{"indices": indicesNamesString}
		} else {
			indicesInfo := conn.GetCatIndexInfo("")
			indicesNamesString = strings.Join(indicesNames(indicesInfo), ",")
		}
		log.Println("Taking snapshot of indices:", indicesNamesString)

		// Take Snapshot
		_, err := conn.TakeSnapshot(*destination, date, nil, query)
		if err != nil {
			fmt.Fprintf(os.Stderr, "take: error %v\n", err)
			os.Exit(1)
		}

		// Poll for Snapshot status until it is done
		log.Println("waiting for snapshot", date, "to be ready...", state)
		for state != "SUCCESS" {
			snapshots, err := conn.GetSnapshotByName(*destination, date, nil)
			if err != nil {
				fmt.Fprintf(os.Stderr, "take: error getting snapshot %s %v\n", *destination, err)
				os.Exit(1)
			}

			if len(snapshots.Snapshots) < 1 {
				break
			}

			state = snapshots.Snapshots[0].State
			log.Println("waiting for snapshot", date, "to be ready...", state)
			time.Sleep(5 * time.Second)
		}

	},
}

func init() {
	RootCmd.AddCommand(takeCmd)

	createRepositoryTake = takeCmd.PersistentFlags().BoolP("create-repository", "r", false, "Create repository")
	allIndices = takeCmd.PersistentFlags().BoolP("all", "a", false,
		"Take snapshot of all indices. Otherwise, only those matching the destination")
}

func indicesNames(catIndexInfo []es.CatIndexInfo) []string {
	var names []string
	for _, cii := range catIndexInfo {
		names = append(names, cii.Name)
	}
	return names
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/Masterminds/cookoo"
)

// Tree prints a tree representing dependencies.
func Tree(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	buildContext, err := GetBuildContext()
	if err != nil {
		return nil, err
	}
	showcore := p.Get("showcore", false).(bool)
	basedir := p.Get("dir", ".").(string)
	myName := guessPackageName(buildContext, basedir)

	if basedir == "." {
		var err error
		basedir, err = os.Getwd()
		if err != nil {
			Error("Could not get working directory")
			return nil, err
		}
	}

	fmt.Println(myName)
	displayTree(buildContext, basedir, myName, 1, showcore)
	return nil, nil
}

// ListDeps lists all of the dependencies of the current project.
//
// Params:
//
// Returns:
//
func ListDeps(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	buildContext, err := GetBuildContext()
	if err != nil {
		return nil, err
	}
	basedir := p.Get("dir", ".").(string)
	myName := guessPackageName(buildContext, basedir)

	basedir, err = filepath.Abs(basedir)
	if err != nil {
		return nil, err
	}

	direct := map[string]*pinfo{}
	d := walkDeps(buildContext, basedir, myName)
	for _, i := range d {
		listDeps(buildContext, direct, i, basedir)
	}

	sortable := make([]string, len(direct))
	i := 0
	for k := range direct {
		sortable[i] = k
		i++
	}

	sort.Strings(sortable)

	for _, k := range sortable {
		dec := "yes"
		if d, ok := direct[k]; ok && d.PType == ptypeUnknown {
			dec = "no"
		}
		vendored := "no"
		if d, ok := direct[k]; ok && d.PType == ptypeVendor {
			vendored = "yes"
		}
		fmt.Printf("%s (Present: %s, Vendored: %s)\n", k, dec, vendored)
	}

	return nil, nil
}

func listDeps(b *BuildCtxt, info map[string]*pinfo, name, path string) {
	found := findPkg(b, name, path)
	switch found.PType {
	case ptypeUnknown:
		info[name] = found
		break
	case ptypeGoroot, ptypeCgo:
		break
	default:
		info[name] = found
		for _, i := range walkDeps(b, found.Path, found.Name) {
			listDeps(b, info, i, found.Path)
		}
	}
}

func displayTree(b *BuildCtxt, basedir, myName string, level int, core bool) {
	deps := walkDeps(b, basedir, myName)
	for _, name := range deps {
		found := findPkg(b, name, basedir)
		if found.PType == ptypeUnknown {
			msg := "glide get " + found.Name
			fmt.Printf("\t%s\t(%s)\n", found.Name, msg)
			continue
		}
		if !core && found.PType == ptypeGoroot || found.PType == ptypeCgo {
			continue
		}
		fmt.Print(strings.Repeat("\t", level))
		fmt.Printf("%s   (%s)\n", found.Name, found.Path)
		displayTree(b, found.Path, found.Name, level+1, core)
	}
}

type ptype int8

const (
	ptypeUnknown ptype = iota
	ptypeLocal
	ptypeVendor
	ptypeGopath
	ptypeGoroot
	ptypeCgo
)

type pinfo struct {
	Name, Path string
	PType      ptype
	Vendored   bool
}

func findPkg(b *BuildCtxt, name, cwd string) *pinfo {
	var fi os.FileInfo
	var err error
	var p string

	info := &pinfo{
		Name: name,
	}

	// Recurse backward to scan other vendor/ directories
	for wd := cwd; wd != "/"; wd = filepath.Dir(wd) {
		p = filepath.Join(wd, "vendor", name)
		if fi, err = os.Stat(p); err == nil && (fi.IsDir() || isLink(fi)) {
			info.Path = p
			info.PType = ptypeVendor
			info.Vendored = true
			return info
		}
	}
	// Check $GOPATH
	for _, r := range strings.Split(b.GOPATH, ":") {
		p = filepath.Join(r, "src", name)
		if fi, err = os.Stat(p); err == nil && (fi.IsDir() || isLink(fi)) {
			info.Path = p
			info.PType = ptypeGopath
			return info
		}
	}

	// Check $GOROOT
	for _, r := range strings.Split(b.GOROOT, ":") {
		p = filepath.Join(r, "src", name)
		if fi, err = os.Stat(p); err == nil && (fi.IsDir() || isLink(fi)) {
			info.Path = p
			info.PType = ptypeGoroot
			return info
		}
	}

	// Finally, if this is "C", we're dealing with cgo
	if name == "C" {
		info.PType = ptypeCgo
	}

	return info
}

func isLink(fi os.FileInfo) bool {
	return fi.Mode()&os.ModeSymlink == os.ModeSymlink
}

func walkDeps(b *BuildCtxt, base, myName string) []string {
	externalDeps := []string{}
	filepath.Walk(base, func(path string, fi os.FileInfo, err error) error {
		if excludeSubtree(path, fi) {
			if fi.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}

		pkg, err := b.ImportDir(path, 0)
		if err != nil {
			return err
		}

		if pkg.Goroot {
			return nil
		}

		for _, imp := range pkg.Imports {
			//if strings.HasPrefix(imp, myName) {
			////Info("Skipping %s because it is a subpackage of %s", imp, myName)
			//continue
			//}
			if imp == myName {
				continue
			}
			externalDeps = append(externalDeps, imp)
		}

		return nil
	})
	return externalDeps
}

func excludeSubtree(path string, fi os.FileInfo) bool {
	top := filepath.Base(path)

	if !fi.IsDir() && !isLink(fi) {
		return true
	}

	// Provisionally, we'll skip vendor. We definitely
	// should skip testdata.
	if top == "vendor" || top == "testdata" {
		return true
	}

	// Skip anything that starts with _
	if strings.HasPrefix(top, "_") || (strings.HasPrefix(top, ".") && top != ".") {
		return true
	}
	return false
}
<commit_msg>Refine the glide list output to show package origin<commit_after>package cmd

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/Masterminds/cookoo"
)

// Tree prints a tree representing dependencies.
func Tree(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	buildContext, err := GetBuildContext()
	if err != nil {
		return nil, err
	}
	showcore := p.Get("showcore", false).(bool)
	basedir := p.Get("dir", ".").(string)
	myName := guessPackageName(buildContext, basedir)

	if basedir == "." {
		var err error
		basedir, err = os.Getwd()
		if err != nil {
			Error("Could not get working directory")
			return nil, err
		}
	}

	fmt.Println(myName)
	displayTree(buildContext, basedir, myName, 1, showcore)
	return nil, nil
}

// ListDeps lists all of the dependencies of the current project.
//
// Params:
//
// Returns:
//
func ListDeps(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	buildContext, err := GetBuildContext()
	if err != nil {
		return nil, err
	}
	basedir := p.Get("dir", ".").(string)
	myName := guessPackageName(buildContext, basedir)

	basedir, err = filepath.Abs(basedir)
	if err != nil {
		return nil, err
	}

	direct := map[string]*pinfo{}
	d := walkDeps(buildContext, basedir, myName)
	for _, i := range d {
		listDeps(buildContext, direct, i, basedir)
	}

	sortable := make([]string, len(direct))
	i := 0
	for k := range direct {
		sortable[i] = k
		i++
	}

	sort.Strings(sortable)

	for _, k := range sortable {
		t := direct[k].PType
		fmt.Printf("%s (Location: %s)\n", k, ptypeString(t))
	}

	return nil, nil
}

func listDeps(b *BuildCtxt, info map[string]*pinfo, name, path string) {
	found := findPkg(b, name, path)
	switch found.PType {
	case ptypeUnknown:
		info[name] = found
		break
	case ptypeGoroot, ptypeCgo:
		break
	default:
		info[name] = found
		for _, i := range walkDeps(b, found.Path, found.Name) {
			listDeps(b, info, i, found.Path)
		}
	}
}

func displayTree(b *BuildCtxt, basedir, myName string, level int, core bool) {
	deps := walkDeps(b, basedir, myName)
	for _, name := range deps {
		found := findPkg(b, name, basedir)
		if found.PType == ptypeUnknown {
			msg := "glide get " + found.Name
			fmt.Printf("\t%s\t(%s)\n", found.Name, msg)
			continue
		}
		if !core && found.PType == ptypeGoroot || found.PType == ptypeCgo {
			continue
		}
		fmt.Print(strings.Repeat("\t", level))
		fmt.Printf("%s   (%s)\n", found.Name, found.Path)
		displayTree(b, found.Path, found.Name, level+1, core)
	}
}

type ptype int8

const (
	ptypeUnknown ptype = iota
	ptypeLocal
	ptypeVendor
	ptypeGopath
	ptypeGoroot
	ptypeCgo
)

func ptypeString(t ptype) string {
	switch t {
	case ptypeLocal:
		return "local"
	case ptypeVendor:
		return "vendored"
	case ptypeGopath:
		return "gopath"
	case ptypeGoroot:
		return "core"
	case ptypeCgo:
		return "cgo"
	default:
		return "missing"
	}
}

type pinfo struct {
	Name, Path string
	PType      ptype
	Vendored   bool
}

func findPkg(b *BuildCtxt, name, cwd string) *pinfo {
	var fi os.FileInfo
	var err error
	var p string

	info := &pinfo{
		Name: name,
	}

	// Recurse backward to scan other vendor/ directories
	for wd := cwd; wd != "/"; wd = filepath.Dir(wd) {
		p = filepath.Join(wd, "vendor", name)
		if fi, err = os.Stat(p); err == nil && (fi.IsDir() || isLink(fi)) {
			info.Path = p
			info.PType = ptypeVendor
			info.Vendored = true
			return info
		}
	}
	// Check $GOPATH
	for _, r := range strings.Split(b.GOPATH, ":") {
		p = filepath.Join(r, "src", name)
		if fi, err = os.Stat(p); err == nil && (fi.IsDir() || isLink(fi)) {
			info.Path = p
			info.PType = ptypeGopath
			return info
		}
	}

	// Check $GOROOT
	for _, r := range strings.Split(b.GOROOT, ":") {
		p = filepath.Join(r, "src", name)
		if fi, err = os.Stat(p); err == nil && (fi.IsDir() || isLink(fi)) {
			info.Path = p
			info.PType = ptypeGoroot
			return info
		}
	}

	// Finally, if this is "C", we're dealing with cgo
	if name == "C" {
		info.PType = ptypeCgo
	}

	return info
}

func isLink(fi os.FileInfo) bool {
	return fi.Mode()&os.ModeSymlink == os.ModeSymlink
}

func walkDeps(b *BuildCtxt, base, myName string) []string {
	externalDeps := []string{}
	filepath.Walk(base, func(path string, fi os.FileInfo, err error) error {
		if excludeSubtree(path, fi) {
			if fi.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}

		pkg, err := b.ImportDir(path, 0)
		if err != nil {
			return err
		}

		if pkg.Goroot {
			return nil
		}

		for _, imp := range pkg.Imports {
			//if strings.HasPrefix(imp, myName) {
			////Info("Skipping %s because it is a subpackage of %s", imp, myName)
			//continue
			//}
			if imp == myName {
				continue
			}
			externalDeps = append(externalDeps, imp)
		}

		return nil
	})
	return externalDeps
}

func excludeSubtree(path string, fi os.FileInfo) bool {
	top := filepath.Base(path)

	if !fi.IsDir() && !isLink(fi) {
		return true
	}

	// Provisionally, we'll skip vendor. We definitely
	// should skip testdata.
	if top == "vendor" || top == "testdata" {
		return true
	}

	// Skip anything that starts with _
	if strings.HasPrefix(top, "_") || (strings.HasPrefix(top, ".") && top != ".") {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package main

import (
	"encoding/json"
	"net/http"
	"os"
	"strings"

	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/xyproto/simpleredis"
)

var (
	masterPool *simpleredis.ConnectionPool
	slavePool  *simpleredis.ConnectionPool
)

func ListRangeHandler(rw http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["key"]
	list := simpleredis.NewList(slavePool, key)
	members := HandleError(list.GetAll()).([]string)
	membersJSON := HandleError(json.MarshalIndent(members, "", "  ")).([]byte)
	rw.Write(membersJSON)
}

func ListPushHandler(rw http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["key"]
	value := mux.Vars(req)["value"]
	list := simpleredis.NewList(masterPool, key)
	HandleError(nil, list.Add(value))
	ListRangeHandler(rw, req)
}

func InfoHandler(rw http.ResponseWriter, req *http.Request) {
	info := HandleError(masterPool.Get(0).Do("INFO")).([]byte)
	rw.Write(info)
}

func EnvHandler(rw http.ResponseWriter, req *http.Request) {
	environment := make(map[string]string)
	for _, item := range os.Environ() {
		splits := strings.Split(item, "=")
		key := splits[0]
		val := strings.Join(splits[1:], "=")
		environment[key] = val
	}

	envJSON := HandleError(json.MarshalIndent(environment, "", "  ")).([]byte)
	rw.Write(envJSON)
}

func HandleError(result interface{}, err error) (r interface{}) {
	if err != nil {
		panic(err)
	}
	return result
}

func main() {
	masterPool = simpleredis.NewConnectionPoolHost("redis-master:6379")
	defer masterPool.Close()
	slavePool = simpleredis.NewConnectionPoolHost("redis-slave:6379")
	defer slavePool.Close()

	r := mux.NewRouter()
	r.Path("/lrange/{key}").Methods("GET").HandlerFunc(ListRangeHandler)
	r.Path("/rpush/{key}/{value}").Methods("GET").HandlerFunc(ListPushHandler)
	r.Path("/info").Methods("GET").HandlerFunc(InfoHandler)
	r.Path("/env").Methods("GET").HandlerFunc(EnvHandler)

	n := negroni.Classic()
	n.UseHandler(r)
	n.Run(":3000")
}
<commit_msg>updates to main.go<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package main

import (
	"encoding/json"
	"net/http"
	"os"
	"strings"

	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/xyproto/simpleredis"
)

var (
	masterPool *simpleredis.ConnectionPool
	replicaPool  *simpleredis.ConnectionPool
)

func ListRangeHandler(rw http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["key"]
	list := simpleredis.NewList(replicaPool, key)
	members := HandleError(list.GetAll()).([]string)
	membersJSON := HandleError(json.MarshalIndent(members, "", "  ")).([]byte)
	rw.Write(membersJSON)
}

func ListPushHandler(rw http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["key"]
	value := mux.Vars(req)["value"]
	list := simpleredis.NewList(masterPool, key)
	HandleError(nil, list.Add(value))
	ListRangeHandler(rw, req)
}

func InfoHandler(rw http.ResponseWriter, req *http.Request) {
	info := HandleError(masterPool.Get(0).Do("INFO")).([]byte)
	rw.Write(info)
}

func EnvHandler(rw http.ResponseWriter, req *http.Request) {
	environment := make(map[string]string)
	for _, item := range os.Environ() {
		splits := strings.Split(item, "=")
		key := splits[0]
		val := strings.Join(splits[1:], "=")
		environment[key] = val
	}

	envJSON := HandleError(json.MarshalIndent(environment, "", "  ")).([]byte)
	rw.Write(envJSON)
}

func HandleError(result interface{}, err error) (r interface{}) {
	if err != nil {
		panic(err)
	}
	return result
}

func main() {
	masterPool = simpleredis.NewConnectionPoolHost("redis-master:6379")
	defer masterPool.Close()
	replicaPool = simpleredis.NewConnectionPoolHost("redis-slave:6379")
	defer replicaPool.Close()

	r := mux.NewRouter()
	r.Path("/lrange/{key}").Methods("GET").HandlerFunc(ListRangeHandler)
	r.Path("/rpush/{key}/{value}").Methods("GET").HandlerFunc(ListPushHandler)
	r.Path("/info").Methods("GET").HandlerFunc(InfoHandler)
	r.Path("/env").Methods("GET").HandlerFunc(EnvHandler)

	n := negroni.Classic()
	n.UseHandler(r)
	n.Run(":3000")
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/absolute8511/go-nsq"
	"github.com/absolute8511/nsq/internal/app"
	"github.com/absolute8511/nsq/internal/version"
)

var (
	showVersion = flag.Bool("version", false, "print version string")

	topic         = flag.String("topic", "", "NSQ topic")
	partition     = flag.Int("partition", -1, "NSQ topic partition")
	channel       = flag.String("channel", "", "NSQ channel")
	maxInFlight   = flag.Int("max-in-flight", 200, "max number of messages to allow in flight")
	totalMessages = flag.Int("n", 0, "total messages to show (will wait if starved)")

	lookupdHTTPAddrs = app.StringArray{}
)

func init() {
	flag.Var(&lookupdHTTPAddrs, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
}

type TailHandler struct {
	totalMessages int
	messagesShown int
}

func (th *TailHandler) HandleMessage(m *nsq.Message) error {
	th.messagesShown++
	_, err := os.Stdout.Write(m.Body)
	if err != nil {
		log.Fatalf("ERROR: failed to write to os.Stdout - %s", err)
	}
	_, err = os.Stdout.WriteString("\n")
	if err != nil {
		log.Fatalf("ERROR: failed to write to os.Stdout - %s", err)
	}
	if th.totalMessages > 0 && th.messagesShown >= th.totalMessages {
		os.Exit(0)
	}
	return nil
}

func main() {
	cfg := nsq.NewConfig()
	// TODO: remove, deprecated
	flag.Var(&nsq.ConfigFlag{cfg}, "reader-opt", "(deprecated) use --consumer-opt")
	flag.Var(&nsq.ConfigFlag{cfg}, "consumer-opt", "option to passthrough to nsq.Consumer (may be given multiple times, http://godoc.org/github.com/absolute8511/go-nsq#Config)")

	flag.Parse()

	if *showVersion {
		fmt.Printf("nsq_tail v%s\n", version.Binary)
		return
	}

	if *channel == "" {
		rand.Seed(time.Now().UnixNano())
		*channel = fmt.Sprintf("tail%06d#ephemeral", rand.Int()%999999)
	}

	if *topic == "" {
		log.Fatal("--topic is required")
	}

	if len(lookupdHTTPAddrs) == 0 {
		log.Fatal("--nsqd-tcp-address or --lookupd-http-address required")
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// Don't ask for more messages than we want
	if *totalMessages > 0 && *totalMessages < *maxInFlight {
		*maxInFlight = *totalMessages
	}

	cfg.UserAgent = fmt.Sprintf("nsq_tail/%s go-nsq/%s", version.Binary, nsq.VERSION)
	cfg.MaxInFlight = *maxInFlight

	consumer, err := nsq.NewPartitionConsumer(*topic, *partition, *channel, cfg)
	if err != nil {
		log.Fatal(err)
	}

	consumer.AddHandler(&TailHandler{totalMessages: *totalMessages})

	err = consumer.ConnectToNSQLookupds(lookupdHTTPAddrs)
	if err != nil {
		log.Fatal(err)
	}

	for {
		select {
		case <-consumer.StopChan:
			return
		case <-sigChan:
			consumer.Stop()
		}
	}
}
<commit_msg>nsq tail support ordered topic<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/absolute8511/go-nsq"
	"github.com/absolute8511/nsq/internal/app"
	"github.com/absolute8511/nsq/internal/version"
)

var (
	showVersion = flag.Bool("version", false, "print version string")

	topic         = flag.String("topic", "", "NSQ topic")
	partition     = flag.Int("partition", -1, "NSQ topic partition")
	channel       = flag.String("channel", "", "NSQ channel")
	ordered       = flag.Bool("ordered", false, "consume in ordered way")
	maxInFlight   = flag.Int("max-in-flight", 200, "max number of messages to allow in flight")
	totalMessages = flag.Int("n", 0, "total messages to show (will wait if starved)")

	lookupdHTTPAddrs = app.StringArray{}
)

func init() {
	flag.Var(&lookupdHTTPAddrs, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
}

type TailHandler struct {
	totalMessages int
	messagesShown int
}

func (th *TailHandler) HandleMessage(m *nsq.Message) error {
	th.messagesShown++
	_, err := os.Stdout.Write(m.Body)
	if err != nil {
		log.Fatalf("ERROR: failed to write to os.Stdout - %s", err)
	}
	_, err = os.Stdout.WriteString("\n")
	if err != nil {
		log.Fatalf("ERROR: failed to write to os.Stdout - %s", err)
	}
	if th.totalMessages > 0 && th.messagesShown >= th.totalMessages {
		os.Exit(0)
	}
	return nil
}

func main() {
	cfg := nsq.NewConfig()
	// TODO: remove, deprecated
	flag.Var(&nsq.ConfigFlag{cfg}, "reader-opt", "(deprecated) use --consumer-opt")
	flag.Var(&nsq.ConfigFlag{cfg}, "consumer-opt", "option to passthrough to nsq.Consumer (may be given multiple times, http://godoc.org/github.com/absolute8511/go-nsq#Config)")

	flag.Parse()

	if *showVersion {
		fmt.Printf("nsq_tail v%s\n", version.Binary)
		return
	}

	if *channel == "" {
		rand.Seed(time.Now().UnixNano())
		*channel = fmt.Sprintf("tail%06d#ephemeral", rand.Int()%999999)
	}

	if *topic == "" {
		log.Fatal("--topic is required")
	}

	if len(lookupdHTTPAddrs) == 0 {
		log.Fatal("--nsqd-tcp-address or --lookupd-http-address required")
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// Don't ask for more messages than we want
	if *totalMessages > 0 && *totalMessages < *maxInFlight {
		*maxInFlight = *totalMessages
	}

	cfg.UserAgent = fmt.Sprintf("nsq_tail/%s go-nsq/%s", version.Binary, nsq.VERSION)
	cfg.MaxInFlight = *maxInFlight
	if *ordered {
		cfg.EnableOrdered = true
	}

	consumer, err := nsq.NewPartitionConsumer(*topic, *partition, *channel, cfg)
	if err != nil {
		log.Fatal(err)
	}

	consumer.AddHandler(&TailHandler{totalMessages: *totalMessages})

	err = consumer.ConnectToNSQLookupds(lookupdHTTPAddrs)
	if err != nil {
		log.Fatal(err)
	}

	for {
		select {
		case <-consumer.StopChan:
			return
		case <-sigChan:
			consumer.Stop()
		}
	}
}
<|endoftext|>


<commit_before>package mastodon

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestAccountUpdate(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, `{"Username": "zzz"}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	a, err := client.AccountUpdate(context.Background(), &Profile{
		DisplayName: String("display_name"),
		Note:        String("note"),
		Avatar:      "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAUoAAADrCAYAAAA...",
		Header:      "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAUoAAADrCAYAAAA...",
	})
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if a.Username != "zzz" {
		t.Fatalf("want %q but %q", "zzz", a.Username)
	}
}

func TestGetBlocks(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, `[{"Username": "foo"}, {"Username": "bar"}]`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	bl, err := client.GetBlocks(context.Background())
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if len(bl) != 2 {
		t.Fatalf("result should be two: %d", len(bl))
	}
	if bl[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", bl[0].Username)
	}
	if bl[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", bl[0].Username)
	}
}

func TestAccountFollow(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/accounts/1234567/follow" {
			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
			return
		}
		fmt.Fprintln(w, `{"id":1234567,"following":true}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	rel, err := client.AccountFollow(context.Background(), 123)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	rel, err = client.AccountFollow(context.Background(), 1234567)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if rel.ID != 1234567 {
		t.Fatalf("want %d but %d", 1234567, rel.ID)
	}
	if !rel.Following {
		t.Fatalf("want %t but %t", true, rel.Following)
	}
}

func TestAccountUnfollow(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/accounts/1234567/unfollow" {
			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
			return
		}
		fmt.Fprintln(w, `{"id":1234567,"following":false}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	rel, err := client.AccountUnfollow(context.Background(), 123)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	rel, err = client.AccountUnfollow(context.Background(), 1234567)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if rel.ID != 1234567 {
		t.Fatalf("want %d but %d", 1234567, rel.ID)
	}
	if rel.Following {
		t.Fatalf("want %t but %t", false, rel.Following)
	}
}

func TestGetFollowRequests(t *testing.T) {
	canErr := true
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if canErr {
			canErr = false
			http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			return
		}
		fmt.Fprintln(w, `[{"Username": "foo"}, {"Username": "bar"}]`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	_, err := client.GetFollowRequests(context.Background())
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	fReqs, err := client.GetFollowRequests(context.Background())
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if len(fReqs) != 2 {
		t.Fatalf("result should be two: %d", len(fReqs))
	}
	if fReqs[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", fReqs[0].Username)
	}
	if fReqs[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", fReqs[0].Username)
	}
}

func TestFollowRequestAuthorize(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/follow_requests/1234567/authorize" {
			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
		}
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	err := client.FollowRequestAuthorize(context.Background(), 123)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	err = client.FollowRequestAuthorize(context.Background(), 1234567)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
}

func TestFollowRequestReject(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/follow_requests/1234567/reject" {
			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
		}
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	err := client.FollowRequestReject(context.Background(), 123)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	err = client.FollowRequestReject(context.Background(), 1234567)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
}

func TestGetMutes(t *testing.T) {
	canErr := true
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if canErr {
			canErr = false
			http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			return
		}
		fmt.Fprintln(w, `[{"Username": "foo"}, {"Username": "bar"}]`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	_, err := client.GetMutes(context.Background())
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	mutes, err := client.GetMutes(context.Background())
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if len(mutes) != 2 {
		t.Fatalf("result should be two: %d", len(mutes))
	}
	if mutes[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", mutes[0].Username)
	}
	if mutes[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", mutes[1].Username)
	}
}
<commit_msg>Fix test error message<commit_after>package mastodon

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestAccountUpdate(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, `{"Username": "zzz"}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	a, err := client.AccountUpdate(context.Background(), &Profile{
		DisplayName: String("display_name"),
		Note:        String("note"),
		Avatar:      "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAUoAAADrCAYAAAA...",
		Header:      "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAUoAAADrCAYAAAA...",
	})
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if a.Username != "zzz" {
		t.Fatalf("want %q but %q", "zzz", a.Username)
	}
}

func TestGetBlocks(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, `[{"Username": "foo"}, {"Username": "bar"}]`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	bl, err := client.GetBlocks(context.Background())
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if len(bl) != 2 {
		t.Fatalf("result should be two: %d", len(bl))
	}
	if bl[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", bl[0].Username)
	}
	if bl[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", bl[0].Username)
	}
}

func TestAccountFollow(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/accounts/1234567/follow" {
			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
			return
		}
		fmt.Fprintln(w, `{"id":1234567,"following":true}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	rel, err := client.AccountFollow(context.Background(), 123)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	rel, err = client.AccountFollow(context.Background(), 1234567)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if rel.ID != 1234567 {
		t.Fatalf("want %d but %d", 1234567, rel.ID)
	}
	if !rel.Following {
		t.Fatalf("want %t but %t", true, rel.Following)
	}
}

func TestAccountUnfollow(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/accounts/1234567/unfollow" {
			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
			return
		}
		fmt.Fprintln(w, `{"id":1234567,"following":false}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	rel, err := client.AccountUnfollow(context.Background(), 123)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	rel, err = client.AccountUnfollow(context.Background(), 1234567)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if rel.ID != 1234567 {
		t.Fatalf("want %d but %d", 1234567, rel.ID)
	}
	if rel.Following {
		t.Fatalf("want %t but %t", false, rel.Following)
	}
}

func TestGetFollowRequests(t *testing.T) {
	canErr := true
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if canErr {
			canErr = false
			http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			return
		}
		fmt.Fprintln(w, `[{"Username": "foo"}, {"Username": "bar"}]`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	_, err := client.GetFollowRequests(context.Background())
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	fReqs, err := client.GetFollowRequests(context.Background())
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if len(fReqs) != 2 {
		t.Fatalf("result should be two: %d", len(fReqs))
	}
	if fReqs[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", fReqs[0].Username)
	}
	if fReqs[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", fReqs[1].Username)
	}
}

func TestFollowRequestAuthorize(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/follow_requests/1234567/authorize" {
			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
		}
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	err := client.FollowRequestAuthorize(context.Background(), 123)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	err = client.FollowRequestAuthorize(context.Background(), 1234567)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
}

func TestFollowRequestReject(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/follow_requests/1234567/reject" {
			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
		}
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	err := client.FollowRequestReject(context.Background(), 123)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	err = client.FollowRequestReject(context.Background(), 1234567)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
}

func TestGetMutes(t *testing.T) {
	canErr := true
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if canErr {
			canErr = false
			http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			return
		}
		fmt.Fprintln(w, `[{"Username": "foo"}, {"Username": "bar"}]`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	_, err := client.GetMutes(context.Background())
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	mutes, err := client.GetMutes(context.Background())
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if len(mutes) != 2 {
		t.Fatalf("result should be two: %d", len(mutes))
	}
	if mutes[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", mutes[0].Username)
	}
	if mutes[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", mutes[1].Username)
	}
}
<|endoftext|>


<commit_before>package action

import (
	"bytes"
	"fmt"
	"html/template"
	"io"
	"os"
	"os/exec"
	"path"
	"regexp"
	"time"

	shellquote "github.com/kballard/go-shellquote"
	homedir "github.com/mitchellh/go-homedir"
	toml "github.com/pelletier/go-toml"
	"github.com/taylorskalyo/stno/datastore"
	cli "gopkg.in/urfave/cli.v1"
)

const defaultTemplate string = `title = ""
datetime = {{.DateTime}}
notes = ""`

type templateData struct {
	DateTime string
}

// Add a new notebook entry.
func Add(c *cli.Context) error {
	// Create temporary file
	tmpfile, err := datastore.TempFile("", "stno", ".toml")
	if err != nil {
		return err
	}
	defer os.Remove(tmpfile.Name())

	// Write template to file
	t, err := template.New("default").Parse(defaultTemplate)
	if err != nil {
		return err
	}
	err = t.Execute(tmpfile, newTemplateData())
	if err != nil {
		return err
	}
	if err := tmpfile.Close(); err != nil {
		return err
	}

	// Open file in editor
	openEditor(tmpfile.Name())

	// Lint file
	rc, err := os.Open(tmpfile.Name())
	if err != nil {
		return err
	}
	defer rc.Close()
	tree, err := toml.LoadReader(rc)
	if err != nil {
		return err
	}
	rc.Seek(0, 0)

	// Copy contents from temp file to entry file
	dir, err := stnoDir(c.GlobalString("notebook"))
	if err != nil {
		return err
	}
	ds, err := datastore.CreateFileStore(dir)
	if err != nil {
		return err
	}
	var buf bytes.Buffer
	datetime, ok := tree.Get("datetime").(time.Time)
	if ok {
		buf.WriteString(fmt.Sprintf("%d", datetime.Unix()))
		buf.WriteString("-")
	}
	title, ok := tree.Get("title").(string)
	if ok {
		r := regexp.MustCompile("[^A-Za-z0-9_-]+")
		buf.WriteString(r.ReplaceAllString(title, "-"))
		buf.WriteString("-")
	}
	_, wc, err := ds.NewUniqueWriteCloser(buf.String())
	if err != nil {
		return err
	}
	defer wc.Close()
	_, err = io.Copy(wc, rc)
	if err != nil {
		return err
	}

	return nil
}

// Query a notebook for a list of entries.
func Query(c *cli.Context) error {
	dir, err := stnoDir(c.GlobalString("notebook"))
	if err != nil {
		return err
	}
	ds, err := datastore.CreateFileStore(dir)
	if err != nil {
		return err
	}

	uuids, err := ds.List()
	if err != nil {
		return err
	}
	for i, uuid := range uuids {
		if i != 0 {
			fmt.Println()
		}
		fmt.Println(uuid)
		rc, err := ds.NewReadCloser(uuid)
		if err != nil {
			return err
		}
		io.Copy(os.Stdout, rc)
		rc.Close()
	}

	return nil
}

func openEditor(path string) error {
	editor := os.Getenv("EDITOR")
	if editor == "" {
		editor = "editor"
	}

	args, err := shellquote.Split(editor)
	if err != nil {
		return err
	}

	editor = args[0]
	args = append(args[1:], path)
	cmd := exec.Command(editor, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}

	return nil
}

// newTemplateData holds values that can be substituted into a template.
func newTemplateData() templateData {
	return templateData{
		DateTime: time.Now().Format(time.RFC3339),
	}
}

func stnoDir(name string) (string, error) {
	return homedir.Expand(path.Join("~/.stno", name))
}
<commit_msg>Abort addition if no changes were made<commit_after>package action

import (
	"bytes"
	"fmt"
	"html/template"
	"io"
	"os"
	"os/exec"
	"path"
	"regexp"
	"time"

	shellquote "github.com/kballard/go-shellquote"
	homedir "github.com/mitchellh/go-homedir"
	toml "github.com/pelletier/go-toml"
	"github.com/taylorskalyo/stno/datastore"
	cli "gopkg.in/urfave/cli.v1"
)

const defaultTemplate string = `title = ""
datetime = {{.DateTime}}
notes = ""`

type templateData struct {
	DateTime string
}

// Add a new notebook entry.
func Add(c *cli.Context) error {
	// Create temporary file
	tmpfile, err := datastore.TempFile("", "stno", ".toml")
	if err != nil {
		return err
	}
	defer os.Remove(tmpfile.Name())

	// Write template to file
	t, err := template.New("default").Parse(defaultTemplate)
	if err != nil {
		return err
	}
	err = t.Execute(tmpfile, newTemplateData())
	if err != nil {
		return err
	}
	fi, err := tmpfile.Stat()
	if err != nil {
		return err
	}
	oldModTime := fi.ModTime()
	if err := tmpfile.Close(); err != nil {
		return err
	}

	// Open file in editor
	openEditor(tmpfile.Name())

	rc, err := os.Open(tmpfile.Name())
	if err != nil {
		return err
	}
	defer rc.Close()

	// Return if there were no changes
	fi, err = rc.Stat()
	if err != nil {
		return err
	}
	if oldModTime == fi.ModTime() {
		fmt.Println("Aborting due to empty entry.")
		return nil
	}

	// Lint file
	tree, err := toml.LoadReader(rc)
	if err != nil {
		return err
	}
	rc.Seek(0, 0)

	// Copy contents from temp file to entry file
	dir, err := stnoDir(c.GlobalString("notebook"))
	if err != nil {
		return err
	}
	ds, err := datastore.CreateFileStore(dir)
	if err != nil {
		return err
	}
	var buf bytes.Buffer
	datetime, ok := tree.Get("datetime").(time.Time)
	if ok {
		buf.WriteString(fmt.Sprintf("%d", datetime.Unix()))
		buf.WriteString("-")
	}
	title, ok := tree.Get("title").(string)
	if ok {
		r := regexp.MustCompile("[^A-Za-z0-9_-]+")
		buf.WriteString(r.ReplaceAllString(title, "-"))
		buf.WriteString("-")
	}
	_, wc, err := ds.NewUniqueWriteCloser(buf.String())
	if err != nil {
		return err
	}
	defer wc.Close()
	_, err = io.Copy(wc, rc)
	if err != nil {
		return err
	}

	return nil
}

// Query a notebook for a list of entries.
func Query(c *cli.Context) error {
	dir, err := stnoDir(c.GlobalString("notebook"))
	if err != nil {
		return err
	}
	ds, err := datastore.CreateFileStore(dir)
	if err != nil {
		return err
	}

	uuids, err := ds.List()
	if err != nil {
		return err
	}
	for i, uuid := range uuids {
		if i != 0 {
			fmt.Println()
		}
		fmt.Println(uuid)
		rc, err := ds.NewReadCloser(uuid)
		if err != nil {
			return err
		}
		io.Copy(os.Stdout, rc)
		rc.Close()
	}

	return nil
}

func openEditor(path string) error {
	editor := os.Getenv("EDITOR")
	if editor == "" {
		editor = "editor"
	}

	args, err := shellquote.Split(editor)
	if err != nil {
		return err
	}

	editor = args[0]
	args = append(args[1:], path)
	cmd := exec.Command(editor, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}

	return nil
}

// newTemplateData holds values that can be substituted into a template.
func newTemplateData() templateData {
	return templateData{
		DateTime: time.Now().Format(time.RFC3339),
	}
}

func stnoDir(name string) (string, error) {
	return homedir.Expand(path.Join("~/.stno", name))
}
<|endoftext|>


<commit_before>// Package violetear - HTTP router
//
// Basic example:
//
//  package main
//
//  import (
//     "fmt"
//     "github.com/nbari/violetear"
//     "log"
//     "net/http"
//  )
//
//  func catchAll(w http.ResponseWriter, r *http.Request) {
//      fmt.Fprintf(w, r.URL.Path[1:])
//  }
//
//  func helloWorld(w http.ResponseWriter, r *http.Request) {
//      fmt.Fprintf(w, r.URL.Path[1:])
//  }
//
//  func handleUUID(w http.ResponseWriter, r *http.Request) {
//      fmt.Fprintf(w, r.URL.Path[1:])
//  }
//
//  func main() {
//      router := violetear.New()
//      router.LogRequests = true
//      router.RequestID = "REQUEST_LOG_ID"
//
//      router.AddRegex(":uuid", `[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}`)
//
//      router.HandleFunc("*", catchAll)
//      router.HandleFunc("/hello/", helloWorld, "GET,HEAD")
//      router.HandleFunc("/root/:uuid/item", handleUUID, "POST,PUT")
//
//      log.Fatal(http.ListenAndServe(":8080", router))
//  }
//
package violetear

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"strings"
	"time"
)

type key int

// ParamsKey used for the context
const (
	ParamsKey        key = 0
	versionSeparator     = "application/vnd."
)

// Params string/interface map used with context
type Params map[string]interface{}

// Router struct
type Router struct {
	// Routes to be matched
	routes *Trie

	// dynamicRoutes map of dynamic routes and regular expresions
	dynamicRoutes dynamicSet

	// LogRequests yes or no
	LogRequests bool

	// NotFoundHandler configurable http.Handler which is called when no matching
	// route is found. If it is not set, http.NotFound is used.
	NotFoundHandler http.Handler

	// NotAllowedHandler configurable http.Handler which is called when method not allowed.
	NotAllowedHandler http.Handler

	// PanicHandler function to handle panics.
	PanicHandler http.HandlerFunc

	// RequestID name of the header to use or create.
	RequestID string

	// Verbose
	Verbose bool
}

var splitPathRx = regexp.MustCompile(`[^/ ]+`)

// New returns a new initialized router.
func New() *Router {
	return &Router{
		routes:        NewTrie(),
		dynamicRoutes: make(dynamicSet),
		Verbose:       true,
	}
}

// Handle registers the handler for the given pattern (path, http.Handler, methods).
func (v *Router) Handle(path string, handler http.Handler, httpMethods ...string) error {
	var version string
	if i := strings.Index(path, "#"); i != -1 {
		version = path[i+1:]
		path = path[:i]
	}
	pathParts := v.splitPath(path)

	// search for dynamic routes
	for _, p := range pathParts {
		if strings.HasPrefix(p, ":") {
			if _, ok := v.dynamicRoutes[p]; !ok {
				return fmt.Errorf("[%s] not found, need to add it using AddRegex(\"%s\", `your regex`)", p, p)
			}
		}
	}

	// if no methods, accept ALL
	methods := "ALL"
	if len(httpMethods) > 0 {
		methods = httpMethods[0]
	}

	if v.Verbose {
		log.Printf("Adding path: %s [%s]", path, methods)
	}

	if err := v.routes.Set(pathParts, handler, methods, version); err != nil {
		return err
	}
	return nil
}

// HandleFunc add a route to the router (path, http.HandlerFunc, methods)
func (v *Router) HandleFunc(path string, handler http.HandlerFunc, httpMethods ...string) error {
	return v.Handle(path, handler, httpMethods...)
}

// AddRegex adds a ":named" regular expression to the dynamicRoutes
func (v *Router) AddRegex(name, regex string) error {
	return v.dynamicRoutes.Set(name, regex)
}

// MethodNotAllowed default handler for 405
func (v *Router) MethodNotAllowed() http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w,
			http.StatusText(http.StatusMethodNotAllowed),
			http.StatusMethodNotAllowed,
		)
	})
}

// ServerHTTP dispatches the handler registered in the matched path
func (v *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	start := time.Now()
	lw := NewResponseWriter(w)
	params := make(Params)

	// panic handler
	defer func() {
		if err := recover(); err != nil {
			if v.PanicHandler != nil {
				v.PanicHandler(w, r)
			} else {
				http.Error(w, http.StatusText(500), http.StatusInternalServerError)
			}
		}
	}()

	// fill the params map
	setParam := func(k, v string) {
		if param, ok := params[k]; ok {
			switch param.(type) {
			case string:
				param = []string{param.(string), v}
			case []string:
				param = append(param.([]string), v)
			}
			params[k] = param
		} else {
			params[k] = v
		}
	}

	// set version based on the value of "Accept: application/vnd.*"
	version := r.Header.Get("Accept")
	if i := strings.LastIndex(version, versionSeparator); i != -1 {
		version = version[len(versionSeparator)+i:]
	} else {
		version = ""
	}

	// _ path never empty, defaults to ("/")
	node, path, leaf, _ := v.routes.Get(v.splitPath(r.URL.Path), version)

	// checkMethod check if method is allowed or not
	checkMethod := func(node *Trie, method string) http.Handler {
		if h, ok := node.Handler[method]; ok {
			return h
		}
		if h, ok := node.Handler["ALL"]; ok {
			return h
		}
		if v.NotAllowedHandler != nil {
			return v.NotAllowedHandler
		}
		return v.MethodNotAllowed()
	}

	var match func(node *Trie, path []string, leaf bool) http.Handler

	// match find a handler for the request
	match = func(node *Trie, path []string, leaf bool) http.Handler {
		catchall := false
		if len(node.Handler) > 0 && leaf {
			return checkMethod(node, r.Method)
		} else if node.HasRegex {
			for _, n := range node.Node {
				if strings.HasPrefix(n.path, ":") {
					rx := v.dynamicRoutes[n.path]
					if rx.MatchString(path[0]) {
						// add param to context
						setParam(n.path, path[0])
						path[0] = n.path
						node, path, leaf, _ := node.Get(path, version)
						return match(node, path, leaf)
					}
				}
			}
			if node.HasCatchall {
				catchall = true
			}
		} else if node.HasCatchall {
			catchall = true
		}
		if catchall {
			for _, n := range node.Node {
				if n.path == "*" {
					// add "*" to context
					setParam("*", path[0])
					return checkMethod(n, r.Method)
				}
			}
		}
		// NotFound
		if v.NotFoundHandler != nil {
			return v.NotFoundHandler
		}
		return http.NotFoundHandler()
	}

	// Request-ID
	if v.RequestID != "" {
		if rid := r.Header.Get(v.RequestID); rid != "" {
			lw.Header().Set(v.RequestID, rid)
		}
	}

	//h http.Handler
	h := match(node, path, leaf)

	// dispatch request
	h.ServeHTTP(lw, r.WithContext(context.WithValue(r.Context(), ParamsKey, params)))

	if v.LogRequests {
		log.Printf("%s [%s] %d %d %v %s",
			r.RemoteAddr,
			r.URL,
			lw.Status(),
			lw.Size(),
			time.Since(start),
			lw.Header().Get(v.RequestID))
	}
	return
}

// splitPath returns an slice of the path
func (v *Router) splitPath(p string) []string {
	pathParts := splitPathRx.FindAllString(p, -1)

	// root (empty slice)
	if len(pathParts) == 0 {
		pathParts = append(pathParts, "/")
	}

	return pathParts
}
<commit_msg>versionHeader instead of versionSeparator 	modified:   violetear.go<commit_after>// Package violetear - HTTP router
//
// Basic example:
//
//  package main
//
//  import (
//     "fmt"
//     "github.com/nbari/violetear"
//     "log"
//     "net/http"
//  )
//
//  func catchAll(w http.ResponseWriter, r *http.Request) {
//      fmt.Fprintf(w, r.URL.Path[1:])
//  }
//
//  func helloWorld(w http.ResponseWriter, r *http.Request) {
//      fmt.Fprintf(w, r.URL.Path[1:])
//  }
//
//  func handleUUID(w http.ResponseWriter, r *http.Request) {
//      fmt.Fprintf(w, r.URL.Path[1:])
//  }
//
//  func main() {
//      router := violetear.New()
//      router.LogRequests = true
//      router.RequestID = "REQUEST_LOG_ID"
//
//      router.AddRegex(":uuid", `[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}`)
//
//      router.HandleFunc("*", catchAll)
//      router.HandleFunc("/hello/", helloWorld, "GET,HEAD")
//      router.HandleFunc("/root/:uuid/item", handleUUID, "POST,PUT")
//
//      log.Fatal(http.ListenAndServe(":8080", router))
//  }
//
package violetear

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"strings"
	"time"
)

type key int

// ParamsKey used for the context
const (
	ParamsKey     key = 0
	versionHeader     = "application/vnd."
)

// Params string/interface map used with context
type Params map[string]interface{}

// Router struct
type Router struct {
	// Routes to be matched
	routes *Trie

	// dynamicRoutes map of dynamic routes and regular expresions
	dynamicRoutes dynamicSet

	// LogRequests yes or no
	LogRequests bool

	// NotFoundHandler configurable http.Handler which is called when no matching
	// route is found. If it is not set, http.NotFound is used.
	NotFoundHandler http.Handler

	// NotAllowedHandler configurable http.Handler which is called when method not allowed.
	NotAllowedHandler http.Handler

	// PanicHandler function to handle panics.
	PanicHandler http.HandlerFunc

	// RequestID name of the header to use or create.
	RequestID string

	// Verbose
	Verbose bool
}

var splitPathRx = regexp.MustCompile(`[^/ ]+`)

// New returns a new initialized router.
func New() *Router {
	return &Router{
		routes:        NewTrie(),
		dynamicRoutes: make(dynamicSet),
		Verbose:       true,
	}
}

// Handle registers the handler for the given pattern (path, http.Handler, methods).
func (v *Router) Handle(path string, handler http.Handler, httpMethods ...string) error {
	var version string
	if i := strings.Index(path, "#"); i != -1 {
		version = path[i+1:]
		path = path[:i]
	}
	pathParts := v.splitPath(path)

	// search for dynamic routes
	for _, p := range pathParts {
		if strings.HasPrefix(p, ":") {
			if _, ok := v.dynamicRoutes[p]; !ok {
				return fmt.Errorf("[%s] not found, need to add it using AddRegex(\"%s\", `your regex`)", p, p)
			}
		}
	}

	// if no methods, accept ALL
	methods := "ALL"
	if len(httpMethods) > 0 {
		methods = httpMethods[0]
	}

	if v.Verbose {
		log.Printf("Adding path: %s [%s]", path, methods)
	}

	if err := v.routes.Set(pathParts, handler, methods, version); err != nil {
		return err
	}
	return nil
}

// HandleFunc add a route to the router (path, http.HandlerFunc, methods)
func (v *Router) HandleFunc(path string, handler http.HandlerFunc, httpMethods ...string) error {
	return v.Handle(path, handler, httpMethods...)
}

// AddRegex adds a ":named" regular expression to the dynamicRoutes
func (v *Router) AddRegex(name, regex string) error {
	return v.dynamicRoutes.Set(name, regex)
}

// MethodNotAllowed default handler for 405
func (v *Router) MethodNotAllowed() http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w,
			http.StatusText(http.StatusMethodNotAllowed),
			http.StatusMethodNotAllowed,
		)
	})
}

// ServerHTTP dispatches the handler registered in the matched path
func (v *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	start := time.Now()
	lw := NewResponseWriter(w)
	params := make(Params)

	// panic handler
	defer func() {
		if err := recover(); err != nil {
			if v.PanicHandler != nil {
				v.PanicHandler(w, r)
			} else {
				http.Error(w, http.StatusText(500), http.StatusInternalServerError)
			}
		}
	}()

	// fill the params map
	setParam := func(k, v string) {
		if param, ok := params[k]; ok {
			switch param.(type) {
			case string:
				param = []string{param.(string), v}
			case []string:
				param = append(param.([]string), v)
			}
			params[k] = param
		} else {
			params[k] = v
		}
	}

	// set version based on the value of "Accept: application/vnd.*"
	version := r.Header.Get("Accept")
	if i := strings.LastIndex(version, versionHeader); i != -1 {
		version = version[len(versionHeader)+i:]
	} else {
		version = ""
	}

	// _ path never empty, defaults to ("/")
	node, path, leaf, _ := v.routes.Get(v.splitPath(r.URL.Path), version)

	// checkMethod check if method is allowed or not
	checkMethod := func(node *Trie, method string) http.Handler {
		if h, ok := node.Handler[method]; ok {
			return h
		}
		if h, ok := node.Handler["ALL"]; ok {
			return h
		}
		if v.NotAllowedHandler != nil {
			return v.NotAllowedHandler
		}
		return v.MethodNotAllowed()
	}

	var match func(node *Trie, path []string, leaf bool) http.Handler

	// match find a handler for the request
	match = func(node *Trie, path []string, leaf bool) http.Handler {
		catchall := false
		if len(node.Handler) > 0 && leaf {
			return checkMethod(node, r.Method)
		} else if node.HasRegex {
			for _, n := range node.Node {
				if strings.HasPrefix(n.path, ":") {
					rx := v.dynamicRoutes[n.path]
					if rx.MatchString(path[0]) {
						// add param to context
						setParam(n.path, path[0])
						path[0] = n.path
						node, path, leaf, _ := node.Get(path, version)
						return match(node, path, leaf)
					}
				}
			}
			if node.HasCatchall {
				catchall = true
			}
		} else if node.HasCatchall {
			catchall = true
		}
		if catchall {
			for _, n := range node.Node {
				if n.path == "*" {
					// add "*" to context
					setParam("*", path[0])
					return checkMethod(n, r.Method)
				}
			}
		}
		// NotFound
		if v.NotFoundHandler != nil {
			return v.NotFoundHandler
		}
		return http.NotFoundHandler()
	}

	// Request-ID
	if v.RequestID != "" {
		if rid := r.Header.Get(v.RequestID); rid != "" {
			lw.Header().Set(v.RequestID, rid)
		}
	}

	//h http.Handler
	h := match(node, path, leaf)

	// dispatch request
	h.ServeHTTP(lw, r.WithContext(context.WithValue(r.Context(), ParamsKey, params)))

	if v.LogRequests {
		log.Printf("%s [%s] %d %d %v %s",
			r.RemoteAddr,
			r.URL,
			lw.Status(),
			lw.Size(),
			time.Since(start),
			lw.Header().Get(v.RequestID))
	}
	return
}

// splitPath returns an slice of the path
func (v *Router) splitPath(p string) []string {
	pathParts := splitPathRx.FindAllString(p, -1)

	// root (empty slice)
	if len(pathParts) == 0 {
		pathParts = append(pathParts, "/")
	}

	return pathParts
}
<|endoftext|>