<commit_before>package main

import  (
        "fmt"
        "os"
        "../go-todotxt"
        "github.com/spf13/cobra"
        "os/user"
        "strings"
        "strconv"
        "github.com/rakyll/globalconf"
        "flag"
        "regexp"
        "os/exec"
        "io/ioutil"
)

func extendedLoader(filename string) (todotxt.TaskList, error) {
        usr, err := user.Current()
        if err != nil {
                return nil, err
        }

        filename = strings.Replace(filename, "~", usr.HomeDir, -1)
        tasks := todotxt.LoadTaskList(filename)

        return tasks, nil
}

type flagValue struct {
        str string
}

func (f *flagValue) String() string {
        return f.str
}

func (f *flagValue) Set(value string) error {
        f.str = value
        return nil
}

func newFlagValue(val string) *flagValue {
        return &flagValue{str: val}
}

func main() {

        conf, _ := globalconf.New("gotodo")

        var numtasks bool
        var sortby string
        var finished bool
        var prettyformat string
        var filename string

        var flagFilename = flag.String("file", "", "Location of the todo.txt file.")

        var flags = make(map[string]string)

        var cmdConfig = &cobra.Command{
            Use:   "config [key] [value]",
            Short: "Show and sets config values",
            Long:  `Config can be used to see and also set configuration variables.`,
            Run: func(cmd *cobra.Command, args []string) {
                    if len(args) == 0{
                            fmt.Printf("Available config variables:\n\n")
                            for k := range flags {
                                    fmt.Printf("%s\n", k)
                            }
                    }
                    if len(args) == 1 {
                            val, exists := flags[args[0]]
                            if exists {
                                    fmt.Printf("%s\n", val)
                            } else {
                                    // otherwise exit with non-zero status
                                    os.Exit(1)
                            }
                    }
                    if len(args) == 2 {
                            _, exists := flags[args[0]]
                            if exists {
                                    f := &flag.Flag{Name: args[0], Value: newFlagValue(args[1])}
                                    conf.Set("", f)
                            } else {
                                    os.Exit(1)
                            }
                    }
            },
        }

        var cmdList = &cobra.Command{
            Use:   "list [keyword]",
            Short: "Lists tasks that contain keyword, if any",
            Long:  `List is the most basic command that is used for listing tasks.
                    You can specify a keyword as well as other options.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if numtasks {
                    fmt.Println(tasks.Len())
                } else {
                    tasks.Sort(sortby)

                    var filteredTasks todotxt.TaskList
                    for _, task := range tasks {
                        if (!task.Finished() && !finished) ||
                           (task.Finished() && finished) {
                           if (len(args) > 0) {
                                   match, _ := regexp.MatchString(args[0], task.Text())
                                   if match {
                                           filteredTasks = append(filteredTasks, task)
                                   }
                           } else {
                                   filteredTasks = append(filteredTasks, task)
                           }
                        }
                    }

                    for _, task := range filteredTasks {
                            task.SetIdPaddingBy(tasks)
                            fmt.Println(task.PrettyPrint(prettyformat))
                    }
                }
            },
        }
        cmdList.Flags().BoolVarP(&numtasks, "num-tasks", "n", false,
                                 "Show the number of tasks")
        cmdList.Flags().BoolVarP(&finished, "finished", "f", false,
                                 "Show finished tasks")
        cmdList.Flags().StringVarP(&sortby, "sort", "s", "prio",
                                   "Sort tasks by parameter (prio|date|len|prio-rev|date-rev|len-rev|id|rand)")
        cmdList.Flags().StringVarP(&prettyformat, "pretty", "", "%i %p %t",
                                   "Pretty print tasks")

        var cmdAdd = &cobra.Command{
            Use:   "add [task]",
            Short: "Adds a task to the todo list.",
            Long:  `Adds a task to the todo list.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                task := strings.Join(args, " ")
                tasks.Add(task)

                tasks.Save(filename)
            },
        }

        var nofinishdate bool
        var cmdDone = &cobra.Command{
            Use:   "done [taskid]",
            Short: "Marks task as done.",
            Long:  `Marks task as done.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if len(args) < 1 {
                        fmt.Println("So what needs to be done?")
                        return
                }

                taskid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                err = tasks.Done(taskid, !nofinishdate)
                if err != nil {
                        fmt.Printf("There was an error %v\n", err)
                }

                tasks.Save(filename)
            },
        }
        cmdDone.Flags().BoolVarP(&nofinishdate, "no-finish-date", "D", false,
                                        "Do not mark finished tasks with date.")

        var cmdArchive = &cobra.Command{
            Use:   "archive [taskid]",
            Short: "Archives task.",
            Long:  `Archives task.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if len(args) < 1 {
                        fmt.Println("So what needs to be done?")
                        return
                }

                taskid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                fmt.Printf("Archiving task %v\n", taskid)

                tasks.Save(filename)
            },
        }

        var setprio string
        var settodo string
        var cmdSet = &cobra.Command{
            Use:   "set [taskid]",
            Short: "Sets given task's parameters.",
            Long:  `Sets given task's parameters.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if len(args) < 1 {
                        fmt.Println("So what do you want to edit?")
                        return
                }

                taskid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                if len(setprio) > 0 {
                        tasks[taskid].SetPriority(setprio[0])
                        tasks[taskid].RebuildRawTodo()
                }

                if len(settodo) > 0 {
                        tasks[taskid].SetTodo(settodo)
                        tasks[taskid].RebuildRawTodo()
                }

              //if len(edittodo) == 0 && len(editprio) == 0 {
              //        // FIXME: nost just vim
              //        out, err := exec.Command("vim", filename).Output()
              //}

                tasks.Save(filename)
            },
        }

        var cmdEdit = &cobra.Command{
            Use:   "edit [taskid]",
            Short: "Edit given task",
            Long:  `Edit given task with your prefered editor.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if len(args) < 1 {
                        fmt.Println("So what do you want to edit?")
                        return
                }

                taskid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                text := tasks[taskid].RawText()

                file, err := ioutil.TempFile(os.TempDir(), "gotodo")
                defer os.Remove(file.Name())

                if err != nil {
                        panic(err)
                }

                e := ioutil.WriteFile(file.Name(), []byte(text), 0644)
                if e != nil {
                        panic(e)
                }

                editor := os.Getenv("EDITOR")
                if len(editor) == 0 {
                        editor = "nano" //FIXME: saner default?
                }

                c := exec.Command(editor, file.Name())

                // nasty hack, see http://stackoverflow.com/a/12089980
                c.Stdin = os.Stdin
                c.Stdout = os.Stdout
                c.Stderr = os.Stderr

                er := c.Run()

                if er != nil {
                        fmt.Println(er.Error())
                        panic(er)
                }

                dat, err := ioutil.ReadFile(file.Name())
                if err != nil {
                        panic(err)
                }

                // FIXME: take just the first line
                tasks[taskid] = todotxt.ParseTask(strings.TrimSpace(string(dat)), taskid)

                tasks.Save(filename)
            },
        }

        cmdSet.PersistentFlags().StringVarP(&setprio, "priority", "p", "",
                                     "Sets task's priority.")
        cmdSet.PersistentFlags().StringVarP(&settodo, "todo", "t", "",
                                     "Sets task's todo.")

        var GotodoCmd = &cobra.Command{
            Use:   "gotodo",
            Short: "Gotodo is a go implementation of todo.txt.",
            Long: `A small, fast and fun implementation of todo.txt`,
            Run: func(cmd *cobra.Command, args []string) {
                cmdList.Run(cmd, nil)
            },
        }

        GotodoCmd.PersistentFlags().StringVarP(&filename, "filename", "", "",
                                     "Load tasks from this file.")

        conf.ParseAll()

        // values here
        flags["file"] = *flagFilename

        // sadly, this is the best we can do right now
        if filename == "" {
                if *flagFilename == "" {
                        filename = "todo.txt"
                } else {
                        filename = *flagFilename
                }
        }

        GotodoCmd.AddCommand(cmdList)
        GotodoCmd.AddCommand(cmdAdd)
        GotodoCmd.AddCommand(cmdDone)
        GotodoCmd.AddCommand(cmdArchive)
        GotodoCmd.AddCommand(cmdSet)
        GotodoCmd.AddCommand(cmdEdit)
        GotodoCmd.AddCommand(cmdConfig)
        GotodoCmd.Execute()
}
<commit_msg>CLeanup<commit_after>package main

import  (
        "fmt"
        "os"
        "../go-todotxt"
        "github.com/spf13/cobra"
        "os/user"
        "strings"
        "strconv"
        "github.com/rakyll/globalconf"
        "flag"
        "regexp"
        "os/exec"
        "io/ioutil"
)

func extendedLoader(filename string) (todotxt.TaskList, error) {
        usr, err := user.Current()
        if err != nil {
                return nil, err
        }

        filename = strings.Replace(filename, "~", usr.HomeDir, -1)
        tasks := todotxt.LoadTaskList(filename)

        return tasks, nil
}

type flagValue struct {
        str string
}

func (f *flagValue) String() string {
        return f.str
}

func (f *flagValue) Set(value string) error {
        f.str = value
        return nil
}

func newFlagValue(val string) *flagValue {
        return &flagValue{str: val}
}

