<commit_before>package main

import (
	"fmt"
	"strings"
	"time"
	"net/http"
	"encoding/json"
	"github.com/rancher/go-rancher/client"
	"os"
	"github.com/ahaynssen/slack-go-webhook"
)

type Rancher struct {
	Url        string `json:"url"`
	AccessKey  string `json:"access_key"`
	SecretKey  string `json:"secret_key"`
	Service    string `json:"service"`
	Image      string `json:"docker_image"`
	StartFirst bool   `json:"start_first"`
	Confirm    bool   `json:"confirm"`
	Timeout    int    `json:"timeout"`
}

func main() {
	fmt.Println("Started upgrade server")
	http.HandleFunc("/upgrade/", requestSuccess)
	http.ListenAndServe(":8080", nil)
}

func requestSuccess(w http.ResponseWriter, r *http.Request){
	vargs := Rancher{
		StartFirst: false,
		Timeout:    30,
	}

	err := json.NewDecoder(r.Body).Decode(&vargs)
	if err != nil {
		fmt.Printf(err.Error())
		Error(w, err.Error(), 400)
		return
	} else {
		if len(vargs.Url) == 0 || len(vargs.AccessKey) == 0 || len(vargs.SecretKey) == 0 {
			Error(w, "Rancher credentials not set", 400)
			return
		}
		go upgradeRancher(vargs)
		w.WriteHeader(200)
		return
	}
}

func upgradeRancher(vargs Rancher) {
	var wantedImage, wantedVer string
	if !strings.HasPrefix(vargs.Image, "docker:") {
		vargs.Image = fmt.Sprintf("docker:%s", vargs.Image)
		parts := strings.Split(vargs.Image, ":")
		wantedImage = parts[1]
		wantedVer = parts[2]
	}
	parts := strings.Split(wantedImage, "/")
	vargs.Service = parts[1]

	rancher, err := client.NewRancherClient(&client.ClientOpts{
		Url:       vargs.Url,
		AccessKey: vargs.AccessKey,
		SecretKey: vargs.SecretKey,
	})

	if err != nil {
		fmt.Printf("Failed to create rancher client: %s\n", err)
		return
	}

	services, err := rancher.Service.List(&client.ListOpts{})
	if err != nil {
		fmt.Printf("Failed to list rancher services: %s\n", err)
		return
	}

	var upgradeLabel = os.Getenv("UPGRADE_LABEL")

	var foundImage, foundVer string
	for services != nil {
		for _, svc := range services.Data {
			if svc.LaunchConfig != nil {
				if _, ok := svc.LaunchConfig.Labels[upgradeLabel]; ok {
					parts := strings.Split(svc.LaunchConfig.ImageUuid, ":")
					foundImage = parts[1]
					foundVer = parts[2]
					if foundImage == wantedImage && ((foundVer < wantedVer) || (wantedVer == "latest")) {
						fmt.Println("Trying to upgrade...")
						err := doUpgrade(vargs, svc, rancher)
						if err != nil {
							fmt.Println(err.Error())
						} else {
							if (vargs.Confirm) {
								fmt.Println("Trying to confirm...")
								err := confirmUpgrade(vargs, svc, rancher)
								url := fmt.Sprintf("https://rancher.connectedfleet.io/env/%s/apps/stacks/%s", svc.AccountId, svc.EnvironmentId)
								if err != nil {
									fmt.Println("Unable to upgrade service %s: %s\n", vargs.Service, err.Error())
									message := fmt.Sprintf("Unable to confirm upgrade to `%s`.\nCheck status at <%[2]s|%[1]s>", vargs.Service, url)
									slackMessage("danger", message)
								} else {
									fmt.Printf("Upgraded %s to %s\n", svc.Name, vargs.Image)
									message := fmt.Sprintf("`%[1]s` has been successfully upgraded to `%[2]s`"+
											       "in Dev\n View in Rancher here: <%[3]s|%[1]s>", vargs.Service, wantedVer, url)
									slackMessage("good", message)

								}
							}
						}
						continue
					}
				}
			}
		}
		services, err = services.Next()
	}
}

func doUpgrade(vargs Rancher, service client.Service, rancher *client.RancherClient) (error) {
	service.LaunchConfig.ImageUuid = vargs.Image
	upgrade := &client.ServiceUpgrade{}
	upgrade.InServiceStrategy = &client.InServiceUpgradeStrategy{
		LaunchConfig:           service.LaunchConfig,
		SecondaryLaunchConfigs: service.SecondaryLaunchConfigs,
		StartFirst:             vargs.StartFirst,
	}
	upgrade.ToServiceStrategy = &client.ToServiceUpgradeStrategy{}
	_, err := rancher.Service.ActionUpgrade(&service, upgrade)
	return err
}

func confirmUpgrade(vargs Rancher, service client.Service, rancher *client.RancherClient) (error) {
	srv, err := retry(func() (interface{}, error) {
		s, e := rancher.Service.ById(service.Id)
		if e != nil {
			return nil, e
		}
		if s.State != "upgraded" {
			return nil, fmt.Errorf("Service not upgraded: %s\n", s.State)
		}
		return s, nil
	}, time.Duration(vargs.Timeout)*time.Second, 3*time.Second)
	if err != nil {
		return err
	}

	_, err = rancher.Service.ActionFinishupgrade(srv.(*client.Service))
	if err != nil {
		return err
	}
	fmt.Printf("Finished upgrade %s\n", vargs.Service)
	return err
}

type retryFunc func() (interface{}, error)

func retry(f retryFunc, timeout time.Duration, interval time.Duration) (interface{}, error) {
	finish := time.After(timeout)
	for {
		result, err := f()
		if err == nil {
			return result, nil
		}
		select {
		case <-finish:
			return nil, err
		case <-time.After(interval):
		}
	}
}

func Error(w http.ResponseWriter, error string, code int) {
	w.Header().Set("Content Type,", "text/plain; charset=UTF-8")
	w.WriteHeader(code)
	fmt.Fprint(w, error)
}

func slackMessage(status string, message string) {
	var webhookUrl = os.Getenv("SLACK_WEBHOOK")

	attachment := slack.Attachment {Color: &status, Text: &message}
	mrkdwn := "text"
	attachment.AddMrkdwn(&mrkdwn)
	payload := slack.Payload {
		Username: "rancher-updater-service",
		Attachments: []slack.Attachment{attachment},
	}
	printable, _ := json.Marshal(payload)
	fmt.Println(string(printable))
	err := slack.Send(webhookUrl, "", payload)
	if len(err) > 0 {
		fmt.Printf("error: %s\n", err)
	}
}<commit_msg>add missing space<commit_after>package main

import (
	"fmt"
	"strings"
	"time"
	"net/http"
	"encoding/json"
	"github.com/rancher/go-rancher/client"
	"os"
	"github.com/ahaynssen/slack-go-webhook"
)

type Rancher struct {
	Url        string `json:"url"`
	AccessKey  string `json:"access_key"`
	SecretKey  string `json:"secret_key"`
	Service    string `json:"service"`
	Image      string `json:"docker_image"`
	StartFirst bool   `json:"start_first"`
	Confirm    bool   `json:"confirm"`
	Timeout    int    `json:"timeout"`
}

func main() {
	fmt.Println("Started upgrade server")
	http.HandleFunc("/upgrade/", requestSuccess)
	http.ListenAndServe(":8080", nil)
}

func requestSuccess(w http.ResponseWriter, r *http.Request){
	vargs := Rancher{
		StartFirst: false,
		Timeout:    30,
	}

	err := json.NewDecoder(r.Body).Decode(&vargs)
	if err != nil {
		fmt.Printf(err.Error())
		Error(w, err.Error(), 400)
		return
	} else {
		if len(vargs.Url) == 0 || len(vargs.AccessKey) == 0 || len(vargs.SecretKey) == 0 {
			Error(w, "Rancher credentials not set", 400)
			return
		}
		go upgradeRancher(vargs)
		w.WriteHeader(200)
		return
	}
}

func upgradeRancher(vargs Rancher) {
	var wantedImage, wantedVer string
	if !strings.HasPrefix(vargs.Image, "docker:") {
		vargs.Image = fmt.Sprintf("docker:%s", vargs.Image)
		parts := strings.Split(vargs.Image, ":")
		wantedImage = parts[1]
		wantedVer = parts[2]
	}
	parts := strings.Split(wantedImage, "/")
	vargs.Service = parts[1]

	rancher, err := client.NewRancherClient(&client.ClientOpts{
		Url:       vargs.Url,
		AccessKey: vargs.AccessKey,
		SecretKey: vargs.SecretKey,
	})

	if err != nil {
		fmt.Printf("Failed to create rancher client: %s\n", err)
		return
	}

	services, err := rancher.Service.List(&client.ListOpts{})
	if err != nil {
		fmt.Printf("Failed to list rancher services: %s\n", err)
		return
	}

	var upgradeLabel = os.Getenv("UPGRADE_LABEL")

	var foundImage, foundVer string
	for services != nil {
		for _, svc := range services.Data {
			if svc.LaunchConfig != nil {
				if _, ok := svc.LaunchConfig.Labels[upgradeLabel]; ok {
					parts := strings.Split(svc.LaunchConfig.ImageUuid, ":")
					foundImage = parts[1]
					foundVer = parts[2]
					if foundImage == wantedImage && ((foundVer < wantedVer) || (wantedVer == "latest")) {
						fmt.Println("Trying to upgrade...")
						err := doUpgrade(vargs, svc, rancher)
						if err != nil {
							fmt.Println(err.Error())
						} else {
							if (vargs.Confirm) {
								fmt.Println("Trying to confirm...")
								err := confirmUpgrade(vargs, svc, rancher)
								url := fmt.Sprintf("https://rancher.connectedfleet.io/env/%s/apps/stacks/%s", svc.AccountId, svc.EnvironmentId)
								if err != nil {
									fmt.Println("Unable to upgrade service %s: %s\n", vargs.Service, err.Error())
									message := fmt.Sprintf("Unable to confirm upgrade to `%s`.\nCheck status at <%[2]s|%[1]s>", vargs.Service, url)
									slackMessage("danger", message)
								} else {
									fmt.Printf("Upgraded %s to %s\n", svc.Name, vargs.Image)
									message := fmt.Sprintf("`%[1]s` has been successfully upgraded to `%[2]s` "+
											       "in Dev\n View in Rancher here: <%[3]s|%[1]s>", vargs.Service, wantedVer, url)
									slackMessage("good", message)

								}
							}
						}
						continue
					}
				}
			}
		}
		services, err = services.Next()
	}
}

func doUpgrade(vargs Rancher, service client.Service, rancher *client.RancherClient) (error) {
	service.LaunchConfig.ImageUuid = vargs.Image
	upgrade := &client.ServiceUpgrade{}
	upgrade.InServiceStrategy = &client.InServiceUpgradeStrategy{
		LaunchConfig:           service.LaunchConfig,
		SecondaryLaunchConfigs: service.SecondaryLaunchConfigs,
		StartFirst:             vargs.StartFirst,
	}
	upgrade.ToServiceStrategy = &client.ToServiceUpgradeStrategy{}
	_, err := rancher.Service.ActionUpgrade(&service, upgrade)
	return err
}

func confirmUpgrade(vargs Rancher, service client.Service, rancher *client.RancherClient) (error) {
	srv, err := retry(func() (interface{}, error) {
		s, e := rancher.Service.ById(service.Id)
		if e != nil {
			return nil, e
		}
		if s.State != "upgraded" {
			return nil, fmt.Errorf("Service not upgraded: %s\n", s.State)
		}
		return s, nil
	}, time.Duration(vargs.Timeout)*time.Second, 3*time.Second)
	if err != nil {
		return err
	}

	_, err = rancher.Service.ActionFinishupgrade(srv.(*client.Service))
	if err != nil {
		return err
	}
	fmt.Printf("Finished upgrade %s\n", vargs.Service)
	return err
}

type retryFunc func() (interface{}, error)

func retry(f retryFunc, timeout time.Duration, interval time.Duration) (interface{}, error) {
	finish := time.After(timeout)
	for {
		result, err := f()
		if err == nil {
			return result, nil
		}
		select {
		case <-finish:
			return nil, err
		case <-time.After(interval):
		}
	}
}

func Error(w http.ResponseWriter, error string, code int) {
	w.Header().Set("Content Type,", "text/plain; charset=UTF-8")
	w.WriteHeader(code)
	fmt.Fprint(w, error)
}

func slackMessage(status string, message string) {
	var webhookUrl = os.Getenv("SLACK_WEBHOOK")

	attachment := slack.Attachment {Color: &status, Text: &message}
	mrkdwn := "text"
	attachment.AddMrkdwn(&mrkdwn)
	payload := slack.Payload {
		Username: "rancher-updater-service",
		Attachments: []slack.Attachment{attachment},
	}
	printable, _ := json.Marshal(payload)
	fmt.Println(string(printable))
	err := slack.Send(webhookUrl, "", payload)
	if len(err) > 0 {
		fmt.Printf("error: %s\n", err)
	}
}<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/docker/docker/opts"
	flag "github.com/docker/docker/pkg/mflag"

	dockerClient "github.com/fsouza/go-dockerclient"
)

