<commit_before>package container

import (
	"crypto/tls"
	"fmt"
	"time"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	defaultStopSignal = "SIGTERM"
	dryRunPrefix      = "DRY: "
)

// A Filter is a prototype for a function that can be used to filter the
// results from a call to the ListContainers() method on the Client.
type Filter func(Container) bool

// A Client is the interface through which Pumba interacts with the Docker API.
type Client interface {
	ListContainers(Filter) ([]Container, error)
	StopContainer(Container, time.Duration, bool) error
	KillContainer(Container, string, bool) error
	StartContainer(Container) error
	RenameContainer(Container, string) error
	RemoveImage(Container, bool, bool) error
	RemoveContainer(Container, bool, bool) error
	DisruptContainer(Container, string, bool) error
}

// NewClient returns a new Client instance which can be used to interact with
// the Docker API.
func NewClient(dockerHost string, tlsConfig *tls.Config, pullImages bool) Client {
	docker, err := dockerclient.NewDockerClient(dockerHost, tlsConfig)

	if err != nil {
		log.Fatalf("Error instantiating Docker client: %s", err)
	}

	return dockerClient{api: docker, pullImages: pullImages}
}

type dockerClient struct {
	api        dockerclient.Client
	pullImages bool
}

func (client dockerClient) ListContainers(fn Filter) ([]Container, error) {
	cs := []Container{}

	log.Debug("Retrieving running containers")

	runningContainers, err := client.api.ListContainers(false, false, "")
	if err != nil {
		return nil, err
	}
	for _, runningContainer := range runningContainers {
		containerInfo, err := client.api.InspectContainer(runningContainer.Id)
		if err != nil {
			return nil, err
		}
		log.Debugf("Running container: %s - (%s)", containerInfo.Name, containerInfo.Id)

		imageInfo, err := client.api.InspectImage(containerInfo.Image)
		if err != nil {
			return nil, err
		}

		c := Container{containerInfo: containerInfo, imageInfo: imageInfo}
		if fn(c) {
			cs = append(cs, c)
		}
	}

	return cs, nil
}

func (client dockerClient) KillContainer(c Container, signal string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sKilling %s (%s) with signal %s", prefix, c.Name(), c.ID(), signal)
	if !dryrun {
		if err := client.api.KillContainer(c.ID(), signal); err != nil {
			return err
		}
	}
	return nil
}

func (client dockerClient) StopContainer(c Container, timeout time.Duration, dryrun bool) error {
	signal := c.StopSignal()
	if signal == "" {
		signal = defaultStopSignal
	}
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sStopping %s (%s) with %s", prefix, c.Name(), c.ID(), signal)

	if !dryrun {
		if err := client.api.KillContainer(c.ID(), signal); err != nil {
			return err
		}

		// Wait for container to exit, but proceed anyway after the timeout elapses
		if err := client.waitForStop(c, timeout); err != nil {
			log.Debugf("Error waiting for container %s (%s) to stop: ''%s'", c.Name(), c.ID(), err.Error())
		}

		log.Debugf("Removing container %s", c.ID())

		if err := client.api.RemoveContainer(c.ID(), true, false); err != nil {
			return err
		}

		// Wait for container to be removed. In this case an error is a good thing
		if err := client.waitForStop(c, timeout); err == nil {
			return fmt.Errorf("Container %s (%s) could not be removed", c.Name(), c.ID())
		}
	}

	return nil
}

func (client dockerClient) StartContainer(c Container) error {
	config := c.runtimeConfig()
	hostConfig := c.hostConfig()
	name := c.Name()

	log.Infof("Starting %s", name)

	newContainerID, err := client.api.CreateContainer(config, name, nil)
	if err != nil {
		return err
	}

	log.Debugf("Starting container %s (%s)", name, newContainerID)

	return client.api.StartContainer(newContainerID, hostConfig)
}

func (client dockerClient) RenameContainer(c Container, newName string) error {
	log.Debugf("Renaming container %s (%s) to %s", c.Name(), c.ID(), newName)
	return client.api.RenameContainer(c.ID(), newName)
}

func (client dockerClient) RemoveImage(c Container, force bool, dryrun bool) error {
	imageID := c.ImageID()
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sRemoving image %s", prefix, imageID)
	if !dryrun {
		_, err := client.api.RemoveImage(imageID, force)
		return err
	}
	return nil
}

func (client dockerClient) RemoveContainer(c Container, force bool, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sRemoving container %s", prefix, c.ID())
	if !dryrun {
		return client.api.RemoveContainer(c.ID(), force, true)
	}
	return nil
}

func (client dockerClient) DisruptContainer(c Container, netemCmd string, dryrun bool) error {
	// find out if we have command, ip or command:ip
	cmd = strings.Split(netemCmd, ":")
	if len(cmd) == 2 {
		return disruptContainerFilterNetwork(c, cmd[0], cmd[1], dryrun)
	}
	else // all network
	{
		return disruptContainerAllNetwork(c, cmd[0], dryrun)
	}
}

func (client dockerClient) disruptContainerAllNetwork(c Container, netemCmd string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sDisrupting container %s with netem cmd %s", prefix, c.ID(), netemCmd)
	if !dryrun {
		// use dockerclient ExecStart to run Traffic Control:
		// 'tc qdisc add dev eth0 root netem delay 100ms'
		// http://www.linuxfoundation.org/collaborate/workgroups/networking/netem
		netemCommand := "tc qdisc add dev eth0 root netem " + strings.ToLower(netemCmd)
		return execOnContainer(c, netemCommand)
	}
	return nil
}

func (client dockerClient) disruptContainerFilterNetwork(c Container, netemCmd string,
	targetIP string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sDisrupting container %s with netem cmd %s on traffic to %s",
		prefix, c.ID(), netemCmd, targetIP)
	if !dryrun {
		// use dockerclient ExecStart to run Traffic Control
		// to filter network, needs to create a priority scheduling, add a low priority
		//  queue, apply netem command on that queue only, then route IP traffic
		//  to the low priority queue
		// 'tc qdisc add dev eth0 root handle 1: prio'
		// 'tc qdisc add dev eth0 parent 1:3 netem delay 3000ms'
		// 'tc filter add dev eth0 protocol ip parent 1:0 prio 3 /
		//  u32 match ip dst 172.19.0.3 flowid 1:3'
		// http://www.linuxfoundation.org/collaborate/workgroups/networking/netem
		handleCommand := "tc qdisc add dev eth0 root handle 1: prio"
		err := execOnContainer(c, handleCommand)
		if err != nil {
				return err
			}

		netemCommand := "tc qdisc add dev eth0 parent 1:3 netem " + strings.ToLower(netemCmd)
		err := execOnContainer(c, netemCommand)
		if err != nil {
				return err
			}

		filterCommand := "tc filter add dev eth0 protocol ip parent 1:0 prio 3 "+
			"u32 match ip dst " + strings.ToLower(targetIP) + " flowid 1:3"
		return execOnContainer(c, filterCommand)
	}
	return nil
}

func (client dockerClient) execOnContainer(c Container, execCmd string) error {
	execConfig := &dockerclient.ExecConfig{
		Cmd: strings.Split(execCmd, " "),
		Container: c.ID(),
	}
	_id, err := client.api.ExecCreate(execConfig)
	if err != nil {
			return err
		}

	log.Debugf("Starting Exec %s (%s)", execCmd, _id)
	return client.api.ExecStart(_id, execConfig)
}

func (client dockerClient) waitForStop(c Container, waitTime time.Duration) error {
	timeout := time.After(waitTime)

	for {
		select {
		case <-timeout:
			return nil
		default:
			if ci, err := client.api.InspectContainer(c.ID()); err != nil {
				return err
			} else if !ci.State.Running {
				return nil
			}
		}

		time.Sleep(1 * time.Second)
	}
}
<commit_msg>filter outgoing traffic so that only certain IP will be affected<commit_after>package container

import (
	"crypto/tls"
	"fmt"
	"time"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	defaultStopSignal = "SIGTERM"
	dryRunPrefix      = "DRY: "
)

// A Filter is a prototype for a function that can be used to filter the
// results from a call to the ListContainers() method on the Client.
type Filter func(Container) bool

// A Client is the interface through which Pumba interacts with the Docker API.
type Client interface {
	ListContainers(Filter) ([]Container, error)
	StopContainer(Container, time.Duration, bool) error
	KillContainer(Container, string, bool) error
	StartContainer(Container) error
	RenameContainer(Container, string) error
	RemoveImage(Container, bool, bool) error
	RemoveContainer(Container, bool, bool) error
	DisruptContainer(Container, string, bool) error
}

// NewClient returns a new Client instance which can be used to interact with
// the Docker API.
func NewClient(dockerHost string, tlsConfig *tls.Config, pullImages bool) Client {
	docker, err := dockerclient.NewDockerClient(dockerHost, tlsConfig)

	if err != nil {
		log.Fatalf("Error instantiating Docker client: %s", err)
	}

	return dockerClient{api: docker, pullImages: pullImages}
}

type dockerClient struct {
	api        dockerclient.Client
	pullImages bool
}

func (client dockerClient) ListContainers(fn Filter) ([]Container, error) {
	cs := []Container{}

	log.Debug("Retrieving running containers")

	runningContainers, err := client.api.ListContainers(false, false, "")
	if err != nil {
		return nil, err
	}
	for _, runningContainer := range runningContainers {
		containerInfo, err := client.api.InspectContainer(runningContainer.Id)
		if err != nil {
			return nil, err
		}
		log.Debugf("Running container: %s - (%s)", containerInfo.Name, containerInfo.Id)

		imageInfo, err := client.api.InspectImage(containerInfo.Image)
		if err != nil {
			return nil, err
		}

		c := Container{containerInfo: containerInfo, imageInfo: imageInfo}
		if fn(c) {
			cs = append(cs, c)
		}
	}

	return cs, nil
}

func (client dockerClient) KillContainer(c Container, signal string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sKilling %s (%s) with signal %s", prefix, c.Name(), c.ID(), signal)
	if !dryrun {
		if err := client.api.KillContainer(c.ID(), signal); err != nil {
			return err
		}
	}
	return nil
}

func (client dockerClient) StopContainer(c Container, timeout time.Duration, dryrun bool) error {
	signal := c.StopSignal()
	if signal == "" {
		signal = defaultStopSignal
	}
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sStopping %s (%s) with %s", prefix, c.Name(), c.ID(), signal)

	if !dryrun {
		if err := client.api.KillContainer(c.ID(), signal); err != nil {
			return err
		}

		// Wait for container to exit, but proceed anyway after the timeout elapses
		if err := client.waitForStop(c, timeout); err != nil {
			log.Debugf("Error waiting for container %s (%s) to stop: ''%s'", c.Name(), c.ID(), err.Error())
		}

		log.Debugf("Removing container %s", c.ID())

		if err := client.api.RemoveContainer(c.ID(), true, false); err != nil {
			return err
		}

		// Wait for container to be removed. In this case an error is a good thing
		if err := client.waitForStop(c, timeout); err == nil {
			return fmt.Errorf("Container %s (%s) could not be removed", c.Name(), c.ID())
		}
	}

	return nil
}

func (client dockerClient) StartContainer(c Container) error {
	config := c.runtimeConfig()
	hostConfig := c.hostConfig()
	name := c.Name()

	log.Infof("Starting %s", name)

	newContainerID, err := client.api.CreateContainer(config, name, nil)
	if err != nil {
		return err
	}

	log.Debugf("Starting container %s (%s)", name, newContainerID)

	return client.api.StartContainer(newContainerID, hostConfig)
}

func (client dockerClient) RenameContainer(c Container, newName string) error {
	log.Debugf("Renaming container %s (%s) to %s", c.Name(), c.ID(), newName)
	return client.api.RenameContainer(c.ID(), newName)
}

func (client dockerClient) RemoveImage(c Container, force bool, dryrun bool) error {
	imageID := c.ImageID()
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sRemoving image %s", prefix, imageID)
	if !dryrun {
		_, err := client.api.RemoveImage(imageID, force)
		return err
	}
	return nil
}

func (client dockerClient) RemoveContainer(c Container, force bool, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sRemoving container %s", prefix, c.ID())
	if !dryrun {
		return client.api.RemoveContainer(c.ID(), force, true)
	}
	return nil
}

func (client dockerClient) DisruptContainer(c Container, netemCmd string, dryrun bool) error {
	// find out if we have command, ip or command:ip
	cmd = strings.Split(netemCmd, ":")
	if len(cmd) == 2 {
		return disruptContainerFilterNetwork(c, cmd[0], cmd[1], dryrun)
	}
	else {// all network
		return disruptContainerAllNetwork(c, cmd[0], dryrun)
	}
}