func main() {

        conf, _ := globalconf.New("gotodo")

        var numtasks bool
        var sortby string
        var finished bool
        var prettyformat string
        var filename string

        var flagFilename = flag.String("file", "", "Location of the todo.txt file.")

        var flags = make(map[string]string)

        var cmdConfig = &cobra.Command{
            Use:   "config [key] [value]",
            Short: "Show and sets config values",
            Long:  `Config can be used to see and also set configuration variables.`,
            Run: func(cmd *cobra.Command, args []string) {
                    if len(args) == 0{
                            fmt.Printf("Available config variables:\n\n")
                            for k := range flags {
                                    fmt.Printf("%s\n", k)
                            }
                    }
                    if len(args) == 1 {
                            val, exists := flags[args[0]]
                            if exists {
                                    fmt.Printf("%s\n", val)
                            } else {
                                    // otherwise exit with non-zero status
                                    os.Exit(1)
                            }
                    }
                    if len(args) == 2 {
                            _, exists := flags[args[0]]
                            if exists {
                                    f := &flag.Flag{Name: args[0], Value: newFlagValue(args[1])}
                                    conf.Set("", f)
                            } else {
                                    os.Exit(1)
                            }
                    }
            },
        }

        var cmdList = &cobra.Command{
            Use:   "list [keyword]",
            Short: "Lists tasks that contain keyword, if any",
            Long:  `List is the most basic command that is used for listing tasks.
                    You can specify a keyword as well as other options.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if numtasks {
                    fmt.Println(tasks.Len())
                } else {
                    tasks.Sort(sortby)

                    var filteredTasks todotxt.TaskList
                    for _, task := range tasks {
                        if (!task.Finished() && !finished) ||
                           (task.Finished() && finished) {
                           if (len(args) > 0) {
                                   match, _ := regexp.MatchString(args[0], task.Text())
                                   if match {
                                           filteredTasks = append(filteredTasks, task)
                                   }
                           } else {
                                   filteredTasks = append(filteredTasks, task)
                           }
                        }
                    }

                    for _, task := range filteredTasks {
                            task.SetIdPaddingBy(tasks)
                            fmt.Println(task.PrettyPrint(prettyformat))
                    }
                }
            },
        }
        cmdList.Flags().BoolVarP(&numtasks, "num-tasks", "n", false,
                                 "Show the number of tasks")
        cmdList.Flags().BoolVarP(&finished, "finished", "f", false,
                                 "Show finished tasks")
        cmdList.Flags().StringVarP(&sortby, "sort", "s", "prio",
                                   "Sort tasks by parameter (prio|date|len|prio-rev|date-rev|len-rev|id|rand)")
        cmdList.Flags().StringVarP(&prettyformat, "pretty", "", "%i %p %t",
                                   "Pretty print tasks")

        var cmdAdd = &cobra.Command{
            Use:   "add [task]",
            Short: "Adds a task to the todo list.",
            Long:  `Adds a task to the todo list.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                task := strings.Join(args, " ")
                tasks.Add(task)

                tasks.Save(filename)
            },
        }

        var nofinishdate bool
        var cmdDone = &cobra.Command{
            Use:   "done [taskid]",
            Short: "Marks task as done.",
            Long:  `Marks task as done.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if len(args) < 1 {
                        fmt.Println("So what needs to be done?")
                        return
                }

                taskid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                err = tasks.Done(taskid, !nofinishdate)
                if err != nil {
                        fmt.Printf("There was an error %v\n", err)
                }

                tasks.Save(filename)
            },
        }
        cmdDone.Flags().BoolVarP(&nofinishdate, "no-finish-date", "D", false,
                                        "Do not mark finished tasks with date.")

        var cmdArchive = &cobra.Command{
            Use:   "archive [taskid]",
            Short: "Archives task.",
            Long:  `Archives task.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if len(args) < 1 {
                        fmt.Println("So what needs to be done?")
                        return
                }

                taskid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                fmt.Printf("Archiving task %v\n", taskid)

                tasks.Save(filename)
            },
        }

        var setprio string
        var settodo string
        var cmdSet = &cobra.Command{
            Use:   "set [taskid]",
            Short: "Sets given task's parameters.",
            Long:  `Sets given task's parameters.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if len(args) < 1 {
                        fmt.Println("So what do you want to edit?")
                        return
                }

                taskid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                if len(setprio) > 0 {
                        tasks[taskid].SetPriority(setprio[0])
                        tasks[taskid].RebuildRawTodo()
                }

                if len(settodo) > 0 {
                        tasks[taskid].SetTodo(settodo)
                        tasks[taskid].RebuildRawTodo()
                }

                tasks.Save(filename)
            },
        }

        var cmdEdit = &cobra.Command{
            Use:   "edit [taskid]",
            Short: "Edit given task",
            Long:  `Edit given task with your prefered editor.`,
            Run: func(cmd *cobra.Command, args []string) {
                tasks, err := extendedLoader(filename)
                if err != nil {
                        fmt.Println(err)
                        return
                }

                if len(args) < 1 {
                        fmt.Println("So what do you want to edit?")
                        return
                }

                taskid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                text := tasks[taskid].RawText()

                file, err := ioutil.TempFile(os.TempDir(), "gotodo")
                defer os.Remove(file.Name())

                if err != nil {
                        panic(err)
                }

                e := ioutil.WriteFile(file.Name(), []byte(text), 0644)
                if e != nil {
                        panic(e)
                }

                editor := os.Getenv("EDITOR")
                if len(editor) == 0 {
                        editor = "nano" //FIXME: saner default?
                }

                c := exec.Command(editor, file.Name())

                // nasty hack, see http://stackoverflow.com/a/12089980
                c.Stdin = os.Stdin
                c.Stdout = os.Stdout
                c.Stderr = os.Stderr

                er := c.Run()

                if er != nil {
                        fmt.Println(er.Error())
                        panic(er)
                }

                dat, err := ioutil.ReadFile(file.Name())
                if err != nil {
                        panic(err)
                }

                // FIXME: take just the first line
                tasks[taskid] = todotxt.ParseTask(strings.TrimSpace(string(dat)), taskid)

                tasks.Save(filename)
            },
        }

        cmdSet.PersistentFlags().StringVarP(&setprio, "priority", "p", "",
                                     "Sets task's priority.")
        cmdSet.PersistentFlags().StringVarP(&settodo, "todo", "t", "",
                                     "Sets task's todo.")

        var GotodoCmd = &cobra.Command{
            Use:   "gotodo",
            Short: "Gotodo is a go implementation of todo.txt.",
            Long: `A small, fast and fun implementation of todo.txt`,
            Run: func(cmd *cobra.Command, args []string) {
                cmdList.Run(cmd, nil)
            },
        }

        GotodoCmd.PersistentFlags().StringVarP(&filename, "filename", "", "",
                                     "Load tasks from this file.")

        conf.ParseAll()

        // values here
        flags["file"] = *flagFilename

        // sadly, this is the best we can do right now
        if filename == "" {
                if *flagFilename == "" {
                        filename = "todo.txt"
                } else {
                        filename = *flagFilename
                }
        }

        GotodoCmd.AddCommand(cmdList)
        GotodoCmd.AddCommand(cmdAdd)
        GotodoCmd.AddCommand(cmdDone)
        GotodoCmd.AddCommand(cmdArchive)
        GotodoCmd.AddCommand(cmdSet)
        GotodoCmd.AddCommand(cmdEdit)
        GotodoCmd.AddCommand(cmdConfig)
        GotodoCmd.Execute()
}
<|endoftext|>


<commit_before>package information

import (
	"fmt"
	"strings"

	"github.com/Necroforger/Fantasia/system"
	"github.com/Necroforger/discordgo"
	"github.com/Necroforger/dream"
	"github.com/Necroforger/ytdl"
)

// Module ...
type Module struct{}

// Build adds this modules commands to the system's router
func (m *Module) Build(s *system.System) {
	r := s.CommandRouter
	r.SetCategory("information")

	// Subrouter must be added before adding commands because it sets the current category
	// For the commands to be added to as the parent routers category.
	test, _ := system.NewSubCommandRouter("^"+r.Prefix+"test", "test")
	r.AddSubrouter(test)

	test.Router.On("help", m.Help).Set("", "Lists the available commands")
	test.Router.On("args", m.Argtest).Set("", "Displays your arguments")
	test.Router.On("depthcharge", m.Depthmap).Set("", "Displays the depth of the command routers commands")
	test.Router.On("play", m.Play).Set("", "plays the given youtube URL")

	sub, _ := system.NewSubCommandRouter("^ meme", "meme")
	test.Router.AddSubrouter(sub)

	sub.Router.On("cachigga", func(ctx *system.Context) { ctx.ReplyStatus(system.StatusWarning, "warning you have been cachiggad") })
}

// Help ...
func (m *Module) Help(ctx *system.Context) {
	ctx.ReplyStatus(system.StatusWarning, "Help command not yet implemented")
}

// Argtest ...
func (m *Module) Argtest(ctx *system.Context) {
	var text string
	for i, v := range ctx.Args {
		text += fmt.Sprintf("%d:\t%s\n", i, v)
	}
	ctx.ReplyStatus(system.StatusNotify, text)
}

// Depthmap maps the depth of subrouters and their commands
func (m *Module) Depthmap(ctx *system.Context) {

	depthString := func(text string, depth int, subrouter bool) string {
		quote := ""
		if subrouter {
			quote = "**"
		}
		return strings.Repeat("  ", depth) + quote + text + quote + "\n"
	}

	var depthcharge func(r *system.CommandRouter, embed *dream.Embed, depth int) *dream.Embed
	depthcharge = func(r *system.CommandRouter, embed *dream.Embed, depth int) *dream.Embed {
		if embed == nil {
			embed = dream.NewEmbed()
		}

		getField := func(name string) *discordgo.MessageEmbedField {
			for _, v := range embed.Fields {
				if v.Name == name {
					return v
				}
			}
			if name == "" {
				name = "undefined"
			}
			field := &discordgo.MessageEmbedField{Name: name}
			embed.Fields = append(embed.Fields, field)
			return field
		}

		for _, v := range r.Routes {
			field := getField(v.Category)
			field.Value += depthString(v.Name, depth, false)
		}

		for _, v := range r.Subrouters {
			field := getField(v.Category())
			field.Value += depthString(v.Name, depth, true)
			embed = depthcharge(v.Router, embed, depth+1)
		}

		return embed
	}

	_, err := ctx.ReplyEmbed(depthcharge(ctx.System.CommandRouter, nil, 0).MessageEmbed)
	if err != nil {
		ctx.Reply(fmt.Sprint(err))
	}
}

// Play plays the given song
func (m *Module) Play(ctx *system.Context) {
	if ctx.Args.After() == "" {
		ctx.ReplyStatus(system.StatusError, "No arguments provided")
		return
	}

	vs, err := ctx.Ses.UserVoiceState(ctx.Msg.Author.ID)
	if err != nil {
		ctx.ReplyStatus(system.StatusError, fmt.Sprint(err))
		return
	}

	vc, err := ctx.Ses.ChannelVoiceJoin(vs.GuildID, vs.ChannelID, false, true)
	if err != nil {
		ctx.ReplyStatus(system.StatusError, fmt.Sprint(err))
		return
	}

	info, err := ytdl.GetVideoInfo(ctx.Args.After())
	if err != nil {
		ctx.ReplyStatus(system.StatusError, fmt.Sprint(err))
		return
	}

	stream, err := system.YoutubeDLFromInfo(info)
	if err != nil {
		ctx.ReplyStatus(system.StatusError, fmt.Sprint(err))
		return
	}

	disp := ctx.Ses.PlayStream(vc, stream)

	ctx.ReplyStatus(system.StatusSuccess, "playing `"+info.Title+"`\ntype 'stop' to stop playing the video")

	for {
		var msg *discordgo.MessageCreate
		for msg = ctx.Ses.NextMessageCreate(); msg.Author.ID != ctx.Msg.Author.ID; msg = ctx.Ses.NextMessageCreate() {
		}
		if msg.Content == "stop" {
			disp.Stop()
			ctx.ReplyStatus(system.StatusNotify, "Video stopped")
			return
		}
	}
}
<commit_msg>Ping test command<commit_after>package information

import (
	"fmt"
	"strings"
	"time"

	"github.com/Necroforger/Fantasia/system"
	"github.com/Necroforger/discordgo"
	"github.com/Necroforger/dream"
	"github.com/Necroforger/ytdl"
)

// Module ...
type Module struct{}

// Build adds this modules commands to the system's router
func (m *Module) Build(s *system.System) {
	r := s.CommandRouter
	r.SetCategory("information")

	// Subrouter must be added before adding commands because it sets the current category
	// For the commands to be added to as the parent routers category.
	test, _ := system.NewSubCommandRouter("^"+r.Prefix+"test", "test")
	r.AddSubrouter(test)

	test.Router.On("help", m.Depthmap).Set("", "Lists the available commands")
	test.Router.On("args", m.Argtest).Set("", "Displays your arguments")
	test.Router.On("play", m.Play).Set("", "plays the given youtube URL")
	test.Router.On("ping", func(ctx *system.Context) {
		start := time.Now()
		m, err := ctx.ReplyStatus(system.StatusNotify, "pinging...")
		if err == nil {
			ctx.Ses.DG.ChannelMessageEditEmbed(m.ChannelID, m.ID, dream.NewEmbed().
				SetDescription(fmt.Sprint(time.Since(start))).
				SetTitle("pong").
				SetColor(system.StatusSuccess).
				MessageEmbed)
		}
	})

	sub, _ := system.NewSubCommandRouter("^ meme", "meme")
	test.Router.AddSubrouter(sub)

	sub.Router.On("cachigga", func(ctx *system.Context) { ctx.ReplyStatus(system.StatusWarning, "warning you have been cachiggad") })
}

// Help ...
func (m *Module) Help(ctx *system.Context) {
	ctx.ReplyStatus(system.StatusWarning, "Help command not yet implemented")
}

// Argtest ...
func (m *Module) Argtest(ctx *system.Context) {
	var text string
	for i, v := range ctx.Args {
		text += fmt.Sprintf("%d:\t[%s]\n", i, v)
	}
	ctx.ReplyStatus(system.StatusNotify, text)
}

// Depthmap maps the depth of subrouters and their commands to an embed
func (m *Module) Depthmap(ctx *system.Context) {

	depthString := func(text string, depth int, subrouter bool) string {
		quote := ""
		if subrouter {
			quote = "`"
		}
		return strings.Repeat("  ", depth) + quote + text + quote + "\n"
	}

	var depthcharge func(r *system.CommandRouter, embed *dream.Embed, depth int) *dream.Embed
	depthcharge = func(r *system.CommandRouter, embed *dream.Embed, depth int) *dream.Embed {
		if embed == nil {
			embed = dream.NewEmbed()
		}

		getField := func(name string) *discordgo.MessageEmbedField {
			for _, v := range embed.Fields {
				if v.Name == name {
					return v
				}
			}
			if name == "" {
				name = "undefined"
			}
			field := &discordgo.MessageEmbedField{Name: name}
			embed.Fields = append(embed.Fields, field)
			return field
		}

		for _, v := range r.Routes {
			field := getField(v.Category)
			field.Value += depthString(v.Name, depth, false)
		}

		for _, v := range r.Subrouters {
			field := getField(v.Category())
			field.Value += depthString(v.Name, depth, true)
			embed = depthcharge(v.Router, embed, depth+1)
		}

		return embed
	}

	_, err := ctx.ReplyEmbed(depthcharge(ctx.System.CommandRouter, nil, 0).SetColor(system.StatusNotify).
		SetThumbnail(ctx.Ses.DG.State.User.AvatarURL("2048")).
		InlineAllFields().
		SetDescription("subcommands are represented by indentation.").
		MessageEmbed)
	if err != nil {
		ctx.Reply(fmt.Sprint(err))
	}
}