var (
	SYSFS       string        = "/sys/fs/cgroup"
	PROCS       string        = "cgroup.procs"
	CGROUP_PROC string        = "/proc/%d/cgroup"
	INTERVAL    time.Duration = 1000
)

type Context struct {
	Args         []string
	Cgroups      []string
	AllCgroups   bool
	Logs         bool
	Notify       bool
	Name         string
	Env          bool
	Rm           bool
	Id           string
	NotifySocket string
	Cmd          *exec.Cmd
	Pid          int
	PidFile      string
	Client       *dockerClient.Client
	Detach       bool
}

func setupEnvironment(c *Context) {
	newArgs := []string{}
	if c.Notify && len(c.NotifySocket) > 0 {
		newArgs = append(newArgs, "-e", fmt.Sprintf("NOTIFY_SOCKET=%s", c.NotifySocket))
		newArgs = append(newArgs, "-v", fmt.Sprintf("%s:%s", c.NotifySocket, c.NotifySocket))
	} else {
		c.Notify = false
	}

	if c.Env {
		for _, val := range os.Environ() {
			if !strings.HasPrefix(val, "HOME=") && !strings.HasPrefix(val, "PATH=") {
				newArgs = append(newArgs, "-e", val)
			}
		}
	}

	if len(newArgs) > 0 {
		c.Args = append(newArgs, c.Args...)
	}
}

func parseContext(args []string) (*Context, error) {
	c := &Context{
		Logs:       true,
		AllCgroups: false,
	}

	flags := flag.NewFlagSet("systemd-docker", flag.ContinueOnError)

	flCgroups := opts.NewListOpts(nil)

	flags.StringVar(&c.PidFile, []string{"p", "-pid-file"}, "", "pipe file")
	flags.BoolVar(&c.Logs, []string{"l", "-logs"}, true, "pipe logs")
	flags.BoolVar(&c.Notify, []string{"n", "-notify"}, false, "setup systemd notify for container")
	flags.BoolVar(&c.Env, []string{"e", "-env"}, false, "inherit environment variable")
	flags.Var(&flCgroups, []string{"c", "-cgroups"}, "cgroups to take ownership of or 'all' for all cgroups available")

	err := flags.Parse(args)
	if err != nil {
		return nil, err
	}

	foundD := false
	var name string

	runArgs := flags.Args()
	if len(runArgs) == 0 || runArgs[0] != "run" {
		log.Println("Args:", runArgs)
		return nil, errors.New("run not found in arguments")
	}

	runArgs = runArgs[1:]
	newArgs := make([]string, 0, len(runArgs))

	for i, arg := range runArgs {
		/* This is tedious, but flag can't ignore unknown flags and I don't want to define them all */
		add := true

		switch {
		case arg == "-rm" || arg == "--rm":
			c.Rm = true
			add = false
		case arg == "-d" || arg == "-detach" || arg == "--detach":
			foundD = true
		case strings.HasPrefix(arg, "-name") || strings.HasPrefix(arg, "--name"):
			if strings.Contains(arg, "=") {
				name = strings.SplitN(arg, "=", 2)[1]
			} else if len(runArgs) > i+1 {
				name = runArgs[i+1]
			}
		}

		if add {
			newArgs = append(newArgs, arg)
		}
	}

	if !foundD {
		newArgs = append([]string{"-d"}, newArgs...)
	}

	c.Name = name
	c.NotifySocket = os.Getenv("NOTIFY_SOCKET")
	c.Args = newArgs
	c.Cgroups = flCgroups.GetAll()
	c.Detach = foundD

	for _, val := range c.Cgroups {
		if val == "all" {
			c.Cgroups = nil
			c.AllCgroups = true
			break
		}
	}

	setupEnvironment(c)

	return c, nil
}

func lookupNamedContainer(c *Context) error {
	client, err := getClient(c)
	if err != nil {
		return err
	}

	container, err := client.InspectContainer(c.Name)
	if _, ok := err.(*dockerClient.NoSuchContainer); ok {
		return nil
	}
	if err != nil || container == nil {
		return err
	}

	if container.State.Running {
		c.Id = container.ID
		c.Pid = container.State.Pid
		return nil
	} else if c.Rm {
		return client.RemoveContainer(dockerClient.RemoveContainerOptions{
			ID:    container.ID,
			Force: true,
		})
	} else {
		client, err := getClient(c)
		err = client.StartContainer(container.ID, nil)
		if err != nil {
			return err
		}

		container, err = client.InspectContainer(c.Name)
		if err != nil {
			return err
		}

		c.Id = container.ID
		c.Pid = container.State.Pid

		return nil
	}
}

func launchContainer(c *Context) error {
	args := append([]string{"run"}, c.Args...)
	c.Cmd = exec.Command("docker", args...)

	errorPipe, err := c.Cmd.StderrPipe()
	if err != nil {
		return err
	}

	outputPipe, err := c.Cmd.StdoutPipe()
	if err != nil {
		return err
	}

	err = c.Cmd.Start()
	if err != nil {
		return err
	}

	go io.Copy(os.Stderr, errorPipe)

	bytes, err := ioutil.ReadAll(outputPipe)
	if err != nil {
		return err
	}

	c.Id = strings.TrimSpace(string(bytes))

	err = c.Cmd.Wait()
	if err != nil {
		return err
	}

	if !c.Cmd.ProcessState.Success() {
		return err
	}

	c.Pid, err = getContainerPid(c)

	return err
}

func runContainer(c *Context) error {
	if len(c.Name) > 0 {
		err := lookupNamedContainer(c)
		if err != nil {
			return err
		}

	}

	if len(c.Id) == 0 {
		err := launchContainer(c)
		if err != nil {
			return err
		}
	}

	if c.Pid == 0 {
		return errors.New("Failed to launch container, pid is 0")
	}

	return nil
}

func getClient(c *Context) (*dockerClient.Client, error) {
	if c.Client != nil {
		return c.Client, nil
	}

	endpoint := os.Getenv("DOCKER_HOST")
	if len(endpoint) == 0 {
		endpoint = "unix:///var/run/docker.sock"
	}

	return dockerClient.NewVersionedClient(endpoint, "1.11")
}

func getContainerPid(c *Context) (int, error) {
	client, err := getClient(c)
	if err != nil {
		return 0, err
	}

	container, err := client.InspectContainer(c.Id)
	if err != nil {
		return 0, err
	}

	if container == nil {
		return 0, errors.New(fmt.Sprintf("Failed to find container %s", c.Id))
	}

	if container.State.Pid <= 0 {
		return 0, errors.New(fmt.Sprintf("Pid is %d for container %s", container.State.Pid, c.Id))
	}

	return container.State.Pid, nil
}

func getCgroupsForPid(pid int) (map[string]string, error) {
	file, err := os.Open(fmt.Sprintf(CGROUP_PROC, pid))
	if err != nil {
		return nil, err
	}

	ret := map[string]string{}

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.SplitN(scanner.Text(), ":", 3)
		if len(line) != 3 {
			continue
		}

		ret[line[1]] = line[2]
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return ret, nil
}

func constructCgroupPath(cgroupName string, cgroupPath string) string {
	return path.Join(SYSFS, strings.TrimPrefix(cgroupName, "name="), cgroupPath, PROCS)
}

func getCgroupPids(cgroupName string, cgroupPath string) ([]string, error) {
	ret := []string{}

	file, err := os.Open(constructCgroupPath(cgroupName, cgroupPath))
	if err != nil {
		return nil, err
	}

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		ret = append(ret, strings.TrimSpace(scanner.Text()))
	}

	if err = scanner.Err(); err != nil {
		return nil, err
	}

	return ret, nil
}

func writePid(pid string, path string) error {
	return ioutil.WriteFile(path, []byte(pid), 0644)
}

func moveCgroups(c *Context) (bool, error) {
	moved := false
	currentCgroups, err := getCgroupsForPid(os.Getpid())
	if err != nil {
		return false, err
	}

	containerCgroups, err := getCgroupsForPid(c.Pid)
	if err != nil {
		return false, err
	}

	var ns []string

	if c.AllCgroups || c.Cgroups == nil || len(c.Cgroups) == 0 {
		ns = make([]string, 0, len(containerCgroups))
		for value, _ := range containerCgroups {
			ns = append(ns, value)
		}
	} else {
		ns = c.Cgroups
	}

	for _, nsName := range ns {
		currentPath, ok := currentCgroups[nsName]
		if !ok {
			continue
		}

		containerPath, ok := containerCgroups[nsName]
		if !ok {
			continue
		}

		if currentPath == containerPath || containerPath == "/" {
			continue
		}

		pids, err := getCgroupPids(nsName, containerPath)
		if err != nil {
			return false, err
		}

		for _, pid := range pids {
			pidInt, err := strconv.Atoi(pid)
			if err != nil {
				continue
			}

			if pidDied(pidInt) {
				continue
			}

			currentFullPath := constructCgroupPath(nsName, currentPath)
			log.Printf("Moving pid %s to %s\n", pid, currentFullPath)
			err = writePid(pid, currentFullPath)
			if err != nil {
				return false, err
			}

			moved = true
		}
	}

	return moved, nil
}

func pidDied(pid int) bool {
	_, err := os.Stat(fmt.Sprintf("/proc/%d", pid))
	return os.IsNotExist(err)
}

func notify(c *Context) error {
	if pidDied(c.Pid) {
		return errors.New("Container exited before we could notify systemd")
	}

	if len(c.NotifySocket) == 0 {
		return nil
	}

	conn, err := net.Dial("unixgram", c.NotifySocket)
	if err != nil {
		return err
	}

	defer conn.Close()

	_, err = conn.Write([]byte(fmt.Sprintf("MAINPID=%d", c.Pid)))
	if err != nil {
		return err
	}

	if pidDied(c.Pid) {
		conn.Write([]byte(fmt.Sprintf("MAINPID=%d", os.Getpid())))
		return errors.New("Container exited before we could notify systemd")
	}

	if !c.Notify {
		_, err = conn.Write([]byte("READY=1"))
		if err != nil {
			return err
		}
	}

	return nil
}

func pidFile(c *Context) error {
	if len(c.PidFile) == 0 || c.Pid <= 0 {
		return nil
	}

	err := ioutil.WriteFile(c.PidFile, []byte(strconv.Itoa(c.Pid)), 0644)
	if err != nil {
		return err
	}

	return nil
}

func pipeLogs(c *Context) error {
	if !c.Logs {
		return nil
	}

	client, err := getClient(c)
	if err != nil {
		return err
	}

	err = client.Logs(dockerClient.LogsOptions{
		Container:    c.Id,
		Follow:       true,
		Stdout:       true,
		Stderr:       true,
		OutputStream: os.Stdout,
		ErrorStream:  os.Stderr,
	})

	return err
}

func keepAlive(c *Context) error {
	if c.Logs || c.Rm {
		client, err := getClient(c)
		if err != nil {
			return err
		}

		/* Good old polling... */
		for true {
			container, err := client.InspectContainer(c.Id)
			if err != nil {
				return err
			}

			if container.State.Running {
				time.Sleep(INTERVAL * time.Millisecond)
			} else {
				return nil
			}
		}
	}

	return nil
}

func rmContainer(c *Context) error {
	if !c.Rm {
		return nil
	}

	client, err := getClient(c)
	if err != nil {
		return err
	}

	return client.RemoveContainer(dockerClient.RemoveContainerOptions{
		ID:    c.Id,
		Force: true,
	})
}

func mainWithArgs(args []string) (*Context, error) {
	c, err := parseContext(args)
	if err != nil {
		return c, err
	}

	err = runContainer(c)
	if err != nil {
		return c, err
	}

	_, err = moveCgroups(c)
	if err != nil {
		return c, err
	}

	err = notify(c)
	if err != nil {
		return c, err
	}

	err = pidFile(c)
	if err != nil {
		return c, err
	}

	if !c.Detach {
		go pipeLogs(c)

		err = keepAlive(c)
		if err != nil {
			return c, err
		}

		err = rmContainer(c)
		if err != nil {
			return c, err
		}
	}

	return c, nil
}

func main() {
	_, err := mainWithArgs(os.Args[1:])
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Detach argument goes to systemd-docker, the Docker detach argument is always true as before<commit_after>package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/docker/docker/opts"
	flag "github.com/docker/docker/pkg/mflag"

	dockerClient "github.com/fsouza/go-dockerclient"
)

var (
	SYSFS       string        = "/sys/fs/cgroup"
	PROCS       string        = "cgroup.procs"
	CGROUP_PROC string        = "/proc/%d/cgroup"
	INTERVAL    time.Duration = 1000
)

type Context struct {
	Args         []string
	Cgroups      []string
	AllCgroups   bool
	Logs         bool
	Notify       bool
	Name         string
	Env          bool
	Rm           bool
	Id           string
	NotifySocket string
	Cmd          *exec.Cmd
	Pid          int
	PidFile      string
	Client       *dockerClient.Client
	Detach       bool
}