func (client dockerClient) disruptContainerAllNetwork(c Container, netemCmd string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sDisrupting container %s with netem cmd %s", prefix, c.ID(), netemCmd)
	if !dryrun {
		// use dockerclient ExecStart to run Traffic Control:
		// 'tc qdisc add dev eth0 root netem delay 100ms'
		// http://www.linuxfoundation.org/collaborate/workgroups/networking/netem
		netemCommand := "tc qdisc add dev eth0 root netem " + strings.ToLower(netemCmd)
		return execOnContainer(c, netemCommand)
	}
	return nil
}

func (client dockerClient) disruptContainerFilterNetwork(c Container, netemCmd string,
	targetIP string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sDisrupting container %s with netem cmd %s on traffic to %s",
		prefix, c.ID(), netemCmd, targetIP)
	if !dryrun {
		// use dockerclient ExecStart to run Traffic Control
		// to filter network, needs to create a priority scheduling, add a low priority
		//  queue, apply netem command on that queue only, then route IP traffic
		//  to the low priority queue
		// 'tc qdisc add dev eth0 root handle 1: prio'
		// 'tc qdisc add dev eth0 parent 1:3 netem delay 3000ms'
		// 'tc filter add dev eth0 protocol ip parent 1:0 prio 3 /
		//  u32 match ip dst 172.19.0.3 flowid 1:3'
		// http://www.linuxfoundation.org/collaborate/workgroups/networking/netem
		handleCommand := "tc qdisc add dev eth0 root handle 1: prio"
		err_handle := execOnContainer(c, handleCommand)
		if err_handle != nil {
				return err_handle
			}

		netemCommand := "tc qdisc add dev eth0 parent 1:3 netem " + strings.ToLower(netemCmd)
		err_netem := execOnContainer(c, netemCommand)
		if err_netem != nil {
				return err_netem
			}

		filterCommand := "tc filter add dev eth0 protocol ip parent 1:0 prio 3 "+
			"u32 match ip dst " + strings.ToLower(targetIP) + " flowid 1:3"
		return execOnContainer(c, filterCommand)
	}
	return nil
}

func (client dockerClient) execOnContainer(c Container, execCmd string) error {
	execConfig := &dockerclient.ExecConfig{
		Cmd: strings.Split(execCmd, " "),
		Container: c.ID(),
	}
	_id, err := client.api.ExecCreate(execConfig)
	if err != nil {
			return err
		}

	log.Debugf("Starting Exec %s (%s)", execCmd, _id)
	return client.api.ExecStart(_id, execConfig)
}

func (client dockerClient) waitForStop(c Container, waitTime time.Duration) error {
	timeout := time.After(waitTime)

	for {
		select {
		case <-timeout:
			return nil
		default:
			if ci, err := client.api.InspectContainer(c.ID()); err != nil {
				return err
			} else if !ci.State.Running {
				return nil
			}
		}

		time.Sleep(1 * time.Second)
	}
}
<|endoftext|>


<commit_before>package _main

import (
	"fmt"
	"github.com/kd5pbo/minimalirc"
	"regexp"
)

func main() {

	/* Make an IRC struct */
	irc := minimalirc.New("irc.freenode.net", 7000, true, "", "jbond",
		"james", "James Bond")

	/* Set some settings */
	irc.IdNick = "jbond"
	irc.IdPass = "iloveturtles"
	irc.Channel = "#mi5"
	irc.Pongs = true

	/* Connect to the server */
	if err := irc.Connect(); nil != err {
		fmt.Printf("Failed to connect to server: %v\n", err)
		return
	}
	fmt.Printf("Connected to server.\n")

	secretMessageRE := regexp.MustCompile(
		`:(\S+) PRIVMSG #mi5 :SECRET MESSAGE: (.*)`)
	/* Wait for a secret message */
	for {
		line, ok := <-irc.C
		if !ok {
			err := <-irc.E
			fmt.Printf("Error reading from server: %v\n", err)
			return
		}
		if g := secretMessageRE.FindStringSubmatch(line); 3 == len(g) {
			fmt.Printf("Secret message from %v: %v\n", g[1], g[2])
		}
	}

}
<commit_msg>Even less mainy<commit_after>//package main

import (
	"fmt"
	"github.com/kd5pbo/minimalirc"
	"regexp"
)

func main() {

	/* Make an IRC struct */
	irc := minimalirc.New("irc.freenode.net", 7000, true, "", "jbond",
		"james", "James Bond")

	/* Set some settings */
	irc.IdNick = "jbond"
	irc.IdPass = "iloveturtles"
	irc.Channel = "#mi5"
	irc.Pongs = true

	/* Connect to the server */
	if err := irc.Connect(); nil != err {
		fmt.Printf("Failed to connect to server: %v\n", err)
		return
	}
	fmt.Printf("Connected to server.\n")

	secretMessageRE := regexp.MustCompile(
		`:(\S+) PRIVMSG #mi5 :SECRET MESSAGE: (.*)`)
	/* Wait for a secret message */
	for {
		line, ok := <-irc.C
		if !ok {
			err := <-irc.E
			fmt.Printf("Error reading from server: %v\n", err)
			return
		}
		if g := secretMessageRE.FindStringSubmatch(line); 3 == len(g) {
			fmt.Printf("Secret message from %v: %v\n", g[1], g[2])
		}
	}

}
<|endoftext|>


<commit_before>package wizard

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"gopkg.in/yaml.v1"

	"github.com/mgutz/ansi"
	"github.com/phrase/phraseapp-go/phraseapp"
)

type WizardData struct {
	AccessToken string        `yaml:"access_token"`
	ProjectId   string        `yaml:"project_id"`
	Format      string        `yaml:"file_format"`
	MainFormat  string        `yaml:"-"`
	Step        string        `yaml:"-"`
	Pull        []*PullConfig `yaml:"pull,omitempty"`
	Push        []*PushConfig `yaml:"push,omitempty"`
}

type WizardWrapper struct {
	Data *WizardData `yaml:"phraseapp"`
}

type PushConfig struct {
	Dir         string      `yaml:"dir,omitempty"`
	File        string      `yaml:"file,omitempty"`
	ProjectId   string      `yaml:"project_id,omitempty"`
	AccessToken string      `yaml:"access_token,omitempty"`
	Params      *PushParams `yaml:"params,omitempty"`
}

type PullConfig struct {
	Dir         string      `yaml:"dir,omitempty"`
	File        string      `yaml:"file,omitempty"`
	ProjectId   string      `yaml:"project_id,omitempty"`
	AccessToken string      `yaml:"access_token,omitempty"`
	Params      *PullParams `yaml:"params,omitempty"`
}

type PullParams struct {
	FileFormat string `yaml:"file_format,omitempty"`
	LocaleId   string `yaml:"locale_id,omitempty"`
}
type PushParams struct {
	FileFormat string `yaml:"file_format,omitempty"`
	LocaleId   string `yaml:"locale_id,omitempty"`
}

func clean() {
	switch runtime.GOOS {
	case "darwin":
		cmd := exec.Command("clear")
		cmd.Stdout = os.Stdout
		cmd.Run()
	case "linux":
		cmd := exec.Command("clear")
		cmd.Stdout = os.Stdout
		cmd.Run()
	case "windows":
		cmd := exec.Command("cls")
		cmd.Stdout = os.Stdout
		cmd.Run()
	default:
		fmt.Printf("%s unsupported", runtime.GOOS)
		panic("Do not know")
	}
}

func spinner(waitMsg string, position int, channelEnd *ChannelEnd, wg *sync.WaitGroup) {
	if channelEnd.closed {
		wg.Done()
		return
	}

	wg.Add(1)
	chars := []string{"⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷", "⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"}
	if position > len(chars)-1 {
		position = 0
	}
	postfix := ""
	prefix := ""
	for counter, str := range chars {
		if counter < position {
			postfix = fmt.Sprint(postfix, str)
		} else {
			prefix = fmt.Sprint(prefix, str)
		}
	}
	clean()
	printWait(fmt.Sprintf("%s %s%s", waitMsg, prefix, postfix))

	time.Sleep(100 * time.Millisecond)

	spinner(waitMsg, position+1, channelEnd, wg)
	wg.Done()
}

func printError(errorMsg string) {
	red := ansi.ColorCode("red+b:black")
	reset := ansi.ColorCode("reset")

	fmt.Println(red, errorMsg, reset)
}

func printWait(msg string) {
	yellow := ansi.ColorCode("yellow+b:black")
	reset := ansi.ColorCode("reset")
	fmt.Print(yellow, msg, reset)
}

func printSuccess(msg string) {
	green := ansi.ColorCode("green+b:black")
	reset := ansi.ColorCode("reset")

	fmt.Println(green, msg, reset)
}

func DisplayWizard(data *WizardData, step string, errorMsg string) {
	clean()

	if errorMsg != "" {
		printError(errorMsg)
	}
	switch {

	case step == "" || data.AccessToken == "":
		data.Step = "token"
		tokenStep(data)
		return
	case step == "newProject":
		data.Step = "newProject"
		newProjectStep(data)
		return
	case step == "selectProject":
		data.Step = "selectProject"
		selectProjectStep(data)
		return
	case step == "selectFormat":
		data.Step = "selectFormat"
		selectFormat(data)
		return
	case step == "pushConfig":
		data.Step = "pushConfig"
		pushConfig(data)
		return
	case step == "pullConfig":
		data.Step = "pullConfig"
		pullConfig(data)
		return
	case step == "finish":
		writeConfig(data, ".phraseapp.yaml")
		return
	}

}

func defaultPushPath(data *WizardData) string {
	switch data.Format {
	case "yml":
		return "config/locales/<locale_name>.yml"
	case "strings":
		return "<locale_name>.lproj/Localizable.strings"
	default:
		return "./"
	}
}

func defaultPullPath(data *WizardData) string {
	defaultPath := ""
	if data.Push[0] != nil {
		if data.Push[0].File != "" {
			defaultPath = filepath.Dir(data.Push[0].File)
		} else {
			defaultPath = data.Push[0].Dir
		}
	}
	return defaultPath
}

func pushConfig(data *WizardData) {
	defaultPath := defaultPushPath(data)
	fmt.Printf("Enter the path to your language files [Press enter to use default: %s]: ", defaultPath)
	var pushPath string
	fmt.Scanln(&pushPath)
	if pushPath == "" {
		pushPath = defaultPath
	}

	data.Push = make([]*PushConfig, 1)
	if strings.HasSuffix(pushPath, "/") || strings.HasSuffix(pushPath, ".") {
		data.Push[0] = &PushConfig{Dir: pushPath}
	} else {
		data.Push[0] = &PushConfig{File: pushPath}
	}
	DisplayWizard(data, next(data), "")
}

func pullConfig(data *WizardData) {
	defaultPath := defaultPullPath(data)

	fmt.Printf("Enter the path you want to put the downlaaded language file from Phrase [Press enter to use default: %s]: ", defaultPath)
	var pullPath string
	fmt.Scanln(&pullPath)
	if pullPath == "" {
		pullPath = defaultPath
	}

	data.Pull = make([]*PullConfig, 1)
	if strings.HasSuffix(pullPath, "/") || strings.HasSuffix(pullPath, ".") {
		data.Pull[0] = &PullConfig{Dir: pullPath}
	} else {
		data.Pull[0] = &PullConfig{File: pullPath}
	}
	DisplayWizard(data, next(data), "")
}

func selectFormat(data *WizardData) {
	auth := phraseapp.AuthCredentials{Token: data.AccessToken}
	phraseapp.RegisterAuthCredentials(&auth, nil)
	formats, err := phraseapp.FormatsList(1, 25)
	if err != nil {
		panic(err.Error())
	}

	for counter, format := range formats {
		fmt.Printf("%2d. %s - %s, File-Extension: %s\n", counter+1, format.ApiName, format.Name, format.Extension)
	}

	var id string
	mainFormatDefault := ""
	if data.MainFormat != "" {
		mainFormatDefault = fmt.Sprintf(" [default: %s]", data.MainFormat)
	}
	fmt.Printf("Select the format you want to use for language files you download from PhraseApp%s: ", mainFormatDefault)
	fmt.Scanln(&id)
	if id == "" && data.MainFormat != "" {
		data.Format = data.MainFormat
		DisplayWizard(data, next(data), "")
		return
	}
	number, err := strconv.Atoi(id)
	if err != nil || number < 1 || number > len(formats)+1 {
		DisplayWizard(data, "selectFormat", fmt.Sprintf("Argument Error: Please select a format from the list by specifying its position in the list."))
		return
	}
	data.Format = formats[number-1].ApiName
	DisplayWizard(data, next(data), "")
}