// Play plays the given song
func (m *Module) Play(ctx *system.Context) {
	if ctx.Args.After() == "" {
		ctx.ReplyStatus(system.StatusError, "No arguments provided")
		return
	}

	vs, err := ctx.Ses.UserVoiceState(ctx.Msg.Author.ID)
	if err != nil {
		ctx.ReplyStatus(system.StatusError, fmt.Sprint(err))
		return
	}

	vc, err := ctx.Ses.ChannelVoiceJoin(vs.GuildID, vs.ChannelID, false, true)
	if err != nil {
		ctx.ReplyStatus(system.StatusError, fmt.Sprint(err))
		return
	}

	info, err := ytdl.GetVideoInfo(ctx.Args.After())
	if err != nil {
		ctx.ReplyStatus(system.StatusError, fmt.Sprint(err))
		return
	}

	stream, err := system.YoutubeDLFromInfo(info)
	if err != nil {
		ctx.ReplyStatus(system.StatusError, fmt.Sprint(err))
		return
	}

	disp := ctx.Ses.PlayStream(vc, stream)

	ctx.ReplyStatus(system.StatusSuccess, "playing `"+info.Title+"`\ntype 'stop' to stop playing the video")

	for {
		var msg *discordgo.MessageCreate
		for msg = ctx.Ses.NextMessageCreate(); msg.Author.ID != ctx.Msg.Author.ID; msg = ctx.Ses.NextMessageCreate() {
		}
		if msg.Content == "stop" {
			disp.Stop()
			ctx.ReplyStatus(system.StatusNotify, "Video stopped")
			return
		}
	}
}
<|endoftext|>


<commit_before>package health

import (
	"encoding/json"
	"net/http"
)

type (
	// Option adds optional parameter for the HealthcheckHandlerFunc
	Option func(*health)

	// Checker checks the status of the dependency and returns error.
	// In case the dependency is working as expected, return nil.
	Checker interface {
		Check() error
	}

	health struct {
		checkers map[string]Checker
	}
)

func (h *health) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	code := http.StatusOK
	var errorMsgs map[string]string
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	for key, checker := range h.checkers {
		if err := checker.Check(); err != nil {
			errorMsgs[key] = err.Error()
			code = http.StatusServiceUnavailable
		}
	}
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(struct {
		Status string            `json:"status,omitempty"`
		Errors map[string]string `json:"errors,omitempty"`
	}{
		Status: http.StatusText(code),
		Errors: errorMsgs,
	})
}

func new(opts ...Option) *health {
	h := &health{
		checkers: make(map[string]Checker),
	}
	for _, opt := range opts {
		opt(h)
	}
	return h
}

// NewHandlerFunc returns a http.Handler
func NewHandlerFunc(opts ...Option) http.Handler {
	return new(opts...)
}

// NewHandler returns a http.Handler
func NewHandler(opts ...Option) http.HandlerFunc {
	return new(opts...).ServeHTTP
}

// WithChecker adds a status checker that needs to be added as part of healthcheck. i.e database, cache or any external dependency
func WithChecker(name string, s Checker) Option {
	return func(h *health) {
		h.checkers[name] = s
	}
}
<commit_msg>add; convinient handler for checker<commit_after>package health

import (
	"encoding/json"
	"net/http"
)

type (
	// Option adds optional parameter for the HealthcheckHandlerFunc
	Option func(*health)

	// Checker checks the status of the dependency and returns error.
	// In case the dependency is working as expected, return nil.
	Checker interface {
		Check() error
	}

	// CheckFunc is a convenience type to create functions that implement
	// the Checker interface
	CheckerFunc func() error

	// Check Implements the Checker interface to allow for any func() error method
	// to be passed as a Checker
	func (c CheckerFunc) Check() error {
		return c()
	}

	health struct {
		checkers map[string]Checker
	}
)

func (h *health) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	code := http.StatusOK
	var errorMsgs map[string]string
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	for key, checker := range h.checkers {
		if err := checker.Check(); err != nil {
			errorMsgs[key] = err.Error()
			code = http.StatusServiceUnavailable
		}
	}
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(struct {
		Status string            `json:"status,omitempty"`
		Errors map[string]string `json:"errors,omitempty"`
	}{
		Status: http.StatusText(code),
		Errors: errorMsgs,
	})
}

func new(opts ...Option) *health {
	h := &health{
		checkers: make(map[string]Checker),
	}
	for _, opt := range opts {
		opt(h)
	}
	return h
}

// NewHandlerFunc returns a http.Handler
func NewHandlerFunc(opts ...Option) http.Handler {
	return new(opts...)
}

// NewHandler returns a http.Handler
func NewHandler(opts ...Option) http.HandlerFunc {
	return new(opts...).ServeHTTP
}

// WithChecker adds a status checker that needs to be added as part of healthcheck. i.e database, cache or any external dependency
func WithChecker(name string, s Checker) Option {
	return func(h *health) {
		h.checkers[name] = s
	}
}
<|endoftext|>


<commit_before>package db

import (
	"crypto/sha512"
	"fmt"
	"emp/objects"
	"time"
)

func AddPubkey(log chan string, pubkey objects.EncryptedPubkey) error {
	mutex.Lock()
	defer mutex.Unlock()

	hash := pubkey.AddrHash.GetBytes()
	payload := append(pubkey.IV[:], pubkey.Payload...)

	if hashList == nil || dbConn == nil {
		return DBError(EUNINIT)
	}
	if Contains(pubkey.AddrHash) == PUBKEY {
		return nil
	}

	err := dbConn.Exec("INSERT INTO pubkey VALUES (?, ?)", hash, payload)
	if err != nil {
		log <- fmt.Sprintf("Error inserting pubkey into db... %s", err)
		return err
	}

	Add(pubkey.AddrHash, PUBKEY)
	return nil
}

func GetPubkey(log chan string, addrHash objects.Hash) *objects.EncryptedPubkey {
	mutex.Lock()
	defer mutex.Unlock()

	hash := addrHash.GetBytes()

	if hashList == nil || dbConn == nil {
		return nil
	}
	if hashList[string(hash)] != PUBKEY {
		return nil
	}

	for s, err := dbConn.Query("SELECT payload FROM pubkey WHERE hash=?", hash); err == nil; err = s.Next() {
		var payload []byte
		s.Scan(&payload) // Assigns 1st column to rowid, the rest to row
		pub := new(objects.EncryptedPubkey)
		pub.AddrHash = addrHash
		copy(pub.IV[:], payload[:16])
		pub.Payload = payload[16:]
		return pub
	}
	// Not Found
	return nil
}

func AddPurge(log chan string, p objects.Purge) error {
	mutex.Lock()
	defer mutex.Unlock()

	txid := p.GetBytes()
	hashArr := sha512.Sum384(txid)
	hash := hashArr[:]

	if hashList == nil || dbConn == nil {
		return DBError(EUNINIT)
	}
	hashObj := new(objects.Hash)
	hashObj.FromBytes(hash)

	if Contains(*hashObj) == PURGE {
		return nil
	}

	err := dbConn.Exec("INSERT INTO purge VALUES (?, ?)", hash, txid)
	if err != nil {
		log <- fmt.Sprintf("Error inserting purge into db... %s", err)
		return err
	}

	Add(*hashObj, PURGE)
	return nil
}

func GetPurge(log chan string, txidHash objects.Hash) *objects.Purge {
	mutex.Lock()
	defer mutex.Unlock()

	hash := txidHash.GetBytes()

	if hashList == nil || dbConn == nil {
		return nil
	}
	if hashList[string(hash)] != PURGE {
		return nil
	}

	for s, err := dbConn.Query("SELECT txid FROM purge WHERE hash=?", hash); err == nil; err = s.Next() {
		var txid []byte
		s.Scan(&txid) // Assigns 1st column to rowid, the rest to row
		p := new(objects.Purge)
		p.FromBytes(txid)
		return p
	}
	// Not Found
	return nil
}

func AddMessage(log chan string, msg *objects.Message) error {
	mutex.Lock()
	defer mutex.Unlock()

	if hashList == nil || dbConn == nil {
		return DBError(EUNINIT)
	}
	if Contains(msg.TxidHash) == MSG {
		return nil
	}

	err := dbConn.Exec("INSERT INTO msg VALUES (?, ?, ?, ?)", msg.TxidHash.GetBytes(), msg.AddrHash.GetBytes(), msg.Timestamp.Unix(), msg.Content.GetBytes())
	if err != nil {
		log <- fmt.Sprintf("Error inserting message into db... %s", err)
		return err
	}

	Add(msg.TxidHash, MSG)
	return nil

}