func setupEnvironment(c *Context) {
	newArgs := []string{}
	if c.Notify && len(c.NotifySocket) > 0 {
		newArgs = append(newArgs, "-e", fmt.Sprintf("NOTIFY_SOCKET=%s", c.NotifySocket))
		newArgs = append(newArgs, "-v", fmt.Sprintf("%s:%s", c.NotifySocket, c.NotifySocket))
	} else {
		c.Notify = false
	}

	if c.Env {
		for _, val := range os.Environ() {
			if !strings.HasPrefix(val, "HOME=") && !strings.HasPrefix(val, "PATH=") {
				newArgs = append(newArgs, "-e", val)
			}
		}
	}

	if len(newArgs) > 0 {
		c.Args = append(newArgs, c.Args...)
	}
}

func parseContext(args []string) (*Context, error) {
	c := &Context{
		Logs:       true,
		AllCgroups: false,
	}

	flags := flag.NewFlagSet("systemd-docker", flag.ContinueOnError)

	flCgroups := opts.NewListOpts(nil)

	flags.StringVar(&c.PidFile, []string{"p", "-pid-file"}, "", "pipe file")
	flags.BoolVar(&c.Logs, []string{"l", "-logs"}, true, "pipe logs")
	flags.BoolVar(&c.Notify, []string{"n", "-notify"}, false, "setup systemd notify for container")
	flags.BoolVar(&c.Env, []string{"e", "-env"}, false, "inherit environment variable")
	flags.Var(&flCgroups, []string{"c", "-cgroups"}, "cgroups to take ownership of or 'all' for all cgroups available")
	flags.BoolVar(&c.Detach, []string{"d", "-detach"}, false, "detach from container")

	err := flags.Parse(args)
	if err != nil {
		return nil, err
	}

	foundD := false
	var name string

	runArgs := flags.Args()
	if len(runArgs) == 0 || runArgs[0] != "run" {
		log.Println("Args:", runArgs)
		return nil, errors.New("run not found in arguments")
	}

	runArgs = runArgs[1:]
	newArgs := make([]string, 0, len(runArgs))

	for i, arg := range runArgs {
		/* This is tedious, but flag can't ignore unknown flags and I don't want to define them all */
		add := true

		switch {
		case arg == "-rm" || arg == "--rm":
			c.Rm = true
			add = false
		case arg == "-d" || arg == "-detach" || arg == "--detach":
			foundD = true
		case strings.HasPrefix(arg, "-name") || strings.HasPrefix(arg, "--name"):
			if strings.Contains(arg, "=") {
				name = strings.SplitN(arg, "=", 2)[1]
			} else if len(runArgs) > i+1 {
				name = runArgs[i+1]
			}
		}

		if add {
			newArgs = append(newArgs, arg)
		}
	}

	if !foundD {
		newArgs = append([]string{"-d"}, newArgs...)
	}

	c.Name = name
	c.NotifySocket = os.Getenv("NOTIFY_SOCKET")
	c.Args = newArgs
	c.Cgroups = flCgroups.GetAll()

	for _, val := range c.Cgroups {
		if val == "all" {
			c.Cgroups = nil
			c.AllCgroups = true
			break
		}
	}

	setupEnvironment(c)

	return c, nil
}

func lookupNamedContainer(c *Context) error {
	client, err := getClient(c)
	if err != nil {
		return err
	}

	container, err := client.InspectContainer(c.Name)
	if _, ok := err.(*dockerClient.NoSuchContainer); ok {
		return nil
	}
	if err != nil || container == nil {
		return err
	}

	if container.State.Running {
		c.Id = container.ID
		c.Pid = container.State.Pid
		return nil
	} else if c.Rm {
		return client.RemoveContainer(dockerClient.RemoveContainerOptions{
			ID:    container.ID,
			Force: true,
		})
	} else {
		client, err := getClient(c)
		err = client.StartContainer(container.ID, nil)
		if err != nil {
			return err
		}

		container, err = client.InspectContainer(c.Name)
		if err != nil {
			return err
		}

		c.Id = container.ID
		c.Pid = container.State.Pid

		return nil
	}
}

func launchContainer(c *Context) error {
	args := append([]string{"run"}, c.Args...)
	c.Cmd = exec.Command("docker", args...)

	errorPipe, err := c.Cmd.StderrPipe()
	if err != nil {
		return err
	}

	outputPipe, err := c.Cmd.StdoutPipe()
	if err != nil {
		return err
	}

	err = c.Cmd.Start()
	if err != nil {
		return err
	}

	go io.Copy(os.Stderr, errorPipe)

	bytes, err := ioutil.ReadAll(outputPipe)
	if err != nil {
		return err
	}

	c.Id = strings.TrimSpace(string(bytes))

	err = c.Cmd.Wait()
	if err != nil {
		return err
	}

	if !c.Cmd.ProcessState.Success() {
		return err
	}

	c.Pid, err = getContainerPid(c)

	return err
}

func runContainer(c *Context) error {
	if len(c.Name) > 0 {
		err := lookupNamedContainer(c)
		if err != nil {
			return err
		}

	}

	if len(c.Id) == 0 {
		err := launchContainer(c)
		if err != nil {
			return err
		}
	}

	if c.Pid == 0 {
		return errors.New("Failed to launch container, pid is 0")
	}

	return nil
}

func getClient(c *Context) (*dockerClient.Client, error) {
	if c.Client != nil {
		return c.Client, nil
	}

	endpoint := os.Getenv("DOCKER_HOST")
	if len(endpoint) == 0 {
		endpoint = "unix:///var/run/docker.sock"
	}

	return dockerClient.NewVersionedClient(endpoint, "1.11")
}

func getContainerPid(c *Context) (int, error) {
	client, err := getClient(c)
	if err != nil {
		return 0, err
	}

	container, err := client.InspectContainer(c.Id)
	if err != nil {
		return 0, err
	}

	if container == nil {
		return 0, errors.New(fmt.Sprintf("Failed to find container %s", c.Id))
	}

	if container.State.Pid <= 0 {
		return 0, errors.New(fmt.Sprintf("Pid is %d for container %s", container.State.Pid, c.Id))
	}

	return container.State.Pid, nil
}

func getCgroupsForPid(pid int) (map[string]string, error) {
	file, err := os.Open(fmt.Sprintf(CGROUP_PROC, pid))
	if err != nil {
		return nil, err
	}

	ret := map[string]string{}

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.SplitN(scanner.Text(), ":", 3)
		if len(line) != 3 {
			continue
		}

		ret[line[1]] = line[2]
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return ret, nil
}

func constructCgroupPath(cgroupName string, cgroupPath string) string {
	return path.Join(SYSFS, strings.TrimPrefix(cgroupName, "name="), cgroupPath, PROCS)
}

func getCgroupPids(cgroupName string, cgroupPath string) ([]string, error) {
	ret := []string{}

	file, err := os.Open(constructCgroupPath(cgroupName, cgroupPath))
	if err != nil {
		return nil, err
	}

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		ret = append(ret, strings.TrimSpace(scanner.Text()))
	}

	if err = scanner.Err(); err != nil {
		return nil, err
	}

	return ret, nil
}

func writePid(pid string, path string) error {
	return ioutil.WriteFile(path, []byte(pid), 0644)
}

func moveCgroups(c *Context) (bool, error) {
	moved := false
	currentCgroups, err := getCgroupsForPid(os.Getpid())
	if err != nil {
		return false, err
	}

	containerCgroups, err := getCgroupsForPid(c.Pid)
	if err != nil {
		return false, err
	}

	var ns []string

	if c.AllCgroups || c.Cgroups == nil || len(c.Cgroups) == 0 {
		ns = make([]string, 0, len(containerCgroups))
		for value, _ := range containerCgroups {
			ns = append(ns, value)
		}
	} else {
		ns = c.Cgroups
	}

	for _, nsName := range ns {
		currentPath, ok := currentCgroups[nsName]
		if !ok {
			continue
		}

		containerPath, ok := containerCgroups[nsName]
		if !ok {
			continue
		}

		if currentPath == containerPath || containerPath == "/" {
			continue
		}

		pids, err := getCgroupPids(nsName, containerPath)
		if err != nil {
			return false, err
		}

		for _, pid := range pids {
			pidInt, err := strconv.Atoi(pid)
			if err != nil {
				continue
			}

			if pidDied(pidInt) {
				continue
			}

			currentFullPath := constructCgroupPath(nsName, currentPath)
			log.Printf("Moving pid %s to %s\n", pid, currentFullPath)
			err = writePid(pid, currentFullPath)
			if err != nil {
				return false, err
			}

			moved = true
		}
	}

	return moved, nil
}

func pidDied(pid int) bool {
	_, err := os.Stat(fmt.Sprintf("/proc/%d", pid))
	return os.IsNotExist(err)
}

func notify(c *Context) error {
	if pidDied(c.Pid) {
		return errors.New("Container exited before we could notify systemd")
	}

	if len(c.NotifySocket) == 0 {
		return nil
	}

	conn, err := net.Dial("unixgram", c.NotifySocket)
	if err != nil {
		return err
	}

	defer conn.Close()

	_, err = conn.Write([]byte(fmt.Sprintf("MAINPID=%d", c.Pid)))
	if err != nil {
		return err
	}

	if pidDied(c.Pid) {
		conn.Write([]byte(fmt.Sprintf("MAINPID=%d", os.Getpid())))
		return errors.New("Container exited before we could notify systemd")
	}

	if !c.Notify {
		_, err = conn.Write([]byte("READY=1"))
		if err != nil {
			return err
		}
	}

	return nil
}

func pidFile(c *Context) error {
	if len(c.PidFile) == 0 || c.Pid <= 0 {
		return nil
	}

	err := ioutil.WriteFile(c.PidFile, []byte(strconv.Itoa(c.Pid)), 0644)
	if err != nil {
		return err
	}

	return nil
}

func pipeLogs(c *Context) error {
	if !c.Logs {
		return nil
	}

	client, err := getClient(c)
	if err != nil {
		return err
	}

	err = client.Logs(dockerClient.LogsOptions{
		Container:    c.Id,
		Follow:       true,
		Stdout:       true,
		Stderr:       true,
		OutputStream: os.Stdout,
		ErrorStream:  os.Stderr,
	})

	return err
}

func keepAlive(c *Context) error {
	if c.Logs || c.Rm {
		client, err := getClient(c)
		if err != nil {
			return err
		}

		/* Good old polling... */
		for true {
			container, err := client.InspectContainer(c.Id)
			if err != nil {
				return err
			}

			if container.State.Running {
				time.Sleep(INTERVAL * time.Millisecond)
			} else {
				return nil
			}
		}
	}

	return nil
}

func rmContainer(c *Context) error {
	if !c.Rm {
		return nil
	}

	client, err := getClient(c)
	if err != nil {
		return err
	}

	return client.RemoveContainer(dockerClient.RemoveContainerOptions{
		ID:    c.Id,
		Force: true,
	})
}

func mainWithArgs(args []string) (*Context, error) {
	c, err := parseContext(args)
	if err != nil {
		return c, err
	}

	err = runContainer(c)
	if err != nil {
		return c, err
	}

	_, err = moveCgroups(c)
	if err != nil {
		return c, err
	}

	err = notify(c)
	if err != nil {
		return c, err
	}

	err = pidFile(c)
	if err != nil {
		return c, err
	}

	if !c.Detach {
		go pipeLogs(c)

		err = keepAlive(c)
		if err != nil {
			return c, err
		}

		err = rmContainer(c)
		if err != nil {
			return c, err
		}
	}

	return c, nil
}

func main() {
	_, err := mainWithArgs(os.Args[1:])
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/vampirewalk/slackbot/Godeps/_workspace/src/github.com/nlopes/slack"
	"log"
	"net/url"
	"os"
	"strings"
)

func main() {
	token := os.Getenv("SLACK_TOKEN")
	if token == "" {
		log.Fatal("SLACK_TOKEN is required")
	}

	api := slack.New(token)
	bot := Bot{}

	rtm := api.NewRTM()
	go rtm.ManageConnection()

Loop:
	for {
		select {
		case msg := <-rtm.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.ConnectedEvent:
				bot.SetMe(ev.Info.User)

			case *slack.MessageEvent:
				// if msg contains github url
				repoURL := bot.ExtractGithubRepoURL(ev.Msg.Text)
				if repoURL != "" && !bot.IsMe(ev.Msg.User) {
					log.Printf("Received message \"%s\" from %s", ev.Msg.Text, ev.Msg.User)
					u, err := url.Parse(repoURL)
					if err != nil {
						rtm.SendMessage(rtm.NewOutgoingMessage("url error", ev.Channel))
					}
					results := strings.Split(u.Path, "/")
					state, err := bot.GetRepoState(results[2])
					if err != nil {
						rtm.SendMessage(rtm.NewOutgoingMessage("Parse error", ev.Channel))
					}
					report := bot.ReportRepoState(state)
					rtm.SendMessage(rtm.NewOutgoingMessage(report, ev.Channel))
				}

			case *slack.RTMError:
				log.Printf("Error: %s\n", ev.Error())

			case *slack.InvalidAuthEvent:
				log.Fatal("Invalid credentials")
				break Loop
			}
		}
	}
}
<commit_msg>Add log<commit_after>package main

import (
	"github.com/vampirewalk/slackbot/Godeps/_workspace/src/github.com/nlopes/slack"
	"log"
	"net/url"
	"os"
	"strings"
)