func writeConfig(data *WizardData, filename string) {
	wrapper := WizardWrapper{Data: data}
	bytes, err := yaml.Marshal(wrapper)
	if err != nil {
		panic(err.Error())
	}
	str := fmt.Sprintf("Success! We have created the config file for you %s:", filename)
	printSuccess(str)
	fmt.Println("")
	fmt.Println(string(bytes))

	printSuccess("You can make changes to this file, see this documentation for more advanced options: http://docs.phraseapp.com/api/v2/config")
	printSuccess("Now start using phraseapp push & pull for your workflow:")
	fmt.Println("")
	fmt.Println("$ phraseapp push")
	fmt.Println("$ phraseapp pull")
	fmt.Println("")
	var initialPush string
	fmt.Print("Enter yes to push your locales now for the first time: ")
	fmt.Scanln(&initialPush)
	if initialPush == "y" {
		fmt.Println("Pushing....")
	}
	fmt.Println("Setup completed!")
}

func next(data *WizardData) string {
	switch data.Step {
	case "", "token":
		return "selectProject"
	case "selectProject":
		return "selectFormat"
	case "newProject":
		return "selectFormat"
	case "selectFormat":
		return "pushConfig"
	case "pushConfig":
		return "pullConfig"
	case "pullConfig":
		return "finish"
	}
	return ""
}

func tokenStep(data *WizardData) {
	fmt.Print("Please enter you API Access Token (Generate one in your profile at phraseapp.com): ")
	fmt.Scanln(&data.AccessToken)
	data.AccessToken = strings.ToLower(data.AccessToken)
	success, err := regexp.MatchString("^[0-9a-f]{64}$", data.AccessToken)
	if err != nil {
		panic(err.Error())
	}
	if success == true {
		DisplayWizard(data, next(data), "")
	} else {
		data.AccessToken = ""
		DisplayWizard(data, "", "Argument Error: AccessToken must be 64 letters long and can only contain a-f, 0-9")
	}
}

func newProjectStep(data *WizardData) {
	fmt.Print("Enter name of new project: ")
	projectParam := &phraseapp.ProjectParams{}
	fmt.Scanln(&projectParam.Name)

	res, err := phraseapp.ProjectCreate(projectParam)
	if err != nil {
		success, match_err := regexp.MatchString("401", err.Error())
		if match_err != nil {
			fmt.Println(err.Error())
			panic(match_err.Error())
		}
		if success {
			data.AccessToken = ""
			DisplayWizard(data, "", fmt.Sprintf("Argument Error: Your AccessToken '%s' has no write scope. Please create a new Access Token with read and write scope.", data.AccessToken))
		} else {
			success, match_err := regexp.MatchString("Validation failed", err.Error())
			if match_err != nil {
				fmt.Println(err.Error())
				panic(match_err.Error())
			}
			if success {
				DisplayWizard(data, "newProject", err.Error())
				return
			} else {
				panic(err.Error())
			}
		}
	} else {
		data.ProjectId = res.Id
		DisplayWizard(data, next(data), "")
		return
	}
}

type ChannelEnd struct {
	closed bool
}

func selectProjectStep(data *WizardData) {
	auth := phraseapp.AuthCredentials{Token: data.AccessToken}
	fmt.Println("Please select your project:")
	phraseapp.RegisterAuthCredentials(&auth, nil)
	var wg sync.WaitGroup
	out := make(chan []phraseapp.Project, 1)
	wg.Add(1)
	var err error
	channelEnd := ChannelEnd{}
	getProjects := func(channelEnd *ChannelEnd) {
		var projects []*phraseapp.Project
		time.Sleep(2000 * time.Millisecond)
		projects, err = phraseapp.ProjectsList(1, 25)
		var array []phraseapp.Project
		for _, res := range projects {
			array = append(array, *res)
		}
		out <- array
		channelEnd.closed = true
		return
	}
	go getProjects(&channelEnd)
	go func(channelEnd *ChannelEnd, wg *sync.WaitGroup) {
		spinner("Loading projects... ", 0, channelEnd, wg)
	}(&channelEnd, &wg)
	var projects []phraseapp.Project

	projects = <-out
	clean()
	wg.Wait()
	close(out)

	if err != nil {
		success, match_err := regexp.MatchString("401", err.Error())
		if match_err != nil {
			fmt.Println(err.Error())
			panic(match_err.Error())
		}
		if success {
			errorMsg := fmt.Sprintf("Argument Error: AccessToken '%s' is invalid. It may be revoked. Please create a new Access Token.", data.AccessToken)
			data.AccessToken = ""
			DisplayWizard(data, "", errorMsg)
		} else {
			panic(err.Error())
		}
	}

	if len(projects) == 1 {
		data.ProjectId = projects[0].Id
		data.MainFormat = projects[0].MainFormat
		fmt.Printf("You've got one project, %s. Answer \"y\" to select this or \"n\" to create a new project: ")
		var answer string
		fmt.Scanln(&answer)
		if answer == "y" {
			DisplayWizard(data, next(data), "")
			return
		} else {
			data.ProjectId = ""
			data.MainFormat = ""
			DisplayWizard(data, "newProject", "")
			return
		}
	}
	for counter, project := range projects {
		fmt.Printf("%2d. %s (Id: %s)\n", counter+1, project.Name, project.Id)
	}
	fmt.Printf("%2d. Create new project\n", len(projects)+1)
	fmt.Print("Select project: ")
	var id string
	fmt.Scanln(&id)
	number, err := strconv.Atoi(id)
	if err != nil || number < 1 || number > len(projects)+1 {
		DisplayWizard(data, "selectProject", fmt.Sprintf("Argument Error: Please select a project from the list by specifying its position in the list, e.g. 2 for the second project."))
		return
	}

	if number == len(projects)+1 {
		DisplayWizard(data, "newProject", "")
		return
	}

	selectedProject := projects[number-1]
	data.ProjectId = selectedProject.Id
	data.MainFormat = selectedProject.MainFormat
	DisplayWizard(data, next(data), "")
}
<commit_msg>adds writing of file to default .phraseapp.yaml for now<commit_after>package wizard

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"gopkg.in/yaml.v1"

	"github.com/mgutz/ansi"
	"github.com/phrase/phraseapp-go/phraseapp"
)

type WizardData struct {
	AccessToken string        `yaml:"access_token"`
	ProjectId   string        `yaml:"project_id"`
	Format      string        `yaml:"file_format"`
	MainFormat  string        `yaml:"-"`
	Step        string        `yaml:"-"`
	Pull        []*PullConfig `yaml:"pull,omitempty"`
	Push        []*PushConfig `yaml:"push,omitempty"`
}

type WizardWrapper struct {
	Data *WizardData `yaml:"phraseapp"`
}

type PushConfig struct {
	Dir         string      `yaml:"dir,omitempty"`
	File        string      `yaml:"file,omitempty"`
	ProjectId   string      `yaml:"project_id,omitempty"`
	AccessToken string      `yaml:"access_token,omitempty"`
	Params      *PushParams `yaml:"params,omitempty"`
}

type PullConfig struct {
	Dir         string      `yaml:"dir,omitempty"`
	File        string      `yaml:"file,omitempty"`
	ProjectId   string      `yaml:"project_id,omitempty"`
	AccessToken string      `yaml:"access_token,omitempty"`
	Params      *PullParams `yaml:"params,omitempty"`
}

type PullParams struct {
	FileFormat string `yaml:"file_format,omitempty"`
	LocaleId   string `yaml:"locale_id,omitempty"`
}
type PushParams struct {
	FileFormat string `yaml:"file_format,omitempty"`
	LocaleId   string `yaml:"locale_id,omitempty"`
}

func clean() {
	switch runtime.GOOS {
	case "darwin":
		cmd := exec.Command("clear")
		cmd.Stdout = os.Stdout
		cmd.Run()
	case "linux":
		cmd := exec.Command("clear")
		cmd.Stdout = os.Stdout
		cmd.Run()
	case "windows":
		cmd := exec.Command("cls")
		cmd.Stdout = os.Stdout
		cmd.Run()
	default:
		fmt.Printf("%s unsupported", runtime.GOOS)
		panic("Do not know")
	}
}

func spinner(waitMsg string, position int, channelEnd *ChannelEnd, wg *sync.WaitGroup) {
	if channelEnd.closed {
		wg.Done()
		return
	}

	wg.Add(1)
	chars := []string{"⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷", "⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"}
	if position > len(chars)-1 {
		position = 0
	}
	postfix := ""
	prefix := ""
	for counter, str := range chars {
		if counter < position {
			postfix = fmt.Sprint(postfix, str)
		} else {
			prefix = fmt.Sprint(prefix, str)
		}
	}
	clean()
	printWait(fmt.Sprintf("%s %s%s", waitMsg, prefix, postfix))

	time.Sleep(100 * time.Millisecond)

	spinner(waitMsg, position+1, channelEnd, wg)
	wg.Done()
}

func printError(errorMsg string) {
	red := ansi.ColorCode("red+b:black")
	reset := ansi.ColorCode("reset")

	fmt.Println(red, errorMsg, reset)
}

func printWait(msg string) {
	yellow := ansi.ColorCode("yellow+b:black")
	reset := ansi.ColorCode("reset")
	fmt.Print(yellow, msg, reset)
}

func printSuccess(msg string) {
	green := ansi.ColorCode("green+b:black")
	reset := ansi.ColorCode("reset")

	fmt.Println(green, msg, reset)
}

func DisplayWizard(data *WizardData, step string, errorMsg string) {
	clean()

	if errorMsg != "" {
		printError(errorMsg)
	}
	switch {

	case step == "" || data.AccessToken == "":
		data.Step = "token"
		tokenStep(data)
		return
	case step == "newProject":
		data.Step = "newProject"
		newProjectStep(data)
		return
	case step == "selectProject":
		data.Step = "selectProject"
		selectProjectStep(data)
		return
	case step == "selectFormat":
		data.Step = "selectFormat"
		selectFormat(data)
		return
	case step == "pushConfig":
		data.Step = "pushConfig"
		pushConfig(data)
		return
	case step == "pullConfig":
		data.Step = "pullConfig"
		pullConfig(data)
		return
	case step == "finish":
		writeConfig(data, ".phraseapp.yaml")
		return
	}

}

func defaultPushPath(data *WizardData) string {
	switch data.Format {
	case "yml":
		return "config/locales/<locale_name>.yml"
	case "strings":
		return "<locale_name>.lproj/Localizable.strings"
	default:
		return "./"
	}
}

func defaultPullPath(data *WizardData) string {
	defaultPath := ""
	if data.Push[0] != nil {
		if data.Push[0].File != "" {
			defaultPath = filepath.Dir(data.Push[0].File)
		} else {
			defaultPath = data.Push[0].Dir
		}
	}
	return defaultPath
}

func pushConfig(data *WizardData) {
	defaultPath := defaultPushPath(data)
	fmt.Printf("Enter the path to your language files [Press enter to use default: %s]: ", defaultPath)
	var pushPath string
	fmt.Scanln(&pushPath)
	if pushPath == "" {
		pushPath = defaultPath
	}

	data.Push = make([]*PushConfig, 1)
	if strings.HasSuffix(pushPath, "/") || strings.HasSuffix(pushPath, ".") {
		data.Push[0] = &PushConfig{Dir: pushPath}
	} else {
		data.Push[0] = &PushConfig{File: pushPath}
	}
	DisplayWizard(data, next(data), "")
}

func pullConfig(data *WizardData) {
	defaultPath := defaultPullPath(data)

	fmt.Printf("Enter the path you want to put the downlaaded language file from Phrase [Press enter to use default: %s]: ", defaultPath)
	var pullPath string
	fmt.Scanln(&pullPath)
	if pullPath == "" {
		pullPath = defaultPath
	}

	data.Pull = make([]*PullConfig, 1)
	if strings.HasSuffix(pullPath, "/") || strings.HasSuffix(pullPath, ".") {
		data.Pull[0] = &PullConfig{Dir: pullPath}
	} else {
		data.Pull[0] = &PullConfig{File: pullPath}
	}
	DisplayWizard(data, next(data), "")
}