func GetMessage(log chan string, txidHash objects.Hash) *objects.Message {
	mutex.Lock()
	defer mutex.Unlock()

	hash := txidHash.GetBytes()

	if hashList == nil || dbConn == nil {
		return nil
	}
	if hashList[string(hash)] != MSG {
		return nil
	}

	msg := new(objects.Message)

	for s, err := dbConn.Query("SELECT * FROM msg WHERE hash=?", hash); err == nil; err = s.Next() {
		var timestamp int64
		var encrypted []byte
		s.Scan(&msg.TxidHash, &msg.AddrHash, &timestamp, &encrypted)

		msg.Timestamp = time.Unix(timestamp, 0)
		msg.Content.FromBytes(encrypted)

		return msg
	}
	// Not Found
	return nil
}

func RemoveHash(log chan string, hashObj objects.Hash) error {
	mutex.Lock()
	defer mutex.Unlock()

	hash := hashObj.GetBytes()

	if hashList == nil || dbConn == nil {
		return DBError(EUNINIT)
	}

	var sql string

	switch Contains(hashObj) {
	case PUBKEY:
		sql = "DELETE FROM pubkey WHERE hash=?"
	case MSG:
		sql = "DELETE FROM msg WHERE hash=?"
	case PURGE:
		sql = "DELETE FROM purge WHERE hash=?"
	default:
		return nil
	}

	err := dbConn.Exec(sql, hash)
	if err != nil {
		log <- fmt.Sprintf("Error deleting hash from db... %s", err)
		return nil
	}

	Delete(hashObj)
	return nil
}
<commit_msg>Fixed message db bug.<commit_after>package db

import (
	"crypto/sha512"
	"fmt"
	"emp/objects"
	"time"
)

func AddPubkey(log chan string, pubkey objects.EncryptedPubkey) error {
	mutex.Lock()
	defer mutex.Unlock()

	hash := pubkey.AddrHash.GetBytes()
	payload := append(pubkey.IV[:], pubkey.Payload...)

	if hashList == nil || dbConn == nil {
		return DBError(EUNINIT)
	}
	if Contains(pubkey.AddrHash) == PUBKEY {
		return nil
	}

	err := dbConn.Exec("INSERT INTO pubkey VALUES (?, ?)", hash, payload)
	if err != nil {
		log <- fmt.Sprintf("Error inserting pubkey into db... %s", err)
		return err
	}

	Add(pubkey.AddrHash, PUBKEY)
	return nil
}

func GetPubkey(log chan string, addrHash objects.Hash) *objects.EncryptedPubkey {
	mutex.Lock()
	defer mutex.Unlock()

	hash := addrHash.GetBytes()

	if hashList == nil || dbConn == nil {
		return nil
	}
	if hashList[string(hash)] != PUBKEY {
		return nil
	}

	for s, err := dbConn.Query("SELECT payload FROM pubkey WHERE hash=?", hash); err == nil; err = s.Next() {
		var payload []byte
		s.Scan(&payload) // Assigns 1st column to rowid, the rest to row
		pub := new(objects.EncryptedPubkey)
		pub.AddrHash = addrHash
		copy(pub.IV[:], payload[:16])
		pub.Payload = payload[16:]
		return pub
	}
	// Not Found
	return nil
}

func AddPurge(log chan string, p objects.Purge) error {
	mutex.Lock()
	defer mutex.Unlock()

	txid := p.GetBytes()
	hashArr := sha512.Sum384(txid)
	hash := hashArr[:]

	if hashList == nil || dbConn == nil {
		return DBError(EUNINIT)
	}
	hashObj := new(objects.Hash)
	hashObj.FromBytes(hash)

	if Contains(*hashObj) == PURGE {
		return nil
	}

	err := dbConn.Exec("INSERT INTO purge VALUES (?, ?)", hash, txid)
	if err != nil {
		log <- fmt.Sprintf("Error inserting purge into db... %s", err)
		return err
	}

	Add(*hashObj, PURGE)
	return nil
}

func GetPurge(log chan string, txidHash objects.Hash) *objects.Purge {
	mutex.Lock()
	defer mutex.Unlock()

	hash := txidHash.GetBytes()

	if hashList == nil || dbConn == nil {
		return nil
	}
	if hashList[string(hash)] != PURGE {
		return nil
	}

	for s, err := dbConn.Query("SELECT txid FROM purge WHERE hash=?", hash); err == nil; err = s.Next() {
		var txid []byte
		s.Scan(&txid) // Assigns 1st column to rowid, the rest to row
		p := new(objects.Purge)
		p.FromBytes(txid)
		return p
	}
	// Not Found
	return nil
}

func AddMessage(log chan string, msg *objects.Message) error {
	mutex.Lock()
	defer mutex.Unlock()

	if hashList == nil || dbConn == nil {
		return DBError(EUNINIT)
	}
	if Contains(msg.TxidHash) == MSG {
		return nil
	}

	err := dbConn.Exec("INSERT INTO msg VALUES (?, ?, ?, ?)", msg.TxidHash.GetBytes(), msg.AddrHash.GetBytes(), msg.Timestamp.Unix(), msg.Content.GetBytes())
	if err != nil {
		log <- fmt.Sprintf("Error inserting message into db... %s", err)
		return err
	}

	Add(msg.TxidHash, MSG)
	return nil

}

func GetMessage(log chan string, txidHash objects.Hash) *objects.Message {
	mutex.Lock()
	defer mutex.Unlock()

	hash := txidHash.GetBytes()

	if hashList == nil || dbConn == nil {
		return nil
	}
	if hashList[string(hash)] != MSG {
		return nil
	}

	msg := new(objects.Message)

	for s, err := dbConn.Query("SELECT * FROM msg WHERE hash=?", hash); err == nil; err = s.Next() {
		var timestamp int64
		encrypted := make([]byte, 0, 0)
		txidhash := make([]byte, 0, 0)
		addrhash := make([]byte, 0, 0)
		s.Scan(&txidhash, &addrhash, &timestamp, &encrypted)

		msg.TxidHash.FromBytes(txidhash)
		msg.AddrHash.FromBytes(addrhash)
		msg.Timestamp = time.Unix(timestamp, 0)
		msg.Content.FromBytes(encrypted)

		fmt.Println("Pulling Message from Database: ", msg)

		return msg
	}
	// Not Found
	return nil
}

func RemoveHash(log chan string, hashObj objects.Hash) error {
	mutex.Lock()
	defer mutex.Unlock()

	hash := hashObj.GetBytes()

	if hashList == nil || dbConn == nil {
		return DBError(EUNINIT)
	}

	var sql string

	switch Contains(hashObj) {
	case PUBKEY:
		sql = "DELETE FROM pubkey WHERE hash=?"
	case MSG:
		sql = "DELETE FROM msg WHERE hash=?"
	case PURGE:
		sql = "DELETE FROM purge WHERE hash=?"
	default:
		return nil
	}

	err := dbConn.Exec(sql, hash)
	if err != nil {
		log <- fmt.Sprintf("Error deleting hash from db... %s", err)
		return nil
	}

	Delete(hashObj)
	return nil
}
<|endoftext|>