func main() {
	token := os.Getenv("SLACK_TOKEN")
	if token == "" {
		log.Fatal("SLACK_TOKEN is required")
	}

	api := slack.New(token)
	bot := Bot{}

	rtm := api.NewRTM()
	go rtm.ManageConnection()

Loop:
	for {
		select {
		case msg := <-rtm.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.ConnectedEvent:
				bot.SetMe(ev.Info.User)

			case *slack.MessageEvent:
				// if msg contains github url
				repoURL := bot.ExtractGithubRepoURL(ev.Msg.Text)
				if repoURL != "" && !bot.IsMe(ev.Msg.User) {
					log.Printf("Received message \"%s\" from %s", ev.Msg.Text, ev.Msg.User)
					u, err := url.Parse(repoURL)
					if err != nil {
						log.Printf("URL error")
						rtm.SendMessage(rtm.NewOutgoingMessage("URL error", ev.Channel))
					}
					results := strings.Split(u.Path, "/")
					state, err := bot.GetRepoState(results[2])
					if err != nil {
						log.Printf("Parse error")
						rtm.SendMessage(rtm.NewOutgoingMessage("Parse error", ev.Channel))
					}
					report := bot.ReportRepoState(state)
					log.Printf(report)
					rtm.SendMessage(rtm.NewOutgoingMessage(report, ev.Channel))
				}

			case *slack.RTMError:
				log.Printf("Error: %s\n", ev.Error())

			case *slack.InvalidAuthEvent:
				log.Fatal("Invalid credentials")
				break Loop
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

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

const (
	description = "horeb: Speaking in tongues via stdout"
)

func usage() {
	fmt.Printf("\n%s\n\n", description)
	flag.PrintDefaults()
	fmt.Println()
}

func main() {
	rand.Seed(time.Now().UnixNano())

	flag.Usage = usage

	color := flag.Bool("c", false, "colorize output")
	dump := flag.Bool("d", false, "print all Blocks")
	list := flag.Bool("l", false, "list all Block names and codepoint ranges")
	nchars := flag.Int("n", 30, "number of characters to print")
	flag.Parse()

	blocks := []string{"all"}
	if flag.NArg() > 0 {
		blocks = flag.Args()
		// special all value means all blocks
		if blocks[0] == "all" {
			// FIXME: remove blocks[0] ("all") to avoid error
			for k := range Blocks {
				blocks = append(blocks, k)
			}
		}
	}

	switch {
	case *color:
		// TODO: implement colorized output
		return
	case *list:
		printBlocks(false)
	case *dump:
		printBlocks(true)
	case len(blocks) == 1:
		if b, ok := Blocks[blocks[0]]; ok {
			b.PrintRandom(*nchars)
		} else {
			log.Fatalf("Unknown block name: %s\n", blocks[0])
		}
	case len(blocks) > 1:
		bm := map[string]*UnicodeBlock{}
		for _, b := range blocks {
			if val, ok := Blocks[b]; ok {
				bm[b] = val
			} else {
				log.Printf("Unknown block name: %s\n", b)
			}
		}
		if len(bm) > 0 {
			for i := 0; i < *nchars; i++ {
				block, err := RandomBlock(bm)
				if err != nil {
					log.Fatal(err)
				}
				fmt.Printf("%c ", block.RandomRune())
			}
		}
		fmt.Println()
	}
}
<commit_msg>Remove "all" value after use<commit_after>package main

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

const (
	description = "horeb: Speaking in tongues via stdout"
)

func usage() {
	fmt.Printf("\n%s\n\n", description)
	flag.PrintDefaults()
	fmt.Println()
}

func main() {
	rand.Seed(time.Now().UnixNano())

	flag.Usage = usage

	color := flag.Bool("c", false, "colorize output")
	dump := flag.Bool("d", false, "print all Blocks")
	list := flag.Bool("l", false, "list all Block names and codepoint ranges")
	nchars := flag.Int("n", 30, "number of characters to print")
	flag.Parse()

	blocks := []string{"all"}
	if flag.NArg() > 0 {
		blocks = flag.Args()
	}
	// special value means all blocks
	if blocks[0] == "all" {
		// remove "all" value after use
		blocks = blocks[:0]
		for k := range Blocks {
			blocks = append(blocks, k)
		}
	}

	switch {
	case *color:
		// TODO: implement colorized output
		return
	case *list:
		printBlocks(false)
	case *dump:
		printBlocks(true)
	case len(blocks) == 1:
		if b, ok := Blocks[blocks[0]]; ok {
			b.PrintRandom(*nchars)
		} else {
			log.Fatalf("Unknown block name: %s\n", blocks[0])
		}
	case len(blocks) > 1:
		bm := map[string]*UnicodeBlock{}
		for _, b := range blocks {
			if val, ok := Blocks[b]; ok {
				bm[b] = val
			} else {
				log.Printf("Unknown block name: %s\n", b)
			}
		}
		if len(bm) > 0 {
			for i := 0; i < *nchars; i++ {
				block, err := RandomBlock(bm)
				if err != nil {
					log.Fatal(err)
				}
				fmt.Printf("%c ", block.RandomRune())
			}
		}
		fmt.Println()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/koding/klient/collaboration"
	"github.com/koding/klient/command"
	"github.com/koding/klient/docker"
	"github.com/koding/klient/fs"
	"github.com/koding/klient/protocol"
	"github.com/koding/klient/terminal"
	"github.com/koding/klient/usage"

	"github.com/koding/klient/Godeps/_workspace/src/github.com/koding/kite"
	"github.com/koding/klient/Godeps/_workspace/src/github.com/koding/kite/config"
)

var (
	flagIP          = flag.String("ip", "", "Change public ip")
	flagPort        = flag.Int("port", 56789, "Change running port")
	flagVersion     = flag.Bool("version", false, "Show version and exit")
	flagEnvironment = flag.String("env", protocol.Environment, "Change environment")
	flagRegion      = flag.String("region", protocol.Region, "Change region")
	flagRegisterURL = flag.String("register-url", "", "Change register URL to kontrol")
	flagDebug       = flag.Bool("debug", false, "Debug mode")

	// update parameters
	flagUpdateInterval = flag.Duration("update-interval", time.Minute*5,
		"Change interval for checking for new updates")
	flagUpdateURL = flag.String("update-url",
		"https://s3.amazonaws.com/koding-klient/"+protocol.Environment+"/latest-version.txt",
		"Change update endpoint for latest version")

	VERSION = protocol.Version
	NAME    = protocol.Name

	// this is our main reference to count and measure metrics for the klient
	// we count only those methods, please add/remove methods here that will
	// reset the timer of a klient.
	usg = usage.NewUsage(map[string]bool{
		"fs.readDirectory":    true,
		"fs.glob":             true,
		"fs.readFile":         true,
		"fs.writeFile":        true,
		"fs.uniquePath":       true,
		"fs.getInfo":          true,
		"fs.setPermissions":   true,
		"fs.remove":           true,
		"fs.rename":           true,
		"fs.createDirectory":  true,
		"fs.move":             true,
		"fs.copy":             true,
		"webterm.getSessions": true,
		"webterm.connect":     true,
		"webterm.killSession": true,
		"exec":                true,
		"klient.share":        true,
		"klient.unshare":      true,
		"klient.shared":       true,
		"docker.create":       true,
		"docker.connect":      true,
		"docker.stop":         true,
		"docker.start":        true,
		"docker.remove":       true,
		"docker.list":         true,
	})

	// this is used to allow other users to call any klient method.
	collab = collaboration.New()

	// we also could use an atomic boolean this is simple for now.
	updating   = false
	updatingMu sync.Mutex // protects updating
)

func main() {
	flag.Parse()
	if *flagVersion {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	k := newKite()

	// Close the klient.db in any case. Corrupt db would be catastrophic
	defer collab.Close()

	k.Log.Info("Running as version %s", VERSION)
	k.Run()
}

func newKite() *kite.Kite {
	k := kite.New(NAME, VERSION)

	if *flagDebug {
		k.SetLogLevel(kite.DEBUG)
	}

	conf := config.MustGet()
	k.Config = conf
	k.Config.Port = *flagPort
	k.Config.Environment = *flagEnvironment
	k.Config.Region = *flagRegion
	k.Id = conf.Id // always boot up with the same id in the kite.key

	if *flagUpdateInterval < time.Minute {
		k.Log.Warning("Update interval can't be less than one minute. Setting to one minute.")
		*flagUpdateInterval = time.Minute
	}

	updater := &Updater{
		Endpoint: *flagUpdateURL,
		Interval: *flagUpdateInterval,
		Log:      k.Log,
	}

	// before we register check for latest update and re-update itself before
	// we continue
	k.Log.Info("Checking for new updates")
	if err := updater.checkAndUpdate(); err != nil {
		k.Log.Warning("Self-update: %s", err)
	}

	go updater.Run()

	userIn := func(user string, users ...string) bool {
		for _, u := range users {
			if u == user {
				return true
			}
		}
		return false
	}

	// don't pass any request if the caller is outside of our scope.
	// don't allow anyone to call a method if we are during an update.
	k.PreHandleFunc(func(r *kite.Request) (interface{}, error) {
		// only authenticated methods have correct username. For example
		// kite.ping has authentication disabled so username can be empty.
		if r.Auth != nil {
			// Koding (kloud) connects to much, don't display it.
			if r.Username != "koding" {
				k.Log.Info("Kite '%s/%s/%s' called method: '%s'",
					r.Username, r.Client.Environment, r.Client.Name, r.Method)
			}

			// Allow these users by default
			allowedUsers := []string{k.Config.Username, "koding"}

			// Allow collaboration users as well
			sharedUsers, err := collab.GetAll()
			if err != nil {
				return nil, fmt.Errorf("Can't read shared users from the storage. Err: %v", err)
			}

			sharedUsernames := make([]string, 0)
			for username := range sharedUsers {
				sharedUsernames = append(sharedUsernames, username)
			}

			allowedUsers = append(allowedUsers, sharedUsernames...)

			if !userIn(r.Username, allowedUsers...) {
				return nil, fmt.Errorf("User '%s' is not allowed to make a call to us.", r.Username)
			}
		}

		updatingMu.Lock()
		defer updatingMu.Unlock()

		if updating {
			return nil, errors.New("Updating klient. Can't accept any method.")
		}

		return true, nil
	})

	// Metrics, is used by Kloud to get usage so Kloud can stop free VMs
	k.PreHandleFunc(usg.Counter) // we measure every incoming request
	k.HandleFunc("klient.usage", usg.Current)

	// Collaboration, is used by our Koding.com browser client.
	k.HandleFunc("klient.share", collab.Share)
	k.HandleFunc("klient.unshare", collab.Unshare)
	k.HandleFunc("klient.shared", collab.Shared)

	// Filesystem
	k.HandleFunc("fs.readDirectory", fs.ReadDirectory)
	k.HandleFunc("fs.glob", fs.Glob)
	k.HandleFunc("fs.readFile", fs.ReadFile)
	k.HandleFunc("fs.writeFile", fs.WriteFile)
	k.HandleFunc("fs.uniquePath", fs.UniquePath)
	k.HandleFunc("fs.getInfo", fs.GetInfo)
	k.HandleFunc("fs.setPermissions", fs.SetPermissions)
	k.HandleFunc("fs.remove", fs.Remove)
	k.HandleFunc("fs.rename", fs.Rename)
	k.HandleFunc("fs.createDirectory", fs.CreateDirectory)
	k.HandleFunc("fs.move", fs.Move)
	k.HandleFunc("fs.copy", fs.Copy)

	// Docker
	dock := docker.New("unix:///var/run/docker.sock", k.Log)
	k.HandleFunc("docker.create", dock.Create)
	k.HandleFunc("docker.connect", dock.Connect)
	k.HandleFunc("docker.stop", dock.Stop)
	k.HandleFunc("docker.start", dock.Start)
	k.HandleFunc("docker.remove", dock.RemoveContainer)
	k.HandleFunc("docker.list", dock.List)

	// Execution
	k.HandleFunc("exec", command.Exec)

	// Terminal
	term := terminal.New(k.Log)
	term.InputHook = usg.Reset
	k.HandleFunc("webterm.getSessions", term.GetSessions)
	k.HandleFunc("webterm.connect", term.Connect)
	k.HandleFunc("webterm.killSession", term.KillSession)
	k.HandleFunc("webterm.killSessions", term.KillSessions)

	var disconnectTimer *time.Timer

	k.OnFirstRequest(func(c *kite.Client) {
		// Koding (kloud) connects to much, don't display it.
		if c.Username != "koding" {
			k.Log.Info("Kite '%s/%s/%s' is connected", c.Username, c.Environment, c.Name)
		}

		if c.Username != k.Config.Username {
			return // we don't care for others
		}

		// it's still not initialized, so don't do anything
		if disconnectTimer != nil {
			// stop previously started disconnect timer.
			k.Log.Info("Disconnection timer is cancelled.")
			disconnectTimer.Stop()
		}

	})

	// Unshare collab users if the klient owner disconnects
	k.OnDisconnect(func(c *kite.Client) {
		// Koding (kloud) connects to much, don't display it.
		if c.Username != "koding" {
			k.Log.Info("Kite '%s/%s/%s' is disconnected", c.Username, c.Environment, c.Name)
		}

		if c.Username != k.Config.Username {
			return // we don't care for others
		}

		// if there is any previously created timers stop them so we don't leak
		// goroutines
		if disconnectTimer != nil {
			disconnectTimer.Stop()
		}

		k.Log.Info("Disconnection timer of 1 minutes is fired.")
		disconnectTimer = time.NewTimer(time.Minute * 1)

		// Close all active sessions of the current. Do not close it
		// immediately, instead of give some time so users can safely exit. If
		// the user reconnects again the timer will be stopped so we don't
		// unshare for network hiccups accidentally.
		go func() {
			select {
			case <-disconnectTimer.C:
				sharedUsers, err := collab.GetAll()
				if err != nil {
					k.Log.Warning("Couldn't unshare users: '%s'", err)
					return
				}

				if len(sharedUsers) == 0 {
					return // nothing to do ...
				}

				k.Log.Info("Unsharing users '%s'", sharedUsers)
				for user, option := range sharedUsers {
					// dont touch permanent users
					if option.Permanent {
						k.Log.Info("User is permanent, avoiding it: '%s'", user)
						continue
					}

					if err := collab.Delete(user); err != nil {
						k.Log.Warning("Couldn't delete user from storage: '%s'", err)
					}
					term.CloseSessions(user)
				}
			}
		}()
	})

	if err := register(k); err != nil {
		panic(err)
	}

	return k
}

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func HasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func AddPort(host, port string) string {
	if ok := HasPort(host); ok {
		return host
	}

	return host + ":" + port
}
<commit_msg>docker: disable until we support officiall Docker<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/koding/klient/collaboration"
	"github.com/koding/klient/command"
	"github.com/koding/klient/fs"
	"github.com/koding/klient/protocol"
	"github.com/koding/klient/terminal"
	"github.com/koding/klient/usage"

	"github.com/koding/klient/Godeps/_workspace/src/github.com/koding/kite"
	"github.com/koding/klient/Godeps/_workspace/src/github.com/koding/kite/config"
)

var (
	flagIP          = flag.String("ip", "", "Change public ip")
	flagPort        = flag.Int("port", 56789, "Change running port")
	flagVersion     = flag.Bool("version", false, "Show version and exit")
	flagEnvironment = flag.String("env", protocol.Environment, "Change environment")
	flagRegion      = flag.String("region", protocol.Region, "Change region")
	flagRegisterURL = flag.String("register-url", "", "Change register URL to kontrol")
	flagDebug       = flag.Bool("debug", false, "Debug mode")

	// update parameters
	flagUpdateInterval = flag.Duration("update-interval", time.Minute*5,
		"Change interval for checking for new updates")
	flagUpdateURL = flag.String("update-url",
		"https://s3.amazonaws.com/koding-klient/"+protocol.Environment+"/latest-version.txt",
		"Change update endpoint for latest version")

	VERSION = protocol.Version
	NAME    = protocol.Name

	// this is our main reference to count and measure metrics for the klient
	// we count only those methods, please add/remove methods here that will
	// reset the timer of a klient.
	usg = usage.NewUsage(map[string]bool{
		"fs.readDirectory":    true,
		"fs.glob":             true,
		"fs.readFile":         true,
		"fs.writeFile":        true,
		"fs.uniquePath":       true,
		"fs.getInfo":          true,
		"fs.setPermissions":   true,
		"fs.remove":           true,
		"fs.rename":           true,
		"fs.createDirectory":  true,
		"fs.move":             true,
		"fs.copy":             true,
		"webterm.getSessions": true,
		"webterm.connect":     true,
		"webterm.killSession": true,
		"exec":                true,
		"klient.share":        true,
		"klient.unshare":      true,
		"klient.shared":       true,
		// Disabled until we have Docker support, no need to bloat the binary :)
		// "docker.create":       true,
		// "docker.connect":      true,
		// "docker.stop":         true,
		// "docker.start":        true,
		// "docker.remove":       true,
		// "docker.list":         true,
	})

	// this is used to allow other users to call any klient method.
	collab = collaboration.New()

	// we also could use an atomic boolean this is simple for now.
	updating   = false
	updatingMu sync.Mutex // protects updating
)

func main() {
	flag.Parse()
	if *flagVersion {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	k := newKite()

	// Close the klient.db in any case. Corrupt db would be catastrophic
	defer collab.Close()

	k.Log.Info("Running as version %s", VERSION)
	k.Run()
}

func newKite() *kite.Kite {
	k := kite.New(NAME, VERSION)

	if *flagDebug {
		k.SetLogLevel(kite.DEBUG)
	}

	conf := config.MustGet()
	k.Config = conf
	k.Config.Port = *flagPort
	k.Config.Environment = *flagEnvironment
	k.Config.Region = *flagRegion
	k.Id = conf.Id // always boot up with the same id in the kite.key

	if *flagUpdateInterval < time.Minute {
		k.Log.Warning("Update interval can't be less than one minute. Setting to one minute.")
		*flagUpdateInterval = time.Minute
	}

	updater := &Updater{
		Endpoint: *flagUpdateURL,
		Interval: *flagUpdateInterval,
		Log:      k.Log,
	}

	// before we register check for latest update and re-update itself before
	// we continue
	k.Log.Info("Checking for new updates")
	if err := updater.checkAndUpdate(); err != nil {
		k.Log.Warning("Self-update: %s", err)
	}

	go updater.Run()

	userIn := func(user string, users ...string) bool {
		for _, u := range users {
			if u == user {
				return true
			}
		}
		return false
	}

	// don't pass any request if the caller is outside of our scope.
	// don't allow anyone to call a method if we are during an update.
	k.PreHandleFunc(func(r *kite.Request) (interface{}, error) {
		// only authenticated methods have correct username. For example
		// kite.ping has authentication disabled so username can be empty.
		if r.Auth != nil {
			// Koding (kloud) connects to much, don't display it.
			if r.Username != "koding" {
				k.Log.Info("Kite '%s/%s/%s' called method: '%s'",
					r.Username, r.Client.Environment, r.Client.Name, r.Method)
			}

			// Allow these users by default
			allowedUsers := []string{k.Config.Username, "koding"}

			// Allow collaboration users as well
			sharedUsers, err := collab.GetAll()
			if err != nil {
				return nil, fmt.Errorf("Can't read shared users from the storage. Err: %v", err)
			}

			sharedUsernames := make([]string, 0)
			for username := range sharedUsers {
				sharedUsernames = append(sharedUsernames, username)
			}

			allowedUsers = append(allowedUsers, sharedUsernames...)

			if !userIn(r.Username, allowedUsers...) {
				return nil, fmt.Errorf("User '%s' is not allowed to make a call to us.", r.Username)
			}
		}

		updatingMu.Lock()
		defer updatingMu.Unlock()

		if updating {
			return nil, errors.New("Updating klient. Can't accept any method.")
		}

		return true, nil
	})

	// Metrics, is used by Kloud to get usage so Kloud can stop free VMs
	k.PreHandleFunc(usg.Counter) // we measure every incoming request
	k.HandleFunc("klient.usage", usg.Current)

	// Collaboration, is used by our Koding.com browser client.
	k.HandleFunc("klient.share", collab.Share)
	k.HandleFunc("klient.unshare", collab.Unshare)
	k.HandleFunc("klient.shared", collab.Shared)

	// Filesystem
	k.HandleFunc("fs.readDirectory", fs.ReadDirectory)
	k.HandleFunc("fs.glob", fs.Glob)
	k.HandleFunc("fs.readFile", fs.ReadFile)
	k.HandleFunc("fs.writeFile", fs.WriteFile)
	k.HandleFunc("fs.uniquePath", fs.UniquePath)
	k.HandleFunc("fs.getInfo", fs.GetInfo)
	k.HandleFunc("fs.setPermissions", fs.SetPermissions)
	k.HandleFunc("fs.remove", fs.Remove)
	k.HandleFunc("fs.rename", fs.Rename)
	k.HandleFunc("fs.createDirectory", fs.CreateDirectory)
	k.HandleFunc("fs.move", fs.Move)
	k.HandleFunc("fs.copy", fs.Copy)

	// // Docker
	// Disabled until we have Docker support, no need to bloat the binary :)
	// dock := docker.New("unix:///var/run/docker.sock", k.Log)
	// k.HandleFunc("docker.create", dock.Create)
	// k.HandleFunc("docker.connect", dock.Connect)
	// k.HandleFunc("docker.stop", dock.Stop)
	// k.HandleFunc("docker.start", dock.Start)
	// k.HandleFunc("docker.remove", dock.RemoveContainer)
	// k.HandleFunc("docker.list", dock.List)

	// Execution
	k.HandleFunc("exec", command.Exec)

	// Terminal
	term := terminal.New(k.Log)
	term.InputHook = usg.Reset
	k.HandleFunc("webterm.getSessions", term.GetSessions)
	k.HandleFunc("webterm.connect", term.Connect)
	k.HandleFunc("webterm.killSession", term.KillSession)
	k.HandleFunc("webterm.killSessions", term.KillSessions)

	var disconnectTimer *time.Timer

	k.OnFirstRequest(func(c *kite.Client) {
		// Koding (kloud) connects to much, don't display it.
		if c.Username != "koding" {
			k.Log.Info("Kite '%s/%s/%s' is connected", c.Username, c.Environment, c.Name)
		}

		if c.Username != k.Config.Username {
			return // we don't care for others
		}

		// it's still not initialized, so don't do anything
		if disconnectTimer != nil {
			// stop previously started disconnect timer.
			k.Log.Info("Disconnection timer is cancelled.")
			disconnectTimer.Stop()
		}

	})

	// Unshare collab users if the klient owner disconnects
	k.OnDisconnect(func(c *kite.Client) {
		// Koding (kloud) connects to much, don't display it.
		if c.Username != "koding" {
			k.Log.Info("Kite '%s/%s/%s' is disconnected", c.Username, c.Environment, c.Name)
		}

		if c.Username != k.Config.Username {
			return // we don't care for others
		}

		// if there is any previously created timers stop them so we don't leak
		// goroutines
		if disconnectTimer != nil {
			disconnectTimer.Stop()
		}

		k.Log.Info("Disconnection timer of 1 minutes is fired.")
		disconnectTimer = time.NewTimer(time.Minute * 1)

		// Close all active sessions of the current. Do not close it
		// immediately, instead of give some time so users can safely exit. If
		// the user reconnects again the timer will be stopped so we don't
		// unshare for network hiccups accidentally.
		go func() {
			select {
			case <-disconnectTimer.C:
				sharedUsers, err := collab.GetAll()
				if err != nil {
					k.Log.Warning("Couldn't unshare users: '%s'", err)
					return
				}

				if len(sharedUsers) == 0 {
					return // nothing to do ...
				}

				k.Log.Info("Unsharing users '%s'", sharedUsers)
				for user, option := range sharedUsers {
					// dont touch permanent users
					if option.Permanent {
						k.Log.Info("User is permanent, avoiding it: '%s'", user)
						continue
					}

					if err := collab.Delete(user); err != nil {
						k.Log.Warning("Couldn't delete user from storage: '%s'", err)
					}
					term.CloseSessions(user)
				}
			}
		}()
	})

	if err := register(k); err != nil {
		panic(err)
	}

	return k
}

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func HasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func AddPort(host, port string) string {
	if ok := HasPort(host); ok {
		return host
	}

	return host + ":" + port
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"

	"github.com/ogier/pflag"
)

var (
	name    = "prod"
	version = "0.0.0"

	flagset   = pflag.NewFlagSet(name, pflag.ContinueOnError)
	isHelp    = flagset.BoolP("help", "h", false, "")
	isVersion = flagset.BoolP("version", "", false, "")
)

func printUsage() {
	fmt.Fprintf(os.Stderr, `
Usage: %s [OPTION]... [FILE]...
Output direct product of lines of each files.

Options:
      --help                display this help text and exit
      --version             display version information and exit
`[1:], name)
}

func printVersion() {
	fmt.Fprintf(os.Stderr, "%s\n", version)
}

func printErr(err interface{}) {
	fmt.Fprintf(os.Stderr, "%s: %s\n", name, err)
}

func guideToHelp() {
	fmt.Fprintf(os.Stderr, "Try '%s --help' for more information.\n", name)
}

func _main() int {
	flagset.SetOutput(ioutil.Discard)
	if err := flagset.Parse(os.Args[1:]); err != nil {
		printErr(err)
		return 2
	}
	if *isHelp {
		printUsage()
		return 0
	}
	if *isVersion {
		printVersion()
		return 0
	}

	var rs []io.Reader
	if flagset.NArg() == 0 {
		rs = append(rs, os.Stdin)
	} else {
		for _, arg := range flagset.Args() {
			f, err := os.Open(arg)
			if err != nil {
				printErr(err)
				return 1
			}
			defer f.Close()
			rs = append(rs, f)
		}
	}

	var aa [][]string
	for _, r := range rs {
		a := make([]string, 0, 64)
		b := bufio.NewScanner(r)
		for b.Scan() {
			a = append(a, b.Text())
		}
		if err := b.Err(); err != nil {
			printErr(err)
			return 1
		}
		aa = append(aa, a)
	}

	ss := make([]string, len(aa))
	for indexes := range Product(aa) {
		for i, index := range indexes {
			ss[i] = aa[i][index]
		}
		fmt.Println(strings.Join(ss, "\t"))
	}
	return 0
}

func main() {
	e := _main()
	os.Exit(e)
}
<commit_msg>Guide to help if flagset.Parse fail<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"

	"github.com/ogier/pflag"
)

var (
	name    = "prod"
	version = "0.0.0"

	flagset   = pflag.NewFlagSet(name, pflag.ContinueOnError)
	isHelp    = flagset.BoolP("help", "h", false, "")
	isVersion = flagset.BoolP("version", "", false, "")
)

func printUsage() {
	fmt.Fprintf(os.Stderr, `
Usage: %s [OPTION]... [FILE]...
Output direct product of lines of each files.

Options:
      --help                display this help text and exit
      --version             display version information and exit
`[1:], name)
}

func printVersion() {
	fmt.Fprintf(os.Stderr, "%s\n", version)
}

func printErr(err interface{}) {
	fmt.Fprintf(os.Stderr, "%s: %s\n", name, err)
}

func guideToHelp() {
	fmt.Fprintf(os.Stderr, "Try '%s --help' for more information.\n", name)
}

func _main() int {
	flagset.SetOutput(ioutil.Discard)
	if err := flagset.Parse(os.Args[1:]); err != nil {
		printErr(err)
		guideToHelp()
		return 2
	}
	if *isHelp {
		printUsage()
		return 0
	}
	if *isVersion {
		printVersion()
		return 0
	}

	var rs []io.Reader
	if flagset.NArg() == 0 {
		rs = append(rs, os.Stdin)
	} else {
		for _, arg := range flagset.Args() {
			f, err := os.Open(arg)
			if err != nil {
				printErr(err)
				return 1
			}
			defer f.Close()
			rs = append(rs, f)
		}
	}

	var aa [][]string
	for _, r := range rs {
		a := make([]string, 0, 64)
		b := bufio.NewScanner(r)
		for b.Scan() {
			a = append(a, b.Text())
		}
		if err := b.Err(); err != nil {
			printErr(err)
			return 1
		}
		aa = append(aa, a)
	}

	ss := make([]string, len(aa))
	for indexes := range Product(aa) {
		for i, index := range indexes {
			ss[i] = aa[i][index]
		}
		fmt.Println(strings.Join(ss, "\t"))
	}
	return 0
}

func main() {
	e := _main()
	os.Exit(e)
}
<|endoftext|>


<commit_before>// Copyright 2015 mokey Authors. All rights reserved.
// Use of this source code is governed by a BSD style
// license that can be found in the LICENSE file.

package main

import (
	"errors"
	"fmt"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"github.com/ubccr/mokey/server"
	"github.com/ubccr/mokey/tools"
	"github.com/urfave/cli"
)

func init() {
	viper.SetConfigName("mokey")
	viper.SetConfigType("yaml")
	viper.AddConfigPath("/etc/mokey/")
}

func main() {
	app := cli.NewApp()
	app.Name = "mokey"
	app.Authors = []cli.Author{cli.Author{Name: "Andrew E. Bruno", Email: "aebruno2@buffalo.edu"}}
	app.Usage = "mokey"
	app.Version = "0.5.2"
	app.Flags = []cli.Flag{
		&cli.StringFlag{Name: "conf,c", Usage: "Path to conf file"},
		&cli.BoolFlag{Name: "debug,d", Usage: "Print debug messages"},
	}
	app.Before = func(c *cli.Context) error {
		if c.GlobalBool("debug") {
			log.SetLevel(log.InfoLevel)
		} else {
			log.SetLevel(log.WarnLevel)
		}

		conf := c.GlobalString("conf")
		if len(conf) > 0 {
			viper.SetConfigFile(conf)
		}

		err := viper.ReadInConfig()
		if err != nil {
			return fmt.Errorf("Failed reading config file - %s", err)
		}

		if !viper.IsSet("enc_key") || !viper.IsSet("auth_key") {
			log.Fatal("Please ensure authentication and encryption keys are set")
		}

		return nil
	}
	app.Commands = []cli.Command{
		{
			Name:  "server",
			Usage: "Run http server",
			Action: func(c *cli.Context) error {
				err := server.Run()
				if err != nil {
					log.Fatal(err)
					return cli.NewExitError(err, 1)
				}

				return nil
			},
		},
		{
			Name:  "resetpw",
			Usage: "Send reset password email",
			Flags: []cli.Flag{
				&cli.StringFlag{Name: "uid, u", Usage: "User id"},
			},
			Action: func(c *cli.Context) error {
				uid := c.String("uid")
				if len(uid) == 0 {
					return cli.NewExitError(errors.New("Please provide a uid"), 1)
				}

				err := tools.SendResetPasswordEmail(uid)
				if err != nil {
					return cli.NewExitError(err, 1)
				}

				return nil
			},
		},
		{
			Name:  "verify-email",
			Usage: "Re-send verify email",
			Flags: []cli.Flag{
				&cli.StringFlag{Name: "uid, u", Usage: "User id"},
			},
			Action: func(c *cli.Context) error {
				uid := c.String("uid")
				if len(uid) == 0 {
					return cli.NewExitError(errors.New("Please provide a uid"), 1)
				}

				err := tools.SendVerifyEmail(uid)
				if err != nil {
					return cli.NewExitError(err, 1)
				}

				return nil
			},
		},
		{
			Name:  "status",
			Usage: "Display token status for user",
			Flags: []cli.Flag{
				&cli.StringFlag{Name: "uid, u", Usage: "User id"},
			},
			Action: func(c *cli.Context) error {
				uid := c.String("uid")
				if len(uid) == 0 {
					return cli.NewExitError(errors.New("Please provide a uid"), 1)
				}

				err := tools.Status(uid)
				if err != nil {
					return cli.NewExitError(err, 1)
				}

				return nil
			},
		}}

	app.RunAndExitOnError()
}
<commit_msg>Bump version<commit_after>// Copyright 2015 mokey Authors. All rights reserved.
// Use of this source code is governed by a BSD style
// license that can be found in the LICENSE file.

package main

import (
	"errors"
	"fmt"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"github.com/ubccr/mokey/server"
	"github.com/ubccr/mokey/tools"
	"github.com/urfave/cli"
)

func init() {
	viper.SetConfigName("mokey")
	viper.SetConfigType("yaml")
	viper.AddConfigPath("/etc/mokey/")
}

func main() {
	app := cli.NewApp()
	app.Name = "mokey"
	app.Authors = []cli.Author{cli.Author{Name: "Andrew E. Bruno", Email: "aebruno2@buffalo.edu"}}
	app.Usage = "mokey"
	app.Version = "0.5.3"
	app.Flags = []cli.Flag{
		&cli.StringFlag{Name: "conf,c", Usage: "Path to conf file"},
		&cli.BoolFlag{Name: "debug,d", Usage: "Print debug messages"},
	}
	app.Before = func(c *cli.Context) error {
		if c.GlobalBool("debug") {
			log.SetLevel(log.InfoLevel)
		} else {
			log.SetLevel(log.WarnLevel)
		}

		conf := c.GlobalString("conf")
		if len(conf) > 0 {
			viper.SetConfigFile(conf)
		}

		err := viper.ReadInConfig()
		if err != nil {
			return fmt.Errorf("Failed reading config file - %s", err)
		}

		if !viper.IsSet("enc_key") || !viper.IsSet("auth_key") {
			log.Fatal("Please ensure authentication and encryption keys are set")
		}

		return nil
	}
	app.Commands = []cli.Command{
		{
			Name:  "server",
			Usage: "Run http server",
			Action: func(c *cli.Context) error {
				err := server.Run()
				if err != nil {
					log.Fatal(err)
					return cli.NewExitError(err, 1)
				}

				return nil
			},
		},
		{
			Name:  "resetpw",
			Usage: "Send reset password email",
			Flags: []cli.Flag{
				&cli.StringFlag{Name: "uid, u", Usage: "User id"},
			},
			Action: func(c *cli.Context) error {
				uid := c.String("uid")
				if len(uid) == 0 {
					return cli.NewExitError(errors.New("Please provide a uid"), 1)
				}

				err := tools.SendResetPasswordEmail(uid)
				if err != nil {
					return cli.NewExitError(err, 1)
				}

				return nil
			},
		},
		{
			Name:  "verify-email",
			Usage: "Re-send verify email",
			Flags: []cli.Flag{
				&cli.StringFlag{Name: "uid, u", Usage: "User id"},
			},
			Action: func(c *cli.Context) error {
				uid := c.String("uid")
				if len(uid) == 0 {
					return cli.NewExitError(errors.New("Please provide a uid"), 1)
				}

				err := tools.SendVerifyEmail(uid)
				if err != nil {
					return cli.NewExitError(err, 1)
				}

				return nil
			},
		},
		{
			Name:  "status",
			Usage: "Display token status for user",
			Flags: []cli.Flag{
				&cli.StringFlag{Name: "uid, u", Usage: "User id"},
			},
			Action: func(c *cli.Context) error {
				uid := c.String("uid")
				if len(uid) == 0 {
					return cli.NewExitError(errors.New("Please provide a uid"), 1)
				}

				err := tools.Status(uid)
				if err != nil {
					return cli.NewExitError(err, 1)
				}

				return nil
			},
		}}

	app.RunAndExitOnError()
}
<|endoftext|>


<commit_before>package main

import (
  "bytes"
  "bufio"
  "fmt"
  "os"
  "io"
  "errors"
  "io/ioutil"
  "log"
  "net/url"
  "os/exec"
  "net/http"
  "mime/multipart"
  "path/filepath"
  "encoding/json"

  "github.com/deiwin/interact"
  "github.com/jhoonb/archivex"
  "github.com/satori/go.uuid"
  "github.com/urfave/cli"
)

func check(e error) {
  if e != nil {
    panic(e)
  }
}

// First let's declare some simple input validators
var (
  checkNotEmpty = func(input string) error {
    // note that the inputs provided to these checks are already trimmed
    if input == "" {
      return errors.New("Input should not be empty!")
    }
    return nil
  }
)

type CaptureConfig struct {
  ID        string
  Recursive string
  Url       string
  Agent     string
  Path      string
  AuthToken string
}

// Configuration is the global configuration object
type EnvConfig struct {
  AuthToken      string `json:"token"`
}

var envConfig = EnvConfig{}

func main() {
  url := "http://code.notable.dev/api/cli/sites"
  app := cli.NewApp()
  app.EnableBashCompletion = true
  app.Name = "notable"
  app.Usage = "Interface with Notable"

  app.Commands = []cli.Command{
    {
      Name:      "code",
      Aliases:   []string{"c"},
      Usage:     "Send local site to Notable",
      Flags: []cli.Flag {
        cli.StringFlag{
          Name: "dest, d",
          Value: ".",
          Usage: "destination",
        },
        cli.StringFlag{
          Name: "token, t",
          Value: "empty",
          Usage: "User Auth Token",
        },
      },
      Action: func(c *cli.Context) error {
        loadAndCheckEnv()

        id := fmt.Sprintf("%s", uuid.NewV4())
        config := CaptureConfig{
          Recursive: "false",
          Agent: "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_2) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/48.0.2564.116 Safari/537.36",
          Url: c.Args().First(),
          Path: c.String("dest"),
          ID: id,
        }

        fmt.Printf("Directory ID: %s", config.ID)
        fetch(config)
        zip(config.ID)
        upload(config, url)
        return nil
      },
    },
    {
      Name:      "login",
      Aliases:   []string{"l"},
      Usage:     "Authenticate the CLI",
      Action: func(c *cli.Context) error {
        actor := interact.NewActor(os.Stdin, os.Stdout)
        message := "Please enter your email address"
        email, err := actor.PromptAndRetry(message, checkNotEmpty)
        if err != nil {
          log.Fatal(err)
        }
        message = "Please enter your password"
        password, err := actor.PromptAndRetry(message, checkNotEmpty)
        if err != nil {
          log.Fatal(err)
        }
        // TODO: post login and get auth token
        fetchToken(email, password)
        return nil
      },
    },
  }

  app.Run(os.Args)
}

func loadAndCheckEnv() {
  config, err := readAuth()

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

  envConfig = config
}

func writeAuth(t string) {
  token_data := []byte(t)
  err := ioutil.WriteFile("._auth", token_data, 0644)
  check(err)
}

func readAuth() (EnvConfig, error) {
  file, err := ioutil.ReadFile("._auth")

  if err != nil {
    return EnvConfig{}, err
  }

  return EnvConfig{
    AuthToken: string(file),
  }, nil
}

func fetchToken(e string, p string) {
  endpoint := "http://notable.dev/api/v5/platform_users/auth_cli"
  v := url.Values{}
  v.Set("email", e)
  v.Add("password", p)
  var err error

  resp, err := http.PostForm(endpoint, v)
  if nil != err {
    panic(err.Error())
  }

  defer resp.Body.Close()
  body, err := ioutil.ReadAll(resp.Body)

  err = json.Unmarshal(body, &envConfig)
  writeAuth(envConfig.AuthToken)
}

func fetch(c CaptureConfig) {
  wGetCheck()
  args := []string{
    fmt.Sprintf("-U '%s'", c.Agent),
    "--no-clobber",
    "--adjust-extension",
    "--span-hosts",
    "--page-requisites",
    "--backup-converted",
    "--html-extension",
    "--convert-links",
    "--no-parent",
    fmt.Sprintf("--directory-prefix=captures/%s", c.ID),
    c.Url,
  }
  cmd := exec.Command("wget", args...)

  cmdReader, err := cmd.StdoutPipe()
  if err != nil {
    fmt.Fprintln(os.Stderr, "Error creating StdoutPipe for Notable", err)
    os.Exit(1)
  }

  scanner := bufio.NewScanner(cmdReader)

  go func() {
    for scanner.Scan() {
      fmt.Printf("Notable capture | %s\n", scanner.Text())
    }
  }()

  err = cmd.Start()
  if err != nil {
    fmt.Fprintln(os.Stderr, "Error starting Cmd", err)
    // os.Exit(1)
  }

  err = cmd.Wait()
  if err != nil {
    fmt.Fprintln(os.Stderr, "Error waiting for Cmd", err)
    // os.Exit(1)
  }

}

func zip(id string) {
  path := fmt.Sprintf("captures/%s", id)
  zip := new(archivex.ZipFile)
  zip.Create(path)
  zip.AddAll(path, true)
  zip.Close()
}

func upload(config CaptureConfig, url string) {
  path := fmt.Sprintf("%s/captures/%s.zip", currentPath(), config.ID)
  post(path, config, url)
}

func wGetCheck() {
  path, err := exec.LookPath("wget")
  if err != nil {
    log.Fatal("\n\n[ERROR] Please install wget using Homebrew:\nbrew up && brew install wget")
  }
  fmt.Printf("wget is available at %s\n", path)
}

func currentPath() string {
  dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
  if err != nil {
    log.Fatal(err)
  }
  return dir
}

func post(path string, config CaptureConfig, url string){
  file, err := os.Open(path)
  if err != nil {
    log.Fatal(err)
  }
  defer file.Close()

  /* Create a buffer to hold this multi-part form */
  body_buf := bytes.NewBufferString("")
  body_writer := multipart.NewWriter(body_buf)
  boundary := body_writer.Boundary()
  fmt.Println(boundary)
  content_type := body_writer.FormDataContentType()
  fmt.Println(content_type)

  /* Create a Form Field in a simpler way */
  body_writer.WriteField("name", config.Url)
  body_writer.WriteField("token", envConfig.AuthToken)

  /* Create a completely custom Form Part (or in this case, a file) */
  // http://golang.org/src/pkg/mime/multipart/writer.go?s=2274:2352#L86
  part, err := body_writer.CreateFormFile("upload", filepath.Base(path))
  if err != nil {
    log.Fatal(err)
  }
  _, err = io.Copy(part, file)

  /* Close the body and send the request */
  body_writer.Close()
  resp, err := http.Post(url, content_type, body_buf)
  if nil != err {
    panic(err.Error())
  }

  /* Handle the response */
  defer resp.Body.Close()
  body, err := ioutil.ReadAll(resp.Body)

  if nil != err {
    fmt.Println("errorination happened reading the body", err)
    return
  }

  fmt.Println(string(body[:]))
}
<commit_msg>Add some color, clean up output.<commit_after>package main

import (
  "bytes"
  "bufio"
  "fmt"
  "os"
  "io"
  "errors"
  "io/ioutil"
  "log"
  "net/url"
  "os/exec"
  "net/http"
  "mime/multipart"
  "path/filepath"
  "encoding/json"

  "github.com/skratchdot/open-golang/open"
  "github.com/deiwin/interact"
  "github.com/jhoonb/archivex"
  "github.com/satori/go.uuid"
  "github.com/fatih/color"
  "github.com/urfave/cli"
)

func check(e error) {
  if e != nil {
    panic(e)
  }
}

// First let's declare some simple input validators
var (
  checkNotEmpty = func(input string) error {
    // note that the inputs provided to these checks are already trimmed
    if input == "" {
      return errors.New("Input should not be empty!")
    }
    return nil
  }
)

type CaptureConfig struct {
  ID        string
  Recursive string
  Url       string
  Agent     string
  Path      string
  AuthToken string
}

// Configuration is the global configuration object
type EnvConfig struct {
  AuthToken      string `json:"token"`
}

var envConfig = EnvConfig{}

func main() {
  url := "http://code.notable.dev/api/cli/sites"
  app := cli.NewApp()
  app.EnableBashCompletion = true
  app.Name = "notable"
  app.Usage = "Interface with Notable"

  app.Commands = []cli.Command{
    {
      Name:      "code",
      Aliases:   []string{"c"},
      Usage:     "Send local site to Notable",
      Flags: []cli.Flag {
        cli.StringFlag{
          Name: "dest, d",
          Value: ".",
          Usage: "destination",
        },
        cli.StringFlag{
          Name: "token, t",
          Value: "empty",
          Usage: "User Auth Token",
        },
      },
      Action: func(c *cli.Context) error {
        loadAndCheckEnv()

        id := fmt.Sprintf("%s", uuid.NewV4())
        config := CaptureConfig{
          Recursive: "false",
          Agent: "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_2) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/48.0.2564.116 Safari/537.36",
          Url: c.Args().First(),
          Path: c.String("dest"),
          ID: id,
        }

        fetch(config)
        zip(config)
        upload(config, url)
        return nil
      },
    },
    {
      Name:      "login",
      Aliases:   []string{"l"},
      Usage:     "Authenticate the CLI",
      Action: func(c *cli.Context) error {
        actor := interact.NewActor(os.Stdin, os.Stdout)
        message := "Please enter your Notable email address"
        email, err := actor.PromptAndRetry(message, checkNotEmpty)
        if err != nil {
          log.Fatal(err)
        }
        message = "Please enter your Notable password"
        password, err := actor.PromptAndRetry(message, checkNotEmpty)
        if err != nil {
          log.Fatal(err)
        }
        // TODO: post login and get auth token
        fetchToken(email, password)
        return nil
      },
    },
  }

  app.Run(os.Args)
}

func loadAndCheckEnv() {
  config, err := readAuth()

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

  envConfig = config
}

func writeAuth(t string) {
  token_data := []byte(t)
  err := ioutil.WriteFile("._auth", token_data, 0644)
  check(err)
}

func readAuth() (EnvConfig, error) {
  file, err := ioutil.ReadFile("._auth")

  if err != nil {
    return EnvConfig{}, err
  }

  return EnvConfig{
    AuthToken: string(file),
  }, nil
}

func fetchToken(e string, p string) {
  endpoint := "http://notable.dev/api/v5/platform_users/auth_cli"
  v := url.Values{}
  v.Set("email", e)
  v.Add("password", p)
  var err error

  resp, err := http.PostForm(endpoint, v)
  if nil != err {
    panic(err.Error())
  }

  defer resp.Body.Close()
  body, err := ioutil.ReadAll(resp.Body)

  err = json.Unmarshal(body, &envConfig)
  writeAuth(envConfig.AuthToken)
}

func fetch(c CaptureConfig) {
  wGetCheck()
  color.Cyan("Capturing url...\n")
  args := []string{
    fmt.Sprintf("-U '%s'", c.Agent),
    "--no-clobber",
    "--adjust-extension",
    "--span-hosts",
    "--page-requisites",
    "--backup-converted",
    "--html-extension",
    "--convert-links",
    "--no-parent",
    fmt.Sprintf("--directory-prefix=captures/%s", c.ID),
    c.Url,
  }
  cmd := exec.Command("wget", args...)

  cmdReader, err := cmd.StdoutPipe()
  if err != nil {
    fmt.Fprintln(os.Stderr, "Error creating StdoutPipe for Notable", err)
    os.Exit(1)
  }

  scanner := bufio.NewScanner(cmdReader)

  go func() {
    for scanner.Scan() {
      fmt.Printf("Notable capture | %s\n", scanner.Text())
    }
  }()

  err = cmd.Start()
  if err != nil {
    fmt.Fprintln(os.Stderr, "Error starting Cmd", err)
    // os.Exit(1)
  }

  err = cmd.Wait()
  if err != nil {
    text := fmt.Sprintf("%s", err)
    if text == "exit status 4" {
      fmt.Fprintln(os.Stderr, "The URL you specified is not accessible.")
      os.Exit(1)
    }
  }

}

func zip(config CaptureConfig) {
  color.Cyan("Compressing...\n")
  path := fmt.Sprintf("captures/%s", config.ID)
  zip := new(archivex.ZipFile)
  zip.Create(path)
  zip.AddAll(path, true)
  zip.Close()
}

func upload(config CaptureConfig, url string) {
  color.Cyan("Uploading...\n")
  path := fmt.Sprintf("%s/captures/%s.zip", currentPath(), config.ID)
  post(path, config, url)
}

func wGetCheck() {
  _, err := exec.LookPath("wget")
  if err != nil {
    log.Fatal("\n\n[ERROR] Please install wget using Homebrew:\nbrew up && brew install wget")
  }
  // fmt.Printf("\nwget is available at %s\n", path)
}

func currentPath() string {
  dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
  if err != nil {
    log.Fatal(err)
  }
  return dir
}

func post(path string, config CaptureConfig, url string){
  file, err := os.Open(path)
  if err != nil {
    log.Fatal(err)
  }
  defer file.Close()

  /* Create a buffer to hold this multi-part form */
  body_buf := bytes.NewBufferString("")
  body_writer := multipart.NewWriter(body_buf)
  // boundary := body_writer.Boundary()
  // fmt.Println(boundary)
  content_type := body_writer.FormDataContentType()
  // fmt.Println(content_type)

  /* Create a Form Field in a simpler way */
  body_writer.WriteField("name", config.Url)
  body_writer.WriteField("token", envConfig.AuthToken)

  /* Create a completely custom Form Part (or in this case, a file) */
  // http://golang.org/src/pkg/mime/multipart/writer.go?s=2274:2352#L86
  part, err := body_writer.CreateFormFile("upload", filepath.Base(path))
  if err != nil {
    log.Fatal(err)
  }
  _, err = io.Copy(part, file)

  /* Close the body and send the request */
  body_writer.Close()
  resp, err := http.Post(url, content_type, body_buf)
  if nil != err {
    panic(err.Error())
  }

  /* Handle the response */
  defer resp.Body.Close()
  body, err := ioutil.ReadAll(resp.Body)

  if nil != err {
    fmt.Println("errorination happened reading the body", err)
    return
  }

  var data map[string]interface{}
  if err := json.Unmarshal(body, &data); err != nil {
    panic(err)
  }

  color.Green("Done! Go give feedback!")

  open.Run(data["url"].(string))
}
<|endoftext|>


<commit_before>package dbfiles

import (
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/juju/errgo"
)

func init() {
	//	log.SetLevel(log.DebugLevel)
}

type DBFiles struct {
	BaseDir string
	Driver
	Structure
	WriteQueue chan (record)

	keysmux *sync.RWMutex
	keys    [][]string
}

type record struct {
	values    []string
	key       []string
	errorChan chan (error)
	basedir   string
}

const DefaultBaseDir = "data"

func New() *DBFiles {
	db := new(DBFiles)
	db.BaseDir = DefaultBaseDir
	db.Driver = CSV{}
	db.Structure = NewFolders()

	db.WriteQueue = make(chan (record), 10000)
	db.keysmux = new(sync.RWMutex)

	go db.runQueue()
	go db.runQueue()
	go db.runQueue()
	go db.runQueue()
	go db.runQueue()
	go db.runQueue()
	go db.runQueue()

	return db
}

func (db DBFiles) runQueue() {
	for {
		record := <-db.WriteQueue
		log.Debug("new record: ", record)
		log.Debug("Basedir: ", record.basedir)

		_, err := os.Stat(record.basedir)
		if os.IsNotExist(err) {
			err := db.Structure.Create(record.basedir)
			if err != nil {
				record.errorChan <- errgo.Notef(err, "can not create structure")
				continue
			}
		}

		file, err := db.Structure.File(record.basedir, db.Driver, record.key)
		if err != nil {
			record.errorChan <- errgo.Notef(err, "can not open file")
			continue
		}

		err = db.Driver.Write(file, record.values)
		if err != nil {
			record.errorChan <- errgo.Notef(err, "can not write values")
			continue
		}

		var data []byte
		io.ReadFull(file, data)
		log.Debug("Data: ", string(data))

		record.errorChan <- nil
		log.Debug("finished writing record: ", record)

		file.Close()
	}
}

func (db *DBFiles) Put(values []string, key ...string) error {
	errorChan := make(chan (error))

	rec := record{
		values:    values,
		key:       key,
		errorChan: errorChan,
		basedir:   db.BaseDir,
	}

	db.WriteQueue <- rec

	err := <-errorChan

	return err
}

func (db DBFiles) Get(key ...string) ([][]string, error) {
	file, err := db.Structure.File(db.BaseDir, db.Driver, key)
	if err != nil {
		return nil, errgo.Notef(err, "can not open file")
	}

	values, err := db.Driver.Read(file)
	if err != nil {
		return nil, errgo.Notef(err, "can not read values")
	}

	return values, nil
}

func (db DBFiles) Keys() ([][]string, error) {
	_, err := os.Stat(db.BaseDir)
	if os.IsNotExist(err) {
		return [][]string{}, nil
	}

	err = filepath.Walk(db.BaseDir, db.walkPopulateKeys)
	if err != nil {
		return nil, errgo.Notef(err, "can not walk through basedir")
	}

	return db.keys, nil
}

func (db *DBFiles) walkPopulateKeys(path string, info os.FileInfo, err error) error {
	if err != nil {
		return errgo.Notef(err, "error is not empty")
	}

	if info == nil {
		return errgo.New("directory info is empty")
	}

	//Skip git folder
	if info.IsDir() && info.Name() == ".git" {
		return filepath.SkipDir
	}

	if info.IsDir() {
		return nil
	}

	// Remove basedir from path
	relpath, err := filepath.Rel(db.BaseDir, path)
	if err != nil {
		return errgo.Notef(err, "can not get relative path")
	}

	// Get driver extention
	driverext := filepath.Ext(relpath)

	// remove driverextention
	nodriverpath := strings.TrimRight(relpath, driverext)

	// Split by path sepperator
	split := strings.Split(nodriverpath, string(os.PathSeparator))

	// Append new key to the db.keys
	db.keysmux.Lock()
	db.keys = append(db.keys, split)
	db.keysmux.Unlock()

	return nil
}

func (db *DBFiles) Destroy() error {
	err := os.RemoveAll(db.BaseDir)
	if err != nil {
		return errgo.Notef(err, "can not remove basedir")
	}

	return nil
}
<commit_msg>Simplified.<commit_after>package dbfiles

import (
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/juju/errgo"
)

func init() {
	//	log.SetLevel(log.DebugLevel)
}

type DBFiles struct {
	BaseDir string
	Driver
	Structure
	WriteQueue chan (record)

	keysmux *sync.RWMutex
	keys    [][]string
}

type record struct {
	values    []string
	key       []string
	errorChan chan (error)
	basedir   string
}

const DefaultBaseDir = "data"

func New() *DBFiles {
	db := new(DBFiles)
	db.BaseDir = DefaultBaseDir
	db.Driver = CSV{}
	db.Structure = NewFolders()

	db.WriteQueue = make(chan (record), 10000)
	db.keysmux = new(sync.RWMutex)

	return db
}

func (db *DBFiles) Put(values []string, key ...string) error {
	record := record{
		values:  values,
		key:     key,
		basedir: db.BaseDir,
	}

	_, err := os.Stat(record.basedir)
	if os.IsNotExist(err) {
		err := db.Structure.Create(record.basedir)
		if err != nil {
			return errgo.Notef(err, "can not create structure")

		}
	}

	file, err := db.Structure.File(record.basedir, db.Driver, record.key)
	if err != nil {
		return errgo.Notef(err, "can not open file")
	}
	defer file.Close()

	err = db.Driver.Write(file, record.values)
	if err != nil {
		return errgo.Notef(err, "can not write values")
	}

	var data []byte
	io.ReadFull(file, data)
	log.Debug("Data: ", string(data))

	log.Debug("finished writing record: ", record)

	return err
}

func (db DBFiles) Get(key ...string) ([][]string, error) {
	file, err := db.Structure.File(db.BaseDir, db.Driver, key)
	if err != nil {
		return nil, errgo.Notef(err, "can not open file")
	}

	values, err := db.Driver.Read(file)
	if err != nil {
		return nil, errgo.Notef(err, "can not read values")
	}

	return values, nil
}

func (db DBFiles) Keys() ([][]string, error) {
	_, err := os.Stat(db.BaseDir)
	if os.IsNotExist(err) {
		return [][]string{}, nil
	}

	err = filepath.Walk(db.BaseDir, db.walkPopulateKeys)
	if err != nil {
		return nil, errgo.Notef(err, "can not walk through basedir")
	}

	return db.keys, nil
}

func (db *DBFiles) walkPopulateKeys(path string, info os.FileInfo, err error) error {
	if err != nil {
		return errgo.Notef(err, "error is not empty")
	}

	if info == nil {
		return errgo.New("directory info is empty")
	}

	//Skip git folder
	if info.IsDir() && info.Name() == ".git" {
		return filepath.SkipDir
	}

	if info.IsDir() {
		return nil
	}

	// Remove basedir from path
	relpath, err := filepath.Rel(db.BaseDir, path)
	if err != nil {
		return errgo.Notef(err, "can not get relative path")
	}

	// Get driver extention
	driverext := filepath.Ext(relpath)

	// remove driverextention
	nodriverpath := relpath[0 : len(relpath)-len(driverext)]

	// Split by path sepperator
	split := strings.Split(nodriverpath, string(os.PathSeparator))

	// Append new key to the db.keys
	db.keysmux.Lock()
	db.keys = append(db.keys, split)
	db.keysmux.Unlock()

	log.Debug("Path: ", path)
	log.Debug("driverext: ", driverext)
	log.Debug("Nodriverpath: ", nodriverpath)
	log.Debug("Split: ", split)

	return nil
}

func (db *DBFiles) Destroy() error {
	err := os.RemoveAll(db.BaseDir)
	if err != nil {
		return errgo.Notef(err, "can not remove basedir")
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"runtime/debug"
	"sort"
	"strings"

	"github.com/crillab/gophersat/bf"
	"github.com/crillab/gophersat/solver"
)

func main() {
	debug.SetGCPercent(300)
	if len(os.Args) > 2 {
		fmt.Fprintf(os.Stderr, "Syntax : %s [file.cnf|file.bf]\n", os.Args[0])
		os.Exit(1)
	}
	f := os.Stdin
	if len(os.Args) == 2 {
		var err error
		f, err = os.Open(os.Args[1])
		if err != nil {
			fmt.Fprintf(os.Stderr, "Could not open %s: %v\n", os.Args[1], err.Error())
			os.Exit(1)
		}
		defer f.Close()
	}
	if err := parseAndSolve(f); err != nil {
		fmt.Fprintf(os.Stdout, "could not solve problem: %v\n", err)
		os.Exit(1)
	}
}

func parseAndSolve(r io.Reader) error {
	content, err := ioutil.ReadAll(r)
	if err != nil {
		return fmt.Errorf("could not read data: %v", err)
	}
	r2 := strings.NewReader(string(content))
	f, errBF := bf.Parse(r2)
	if errBF == nil {
		return solveBF(f)
	}
	r2 = strings.NewReader(string(content))
	pb, errCNF := solver.ParseCNF(r2)
	if errCNF == nil {
		return solveCNF(pb)
	}
	r2 = strings.NewReader(string(content))
	pb, errOPB := solver.ParseOPB(r2)
	if errOPB != nil {
		return fmt.Errorf("could not parse content as DIMACS (%v), as boolean formula (%v), nor as a pseudo-boolean problem (%v)", errCNF, errBF, errOPB)
	}
	return solveCNF(pb)
}

func solveCNF(pb *solver.Problem) error {
	fmt.Printf("c ======================================================================================\n")
	fmt.Printf("c | Number of clauses   : %9d                                                    |\n", len(pb.Clauses))
	fmt.Printf("c | Number of variables : %9d                                                    |\n", pb.NbVars)
	s := solver.New(pb)
	s.Verbose = true
	s.Solve()
	fmt.Printf("c nb conflicts: %d\nc nb restarts: %d\nc nb decisions: %d\n", s.Stats.NbConflicts, s.Stats.NbRestarts, s.Stats.NbDecisions)
	fmt.Printf("c nb unit learned: %d\nc nb binary learned: %d\nc nb learned: %d\n", s.Stats.NbUnitLearned, s.Stats.NbBinaryLearned, s.Stats.NbLearned)
	fmt.Printf("c nb clauses deleted: %d\n", s.Stats.NbDeleted)
	s.OutputModel()
	return nil
}

func countModels(pb *solver.Problem) error {
	fmt.Printf("c ======================================================================================\n")
	fmt.Printf("c | Number of clauses   : %9d                                                    |\n", len(pb.Clauses))
	fmt.Printf("c | Number of variables : %9d                                                    |\n", pb.NbVars)
	s := solver.New(pb)
	s.Verbose = true
	fmt.Printf("%d models found\n", s.CountModels())
	return nil
}

func solveBF(f bf.Formula) error {
	if model := bf.Solve(f); model == nil {
		fmt.Println("UNSATISFIABLE")
	} else {
		fmt.Println("SATISFIABLE")
		keys := make(sort.StringSlice, 0, len(model))
		for k := range model {
			keys = append(keys, k)
		}
		sort.Sort(keys)
		for _, k := range keys {
			fmt.Printf("%s: %t\n", k, model[k])
		}
	}
	return nil
}
<commit_msg>main executable now accepts optimization problems as input<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"runtime/debug"
	"sort"
	"strings"

	"github.com/crillab/gophersat/bf"
	"github.com/crillab/gophersat/solver"
)

func main() {
	debug.SetGCPercent(300)
	if len(os.Args) > 2 {
		fmt.Fprintf(os.Stderr, "Syntax : %s [file.cnf|file.bf]\n", os.Args[0])
		os.Exit(1)
	}
	f := os.Stdin
	if len(os.Args) == 2 {
		var err error
		f, err = os.Open(os.Args[1])
		if err != nil {
			fmt.Fprintf(os.Stderr, "Could not open %s: %v\n", os.Args[1], err.Error())
			os.Exit(1)
		}
		defer f.Close()
	}
	if err := parseAndSolve(f); err != nil {
		fmt.Fprintf(os.Stdout, "could not solve problem: %v\n", err)
		os.Exit(1)
	}
}

func parseAndSolve(r io.Reader) error {
	content, err := ioutil.ReadAll(r)
	if err != nil {
		return fmt.Errorf("could not read data: %v", err)
	}
	r2 := strings.NewReader(string(content))
	f, errBF := bf.Parse(r2)
	if errBF == nil {
		return solveBF(f)
	}
	r2 = strings.NewReader(string(content))
	pb, errCNF := solver.ParseCNF(r2)
	if errCNF == nil {
		return solveCNF(pb)
	}
	r2 = strings.NewReader(string(content))
	pb, errOPB := solver.ParseOPB(r2)
	if errOPB != nil {
		return fmt.Errorf("could not parse content as DIMACS (%v), as boolean formula (%v), nor as a pseudo-boolean problem (%v)", errCNF, errBF, errOPB)
	}
	return minimize(pb)
}

func minimize(pb *solver.Problem) error {
	fmt.Printf("c ======================================================================================\n")
	fmt.Printf("c | Number of clauses   : %9d                                                    |\n", len(pb.Clauses))
	fmt.Printf("c | Number of variables : %9d                                                    |\n", pb.NbVars)
	s := solver.New(pb)
	s.Verbose = true
	s.Minimize()
	fmt.Printf("c nb conflicts: %d\nc nb restarts: %d\nc nb decisions: %d\n", s.Stats.NbConflicts, s.Stats.NbRestarts, s.Stats.NbDecisions)
	fmt.Printf("c nb unit learned: %d\nc nb binary learned: %d\nc nb learned: %d\n", s.Stats.NbUnitLearned, s.Stats.NbBinaryLearned, s.Stats.NbLearned)
	fmt.Printf("c nb clauses deleted: %d\n", s.Stats.NbDeleted)
	s.OutputModel()
	return nil
}

func solveCNF(pb *solver.Problem) error {
	fmt.Printf("c ======================================================================================\n")
	fmt.Printf("c | Number of clauses   : %9d                                                    |\n", len(pb.Clauses))
	fmt.Printf("c | Number of variables : %9d                                                    |\n", pb.NbVars)
	s := solver.New(pb)
	s.Verbose = true
	s.Solve()
	fmt.Printf("c nb conflicts: %d\nc nb restarts: %d\nc nb decisions: %d\n", s.Stats.NbConflicts, s.Stats.NbRestarts, s.Stats.NbDecisions)
	fmt.Printf("c nb unit learned: %d\nc nb binary learned: %d\nc nb learned: %d\n", s.Stats.NbUnitLearned, s.Stats.NbBinaryLearned, s.Stats.NbLearned)
	fmt.Printf("c nb clauses deleted: %d\n", s.Stats.NbDeleted)
	s.OutputModel()
	return nil
}

func countModels(pb *solver.Problem) error {
	fmt.Printf("c ======================================================================================\n")
	fmt.Printf("c | Number of clauses   : %9d                                                    |\n", len(pb.Clauses))
	fmt.Printf("c | Number of variables : %9d                                                    |\n", pb.NbVars)
	s := solver.New(pb)
	s.Verbose = true
	fmt.Printf("%d models found\n", s.CountModels())
	return nil
}

func solveBF(f bf.Formula) error {
	if model := bf.Solve(f); model == nil {
		fmt.Println("UNSATISFIABLE")
	} else {
		fmt.Println("SATISFIABLE")
		keys := make(sort.StringSlice, 0, len(model))
		for k := range model {
			keys = append(keys, k)
		}
		sort.Sort(keys)
		for _, k := range keys {
			fmt.Printf("%s: %t\n", k, model[k])
		}
	}
	return nil
}
<|endoftext|>