func selectFormat(data *WizardData) {
	auth := phraseapp.AuthCredentials{Token: data.AccessToken}
	phraseapp.RegisterAuthCredentials(&auth, nil)
	formats, err := phraseapp.FormatsList(1, 25)
	if err != nil {
		panic(err.Error())
	}

	for counter, format := range formats {
		fmt.Printf("%2d. %s - %s, File-Extension: %s\n", counter+1, format.ApiName, format.Name, format.Extension)
	}

	var id string
	mainFormatDefault := ""
	if data.MainFormat != "" {
		mainFormatDefault = fmt.Sprintf(" [Press enter for default: %s]", data.MainFormat)
	}
	fmt.Printf("Select the format you want to use for language files you download from PhraseApp%s: ", mainFormatDefault)
	fmt.Scanln(&id)
	if id == "" && data.MainFormat != "" {
		data.Format = data.MainFormat
		DisplayWizard(data, next(data), "")
		return
	}
	number, err := strconv.Atoi(id)
	if err != nil || number < 1 || number > len(formats)+1 {
		DisplayWizard(data, "selectFormat", fmt.Sprintf("Argument Error: Please select a format from the list by specifying its position in the list."))
		return
	}
	data.Format = formats[number-1].ApiName
	DisplayWizard(data, next(data), "")
}

func writeConfig(data *WizardData, filename string) {
	wrapper := WizardWrapper{Data: data}
	bytes, err := yaml.Marshal(wrapper)
	if err != nil {
		panic(err.Error())
	}
	err = ioutil.WriteFile(filename, bytes, 0655)
	if err != nil {
		panic(err.Error())
	}
	str := fmt.Sprintf("Success! We have created the config file for you %s:", filename)
	printSuccess(str)
	fmt.Println("")
	fmt.Println(string(bytes))

	printSuccess("You can make changes to this file, see this documentation for more advanced options: http://docs.phraseapp.com/api/v2/config")
	printSuccess("Now start using phraseapp push & pull for your workflow:")
	fmt.Println("")
	fmt.Println("$ phraseapp push")
	fmt.Println("$ phraseapp pull")
	fmt.Println("")
	var initialPush string
	fmt.Print("Enter yes to push your locales now for the first time: ")
	fmt.Scanln(&initialPush)
	if initialPush == "y" {
		fmt.Println("Pushing....")
	}
	fmt.Println("Setup completed!")
}

func next(data *WizardData) string {
	switch data.Step {
	case "", "token":
		return "selectProject"
	case "selectProject":
		return "selectFormat"
	case "newProject":
		return "selectFormat"
	case "selectFormat":
		return "pushConfig"
	case "pushConfig":
		return "pullConfig"
	case "pullConfig":
		return "finish"
	}
	return ""
}

func tokenStep(data *WizardData) {
	fmt.Print("Please enter you API Access Token (Generate one in your profile at phraseapp.com): ")
	fmt.Scanln(&data.AccessToken)
	data.AccessToken = strings.ToLower(data.AccessToken)
	success, err := regexp.MatchString("^[0-9a-f]{64}$", data.AccessToken)
	if err != nil {
		panic(err.Error())
	}
	if success == true {
		DisplayWizard(data, next(data), "")
	} else {
		data.AccessToken = ""
		DisplayWizard(data, "", "Argument Error: AccessToken must be 64 letters long and can only contain a-f, 0-9")
	}
}

func newProjectStep(data *WizardData) {
	fmt.Print("Enter name of new project: ")
	projectParam := &phraseapp.ProjectParams{}
	fmt.Scanln(&projectParam.Name)

	res, err := phraseapp.ProjectCreate(projectParam)
	if err != nil {
		success, match_err := regexp.MatchString("401", err.Error())
		if match_err != nil {
			fmt.Println(err.Error())
			panic(match_err.Error())
		}
		if success {
			data.AccessToken = ""
			DisplayWizard(data, "", fmt.Sprintf("Argument Error: Your AccessToken '%s' has no write scope. Please create a new Access Token with read and write scope.", data.AccessToken))
		} else {
			success, match_err := regexp.MatchString("Validation failed", err.Error())
			if match_err != nil {
				fmt.Println(err.Error())
				panic(match_err.Error())
			}
			if success {
				DisplayWizard(data, "newProject", err.Error())
				return
			} else {
				panic(err.Error())
			}
		}
	} else {
		data.ProjectId = res.Id
		DisplayWizard(data, next(data), "")
		return
	}
}

type ChannelEnd struct {
	closed bool
}

func selectProjectStep(data *WizardData) {
	auth := phraseapp.AuthCredentials{Token: data.AccessToken}
	fmt.Println("Please select your project:")
	phraseapp.RegisterAuthCredentials(&auth, nil)
	var wg sync.WaitGroup
	out := make(chan []phraseapp.Project, 1)
	wg.Add(1)
	var err error
	channelEnd := ChannelEnd{}
	getProjects := func(channelEnd *ChannelEnd) {
		var projects []*phraseapp.Project
		// time.Sleep(500 * time.Millisecond)
		projects, err = phraseapp.ProjectsList(1, 25)
		var array []phraseapp.Project
		for _, res := range projects {
			array = append(array, *res)
		}
		out <- array
		channelEnd.closed = true
		return
	}
	go getProjects(&channelEnd)
	go func(channelEnd *ChannelEnd, wg *sync.WaitGroup) {
		spinner("Loading projects... ", 0, channelEnd, wg)
	}(&channelEnd, &wg)
	var projects []phraseapp.Project

	projects = <-out
	clean()
	wg.Wait()
	close(out)

	if err != nil {
		success, match_err := regexp.MatchString("401", err.Error())
		if match_err != nil {
			fmt.Println(err.Error())
			panic(match_err.Error())
		}
		if success {
			errorMsg := fmt.Sprintf("Argument Error: AccessToken '%s' is invalid. It may be revoked. Please create a new Access Token.", data.AccessToken)
			data.AccessToken = ""
			DisplayWizard(data, "", errorMsg)
		} else {
			panic(err.Error())
		}
	}

	if len(projects) == 1 {
		data.ProjectId = projects[0].Id
		data.MainFormat = projects[0].MainFormat
		fmt.Printf("You've got one project, %s. Answer \"y\" to select this or \"n\" to create a new project: ")
		var answer string
		fmt.Scanln(&answer)
		if answer == "y" {
			DisplayWizard(data, next(data), "")
			return
		} else {
			data.ProjectId = ""
			data.MainFormat = ""
			DisplayWizard(data, "newProject", "")
			return
		}
	}
	for counter, project := range projects {
		fmt.Printf("%2d. %s (Id: %s)\n", counter+1, project.Name, project.Id)
	}
	fmt.Printf("%2d. Create new project\n", len(projects)+1)
	fmt.Print("Select project: ")
	var id string
	fmt.Scanln(&id)
	number, err := strconv.Atoi(id)
	if err != nil || number < 1 || number > len(projects)+1 {
		DisplayWizard(data, "selectProject", fmt.Sprintf("Argument Error: Please select a project from the list by specifying its position in the list, e.g. 2 for the second project."))
		return
	}

	if number == len(projects)+1 {
		DisplayWizard(data, "newProject", "")
		return
	}

	selectedProject := projects[number-1]
	data.ProjectId = selectedProject.Id
	data.MainFormat = selectedProject.MainFormat
	DisplayWizard(data, next(data), "")
}
<|endoftext|>


<commit_before>package worker

import (
	"encoding/json"
	"fmt"
	"github.com/go-redis/redis"
	"github.com/stitchfix/flotilla-os/clients/metrics"
	"github.com/stitchfix/flotilla-os/queue"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/engine"
	flotillaLog "github.com/stitchfix/flotilla-os/log"
	"github.com/stitchfix/flotilla-os/state"
	"gopkg.in/tomb.v2"
)

type statusWorker struct {
	sm                       state.Manager
	ee                       engine.Engine
	conf                     config.Config
	log                      flotillaLog.Logger
	pollInterval             time.Duration
	t                        tomb.Tomb
	engine                   *string
	redisClient              *redis.Client
	workerId                 string
	exceptionExtractorClient *http.Client
	exceptionExtractorUrl    string
}

func (sw *statusWorker) Initialize(conf config.Config, sm state.Manager, ee engine.Engine, log flotillaLog.Logger, pollInterval time.Duration, engine *string, qm queue.Manager) error {
	sw.pollInterval = pollInterval
	sw.conf = conf
	sw.sm = sm
	sw.ee = ee
	sw.log = log
	sw.engine = engine
	sw.workerId = fmt.Sprintf("workerid:%d", rand.Int())

	if sw.conf.IsSet("eks.exception_extractor_url") {
		sw.exceptionExtractorClient = &http.Client{
			Timeout: time.Second * 5,
		}
		sw.exceptionExtractorUrl = sw.conf.GetString("eks.exception_extractor_url")
	}
	sw.setupRedisClient(conf)
	_ = sw.log.Log("message", "initialized a status worker", "engine", *engine)
	return nil
}

func (sw *statusWorker) setupRedisClient(conf config.Config) {
	if *sw.engine == state.EKSEngine {
		sw.redisClient = redis.NewClient(&redis.Options{Addr: conf.GetString("redis_address"), DB: conf.GetInt("redis_db")})
	}
}

func (sw *statusWorker) GetTomb() *tomb.Tomb {
	return &sw.t
}

//
// Run updates status of tasks
//
func (sw *statusWorker) Run() error {
	for {
		select {
		case <-sw.t.Dying():
			sw.log.Log("message", "A status worker was terminated")
			return nil
		default:
			if *sw.engine == state.ECSEngine {
				sw.runOnceECS()
				time.Sleep(sw.pollInterval)
			}

			if *sw.engine == state.EKSEngine {
				sw.runOnceEKS()
				time.Sleep(sw.pollInterval)
			}
		}
	}
}

func (sw *statusWorker) runOnceEKS() {
	rl, err := sw.sm.ListRuns(1000, 0, "started_at", "asc", map[string][]string{
		"queued_at_since": {
			time.Now().AddDate(0, 0, -3).Format(time.RFC3339),
		},
		"status": {state.StatusNeedsRetry, state.StatusRunning, state.StatusQueued, state.StatusPending},
	}, nil, []string{state.EKSEngine})

	if err != nil {
		_ = sw.log.Log("message", "unable to receive runs", "error", fmt.Sprintf("%+v", err))
		return
	}
	runs := rl.Runs
	sw.processEKSRuns(runs)
}

func (sw *statusWorker) processEKSRuns(runs []state.Run) {
	var lockedRuns []state.Run
	for _, run := range runs {
		duration := time.Duration((rand.Intn(15))+25) * time.Second
		lock := sw.acquireLock(run, "status", duration)
		if lock {
			lockedRuns = append(lockedRuns, run)
		}
	}
	_ = metrics.Increment(metrics.StatusWorkerLockedRuns, []string{sw.workerId}, float64(len(lockedRuns)))
	for _, run := range lockedRuns {
		start := time.Now()
		sw.processEKSRun(run)
		_ = metrics.Timing(metrics.StatusWorkerProcessEKSRun, time.Since(start), []string{sw.workerId}, 1)
	}
}
func (sw *statusWorker) acquireLock(run state.Run, purpose string, expiration time.Duration) bool {
	start := time.Now()
	set, err := sw.redisClient.SetNX(fmt.Sprintf("%s-%s", run.RunID, purpose), sw.workerId, expiration).Result()
	if err != nil {
		_ = sw.log.Log("message", "unable to set lock", "error", fmt.Sprintf("%+v", err))
		return true
	}
	_ = metrics.Timing(metrics.StatusWorkerAcquireLock, time.Since(start), []string{sw.workerId}, 1)
	return set
}