<commit_before>/*
Open Source Initiative OSI - The MIT License (MIT):Licensing
The MIT License (MIT)
Copyright (c) 2017 Ralph Caraveo (deckarep@gmail.com)
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 ssh

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/deckarep/blade/lib/recipe"
	"github.com/fatih/color"
	"golang.org/x/crypto/ssh"
)

var (
	concurrencySem        chan int
	hostQueue             = make(chan string)
	hostWg                sync.WaitGroup
	successfullyCompleted int32
	failedCompleted       int32
	sessionLogger         = log.New(os.Stdout, "", 0)
)

// StartSession kicks off a Blade Recipe as a session of work to be completed.
func StartSession(recipe *recipe.BladeRecipe, modifier *SessionModifier) {
	// Assumme root.
	if recipe.Overrides.User == "" {
		recipe.Overrides.User = "root"
	}

	sshConfig := &ssh.ClientConfig{
		User: recipe.Overrides.User,
		Auth: []ssh.AuthMethod{
			SSHAgent(),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	// Apply recipe will apply the recipe arguments to the commands
	// assumming they're defined.
	// TODO: don't loop twice here, then later for each cmd processing.
	sshCmds, err := applyRecipeArgs(recipe.Argument.Set, recipe.Required.Commands)
	if err != nil {
		log.Fatal("Failed to apply recipe arguments to commands with err:", err.Error())
	}

	// Concurrency must be at least 1 to make progress.
	if recipe.Overrides.Concurrency == 0 {
		recipe.Overrides.Concurrency = 1
	}

	concurrencySem = make(chan int, recipe.Overrides.Concurrency)
	go consumeAndLimitConcurrency(sshConfig, sshCmds)

	// If Hosts is defined, just use that as a discrete list.
	allHosts := recipe.Required.Hosts

	if len(allHosts) == 0 {
		// Otherwise do dynamic host lookup here.
		commandSlice := strings.Split(recipe.Required.HostLookupCommand, " ")
		out, err := exec.Command(commandSlice[0], commandSlice[1:]...).Output()
		if err != nil {
			fmt.Println("Couldn't execute command:", err.Error())
			return
		}

		allHosts = strings.Split(string(out), ",")
	}

	log.Print(color.GreenString(fmt.Sprintf("Starting recipe: %s", recipe.Meta.Name)))

	totalHosts := len(allHosts)
	for _, h := range allHosts {
		enqueueHost(h, recipe.Overrides.Port)
	}

	hostWg.Wait()
	log.Print(color.GreenString(fmt.Sprintf("Completed recipe: %s - %d sucess | %d failed | %d total",
		recipe.Meta.Name,
		atomic.LoadInt32(&successfullyCompleted),
		atomic.LoadInt32(&failedCompleted),
		totalHosts)))
}

func executeSession(sshConfig *ssh.ClientConfig, hostname string, commands []string) {

	// inline user overrides sshConfig.User
	userHost := strings.Split(hostname, "@")
	if len(userHost) == 2 {
		newSSHConfig := *sshConfig
		sshConfig = &newSSHConfig
		sshConfig.User = userHost[0]
		hostname = userHost[1]
	}

	backoff.RetryNotify(func() error {
		return startSSHSession(sshConfig, hostname, commands)
	}, backoff.WithMaxTries(backoff.NewExponentialBackOff(), 3),
		func(err error, dur time.Duration) {
			// TODO: handle this better.
			log.Println("Retry notify callback: ", err.Error())
		},
	)
}

func startSSHSession(sshConfig *ssh.ClientConfig, hostname string, commands []string) error {
	var finalError error
	defer func() {
		if finalError != nil {
			log.Println(color.YellowString(hostname) + fmt.Sprintf(" error %s", finalError.Error()))
		}
	}()

	client, err := ssh.Dial("tcp", hostname, sshConfig)
	if err != nil {
		finalError = fmt.Errorf("Failed to dial remote host: %s", err.Error())
		return finalError
	}

	// Since we can run multiple commands, we need to keep track of intermediate failures
	// and log accordingly or do some type of aggregate report.
	// Commands within a single session are executed in serial by design.
	for i, cmd := range commands {
		se := newSingleExecution(client, hostname, cmd, i+1)
		se.execute()
	}

	// Technically this is only successful when errors didn't occur above.
	atomic.AddInt32(&successfullyCompleted, 1)
	return nil
}

func applyRecipeArgs(args []*recipe.Arg, commands []string) ([]string, error) {
	if len(args) == 0 {
		return commands, nil
	}

	// TODO: ensure at least all args are used at least once to minimize end-user errors.
	// TODO: also allow the args to become flags so you can override at the command line.

	var appliedResults []string
	for _, cmd := range commands {
		replacedCmd := cmd
		for _, arg := range args {
			argToken := fmt.Sprintf("{{%s}}", arg.Arg)
			appliedFlagValue := arg.FlagValue()
			if arg.FlagValue() != "" {
				replacedCmd = strings.Replace(replacedCmd, argToken, appliedFlagValue, -1)
			} else {
				replacedCmd = strings.Replace(replacedCmd, argToken, arg.Value, -1)
			}
		}
		appliedResults = append(appliedResults, replacedCmd)
	}

	return appliedResults, nil
}
<commit_msg>spelling<commit_after>/*
Open Source Initiative OSI - The MIT License (MIT):Licensing
The MIT License (MIT)
Copyright (c) 2017 Ralph Caraveo (deckarep@gmail.com)
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 ssh

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/deckarep/blade/lib/recipe"
	"github.com/fatih/color"
	"golang.org/x/crypto/ssh"
)

var (
	concurrencySem        chan int
	hostQueue             = make(chan string)
	hostWg                sync.WaitGroup
	successfullyCompleted int32
	failedCompleted       int32
	sessionLogger         = log.New(os.Stdout, "", 0)
)

// StartSession kicks off a Blade Recipe as a session of work to be completed.
func StartSession(recipe *recipe.BladeRecipe, modifier *SessionModifier) {
	// Assumme root.
	if recipe.Overrides.User == "" {
		recipe.Overrides.User = "root"
	}

	sshConfig := &ssh.ClientConfig{
		User: recipe.Overrides.User,
		Auth: []ssh.AuthMethod{
			SSHAgent(),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	// Apply recipe will apply the recipe arguments to the commands
	// assumming they're defined.
	// TODO: don't loop twice here, then later for each cmd processing.
	sshCmds, err := applyRecipeArgs(recipe.Argument.Set, recipe.Required.Commands)
	if err != nil {
		log.Fatal("Failed to apply recipe arguments to commands with err:", err.Error())
	}

	// Concurrency must be at least 1 to make progress.
	if recipe.Overrides.Concurrency == 0 {
		recipe.Overrides.Concurrency = 1
	}

	concurrencySem = make(chan int, recipe.Overrides.Concurrency)
	go consumeAndLimitConcurrency(sshConfig, sshCmds)

	// If Hosts is defined, just use that as a discrete list.
	allHosts := recipe.Required.Hosts

	if len(allHosts) == 0 {
		// Otherwise do dynamic host lookup here.
		commandSlice := strings.Split(recipe.Required.HostLookupCommand, " ")
		out, err := exec.Command(commandSlice[0], commandSlice[1:]...).Output()
		if err != nil {
			fmt.Println("Couldn't execute command:", err.Error())
			return
		}

		allHosts = strings.Split(string(out), ",")
	}

	log.Print(color.GreenString(fmt.Sprintf("Starting recipe: %s", recipe.Meta.Name)))

	totalHosts := len(allHosts)
	for _, h := range allHosts {
		enqueueHost(h, recipe.Overrides.Port)
	}

	hostWg.Wait()
	log.Print(color.GreenString(fmt.Sprintf("Completed recipe: %s - %d success | %d failed | %d total",
		recipe.Meta.Name,
		atomic.LoadInt32(&successfullyCompleted),
		atomic.LoadInt32(&failedCompleted),
		totalHosts)))
}

func executeSession(sshConfig *ssh.ClientConfig, hostname string, commands []string) {

	// inline user overrides sshConfig.User
	userHost := strings.Split(hostname, "@")
	if len(userHost) == 2 {
		newSSHConfig := *sshConfig
		sshConfig = &newSSHConfig
		sshConfig.User = userHost[0]
		hostname = userHost[1]
	}

	backoff.RetryNotify(func() error {
		return startSSHSession(sshConfig, hostname, commands)
	}, backoff.WithMaxTries(backoff.NewExponentialBackOff(), 3),
		func(err error, dur time.Duration) {
			// TODO: handle this better.
			log.Println("Retry notify callback: ", err.Error())
		},
	)
}

func startSSHSession(sshConfig *ssh.ClientConfig, hostname string, commands []string) error {
	var finalError error
	defer func() {
		if finalError != nil {
			log.Println(color.YellowString(hostname) + fmt.Sprintf(" error %s", finalError.Error()))
		}
	}()

	client, err := ssh.Dial("tcp", hostname, sshConfig)
	if err != nil {
		finalError = fmt.Errorf("Failed to dial remote host: %s", err.Error())
		return finalError
	}

	// Since we can run multiple commands, we need to keep track of intermediate failures
	// and log accordingly or do some type of aggregate report.
	// Commands within a single session are executed in serial by design.
	for i, cmd := range commands {
		se := newSingleExecution(client, hostname, cmd, i+1)
		se.execute()
	}

	// Technically this is only successful when errors didn't occur above.
	atomic.AddInt32(&successfullyCompleted, 1)
	return nil
}

func applyRecipeArgs(args []*recipe.Arg, commands []string) ([]string, error) {
	if len(args) == 0 {
		return commands, nil
	}

	// TODO: ensure at least all args are used at least once to minimize end-user errors.
	// TODO: also allow the args to become flags so you can override at the command line.

	var appliedResults []string
	for _, cmd := range commands {
		replacedCmd := cmd
		for _, arg := range args {
			argToken := fmt.Sprintf("{{%s}}", arg.Arg)
			appliedFlagValue := arg.FlagValue()
			if arg.FlagValue() != "" {
				replacedCmd = strings.Replace(replacedCmd, argToken, appliedFlagValue, -1)
			} else {
				replacedCmd = strings.Replace(replacedCmd, argToken, arg.Value, -1)
			}
		}
		appliedResults = append(appliedResults, replacedCmd)
	}

	return appliedResults, nil
}
<|endoftext|>


<commit_before>package enmime

import (
	"bufio"
	"bytes"
	"fmt"
	"mime"
	"net/textproto"
	"strings"

	"github.com/jhillyerd/enmime/internal/coding"
	"github.com/pkg/errors"
)

const (
	// Standard MIME content dispositions
	cdAttachment = "attachment"
	cdInline     = "inline"

	// Standard MIME content types
	ctAppPrefix        = "application/"
	ctAppOctetStream   = "application/octet-stream"
	ctMultipartAltern  = "multipart/alternative"
	ctMultipartMixed   = "multipart/mixed"
	ctMultipartPrefix  = "multipart/"
	ctMultipartRelated = "multipart/related"
	ctTextPrefix       = "text/"
	ctTextPlain        = "text/plain"
	ctTextHTML         = "text/html"

	// Standard Transfer encodings
	cte7Bit            = "7bit"
	cte8Bit            = "8bit"
	cteBase64          = "base64"
	cteBinary          = "binary"
	cteQuotedPrintable = "quoted-printable"

	// Standard MIME header names
	hnContentDisposition = "Content-Disposition"
	hnContentEncoding    = "Content-Transfer-Encoding"
	hnContentID          = "Content-ID"
	hnContentType        = "Content-Type"
	hnMIMEVersion        = "MIME-Version"

	// Standard MIME header parameters
	hpBoundary = "boundary"
	hpCharset  = "charset"
	hpFile     = "file"
	hpFilename = "filename"
	hpName     = "name"
	hpModDate  = "modification-date"

	utf8 = "utf-8"
)

// AddressHeaders is the set of SMTP headers that contain email addresses, used by
// Envelope.AddressList().  Key characters must be all lowercase.
var AddressHeaders = map[string]bool{
	"bcc":             true,
	"cc":              true,
	"delivered-to":    true,
	"from":            true,
	"reply-to":        true,
	"to":              true,
	"sender":          true,
	"resent-bcc":      true,
	"resent-cc":       true,
	"resent-from":     true,
	"resent-reply-to": true,
	"resent-to":       true,
	"resent-sender":   true,
}

// Terminology from RFC 2047:
//  encoded-word: the entire =?charset?encoding?encoded-text?= string
//  charset: the character set portion of the encoded word
//  encoding: the character encoding type used for the encoded-text
//  encoded-text: the text we are decoding

// readHeader reads a block of SMTP or MIME headers and returns a textproto.MIMEHeader.
// Header parse warnings & errors will be added to p.Errors, io errors will be returned directly.
func readHeader(r *bufio.Reader, p *Part) (textproto.MIMEHeader, error) {
	// buf holds the massaged output for textproto.Reader.ReadMIMEHeader()
	buf := &bytes.Buffer{}
	tp := textproto.NewReader(r)
	firstHeader := true
	for {
		// Pull out each line of the headers as a temporary slice s
		s, err := tp.ReadLineBytes()
		if err != nil {
			buf.Write([]byte{'\r', '\n'})
			break
		}

		firstColon := bytes.IndexByte(s, ':')
		firstSpace := bytes.IndexAny(s, " \t\n\r")
		if firstSpace == 0 {
			// Starts with space: continuation
			buf.WriteByte(' ')
			buf.Write(textproto.TrimBytes(s))
			continue
		}
		if firstColon == 0 {
			// Can't parse line starting with colon: skip
			p.addError(ErrorMalformedHeader, "Header line %q started with a colon", s)
			continue
		}
		if firstColon > 0 {
			// Contains a colon, treat as a new header line
			if !firstHeader {
				// New Header line, end the previous
				buf.Write([]byte{'\r', '\n'})
			}

			// Behavior change in net/textproto package in Golang 1.12.10 and 1.13.1:
			// A space preceding the first colon in a header line is no longer handled
			// automatically due to CVE-2019-16276 which takes advantage of this
			// particular violation of RFC-7230 to exploit HTTP/1.1
			if bytes.Contains(s[:firstColon+1], []byte{' ', ':'}) {
				s = bytes.Replace(s, []byte{' ', ':'}, []byte{':'}, 1)
			}

			s = textproto.TrimBytes(s)
			buf.Write(s)
			firstHeader = false
		} else {
			// No colon: potential non-indented continuation
			if len(s) > 0 {
				// Attempt to detect and repair a non-indented continuation of previous line
				buf.WriteByte(' ')
				buf.Write(s)
				p.addWarning(ErrorMalformedHeader, "Continued line %q was not indented", s)
			} else {
				// Empty line, finish header parsing
				buf.Write([]byte{'\r', '\n'})
				break
			}
		}
	}

	buf.Write([]byte{'\r', '\n'})
	tr := textproto.NewReader(bufio.NewReader(buf))
	header, err := tr.ReadMIMEHeader()
	return header, errors.WithStack(err)
}

// decodeToUTF8Base64Header decodes a MIME header per RFC 2047, reencoding to =?utf-8b?
func decodeToUTF8Base64Header(input string) string {
	if !strings.Contains(input, "=?") {
		// Don't scan if there is nothing to do here
		return input
	}

	// The standard lib performs an incremental inspection of this string, where the
	// "skipSpace" method only strings.trimLeft for spaces and tabs. Here we have a
	// hard dependency on space existing and not on next expected rune.
	//
	// For resolving #112 with the least change, I will implement the
	// "quoted display-name" detector, which will resolve the case specific
	// issue stated in #112, but only in the case of a quoted display-name
	// followed, without whitespace, by addr-spec.
	tokens := strings.FieldsFunc(quotedDisplayName(input), whiteSpaceRune)
	output := make([]string, len(tokens))

	for i, token := range tokens {
		if len(token) > 4 && strings.Contains(token, "=?") {
			// Stash parenthesis, they should not be encoded
			prefix := ""
			suffix := ""
			if token[0] == '(' {
				prefix = "("
				token = token[1:]
			}
			if token[len(token)-1] == ')' {
				suffix = ")"
				token = token[:len(token)-1]
			}
			// Base64 encode token
			output[i] = prefix +
				mime.BEncoding.Encode("UTF-8", coding.DecodeExtHeader(token)) +
				suffix
		} else {
			output[i] = token
		}
	}

	// Return space separated tokens
	return strings.Join(output, " ")
}

func quotedDisplayName(s string) string {
	if !strings.HasPrefix(s, "\"") {
		return s
	}
	idx := strings.LastIndex(s, "\"")
	return fmt.Sprintf("%s %s", s[:idx+1], s[idx+1:])
}

// Detects a RFC-822 linear-white-space, passed to strings.FieldsFunc.
func whiteSpaceRune(r rune) bool {
	return r == ' ' || r == '\t' || r == '\r' || r == '\n'
}
<commit_msg>re-export ParseMediaType (#214)<commit_after>package enmime

import (
	"bufio"
	"bytes"
	"fmt"
	"mime"
	"net/textproto"
	"strings"

	"github.com/jhillyerd/enmime/internal/coding"
	"github.com/jhillyerd/enmime/internal/mediatype"
	"github.com/pkg/errors"
)

const (
	// Standard MIME content dispositions
	cdAttachment = "attachment"
	cdInline     = "inline"

	// Standard MIME content types
	ctAppPrefix        = "application/"
	ctAppOctetStream   = "application/octet-stream"
	ctMultipartAltern  = "multipart/alternative"
	ctMultipartMixed   = "multipart/mixed"
	ctMultipartPrefix  = "multipart/"
	ctMultipartRelated = "multipart/related"
	ctTextPrefix       = "text/"
	ctTextPlain        = "text/plain"
	ctTextHTML         = "text/html"

	// Standard Transfer encodings
	cte7Bit            = "7bit"
	cte8Bit            = "8bit"
	cteBase64          = "base64"
	cteBinary          = "binary"
	cteQuotedPrintable = "quoted-printable"

	// Standard MIME header names
	hnContentDisposition = "Content-Disposition"
	hnContentEncoding    = "Content-Transfer-Encoding"
	hnContentID          = "Content-ID"
	hnContentType        = "Content-Type"
	hnMIMEVersion        = "MIME-Version"

	// Standard MIME header parameters
	hpBoundary = "boundary"
	hpCharset  = "charset"
	hpFile     = "file"
	hpFilename = "filename"
	hpName     = "name"
	hpModDate  = "modification-date"

	utf8 = "utf-8"
)

// AddressHeaders is the set of SMTP headers that contain email addresses, used by
// Envelope.AddressList().  Key characters must be all lowercase.
var AddressHeaders = map[string]bool{
	"bcc":             true,
	"cc":              true,
	"delivered-to":    true,
	"from":            true,
	"reply-to":        true,
	"to":              true,
	"sender":          true,
	"resent-bcc":      true,
	"resent-cc":       true,
	"resent-from":     true,
	"resent-reply-to": true,
	"resent-to":       true,
	"resent-sender":   true,
}

// Terminology from RFC 2047:
//  encoded-word: the entire =?charset?encoding?encoded-text?= string
//  charset: the character set portion of the encoded word
//  encoding: the character encoding type used for the encoded-text
//  encoded-text: the text we are decoding

// ParseMediaType is a more tolerant implementation of Go's mime.ParseMediaType function.
//
// Tolerances accounted for:
//   * Missing ';' between content-type and media parameters
//   * Repeating media parameters
//   * Unquoted values in media parameters containing 'tspecials' characters
func ParseMediaType(ctype string) (mtype string, params map[string]string, invalidParams []string,
	err error) {
	// Export of internal function.
	return mediatype.Parse(ctype)
}

// readHeader reads a block of SMTP or MIME headers and returns a textproto.MIMEHeader.
// Header parse warnings & errors will be added to p.Errors, io errors will be returned directly.
func readHeader(r *bufio.Reader, p *Part) (textproto.MIMEHeader, error) {
	// buf holds the massaged output for textproto.Reader.ReadMIMEHeader()
	buf := &bytes.Buffer{}
	tp := textproto.NewReader(r)
	firstHeader := true
	for {
		// Pull out each line of the headers as a temporary slice s
		s, err := tp.ReadLineBytes()
		if err != nil {
			buf.Write([]byte{'\r', '\n'})
			break
		}

		firstColon := bytes.IndexByte(s, ':')
		firstSpace := bytes.IndexAny(s, " \t\n\r")
		if firstSpace == 0 {
			// Starts with space: continuation
			buf.WriteByte(' ')
			buf.Write(textproto.TrimBytes(s))
			continue
		}
		if firstColon == 0 {
			// Can't parse line starting with colon: skip
			p.addError(ErrorMalformedHeader, "Header line %q started with a colon", s)
			continue
		}
		if firstColon > 0 {
			// Contains a colon, treat as a new header line
			if !firstHeader {
				// New Header line, end the previous
				buf.Write([]byte{'\r', '\n'})
			}

			// Behavior change in net/textproto package in Golang 1.12.10 and 1.13.1:
			// A space preceding the first colon in a header line is no longer handled
			// automatically due to CVE-2019-16276 which takes advantage of this
			// particular violation of RFC-7230 to exploit HTTP/1.1
			if bytes.Contains(s[:firstColon+1], []byte{' ', ':'}) {
				s = bytes.Replace(s, []byte{' ', ':'}, []byte{':'}, 1)
			}

			s = textproto.TrimBytes(s)
			buf.Write(s)
			firstHeader = false
		} else {
			// No colon: potential non-indented continuation
			if len(s) > 0 {
				// Attempt to detect and repair a non-indented continuation of previous line
				buf.WriteByte(' ')
				buf.Write(s)
				p.addWarning(ErrorMalformedHeader, "Continued line %q was not indented", s)
			} else {
				// Empty line, finish header parsing
				buf.Write([]byte{'\r', '\n'})
				break
			}
		}
	}

	buf.Write([]byte{'\r', '\n'})
	tr := textproto.NewReader(bufio.NewReader(buf))
	header, err := tr.ReadMIMEHeader()
	return header, errors.WithStack(err)
}

// decodeToUTF8Base64Header decodes a MIME header per RFC 2047, reencoding to =?utf-8b?
func decodeToUTF8Base64Header(input string) string {
	if !strings.Contains(input, "=?") {
		// Don't scan if there is nothing to do here
		return input
	}

	// The standard lib performs an incremental inspection of this string, where the
	// "skipSpace" method only strings.trimLeft for spaces and tabs. Here we have a
	// hard dependency on space existing and not on next expected rune.
	//
	// For resolving #112 with the least change, I will implement the
	// "quoted display-name" detector, which will resolve the case specific
	// issue stated in #112, but only in the case of a quoted display-name
	// followed, without whitespace, by addr-spec.
	tokens := strings.FieldsFunc(quotedDisplayName(input), whiteSpaceRune)
	output := make([]string, len(tokens))

	for i, token := range tokens {
		if len(token) > 4 && strings.Contains(token, "=?") {
			// Stash parenthesis, they should not be encoded
			prefix := ""
			suffix := ""
			if token[0] == '(' {
				prefix = "("
				token = token[1:]
			}
			if token[len(token)-1] == ')' {
				suffix = ")"
				token = token[:len(token)-1]
			}
			// Base64 encode token
			output[i] = prefix +
				mime.BEncoding.Encode("UTF-8", coding.DecodeExtHeader(token)) +
				suffix
		} else {
			output[i] = token
		}
	}

	// Return space separated tokens
	return strings.Join(output, " ")
}

func quotedDisplayName(s string) string {
	if !strings.HasPrefix(s, "\"") {
		return s
	}
	idx := strings.LastIndex(s, "\"")
	return fmt.Sprintf("%s %s", s[:idx+1], s[idx+1:])
}

// Detects a RFC-822 linear-white-space, passed to strings.FieldsFunc.
func whiteSpaceRune(r rune) bool {
	return r == ' ' || r == '\t' || r == '\r' || r == '\n'
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 Blake Gentry. All rights reserved. Use of
// this source code is governed by an MIT license that can be
// found in the LICENSE file.

package heroku

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"reflect"
	"runtime"
	"strings"

	"code.google.com/p/go-uuid/uuid"
)

const (
	Version          = "0.6.0"
	DefaultAPIURL    = "https://api.heroku.com"
	DefaultUserAgent = "heroku-go/" + Version + " (" + runtime.GOOS + "; " + runtime.GOARCH + ")"
)

// A Client is a Heroku API client. Its zero value is a usable client that uses
// default settings for the Heroku API. The Client has an internal HTTP client
// (HTTP) which defaults to http.DefaultClient.
//
// As with all http.Clients, this Client's Transport has internal state (cached
// HTTP connections), so Clients should be reused instead of created as needed.
// Clients are safe for use by multiple goroutines.
type Client struct {
	// HTTP is the Client's internal http.Client, handling HTTP requests to the
	// Heroku API.
	HTTP *http.Client

	// The URL of the Heroku API to communicate with. Defaults to
	// "https://api.heroku.com".
	URL string

	// Username is the HTTP basic auth username for API calls made by this Client.
	Username string

	// Password is the HTTP basic auth password for API calls made by this Client.
	Password string

	// UserAgent to be provided in API requests. Set to DefaultUserAgent if not
	// specified.
	UserAgent string

	// Debug mode can be used to dump the full request and response to stdout.
	Debug bool

	// AdditionalHeaders are extra headers to add to each HTTP request sent by
	// this Client.
	AdditionalHeaders http.Header
}

func (c *Client) Get(v interface{}, path string) error {
	return c.APIReq(v, "GET", path, nil)
}

func (c *Client) Patch(v interface{}, path string, body interface{}) error {
	return c.APIReq(v, "PATCH", path, body)
}

func (c *Client) Post(v interface{}, path string, body interface{}) error {
	return c.APIReq(v, "POST", path, body)
}

func (c *Client) Put(v interface{}, path string, body interface{}) error {
	return c.APIReq(v, "PUT", path, body)
}

func (c *Client) Delete(path string) error {
	return c.APIReq(nil, "DELETE", path, nil)
}

// Generates an HTTP request for the Heroku API, but does not
// perform the request. The request's Accept header field will be
// set to:
//
//   Accept: application/vnd.heroku+json; version=3
//
// The Request-Id header will be set to a random UUID. The User-Agent header
// will be set to the Client's UserAgent, or DefaultUserAgent if UserAgent is
// not set.
//
// The type of body determines how to encode the request:
//
//   nil         no body
//   io.Reader   body is sent verbatim
//   else        body is encoded as application/json
func (c *Client) NewRequest(method, path string, body interface{}) (*http.Request, error) {
	var ctype string
	var rbody io.Reader

	switch t := body.(type) {
	case nil:
	case string:
		rbody = bytes.NewBufferString(t)
	case io.Reader:
		rbody = t
	default:
		v := reflect.ValueOf(body)
		if !v.IsValid() {
			break
		}
		if v.Type().Kind() == reflect.Ptr {
			v = reflect.Indirect(v)
			if !v.IsValid() {
				break
			}
		}

		j, err := json.Marshal(body)
		if err != nil {
			log.Fatal(err)
		}
		rbody = bytes.NewReader(j)
		ctype = "application/json"
	}
	apiURL := strings.TrimRight(c.URL, "/")
	if apiURL == "" {
		apiURL = DefaultAPIURL
	}
	req, err := http.NewRequest(method, apiURL+path, rbody)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/vnd.heroku+json; version=3")
	req.Header.Set("Request-Id", uuid.New())
	useragent := c.UserAgent
	if useragent == "" {
		useragent = DefaultUserAgent
	}
	req.Header.Set("User-Agent", useragent)
	if ctype != "" {
		req.Header.Set("Content-Type", ctype)
	}
	req.SetBasicAuth(c.Username, c.Password)
	for k, v := range c.AdditionalHeaders {
		req.Header[k] = v
	}
	return req, nil
}

// Sends a Heroku API request and decodes the response into v. As
// described in NewRequest(), the type of body determines how to
// encode the request body. As described in DoReq(), the type of
// v determines how to handle the response body.
func (c *Client) APIReq(v interface{}, meth, path string, body interface{}) error {
	req, err := c.NewRequest(meth, path, body)
	if err != nil {
		return err
	}
	return c.DoReq(req, v)
}

// Submits an HTTP request, checks its response, and deserializes
// the response into v. The type of v determines how to handle
// the response body:
//
//   nil        body is discarded
//   io.Writer  body is copied directly into v
//   else       body is decoded into v as json
//
func (c *Client) DoReq(req *http.Request, v interface{}) error {
	if c.Debug {
		dump, err := httputil.DumpRequestOut(req, true)
		if err != nil {
			log.Println(err)
		} else {
			os.Stderr.Write(dump)
			os.Stderr.Write([]byte{'\n', '\n'})
		}
	}

	httpClient := c.HTTP
	if httpClient == nil {
		httpClient = http.DefaultClient
	}

	res, err := httpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if c.Debug {
		dump, err := httputil.DumpResponse(res, true)
		if err != nil {
			log.Println(err)
		} else {
			os.Stderr.Write(dump)
			os.Stderr.Write([]byte{'\n'})
		}
	}
	if err = checkResp(res); err != nil {
		return err
	}
	switch t := v.(type) {
	case nil:
	case io.Writer:
		_, err = io.Copy(t, res.Body)
	default:
		err = json.NewDecoder(res.Body).Decode(v)
	}
	return err
}

// An Error represents a Heroku API error.
type Error struct {
	error
	Id  string
	URL string
}

type errorResp struct {
	Message string
	Id      string
	URL     string `json:"url"`
}

func checkResp(res *http.Response) error {
	if res.StatusCode/100 != 2 { // 200, 201, 202, etc
		var e errorResp
		err := json.NewDecoder(res.Body).Decode(&e)
		if err != nil {
			return errors.New("Unexpected error: " + res.Status)
		}
		return Error{error: errors.New(e.Message), Id: e.Id, URL: e.URL}
	}
	if msg := res.Header.Get("X-Heroku-Warning"); msg != "" {
		fmt.Fprintln(os.Stderr, strings.TrimSpace(msg))
	}
	return nil
}

type ListRange struct {
	Field      string
	Max        int
	Descending bool
	FirstId    string
	LastId     string
}

func (lr *ListRange) SetHeader(req *http.Request) {
	var hdrval string
	if lr.Field != "" {
		hdrval += lr.Field + " "
	}
	hdrval += lr.FirstId + ".." + lr.LastId
	if lr.Max != 0 {
		hdrval += fmt.Sprintf("; max=%d", lr.Max)
		if lr.Descending {
			hdrval += ", "
		}
	}

	if lr.Descending {
		hdrval += ", order=desc"
	}

	req.Header.Set("Range", hdrval)
	return
}
<commit_msg>version 0.7.0<commit_after>// Copyright (c) 2013 Blake Gentry. All rights reserved. Use of
// this source code is governed by an MIT license that can be
// found in the LICENSE file.

package heroku

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"reflect"
	"runtime"
	"strings"

	"code.google.com/p/go-uuid/uuid"
)

const (
	Version          = "0.7.0"
	DefaultAPIURL    = "https://api.heroku.com"
	DefaultUserAgent = "heroku-go/" + Version + " (" + runtime.GOOS + "; " + runtime.GOARCH + ")"
)

// A Client is a Heroku API client. Its zero value is a usable client that uses
// default settings for the Heroku API. The Client has an internal HTTP client
// (HTTP) which defaults to http.DefaultClient.
//
// As with all http.Clients, this Client's Transport has internal state (cached
// HTTP connections), so Clients should be reused instead of created as needed.
// Clients are safe for use by multiple goroutines.
type Client struct {
	// HTTP is the Client's internal http.Client, handling HTTP requests to the
	// Heroku API.
	HTTP *http.Client

	// The URL of the Heroku API to communicate with. Defaults to
	// "https://api.heroku.com".
	URL string

	// Username is the HTTP basic auth username for API calls made by this Client.
	Username string

	// Password is the HTTP basic auth password for API calls made by this Client.
	Password string

	// UserAgent to be provided in API requests. Set to DefaultUserAgent if not
	// specified.
	UserAgent string

	// Debug mode can be used to dump the full request and response to stdout.
	Debug bool

	// AdditionalHeaders are extra headers to add to each HTTP request sent by
	// this Client.
	AdditionalHeaders http.Header
}

func (c *Client) Get(v interface{}, path string) error {
	return c.APIReq(v, "GET", path, nil)
}

func (c *Client) Patch(v interface{}, path string, body interface{}) error {
	return c.APIReq(v, "PATCH", path, body)
}

func (c *Client) Post(v interface{}, path string, body interface{}) error {
	return c.APIReq(v, "POST", path, body)
}

func (c *Client) Put(v interface{}, path string, body interface{}) error {
	return c.APIReq(v, "PUT", path, body)
}

func (c *Client) Delete(path string) error {
	return c.APIReq(nil, "DELETE", path, nil)
}

// Generates an HTTP request for the Heroku API, but does not
// perform the request. The request's Accept header field will be
// set to:
//
//   Accept: application/vnd.heroku+json; version=3
//
// The Request-Id header will be set to a random UUID. The User-Agent header
// will be set to the Client's UserAgent, or DefaultUserAgent if UserAgent is
// not set.
//
// The type of body determines how to encode the request:
//
//   nil         no body
//   io.Reader   body is sent verbatim
//   else        body is encoded as application/json
func (c *Client) NewRequest(method, path string, body interface{}) (*http.Request, error) {
	var ctype string
	var rbody io.Reader

	switch t := body.(type) {
	case nil:
	case string:
		rbody = bytes.NewBufferString(t)
	case io.Reader:
		rbody = t
	default:
		v := reflect.ValueOf(body)
		if !v.IsValid() {
			break
		}
		if v.Type().Kind() == reflect.Ptr {
			v = reflect.Indirect(v)
			if !v.IsValid() {
				break
			}
		}

		j, err := json.Marshal(body)
		if err != nil {
			log.Fatal(err)
		}
		rbody = bytes.NewReader(j)
		ctype = "application/json"
	}
	apiURL := strings.TrimRight(c.URL, "/")
	if apiURL == "" {
		apiURL = DefaultAPIURL
	}
	req, err := http.NewRequest(method, apiURL+path, rbody)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/vnd.heroku+json; version=3")
	req.Header.Set("Request-Id", uuid.New())
	useragent := c.UserAgent
	if useragent == "" {
		useragent = DefaultUserAgent
	}
	req.Header.Set("User-Agent", useragent)
	if ctype != "" {
		req.Header.Set("Content-Type", ctype)
	}
	req.SetBasicAuth(c.Username, c.Password)
	for k, v := range c.AdditionalHeaders {
		req.Header[k] = v
	}
	return req, nil
}

// Sends a Heroku API request and decodes the response into v. As
// described in NewRequest(), the type of body determines how to
// encode the request body. As described in DoReq(), the type of
// v determines how to handle the response body.
func (c *Client) APIReq(v interface{}, meth, path string, body interface{}) error {
	req, err := c.NewRequest(meth, path, body)
	if err != nil {
		return err
	}
	return c.DoReq(req, v)
}

// Submits an HTTP request, checks its response, and deserializes
// the response into v. The type of v determines how to handle
// the response body:
//
//   nil        body is discarded
//   io.Writer  body is copied directly into v
//   else       body is decoded into v as json
//
func (c *Client) DoReq(req *http.Request, v interface{}) error {
	if c.Debug {
		dump, err := httputil.DumpRequestOut(req, true)
		if err != nil {
			log.Println(err)
		} else {
			os.Stderr.Write(dump)
			os.Stderr.Write([]byte{'\n', '\n'})
		}
	}

	httpClient := c.HTTP
	if httpClient == nil {
		httpClient = http.DefaultClient
	}

	res, err := httpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if c.Debug {
		dump, err := httputil.DumpResponse(res, true)
		if err != nil {
			log.Println(err)
		} else {
			os.Stderr.Write(dump)
			os.Stderr.Write([]byte{'\n'})
		}
	}
	if err = checkResp(res); err != nil {
		return err
	}
	switch t := v.(type) {
	case nil:
	case io.Writer:
		_, err = io.Copy(t, res.Body)
	default:
		err = json.NewDecoder(res.Body).Decode(v)
	}
	return err
}

// An Error represents a Heroku API error.
type Error struct {
	error
	Id  string
	URL string
}

type errorResp struct {
	Message string
	Id      string
	URL     string `json:"url"`
}

func checkResp(res *http.Response) error {
	if res.StatusCode/100 != 2 { // 200, 201, 202, etc
		var e errorResp
		err := json.NewDecoder(res.Body).Decode(&e)
		if err != nil {
			return errors.New("Unexpected error: " + res.Status)
		}
		return Error{error: errors.New(e.Message), Id: e.Id, URL: e.URL}
	}
	if msg := res.Header.Get("X-Heroku-Warning"); msg != "" {
		fmt.Fprintln(os.Stderr, strings.TrimSpace(msg))
	}
	return nil
}

type ListRange struct {
	Field      string
	Max        int
	Descending bool
	FirstId    string
	LastId     string
}

func (lr *ListRange) SetHeader(req *http.Request) {
	var hdrval string
	if lr.Field != "" {
		hdrval += lr.Field + " "
	}
	hdrval += lr.FirstId + ".." + lr.LastId
	if lr.Max != 0 {
		hdrval += fmt.Sprintf("; max=%d", lr.Max)
		if lr.Descending {
			hdrval += ", "
		}
	}

	if lr.Descending {
		hdrval += ", order=desc"
	}

	req.Header.Set("Range", hdrval)
	return
}
<|endoftext|>


<commit_before>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"database/sql/driver"
	"fmt"
	"io"
	"os"
	"strings"
)

var (
	fileRegister   map[string]bool
	readerRegister map[string]func() io.Reader
)

func init() {
	fileRegister = make(map[string]bool)
	readerRegister = make(map[string]func() io.Reader)
}

// RegisterLocalFile adds the given file to the file whitelist,
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
// Alternatively you can allow the use of all local files with
// the DSN parameter 'allowAllFiles=true'
//
//  filePath := "/home/gopher/data.csv"
//  mysql.RegisterLocalFile(filePath)
//  err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterLocalFile(filePath string) {
	fileRegister[strings.Trim(filePath, `"`)] = true
}

// DeregisterLocalFile removes the given filepath from the whitelist.
func DeregisterLocalFile(filePath string) {
	delete(fileRegister, strings.Trim(filePath, `"`))
}

// RegisterReaderHandler registers a handler function which is used
// to receive a io.Reader.
// The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
// If the handler returns a io.ReadCloser Close() is called when the
// request is finished.
//
//  mysql.RegisterReaderHandler("data", func() io.Reader {
//  	var csvReader io.Reader // Some Reader that returns CSV data
//  	... // Open Reader here
//  	return csvReader
//  })
//  err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterReaderHandler(name string, handler func() io.Reader) {
	readerRegister[name] = handler
}

// DeregisterReaderHandler removes the ReaderHandler function with
// the given name from the registry.
func DeregisterReaderHandler(name string) {
	delete(readerRegister, name)
}

func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
	var rdr io.Reader
	data := make([]byte, 4+mc.maxWriteSize)

	if strings.HasPrefix(name, "Reader::") { // io.Reader
		name = name[8:]
		handler, inMap := readerRegister[name]
		if handler != nil {
			rdr = handler()
		}
		if rdr == nil {
			if !inMap {
				err = fmt.Errorf("Reader '%s' is not registered", name)
			} else {
				err = fmt.Errorf("Reader '%s' is <nil>", name)
			}
		}
	} else { // File
		name = strings.Trim(name, `"`)
		if mc.cfg.allowAllFiles || fileRegister[name] {
			rdr, err = os.Open(name)
		} else {
			err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name)
		}
	}

	if rdc, ok := rdr.(io.ReadCloser); ok {
		defer func() {
			if err == nil {
				err = rdc.Close()
			} else {
				rdc.Close()
			}
		}()
	}

	// send content packets
	var ioErr error
	if err == nil {
		var n int
		for err == nil && ioErr == nil {
			n, err = rdr.Read(data[4:])
			if n > 0 {
				ioErr = mc.writePacket(data[:4+n])
			}
		}
		if err == io.EOF {
			err = nil
		}
		if ioErr != nil {
			errLog.Print(ioErr.Error())
			return driver.ErrBadConn
		}
	}

	// send empty packet (termination)
	ioErr = mc.writePacket([]byte{
		0x00,
		0x00,
		0x00,
		mc.sequence,
	})
	if ioErr != nil {
		errLog.Print(ioErr.Error())
		return driver.ErrBadConn
	}

	// read OK packet
	if err == nil {
		return mc.readResultOK()
	} else {
		mc.readPacket()
	}
	return err
}
<commit_msg>infile: lazy map init<commit_after>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"database/sql/driver"
	"fmt"
	"io"
	"os"
	"strings"
)

var (
	fileRegister   map[string]bool
	readerRegister map[string]func() io.Reader
)

// RegisterLocalFile adds the given file to the file whitelist,
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
// Alternatively you can allow the use of all local files with
// the DSN parameter 'allowAllFiles=true'
//
//  filePath := "/home/gopher/data.csv"
//  mysql.RegisterLocalFile(filePath)
//  err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterLocalFile(filePath string) {
	// lazy map init
	if fileRegister == nil {
		fileRegister = make(map[string]bool)
	}

	fileRegister[strings.Trim(filePath, `"`)] = true
}

// DeregisterLocalFile removes the given filepath from the whitelist.
func DeregisterLocalFile(filePath string) {
	delete(fileRegister, strings.Trim(filePath, `"`))
}

// RegisterReaderHandler registers a handler function which is used
// to receive a io.Reader.
// The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
// If the handler returns a io.ReadCloser Close() is called when the
// request is finished.
//
//  mysql.RegisterReaderHandler("data", func() io.Reader {
//  	var csvReader io.Reader // Some Reader that returns CSV data
//  	... // Open Reader here
//  	return csvReader
//  })
//  err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterReaderHandler(name string, handler func() io.Reader) {
	// lazy map init
	if readerRegister == nil {
		readerRegister = make(map[string]func() io.Reader)
	}

	readerRegister[name] = handler
}

// DeregisterReaderHandler removes the ReaderHandler function with
// the given name from the registry.
func DeregisterReaderHandler(name string) {
	delete(readerRegister, name)
}

func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
	var rdr io.Reader
	data := make([]byte, 4+mc.maxWriteSize)

	if strings.HasPrefix(name, "Reader::") { // io.Reader
		name = name[8:]
		handler, inMap := readerRegister[name]
		if handler != nil {
			rdr = handler()
		}
		if rdr == nil {
			if !inMap {
				err = fmt.Errorf("Reader '%s' is not registered", name)
			} else {
				err = fmt.Errorf("Reader '%s' is <nil>", name)
			}
		}
	} else { // File
		name = strings.Trim(name, `"`)
		if mc.cfg.allowAllFiles || fileRegister[name] {
			rdr, err = os.Open(name)
		} else {
			err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name)
		}
	}

	if rdc, ok := rdr.(io.ReadCloser); ok {
		defer func() {
			if err == nil {
				err = rdc.Close()
			} else {
				rdc.Close()
			}
		}()
	}

	// send content packets
	var ioErr error
	if err == nil {
		var n int
		for err == nil && ioErr == nil {
			n, err = rdr.Read(data[4:])
			if n > 0 {
				ioErr = mc.writePacket(data[:4+n])
			}
		}
		if err == io.EOF {
			err = nil
		}
		if ioErr != nil {
			errLog.Print(ioErr.Error())
			return driver.ErrBadConn
		}
	}

	// send empty packet (termination)
	ioErr = mc.writePacket([]byte{
		0x00,
		0x00,
		0x00,
		mc.sequence,
	})
	if ioErr != nil {
		errLog.Print(ioErr.Error())
		return driver.ErrBadConn
	}

	// read OK packet
	if err == nil {
		return mc.readResultOK()
	} else {
		mc.readPacket()
	}
	return err
}
<|endoftext|>


<commit_before>package inject

import (
	"fmt"
	"reflect"
)

type Injector interface {
	Applicator
	Invoker
	TypeMapper
	SetParent(Injector)
}

type Applicator interface {
	Apply(interface{}) error
}

type Invoker interface {
	Invoke(interface{}) ([]reflect.Value, error)
}

type TypeMapper interface {
	Map(interface{})
	MapTo(interface{}, interface{})
	Get(reflect.Type) reflect.Value
}

type injector struct {
	values map[reflect.Type]reflect.Value
	parent Injector
}

func InterfaceOf(value interface{}) reflect.Type {
	t := reflect.TypeOf(value)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()

		if t.Kind() == reflect.Interface {
			return t
		}
	}

	panic("Called inject.InterfaceOf with a value that is not a pointer to an interface. (*MyInterface)(nil)")
	return nil
}

func New() Injector {
	return &injector{
		values: make(map[reflect.Type]reflect.Value),
	}
}

func (inj *injector) Invoke(f interface{}) ([]reflect.Value, error) {
	t := reflect.TypeOf(f)

	var in = make([]reflect.Value, t.NumIn())
	for i := 0; i < t.NumIn(); i++ {
		argType := t.In(i)
		val := inj.Get(argType)
		if !val.IsValid() {
			return nil, fmt.Errorf("Value not found for type %v", argType)
		}

		in[i] = val
	}

	return reflect.ValueOf(f).Call(in), nil
}

func (inj *injector) Apply(val interface{}) error {
	v := reflect.ValueOf(val)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	if v.Kind() != reflect.Struct {
		return nil
	}

	t := v.Type()

	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		structField := t.Field(i)
		if f.CanSet() && structField.Tag == "inject" {
			ft := f.Type()
			v := inj.Get(ft)
			if !v.IsValid() {
				return fmt.Errorf("Value not found for type %v", ft)
			}

			f.Set(v)
		}

	}

	return nil
}

func (i *injector) Map(val interface{}) {
	i.values[reflect.TypeOf(val)] = reflect.ValueOf(val)
}

func (i *injector) MapTo(val interface{}, ifacePtr interface{}) {
	i.values[InterfaceOf(ifacePtr)] = reflect.ValueOf(val)
}

func (i *injector) Get(t reflect.Type) reflect.Value {
	val := i.values[t]
	if !val.IsValid() && i.parent != nil {
		val = i.parent.Get(t)
	}
	return val
}

func (i *injector) SetParent(parent Injector) {
	i.parent = parent
}
<commit_msg>Update inject.go<commit_after>package inject

import (
	"fmt"
	"reflect"
)

type Injector interface {
	Applicator
	Invoker
	TypeMapper
	SetParent(Injector)
}

type Applicator interface {
	Apply(interface{}) error
}

type Invoker interface {
	Invoke(interface{}) ([]reflect.Value, error)
}

type TypeMapper interface {
	Map(interface{}) Injector
	MapTo(interface{}, interface{}) Injector
	Get(reflect.Type) reflect.Value
}

type injector struct {
	values map[reflect.Type]reflect.Value
	parent Injector
}

func InterfaceOf(value interface{}) reflect.Type {
	t := reflect.TypeOf(value)

	for t.Kind() == reflect.Ptr {
		t = t.Elem()
	}

	if t.Kind() != reflect.Interface {
		panic("Called inject.InterfaceOf with a value that is not a pointer to an interface. (*MyInterface)(nil)")
	}

	return t
}

func New() Injector {
	return &injector{
		values: make(map[reflect.Type]reflect.Value),
	}
}

func (inj *injector) Invoke(f interface{}) ([]reflect.Value, error) {
	t := reflect.TypeOf(f)

	var in = make([]reflect.Value, t.NumIn())
	for i := 0; i < t.NumIn(); i++ {
		argType := t.In(i)
		val := inj.Get(argType)
		if !val.IsValid() {
			return nil, fmt.Errorf("Value not found for type %v", argType)
		}

		in[i] = val
	}

	return reflect.ValueOf(f).Call(in), nil
}

func (inj *injector) Apply(val interface{}) error {
	v := reflect.ValueOf(val)

	for v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	if v.Kind() != reflect.Struct {
		return nil
	}

	t := v.Type()

	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		structField := t.Field(i)
		if f.CanSet() && structField.Tag == "inject" {
			ft := f.Type()
			v := inj.Get(ft)
			if !v.IsValid() {
				return fmt.Errorf("Value not found for type %v", ft)
			}

			f.Set(v)
		}

	}

	return nil
}

func (i *injector) Map(val interface{}) Injector {
	i.values[reflect.TypeOf(val)] = reflect.ValueOf(val)
	return i
}

func (i *injector) MapTo(val interface{}, ifacePtr interface{}) Injector {
	i.values[InterfaceOf(ifacePtr)] = reflect.ValueOf(val)
	return i
}

func (i *injector) Get(t reflect.Type) reflect.Value {
	val := i.values[t]
	if !val.IsValid() && i.parent != nil {
		val = i.parent.Get(t)
	}
	return val
}

func (i *injector) SetParent(parent Injector) {
	i.parent = parent
}
<|endoftext|>