func (sw *statusWorker) processEKSRun(run state.Run) {
	reloadRun, err := sw.sm.GetRun(run.RunID)
	if err == nil && reloadRun.Status == state.StatusStopped {
		// Run was updated by another worker process.
		return
	}
	start := time.Now()
	updatedRunWithMetrics, _ := sw.ee.FetchPodMetrics(run)
	_ = metrics.Timing(metrics.StatusWorkerFetchPodMetrics, time.Since(start), []string{sw.workerId}, 1)

	start = time.Now()
	updatedRun, err := sw.ee.FetchUpdateStatus(updatedRunWithMetrics)
	_ = metrics.Timing(metrics.StatusWorkerFetchUpdateStatus, time.Since(start), []string{sw.workerId}, 1)

	if err == nil {
		subRuns, err := sw.sm.ListRuns(1000, 0, "status", "desc", nil, map[string]string{"PARENT_FLOTILLA_RUN_ID": run.RunID}, state.Engines)
		if err == nil && subRuns.Total > 0 {
			var spawnedRuns state.SpawnedRuns
			for _, subRun := range subRuns.Runs {
				spawnedRuns = append(spawnedRuns, state.SpawnedRun{RunID: subRun.RunID})
			}
			updatedRun.SpawnedRuns = &spawnedRuns
		}
	}
	if err != nil {
		message := fmt.Sprintf("%+v", err)
		_ = sw.log.Log("message", "unable to receive eks runs", "error", message)

		minutesInQueue := time.Now().Sub(*run.QueuedAt).Minutes()
		if strings.Contains(message, "not found") && minutesInQueue > float64(30) {
			stoppedAt := time.Now()
			reason := "Job either timed out or not found on the EKS cluster."
			updatedRun.Status = state.StatusStopped
			updatedRun.FinishedAt = &stoppedAt
			updatedRun.ExitReason = &reason
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}

	} else {
		if run.Status != updatedRun.Status && (updatedRun.PodName == run.PodName) {
			_ = sw.log.Log("message", "updating eks run status", "pod", updatedRun.PodName, "status", updatedRun.Status, "exit_code", updatedRun.ExitCode)
			if sw.exceptionExtractorClient != nil {
				jobUrl := fmt.Sprintf("%s/extract/%s", sw.exceptionExtractorUrl, run.RunID)
				res, err := sw.exceptionExtractorClient.Get(jobUrl)
				if err == nil {
					body, err := ioutil.ReadAll(res.Body)
					if body != nil {
						defer res.Body.Close()
						runExceptions := state.RunExceptions{}
						err = json.Unmarshal(body, &runExceptions)
						if err == nil {
							updatedRun.RunExceptions = &runExceptions
						}
					}
				}
			}
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			if err != nil {
				_ = sw.log.Log("message", "unable to save eks runs", "error", fmt.Sprintf("%+v", err))
			}

			if updatedRun.Status == state.StatusStopped {
				//TODO - move to a separate worker.
				//_ = sw.ee.Terminate(run)
			}
		} else {
			if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
				updatedRun.MaxCpuUsed != run.MaxCpuUsed ||
				updatedRun.Cpu != run.Cpu ||
				updatedRun.Memory != run.Memory ||
				updatedRun.PodEvents != run.PodEvents ||
				updatedRun.SpawnedRuns != run.SpawnedRuns {
				_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			}
		}
	}
}

func (sw *statusWorker) processEKSRunMetrics(run state.Run) {
	updatedRun, err := sw.ee.FetchPodMetrics(run)
	if err == nil {
		if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
			updatedRun.MaxCpuUsed != run.MaxCpuUsed {
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}
	}
}

func (sw *statusWorker) runOnceECS() {
	runReceipt, err := sw.ee.PollStatus()

	if err != nil {
		sw.log.Log("message", "unable to receive status message", "error", fmt.Sprintf("%+v", err))
		return
	}

	// Ensure update is in the env required, otherwise, ack without taking action
	update := runReceipt.Run
	if update != nil {
		//
		// Relies on the reserved env var, FLOTILLA_SERVER_MODE to ensure update
		// belongs to -this- mode of Flotilla
		//
		var serverMode string
		if update.Env != nil {
			for _, kv := range *update.Env {
				if kv.Name == "FLOTILLA_SERVER_MODE" {
					serverMode = kv.Value
				}
			}
		}

		shouldProcess := len(serverMode) > 0 && serverMode == sw.conf.GetString("flotilla_mode")
		if shouldProcess {
			run, err := sw.findRun(update.TaskArn)
			if err != nil {
				sw.log.Log("message", "unable to find run to apply update to", "error", fmt.Sprintf("%+v", err))
				return
			}

			_, err = sw.sm.UpdateRun(run.RunID, *update)
			if err != nil {
				sw.log.Log("message", "error applying status update", "run", run.RunID, "error", fmt.Sprintf("%+v", err))
				return
			}

			// emit status update event
			sw.logStatusUpdate(*update)
		}

		if sw.conf.GetString("flotilla_mode") != "dev" {
			_ = sw.log.Log("message", "Acking status update", "arn", update.TaskArn)
		}
		if err = runReceipt.Done(); err != nil {
			sw.log.Log("message", "Acking status update failed", "arn", update.TaskArn, "error", fmt.Sprintf("%+v", err))
		}
	}
}

func (sw *statusWorker) logStatusUpdate(update state.Run) {
	var err error
	var startedAt, finishedAt time.Time
	var env state.EnvList

	if update.StartedAt != nil {
		startedAt = *update.StartedAt
	}

	if update.FinishedAt != nil {
		finishedAt = *update.FinishedAt
	}

	if update.Env != nil {
		env = *update.Env
	}

	if update.ExitCode != nil {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"exit_code", *update.ExitCode,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	} else {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	}

	if err != nil {
		sw.log.Log("message", "Failed to emit status event", "run_id", update.RunID, "error", err.Error())
	}
}

func (sw *statusWorker) findRun(taskArn string) (state.Run, error) {
	var engines []string
	if sw.engine != nil {
		engines = []string{*sw.engine}
	} else {
		engines = nil
	}

	runs, err := sw.sm.ListRuns(1, 0, "started_at", "asc", map[string][]string{
		"task_arn": {taskArn},
	}, nil, engines)
	if err != nil {
		return state.Run{}, errors.Wrapf(err, "problem finding run by task arn [%s]", taskArn)
	}
	if runs.Total > 0 && len(runs.Runs) > 0 {
		return runs.Runs[0], nil
	}
	return state.Run{}, errors.Errorf("no run found for [%s]", taskArn)
}
<commit_msg>only extract exceptions when exit code is not 0<commit_after>package worker

import (
	"encoding/json"
	"fmt"
	"github.com/go-redis/redis"
	"github.com/stitchfix/flotilla-os/clients/metrics"
	"github.com/stitchfix/flotilla-os/queue"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/engine"
	flotillaLog "github.com/stitchfix/flotilla-os/log"
	"github.com/stitchfix/flotilla-os/state"
	"gopkg.in/tomb.v2"
)

type statusWorker struct {
	sm                       state.Manager
	ee                       engine.Engine
	conf                     config.Config
	log                      flotillaLog.Logger
	pollInterval             time.Duration
	t                        tomb.Tomb
	engine                   *string
	redisClient              *redis.Client
	workerId                 string
	exceptionExtractorClient *http.Client
	exceptionExtractorUrl    string
}

func (sw *statusWorker) Initialize(conf config.Config, sm state.Manager, ee engine.Engine, log flotillaLog.Logger, pollInterval time.Duration, engine *string, qm queue.Manager) error {
	sw.pollInterval = pollInterval
	sw.conf = conf
	sw.sm = sm
	sw.ee = ee
	sw.log = log
	sw.engine = engine
	sw.workerId = fmt.Sprintf("workerid:%d", rand.Int())

	if sw.conf.IsSet("eks.exception_extractor_url") {
		sw.exceptionExtractorClient = &http.Client{
			Timeout: time.Second * 5,
		}
		sw.exceptionExtractorUrl = sw.conf.GetString("eks.exception_extractor_url")
	}
	sw.setupRedisClient(conf)
	_ = sw.log.Log("message", "initialized a status worker", "engine", *engine)
	return nil
}

func (sw *statusWorker) setupRedisClient(conf config.Config) {
	if *sw.engine == state.EKSEngine {
		sw.redisClient = redis.NewClient(&redis.Options{Addr: conf.GetString("redis_address"), DB: conf.GetInt("redis_db")})
	}
}

func (sw *statusWorker) GetTomb() *tomb.Tomb {
	return &sw.t
}

//
// Run updates status of tasks
//
func (sw *statusWorker) Run() error {
	for {
		select {
		case <-sw.t.Dying():
			sw.log.Log("message", "A status worker was terminated")
			return nil
		default:
			if *sw.engine == state.ECSEngine {
				sw.runOnceECS()
				time.Sleep(sw.pollInterval)
			}

			if *sw.engine == state.EKSEngine {
				sw.runOnceEKS()
				time.Sleep(sw.pollInterval)
			}
		}
	}
}

func (sw *statusWorker) runOnceEKS() {
	rl, err := sw.sm.ListRuns(1000, 0, "started_at", "asc", map[string][]string{
		"queued_at_since": {
			time.Now().AddDate(0, 0, -3).Format(time.RFC3339),
		},
		"status": {state.StatusNeedsRetry, state.StatusRunning, state.StatusQueued, state.StatusPending},
	}, nil, []string{state.EKSEngine})

	if err != nil {
		_ = sw.log.Log("message", "unable to receive runs", "error", fmt.Sprintf("%+v", err))
		return
	}
	runs := rl.Runs
	sw.processEKSRuns(runs)
}

func (sw *statusWorker) processEKSRuns(runs []state.Run) {
	var lockedRuns []state.Run
	for _, run := range runs {
		duration := time.Duration((rand.Intn(15))+25) * time.Second
		lock := sw.acquireLock(run, "status", duration)
		if lock {
			lockedRuns = append(lockedRuns, run)
		}
	}
	_ = metrics.Increment(metrics.StatusWorkerLockedRuns, []string{sw.workerId}, float64(len(lockedRuns)))
	for _, run := range lockedRuns {
		start := time.Now()
		sw.processEKSRun(run)
		_ = metrics.Timing(metrics.StatusWorkerProcessEKSRun, time.Since(start), []string{sw.workerId}, 1)
	}
}
func (sw *statusWorker) acquireLock(run state.Run, purpose string, expiration time.Duration) bool {
	start := time.Now()
	set, err := sw.redisClient.SetNX(fmt.Sprintf("%s-%s", run.RunID, purpose), sw.workerId, expiration).Result()
	if err != nil {
		_ = sw.log.Log("message", "unable to set lock", "error", fmt.Sprintf("%+v", err))
		return true
	}
	_ = metrics.Timing(metrics.StatusWorkerAcquireLock, time.Since(start), []string{sw.workerId}, 1)
	return set
}

func (sw *statusWorker) processEKSRun(run state.Run) {
	reloadRun, err := sw.sm.GetRun(run.RunID)
	if err == nil && reloadRun.Status == state.StatusStopped {
		// Run was updated by another worker process.
		return
	}
	start := time.Now()
	updatedRunWithMetrics, _ := sw.ee.FetchPodMetrics(run)
	_ = metrics.Timing(metrics.StatusWorkerFetchPodMetrics, time.Since(start), []string{sw.workerId}, 1)

	start = time.Now()
	updatedRun, err := sw.ee.FetchUpdateStatus(updatedRunWithMetrics)
	_ = metrics.Timing(metrics.StatusWorkerFetchUpdateStatus, time.Since(start), []string{sw.workerId}, 1)

	if err == nil {
		subRuns, err := sw.sm.ListRuns(1000, 0, "status", "desc", nil, map[string]string{"PARENT_FLOTILLA_RUN_ID": run.RunID}, state.Engines)
		if err == nil && subRuns.Total > 0 {
			var spawnedRuns state.SpawnedRuns
			for _, subRun := range subRuns.Runs {
				spawnedRuns = append(spawnedRuns, state.SpawnedRun{RunID: subRun.RunID})
			}
			updatedRun.SpawnedRuns = &spawnedRuns
		}
	}
	if err != nil {
		message := fmt.Sprintf("%+v", err)
		_ = sw.log.Log("message", "unable to receive eks runs", "error", message)

		minutesInQueue := time.Now().Sub(*run.QueuedAt).Minutes()
		if strings.Contains(message, "not found") && minutesInQueue > float64(30) {
			stoppedAt := time.Now()
			reason := "Job either timed out or not found on the EKS cluster."
			updatedRun.Status = state.StatusStopped
			updatedRun.FinishedAt = &stoppedAt
			updatedRun.ExitReason = &reason
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}

	} else {
		if run.Status != updatedRun.Status && (updatedRun.PodName == run.PodName) {
			_ = sw.log.Log("message", "updating eks run status", "pod", updatedRun.PodName, "status", updatedRun.Status, "exit_code", updatedRun.ExitCode)

			if sw.exceptionExtractorClient != nil && updatedRun.ExitCode != nil && *(updatedRun.ExitCode) != 0 {
				jobUrl := fmt.Sprintf("%s/extract/%s", sw.exceptionExtractorUrl, run.RunID)
				res, err := sw.exceptionExtractorClient.Get(jobUrl)
				if err == nil {
					body, err := ioutil.ReadAll(res.Body)
					if body != nil {
						defer res.Body.Close()
						runExceptions := state.RunExceptions{}
						err = json.Unmarshal(body, &runExceptions)
						if err == nil {
							updatedRun.RunExceptions = &runExceptions
						}
					}
				}
			}
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			if err != nil {
				_ = sw.log.Log("message", "unable to save eks runs", "error", fmt.Sprintf("%+v", err))
			}

			if updatedRun.Status == state.StatusStopped {
				//TODO - move to a separate worker.
				//_ = sw.ee.Terminate(run)
			}
		} else {
			if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
				updatedRun.MaxCpuUsed != run.MaxCpuUsed ||
				updatedRun.Cpu != run.Cpu ||
				updatedRun.Memory != run.Memory ||
				updatedRun.PodEvents != run.PodEvents ||
				updatedRun.SpawnedRuns != run.SpawnedRuns {
				_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			}
		}
	}
}

func (sw *statusWorker) processEKSRunMetrics(run state.Run) {
	updatedRun, err := sw.ee.FetchPodMetrics(run)
	if err == nil {
		if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
			updatedRun.MaxCpuUsed != run.MaxCpuUsed {
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}
	}
}

func (sw *statusWorker) runOnceECS() {
	runReceipt, err := sw.ee.PollStatus()

	if err != nil {
		sw.log.Log("message", "unable to receive status message", "error", fmt.Sprintf("%+v", err))
		return
	}

	// Ensure update is in the env required, otherwise, ack without taking action
	update := runReceipt.Run
	if update != nil {
		//
		// Relies on the reserved env var, FLOTILLA_SERVER_MODE to ensure update
		// belongs to -this- mode of Flotilla
		//
		var serverMode string
		if update.Env != nil {
			for _, kv := range *update.Env {
				if kv.Name == "FLOTILLA_SERVER_MODE" {
					serverMode = kv.Value
				}
			}
		}

		shouldProcess := len(serverMode) > 0 && serverMode == sw.conf.GetString("flotilla_mode")
		if shouldProcess {
			run, err := sw.findRun(update.TaskArn)
			if err != nil {
				sw.log.Log("message", "unable to find run to apply update to", "error", fmt.Sprintf("%+v", err))
				return
			}

			_, err = sw.sm.UpdateRun(run.RunID, *update)
			if err != nil {
				sw.log.Log("message", "error applying status update", "run", run.RunID, "error", fmt.Sprintf("%+v", err))
				return
			}

			// emit status update event
			sw.logStatusUpdate(*update)
		}

		if sw.conf.GetString("flotilla_mode") != "dev" {
			_ = sw.log.Log("message", "Acking status update", "arn", update.TaskArn)
		}
		if err = runReceipt.Done(); err != nil {
			sw.log.Log("message", "Acking status update failed", "arn", update.TaskArn, "error", fmt.Sprintf("%+v", err))
		}
	}
}

func (sw *statusWorker) logStatusUpdate(update state.Run) {
	var err error
	var startedAt, finishedAt time.Time
	var env state.EnvList

	if update.StartedAt != nil {
		startedAt = *update.StartedAt
	}

	if update.FinishedAt != nil {
		finishedAt = *update.FinishedAt
	}

	if update.Env != nil {
		env = *update.Env
	}

	if update.ExitCode != nil {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"exit_code", *update.ExitCode,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	} else {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	}

	if err != nil {
		sw.log.Log("message", "Failed to emit status event", "run_id", update.RunID, "error", err.Error())
	}
}

func (sw *statusWorker) findRun(taskArn string) (state.Run, error) {
	var engines []string
	if sw.engine != nil {
		engines = []string{*sw.engine}
	} else {
		engines = nil
	}

	runs, err := sw.sm.ListRuns(1, 0, "started_at", "asc", map[string][]string{
		"task_arn": {taskArn},
	}, nil, engines)
	if err != nil {
		return state.Run{}, errors.Wrapf(err, "problem finding run by task arn [%s]", taskArn)
	}
	if runs.Total > 0 && len(runs.Runs) > 0 {
		return runs.Runs[0], nil
	}
	return state.Run{}, errors.Errorf("no run found for [%s]", taskArn)
}
<|endoftext|>


<commit_before>package controllers

import (
    "encoding/json"
    "errors"
    "go-task/models"
    "strconv"
    "strings"

    "github.com/astaxie/beego"
    "github.com/astaxie/beego/validation"
)

// Task
type TaskController struct {
    beego.Controller
}

func (this *TaskController) URLMapping() {
    this.Mapping("Post", this.Post)
    this.Mapping("GetOne", this.GetOne)
    this.Mapping("GetAll", this.GetAll)
    this.Mapping("Put", this.Put)
    this.Mapping("Delete", this.Delete)
}

// @Title Create Task
// @Description Create Task
// @Param body    body  models.Task true    "body for Task content"
// @Success 200 {int} models.Task.Id
// @Failure 403 body is empty
// @router / [post]
func (this *TaskController) Post() {
    var v models.Task
    json.Unmarshal(this.Ctx.Input.RequestBody, &v)

    valid := validation.Validation{}
    valid.Valid(&v)
    if valid.HasErrors() {
        errorMessages := make(map[string]string)
        for _, err := range valid.Errors {
            errorMessages[err.Key] = err.Message
        }
        this.Data["json"] = errorMessages
        this.Ctx.Output.SetStatus(422)
        this.ServeJSON()
        return
    }

    id, err := models.StoreTask(&v)
    if err == nil {
        this.Data["json"] = map[string]int64{"id": id}
        this.Ctx.Output.SetStatus(201)
    } else {
        this.Data["json"] = err.Error()
        this.Ctx.Output.SetStatus(422)
    }
    this.ServeJSON()
}

// @Title Get
// @Description get Task by id
// @Param id    path  string  true    "The key for staticblock"
// @Success 200 {object} models.Task
// @Failure 403 :id is empty
// @router /:id [get]
func (this *TaskController) GetOne() {
    idStr := this.Ctx.Input.Param(":id")
    id, _ := strconv.Atoi(idStr)
    t, err := models.ShowTask(id)
    if err != nil {
        this.Data["json"] = map[string]string{"error": "Task Not Found."}
        this.Ctx.Output.SetStatus(404)
    } else {
        this.Data["json"] = t
    }
    this.ServeJSON()
}

// @Title Get All
// @Description get Task
// @Param query query string  false "Filter. e.g. col1:v1,col2:v2 ..."
// @Param fields  query string  false "Fields returned. e.g. col1,col2 ..."
// @Param sortby  query string  false "Sorted-by fields. e.g. col1,col2 ..."
// @Param order query string  false "Order corresponding to each sortby field, if single value, apply to all sortby fields. e.g. desc,asc ..."
// @Param limit query string  false "Limit the size of result set. Must be an integer"
// @Param offset  query string  false "Start position of result set. Must be an integer"
// @Success 200 {object} models.Task
// @Failure 403
// @router / [get]
func (this *TaskController) GetAll() {
    var fields []string
    var sortby []string
    var order []string
    var query map[string]string = make(map[string]string)
    var limit int64 = 10
    var offset int64 = 0

    // fields: col1,col2,entity.col3
    if v := this.GetString("fields"); v != "" {
        fields = strings.Split(v, ",")
    }
    // limit: 10 (default is 10)
    // if v, err := this.GetInt("limit"); err == nil {
    //  limit = v
    // }
    // // offset: 0 (default is 0)
    // if v, err := this.GetInt("offset"); err == nil {
    //  offset = v
    // }
    // sortby: col1,col2
    if v := this.GetString("sortby"); v != "" {
        sortby = strings.Split(v, ",")
    }
    // order: desc,asc
    if v := this.GetString("order"); v != "" {
        order = strings.Split(v, ",")
    }
    // query: k:v,k:v
    if v := this.GetString("query"); v != "" {
        for _, cond := range strings.Split(v, ",") {
            kv := strings.Split(cond, ":")
            if len(kv) != 2 {
                this.Data["json"] = errors.New("Error: invalid query key/value pair")
                this.ServeJSON()
                return
            }
            k, v := kv[0], kv[1]
            query[k] = v
        }
    }

    l, err := models.GetAllTask(query, fields, sortby, order, offset, limit)
    if err != nil {
        this.Data["json"] = err.Error()
    } else {
        this.Data["json"] = l
    }
    this.ServeJSON()
}

// @Title Update
// @Description update the Task
// @Param id    path  string  true    "The id you want to update"
// @Param body    body  models.Task true    "body for Task content"
// @Success 200 {object} models.Task
// @Failure 403 :id is not int
// @router /:id [put]
func (this *TaskController) Put() {
    idStr := this.Ctx.Input.Param(":id")
    id, _ := strconv.Atoi(idStr)
    v := models.Task{Id: id}
    json.Unmarshal(this.Ctx.Input.RequestBody, &v)
    if err := models.UpdateTask(&v); err == nil {
        this.Data["json"] = "OK"
    } else {
        this.Data["json"] = err.Error()
    }
    this.ServeJSON()
}

// @Title Delete
// @Description delete the Task
// @Param id    path  string  true    "The id you want to delete"
// @Success 200 {string} delete success!
// @Failure 403 id is empty
// @router /:id [delete]
func (this *TaskController) Delete() {
    idStr := this.Ctx.Input.Param(":id")
    id, _ := strconv.Atoi(idStr)
    if err := models.DestroyTask(id); err == nil {
        this.Data["json"] = "OK"
    } else {
        this.Data["json"] = err.Error()
    }
    this.ServeJSON()
}
<commit_msg>join with User & Status<commit_after>package controllers

import (
    "encoding/json"
    "errors"
    "go-task/models"
    "strconv"
    "strings"

    "github.com/astaxie/beego"
    "github.com/astaxie/beego/validation"
)

// Task
type TaskController struct {
    beego.Controller
}

func (this *TaskController) URLMapping() {
    this.Mapping("Post", this.Post)
    this.Mapping("GetOne", this.GetOne)
    this.Mapping("GetAll", this.GetAll)
    this.Mapping("Put", this.Put)
    this.Mapping("Delete", this.Delete)
}

// @Title Create Task
// @Description Create Task
// @Param body    body  models.Task true    "body for Task content"
// @Success 200 {int} models.Task.Id
// @Failure 403 body is empty
// @router / [post]
func (this *TaskController) Post() {
    var v models.Task
    json.Unmarshal(this.Ctx.Input.RequestBody, &v)

    valid := validation.Validation{}
    valid.Valid(&v)
    if valid.HasErrors() {
        errorMessages := make(map[string]string)
        for _, err := range valid.Errors {
            errorMessages[err.Key] = err.Message
        }
        this.Data["json"] = errorMessages
        this.Ctx.Output.SetStatus(422)
        this.ServeJSON()
        return
    }
    beego.Info(&v)
    id, err := models.StoreTask(&v)
    if err == nil {
        this.Data["json"] = map[string]int64{"id": id}
        this.Ctx.Output.SetStatus(201)
    } else {
        this.Data["json"] = err.Error()
        this.Ctx.Output.SetStatus(422)
    }
    this.ServeJSON()
}

// @Title Get
// @Description get Task by id
// @Param id    path  string  true    "The key for staticblock"
// @Success 200 {object} models.Task
// @Failure 403 :id is empty
// @router /:id [get]
func (this *TaskController) GetOne() {
    idStr := this.Ctx.Input.Param(":id")
    id, _ := strconv.Atoi(idStr)
    t, err := models.ShowTask(id)
    if err != nil {
        this.Data["json"] = map[string]string{"error": "Task Not Found."}
        this.Ctx.Output.SetStatus(404)
    } else {
        this.Data["json"] = t
    }
    this.ServeJSON()
}

// @Title Get All
// @Description get Task
// @Param query query string  false "Filter. e.g. col1:v1,col2:v2 ..."
// @Param fields  query string  false "Fields returned. e.g. col1,col2 ..."
// @Param sortby  query string  false "Sorted-by fields. e.g. col1,col2 ..."
// @Param order query string  false "Order corresponding to each sortby field, if single value, apply to all sortby fields. e.g. desc,asc ..."
// @Param limit query string  false "Limit the size of result set. Must be an integer"
// @Param offset  query string  false "Start position of result set. Must be an integer"
// @Success 200 {object} models.Task
// @Failure 403
// @router / [get]
func (this *TaskController) GetAll() {
    var fields []string
    var sortby []string
    var order []string
    var query map[string]string = make(map[string]string)
    var limit int64 = 10
    var offset int64 = 0

    // fields: col1,col2,entity.col3
    if v := this.GetString("fields"); v != "" {
        fields = strings.Split(v, ",")
    }
    // limit: 10 (default is 10)
    // if v, err := this.GetInt("limit"); err == nil {
    //  limit = v
    // }
    // // offset: 0 (default is 0)
    // if v, err := this.GetInt("offset"); err == nil {
    //  offset = v
    // }
    // sortby: col1,col2
    if v := this.GetString("sortby"); v != "" {
        sortby = strings.Split(v, ",")
    }
    // order: desc,asc
    if v := this.GetString("order"); v != "" {
        order = strings.Split(v, ",")
    }
    // query: k:v,k:v
    if v := this.GetString("query"); v != "" {
        for _, cond := range strings.Split(v, ",") {
            kv := strings.Split(cond, ":")
            if len(kv) != 2 {
                this.Data["json"] = errors.New("Error: invalid query key/value pair")
                this.ServeJSON()
                return
            }
            k, v := kv[0], kv[1]
            query[k] = v
        }
    }

    l, err := models.GetAllTask(query, fields, sortby, order, offset, limit)
    if err != nil {
        this.Data["json"] = err.Error()
    } else {
        this.Data["json"] = l
    }
    this.ServeJSON()
}

// @Title Update
// @Description update the Task
// @Param id    path  string  true    "The id you want to update"
// @Param body    body  models.Task true    "body for Task content"
// @Success 200 {object} models.Task
// @Failure 403 :id is not int
// @router /:id [put]
func (this *TaskController) Put() {
    idStr := this.Ctx.Input.Param(":id")
    id, _ := strconv.Atoi(idStr)
    v := models.Task{Id: id}
    json.Unmarshal(this.Ctx.Input.RequestBody, &v)
    if err := models.UpdateTask(&v); err == nil {
        this.Data["json"] = "OK"
    } else {
        this.Data["json"] = err.Error()
    }
    this.ServeJSON()
}

// @Title Delete
// @Description delete the Task
// @Param id    path  string  true    "The id you want to delete"
// @Success 200 {string} delete success!
// @Failure 403 id is empty
// @router /:id [delete]
func (this *TaskController) Delete() {
    idStr := this.Ctx.Input.Param(":id")
    id, _ := strconv.Atoi(idStr)
    if err := models.DestroyTask(id); err == nil {
        this.Data["json"] = "OK"
    } else {
        this.Data["json"] = err.Error()
    }
    this.ServeJSON()
}
<|endoftext|>


<commit_before>package conversation

import (
	"fmt"
	messagebird "github.com/messagebird/go-rest-api/v8"
	"net/url"
	"strconv"
)

const (
	// apiRoot is the absolute URL of the Converstations API. All paths are
	// relative to apiRoot (e.g.
	// https://conversations.messagebird.com/v1/webhooks).
	apiRoot = "https://conversations.messagebird.com/v1"

	whatsappSandboxAPIRoot = "https://whatsapp-sandbox.messagebird.com/v1"

	// path is the path for the Conversation resource, relative to apiRoot.
	path = "conversations"

	// startConversationPath is the path for starting new conversation
	startConversationPath = "start"

	// contactPath is the path for fetching a collection of conversations by contact ID
	contactPath = "contact"

	// messagesPath is the path for the Message resource, relative to apiRoot
	// and path.
	messagesPath = "messages"

	// sendMessagePath is the path for creating the Message resource relative to apiRoot
	sendMessagePath = "send"

	// webhooksPath is the path for the Webhook resource, relative to apiRoot.
	webhooksPath = "webhooks"
)

type GetRequest interface {
	QueryParams() string
}

// PaginationRequest can be used to set pagination options in List().
type PaginationRequest struct {
	Limit, Offset int
}

func (lro *PaginationRequest) QueryParams() string {
	if lro == nil {
		return ""
	}

	query := url.Values{}
	query.Set("limit", strconv.Itoa(lro.Limit))
	query.Set("offset", strconv.Itoa(lro.Offset))

	return query.Encode()
}

// request does the exact same thing as Client.Request. It does, however,
// prefix the path with the Conversation API's root. This ensures the client
// doesn't "handle" this for us: by default, it uses the REST API.
func request(c messagebird.MessageBirdClient, v interface{}, method, path string, data interface{}) error {
	var root string
	if c.IsFeatureEnabled(messagebird.FeatureConversationsAPIWhatsAppSandbox) {
		root = whatsappSandboxAPIRoot
	} else {
		root = apiRoot
	}
	return c.Request(v, method, fmt.Sprintf("%s/%s", root, path), data)
}
<commit_msg>remove GetRequest interface<commit_after>package conversation

import (
	"fmt"
	messagebird "github.com/messagebird/go-rest-api/v8"
	"net/url"
	"strconv"
)

const (
	// apiRoot is the absolute URL of the Converstations API. All paths are
	// relative to apiRoot (e.g.
	// https://conversations.messagebird.com/v1/webhooks).
	apiRoot = "https://conversations.messagebird.com/v1"

	whatsappSandboxAPIRoot = "https://whatsapp-sandbox.messagebird.com/v1"

	// path is the path for the Conversation resource, relative to apiRoot.
	path = "conversations"

	// startConversationPath is the path for starting new conversation
	startConversationPath = "start"

	// contactPath is the path for fetching a collection of conversations by contact ID
	contactPath = "contact"

	// messagesPath is the path for the Message resource, relative to apiRoot
	// and path.
	messagesPath = "messages"

	// sendMessagePath is the path for creating the Message resource relative to apiRoot
	sendMessagePath = "send"

	// webhooksPath is the path for the Webhook resource, relative to apiRoot.
	webhooksPath = "webhooks"
)

// PaginationRequest can be used to set pagination options in List().
type PaginationRequest struct {
	Limit, Offset int
}

func (lro *PaginationRequest) QueryParams() string {
	if lro == nil {
		return ""
	}

	query := url.Values{}
	query.Set("limit", strconv.Itoa(lro.Limit))
	query.Set("offset", strconv.Itoa(lro.Offset))

	return query.Encode()
}

// request does the exact same thing as Client.Request. It does, however,
// prefix the path with the Conversation API's root. This ensures the client
// doesn't "handle" this for us: by default, it uses the REST API.
func request(c messagebird.MessageBirdClient, v interface{}, method, path string, data interface{}) error {
	var root string
	if c.IsFeatureEnabled(messagebird.FeatureConversationsAPIWhatsAppSandbox) {
		root = whatsappSandboxAPIRoot
	} else {
		root = apiRoot
	}
	return c.Request(v, method, fmt.Sprintf("%s/%s", root, path), data)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
	"github.com/coreos/etcd/etcdserver/etcdserverpb"
	"github.com/gtfierro/cs262-project/common"
	"github.com/pkg/profile"
	"golang.org/x/net/context"
	"os"
	"runtime/trace"
	"time"
)

// config flags
var configfile = flag.String("c", "config.ini", "Path to configuration file")

func init() {
	// set up logging
	log.SetOutput(os.Stderr)
}

func main() {
	flag.Parse()
	// configure logging instance
	config, configLogmsg := common.LoadConfig(*configfile)
	common.SetupLogging(config)
	log.Info(configLogmsg)

	if config.Debug.Enable {
		var p interface {
			Stop()
		}
		switch config.Debug.ProfileType {
		case "cpu", "profile":
			p = profile.Start(profile.CPUProfile, profile.ProfilePath("."))
		case "block":
			p = profile.Start(profile.BlockProfile, profile.ProfilePath("."))
		case "mem":
			p = profile.Start(profile.MemProfile, profile.ProfilePath("."))
		case "trace":
			f, err := os.Create("trace.out")
			if err != nil {
				log.Fatal(err)
			}
			trace.Start(f)
		}
		time.AfterFunc(time.Duration(config.Debug.ProfileLength)*time.Second, func() {
			if p != nil {
				p.Stop()
			}
			if config.Debug.ProfileType == "trace" {
				trace.Stop()
			}

			os.Exit(0)
		})
	}

	server := NewServer(config)
	logStartKey := server.rebuildIfNecessary()
	go server.handleLeadership()
	go server.handleBrokerMessages()
	go server.monitorLog(logStartKey)
	go server.monitorGeneralConnections()
	server.listenAndDispatch()
}
<commit_msg>remove unused imports<commit_after>package main

import (
	"flag"
	log "github.com/Sirupsen/logrus"
	"github.com/gtfierro/cs262-project/common"
	"github.com/pkg/profile"
	"os"
	"runtime/trace"
	"time"
)

// config flags
var configfile = flag.String("c", "config.ini", "Path to configuration file")

func init() {
	// set up logging
	log.SetOutput(os.Stderr)
}

func main() {
	flag.Parse()
	// configure logging instance
	config, configLogmsg := common.LoadConfig(*configfile)
	common.SetupLogging(config)
	log.Info(configLogmsg)

	if config.Debug.Enable {
		var p interface {
			Stop()
		}
		switch config.Debug.ProfileType {
		case "cpu", "profile":
			p = profile.Start(profile.CPUProfile, profile.ProfilePath("."))
		case "block":
			p = profile.Start(profile.BlockProfile, profile.ProfilePath("."))
		case "mem":
			p = profile.Start(profile.MemProfile, profile.ProfilePath("."))
		case "trace":
			f, err := os.Create("trace.out")
			if err != nil {
				log.Fatal(err)
			}
			trace.Start(f)
		}
		time.AfterFunc(time.Duration(config.Debug.ProfileLength)*time.Second, func() {
			if p != nil {
				p.Stop()
			}
			if config.Debug.ProfileType == "trace" {
				trace.Stop()
			}

			os.Exit(0)
		})
	}

	server := NewServer(config)
	logStartKey := server.rebuildIfNecessary()
	go server.handleLeadership()
	go server.handleBrokerMessages()
	go server.monitorLog(logStartKey)
	go server.monitorGeneralConnections()
	server.listenAndDispatch()
}
<|endoftext|>


<commit_before>package paillier

import (
	"crypto/rand"
	"errors"
	"math/big"
	"reflect"
	"testing"
)

var MockGenerateSafePrimes = func() (*big.Int, *big.Int, error) {
	return big.NewInt(887), big.NewInt(443), nil
}

func TestCreateThresholdKeyGenerator(t *testing.T) {
	var tests = map[string]struct {
		keyLength     int
		expectedError error
	}{
		"generator successfully created for 20 bit key length": {
			keyLength: 20,
		},
		"generator can't be created for 19 bit key length": {
			keyLength:     19,
			expectedError: errors.New("Public key bit length must be an even number"),
		},
		"generator successfully created for 18 bit key length": {
			keyLength: 18,
		},
		"generator can't be created for 17 bit key length": {
			keyLength:     17,
			expectedError: errors.New("Public key bit length must be an even number"),
		},
		"generator can't be created for 16 bit key length": {
			keyLength:     16,
			expectedError: errors.New("Public key bit length must be at least 18 bits"),
		},
	}

	for testName, test := range tests {
		t.Run(testName, func(t *testing.T) {
			gen, err := GetThresholdKeyGenerator(test.keyLength, 4, 3, rand.Reader)

			if !reflect.DeepEqual(test.expectedError, err) {
				t.Fatalf(
					"Unexpected error\nActual: %v\nExpected: %v",
					err,
					test.expectedError,
				)
			}

			if test.expectedError == nil && gen == nil {
				t.Fatal("Got nil generator, it should be successfully created")
			}
		})
	}
}

func TestGenerateSafePrimesOfThresholdKeyGenerator(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 4, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	p, q, err := tkh.generateSafePrimes()
	if err != nil {
		t.Error(err)
	}
	IsSafePrime(p, q, 16, t)
}

func TestInitPandP1(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 4, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.initPandP1()
	IsSafePrime(tkh.p, tkh.p1, 16, t)
}

func TestInitQandQ1(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 4, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.initQandQ1()
	IsSafePrime(tkh.q, tkh.q1, 16, t)
}

func TestInitPsAndQs(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 4, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.initPsAndQs()

	IsSafePrime(tkh.p, tkh.p1, 16, t)
	IsSafePrime(tkh.q, tkh.q1, 16, t)
}

func TestArePsAndQsGood(t *testing.T) {
	tkh := new(ThresholdKeyGenerator)
	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(887), b(443), b(839), b(419)
	if !tkh.arePsAndQsGood() {
		t.Fail()
	}

	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(887), b(443), b(887), b(443)
	if tkh.arePsAndQsGood() {
		t.Fail()
	}

	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(887), b(443), b(443), b(221)
	if tkh.arePsAndQsGood() {
		t.Fail()
	}
}

func TestInitShortcuts(t *testing.T) {
	tkh := new(ThresholdKeyGenerator)
	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(839), b(419), b(887), b(443)
	tkh.initShortcuts()

	if !reflect.DeepEqual(tkh.n, b(744193)) {
		t.Error("wrong n", tkh.n)
	}
	if !reflect.DeepEqual(tkh.m, b(185617)) {
		t.Error("wrong m", tkh.m)
	}
	if !reflect.DeepEqual(tkh.nm, new(big.Int).Mul(b(744193), b(185617))) {
		t.Error("wrong nm", tkh.nm)
	}
	if !reflect.DeepEqual(tkh.nSquare, new(big.Int).Mul(b(744193), b(744193))) {
		t.Error("wrong nSquare", tkh.nSquare)
	}
}

func TestInitD(t *testing.T) {
	tkh := new(ThresholdKeyGenerator)
	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(863), b(431), b(839), b(419)
	tkh.initShortcuts()
	tkh.initD()
	if n(tkh.d)%n(tkh.m) != 0 {
		t.Fail()
	}
	if n(tkh.d)%n(tkh.n) != 1 {
		t.Fail()
	}
}

func TestInitNumerialValues(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 4, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	if err := tkh.initNumerialValues(); err != nil {
		t.Error(err)
	}
}

func TestGenerateHidingPolynomial(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 15, 10, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	if err := tkh.initNumerialValues(); err != nil {
		t.Error(err)
	}
	if err := tkh.generateHidingPolynomial(); err != nil {
		t.Error(err)
	}
	p := tkh.polynomialCoefficients
	if len(p) != tkh.Threshold {
		t.Fail()
	}
	if n(p[0]) != n(tkh.d) {
		t.Fail()
	}
	for i := 1; i < len(p); i++ {
		if j := n(p[i]); j < 0 || j >= n(tkh.nm) {
			t.Fail()
		}
	}
}

func TestComputeShare(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 5, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.nm = b(103)
	tkh.polynomialCoefficients = []*big.Int{b(29), b(88), b(51)}
	share := tkh.computeShare(2)
	if n(share) != 31 {
		t.Error("error computing a share.  ", share)
	}
}

func TestCreateShares(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 100, 10, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	if err := tkh.initNumerialValues(); err != nil {
		t.Error(err)
	}
	if err := tkh.generateHidingPolynomial(); err != nil {
		t.Error(err)
	}

	if shares := tkh.createShares(); len(shares) != 100 {
		t.Fail()
	}
}

func TestCreateViArray(t *testing.T) {
	tkh := new(ThresholdKeyGenerator)
	tkh.TotalNumberOfDecryptionServers = 10
	tkh.v = b(54)
	tkh.nSquare = b(101 * 101)
	vArr := tkh.createViArray([]*big.Int{b(12), b(90), b(103)})
	exp := []*big.Int{b(6162), b(304), b(2728)}
	if !reflect.DeepEqual(vArr, exp) {
		t.Fail()
	}
}

func TestGetThresholdKeyGenerator(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(50, 10, 6, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	if err := tkh.initNumerialValues(); err != nil {
		t.Error(nil)
	}
}

func TestGenerate(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 10, 6, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tpks, err := tkh.Generate()
	if err != nil {
		t.Error(err)
		return
	}
	if len(tpks) != 10 {
		t.Fail()
	}
	for i, tpk := range tpks {
		if tpk.Id != i+1 {
			t.Fail()
		}
		if len(tpk.Vi) != 10 {
			t.Fail()
		}
		if tpk.N == nil {
			t.Fail()
		}
		if tpk.Threshold != 6 || tpk.TotalNumberOfDecryptionServers != 10 {
			t.Fail()
		}
	}
}

func TestComputeV(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 10, 6, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.n = b(1907 * 1823)
	tkh.nSquare = new(big.Int).Mul(tkh.n, tkh.n)
	for i := 0; i < 100; i++ {
		if err := tkh.computeV(); err != nil {
			t.Error(err)
		}
		if tkh.v.Cmp(tkh.nSquare) > 0 {
			t.Error("v is too big")
		}
		if tkh.v.Cmp(tkh.n) > 0 {
			return
		}
	}
	t.Error(`v has never been bigger than n.  It is suspicious in the sense<
	than it was taken in the range 0...n**2 -1
	`)
}
<commit_msg>Added test ensuring correct n,p,q bit length for threshold key-gen<commit_after>package paillier

import (
	"crypto/rand"
	"errors"
	"math/big"
	"reflect"
	"testing"
)

var MockGenerateSafePrimes = func() (*big.Int, *big.Int, error) {
	return big.NewInt(887), big.NewInt(443), nil
}

func TestCreateThresholdKeyGenerator(t *testing.T) {
	var tests = map[string]struct {
		keyLength     int
		expectedError error
	}{
		"generator successfully created for 20 bit key length": {
			keyLength: 20,
		},
		"generator can't be created for 19 bit key length": {
			keyLength:     19,
			expectedError: errors.New("Public key bit length must be an even number"),
		},
		"generator successfully created for 18 bit key length": {
			keyLength: 18,
		},
		"generator can't be created for 17 bit key length": {
			keyLength:     17,
			expectedError: errors.New("Public key bit length must be an even number"),
		},
		"generator can't be created for 16 bit key length": {
			keyLength:     16,
			expectedError: errors.New("Public key bit length must be at least 18 bits"),
		},
	}

	for testName, test := range tests {
		t.Run(testName, func(t *testing.T) {
			gen, err := GetThresholdKeyGenerator(test.keyLength, 4, 3, rand.Reader)

			if !reflect.DeepEqual(test.expectedError, err) {
				t.Fatalf(
					"Unexpected error\nActual: %v\nExpected: %v",
					err,
					test.expectedError,
				)
			}

			if test.expectedError == nil && gen == nil {
				t.Fatal("Got nil generator, it should be successfully created")
			}
		})
	}
}

func TestGenerateNumbersOfCorrectBitLength(t *testing.T) {
	var tests = map[string]struct {
		keyLength             int
		expectedPrimeLength   int
		expectedModulusLength int
	}{
		"key bit length = 32, prime bit length = 16, modulus bit length = 32": {
			keyLength:             32,
			expectedPrimeLength:   16,
			expectedModulusLength: 32,
		},
		"key bit length = 64, prime bit length = 32, modulus bit length = 64": {
			keyLength:             64,
			expectedPrimeLength:   32,
			expectedModulusLength: 64,
		},
		"key bit length = 128, prime bit length = 64, modulus bit length = 128": {
			keyLength:             128,
			expectedPrimeLength:   64,
			expectedModulusLength: 128,
		},
	}

	for testName, test := range tests {
		t.Run(testName, func(t *testing.T) {
			gen, err := GetThresholdKeyGenerator(test.keyLength, 10, 6, rand.Reader)
			if err != nil {
				t.Fatal(err)
			}

			err = gen.initNumerialValues()
			if err != nil {
				t.Fatal(err)
			}

			if gen.p.BitLen() != test.expectedPrimeLength {
				t.Fatalf(
					"Unexpected prime bit length\nExpected %v\n Actual %v",
					test.expectedModulusLength,
					gen.p.BitLen(),
				)
			}

			if gen.q.BitLen() != test.expectedPrimeLength {
				t.Fatalf(
					"Unexpected prime bit length\nExpected %v\n Actual %v",
					test.expectedModulusLength,
					gen.q.BitLen(),
				)
			}

			if gen.n.BitLen() != test.expectedModulusLength {
				t.Fatalf(
					"Unexpected modulus bit length\nExpected %v\n Actual %v",
					test.expectedModulusLength,
					gen.n.BitLen(),
				)
			}

			if new(big.Int).Mul(gen.p, gen.q).Cmp(gen.n) != 0 {
				t.Fatal("n != pq")
			}
		})
	}
}

func TestInitPandP1(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 4, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.initPandP1()
	IsSafePrime(tkh.p, tkh.p1, 16, t)
}

func TestInitQandQ1(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 4, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.initQandQ1()
	IsSafePrime(tkh.q, tkh.q1, 16, t)
}

func TestInitPsAndQs(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 4, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.initPsAndQs()

	IsSafePrime(tkh.p, tkh.p1, 16, t)
	IsSafePrime(tkh.q, tkh.q1, 16, t)
}

func TestArePsAndQsGood(t *testing.T) {
	tkh := new(ThresholdKeyGenerator)
	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(887), b(443), b(839), b(419)
	if !tkh.arePsAndQsGood() {
		t.Fail()
	}

	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(887), b(443), b(887), b(443)
	if tkh.arePsAndQsGood() {
		t.Fail()
	}

	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(887), b(443), b(443), b(221)
	if tkh.arePsAndQsGood() {
		t.Fail()
	}
}

func TestInitShortcuts(t *testing.T) {
	tkh := new(ThresholdKeyGenerator)
	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(839), b(419), b(887), b(443)
	tkh.initShortcuts()

	if !reflect.DeepEqual(tkh.n, b(744193)) {
		t.Error("wrong n", tkh.n)
	}
	if !reflect.DeepEqual(tkh.m, b(185617)) {
		t.Error("wrong m", tkh.m)
	}
	if !reflect.DeepEqual(tkh.nm, new(big.Int).Mul(b(744193), b(185617))) {
		t.Error("wrong nm", tkh.nm)
	}
	if !reflect.DeepEqual(tkh.nSquare, new(big.Int).Mul(b(744193), b(744193))) {
		t.Error("wrong nSquare", tkh.nSquare)
	}
}

func TestInitD(t *testing.T) {
	tkh := new(ThresholdKeyGenerator)
	tkh.p, tkh.p1, tkh.q, tkh.q1 = b(863), b(431), b(839), b(419)
	tkh.initShortcuts()
	tkh.initD()
	if n(tkh.d)%n(tkh.m) != 0 {
		t.Fail()
	}
	if n(tkh.d)%n(tkh.n) != 1 {
		t.Fail()
	}
}

func TestInitNumerialValues(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 4, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	if err := tkh.initNumerialValues(); err != nil {
		t.Error(err)
	}
}

func TestGenerateHidingPolynomial(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 15, 10, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	if err := tkh.initNumerialValues(); err != nil {
		t.Error(err)
	}
	if err := tkh.generateHidingPolynomial(); err != nil {
		t.Error(err)
	}
	p := tkh.polynomialCoefficients
	if len(p) != tkh.Threshold {
		t.Fail()
	}
	if n(p[0]) != n(tkh.d) {
		t.Fail()
	}
	for i := 1; i < len(p); i++ {
		if j := n(p[i]); j < 0 || j >= n(tkh.nm) {
			t.Fail()
		}
	}
}

func TestComputeShare(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 5, 3, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.nm = b(103)
	tkh.polynomialCoefficients = []*big.Int{b(29), b(88), b(51)}
	share := tkh.computeShare(2)
	if n(share) != 31 {
		t.Error("error computing a share.  ", share)
	}
}

func TestCreateShares(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 100, 10, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	if err := tkh.initNumerialValues(); err != nil {
		t.Error(err)
	}
	if err := tkh.generateHidingPolynomial(); err != nil {
		t.Error(err)
	}

	if shares := tkh.createShares(); len(shares) != 100 {
		t.Fail()
	}
}

func TestCreateViArray(t *testing.T) {
	tkh := new(ThresholdKeyGenerator)
	tkh.TotalNumberOfDecryptionServers = 10
	tkh.v = b(54)
	tkh.nSquare = b(101 * 101)
	vArr := tkh.createViArray([]*big.Int{b(12), b(90), b(103)})
	exp := []*big.Int{b(6162), b(304), b(2728)}
	if !reflect.DeepEqual(vArr, exp) {
		t.Fail()
	}
}

func TestGetThresholdKeyGenerator(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(50, 10, 6, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	if err := tkh.initNumerialValues(); err != nil {
		t.Error(nil)
	}
}

func TestGenerate(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 10, 6, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tpks, err := tkh.Generate()
	if err != nil {
		t.Error(err)
		return
	}
	if len(tpks) != 10 {
		t.Fail()
	}
	for i, tpk := range tpks {
		if tpk.Id != i+1 {
			t.Fail()
		}
		if len(tpk.Vi) != 10 {
			t.Fail()
		}
		if tpk.N == nil {
			t.Fail()
		}
		if tpk.Threshold != 6 || tpk.TotalNumberOfDecryptionServers != 10 {
			t.Fail()
		}
	}
}

func TestComputeV(t *testing.T) {
	tkh, err := GetThresholdKeyGenerator(32, 10, 6, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	tkh.n = b(1907 * 1823)
	tkh.nSquare = new(big.Int).Mul(tkh.n, tkh.n)
	for i := 0; i < 100; i++ {
		if err := tkh.computeV(); err != nil {
			t.Error(err)
		}
		if tkh.v.Cmp(tkh.nSquare) > 0 {
			t.Error("v is too big")
		}
		if tkh.v.Cmp(tkh.n) > 0 {
			return
		}
	}
	t.Error(`v has never been bigger than n.  It is suspicious in the sense<
	than it was taken in the range 0...n**2 -1
	`)
}
<|endoftext|>