<commit_before>package main

import (
	"context"
	"os"

	"github.com/Orientsoft/kubekit/utils"

	"github.com/Orientsoft/kubekit/controllers"

	"github.com/fatih/color"

	cli "gopkg.in/urfave/cli.v1"
)

const (
	FilePort = iota
	ToolkitPort
	VERSION = "0.2"
)

func initialize() {
	//Remove the install log file
	os.Remove("install.log")
	utils.DisplayLogo(VERSION)
}

func main() {
	initialize()

	app := cli.NewApp()
	app.Name = "KubeKit"
	app.Usage = "A toolkit for Kubernetes & apps offline deployment."
	app.Version = VERSION
	app.Action = func(c *cli.Context) error {
		return nil
	}

	app.Commands = []cli.Command{
		{
			Name:      "init",
			Aliases:   []string{"i"},
			Usage:     "Initialize current server with Docker engine & Kubernetes master.",
			ArgsUsage: "K8S_MASTER_IP FILE_SERVER_PORT TOOLKIT_SERVER_PORT",
			Action: func(c *cli.Context) error {

				masterIP := c.Args().Get(0)

				if masterIP == "" {
					color.Red("Please run kubekit with master IP: kubekit i K8S_MASTER_IP")
					os.Exit(0)
				}

				color.Blue("Initialization process started, with kubernetes master IP: %s\r\n", masterIP)
				utils.SaveMasterIP(masterIP)

				filePort := getServerPort(c, FilePort, 1)
				toolkitPort := getServerPort(c, ToolkitPort, 2)

				masterAddr := masterIP + filePort

				srv := utils.StartServer(filePort)
				defer srv.Shutdown(context.Background())

				if !utils.SetupDocker(masterAddr) {
					color.Red("%sProgram terminated...", utils.CrossSymbol)
					os.Exit(1)
				}

				if utils.SetupMaster(masterAddr) {
					// Launch toolkit server
					controllers.StartToolkitServer(filePort, toolkitPort)
				}

				return nil
			},
		},
		{
			Name:      "server",
			Aliases:   []string{"s"},
			Usage:     "Start kubekit file server & toolkit server.",
			ArgsUsage: "FILE_SERVER_PORT TOOLKIT_SERVER_PORT",
			Action: func(c *cli.Context) error {
				filePort := getServerPort(c, FilePort, 0)
				toolkitPort := getServerPort(c, ToolkitPort, 1)

				srv := utils.StartServer(filePort)
				defer srv.Shutdown(context.Background())
				// Launch toolkit server only
				controllers.StartToolkitServer(filePort, toolkitPort)
				return nil
			},
		},
	}

	app.Run(os.Args)
}

// getServerPort parse and get server port by port type and argument position
// portType:  1.File server 2.Toolkit server
func getServerPort(c *cli.Context, portType, argPos int) string {
	var port string

	if len(c.Args().Get(argPos)) > 0 {
		if utils.IsValidPort(c.Args().Get(argPos)) {
			port = ":" + c.Args().Get(argPos)
		} else {
			if portType == FilePort {
				color.Red("%sPlease input a valid file server port!", utils.CrossSymbol)
			} else if portType == ToolkitPort {
				color.Red("%sPlease input a valid toolkit server port!", utils.CrossSymbol)
			}
			os.Exit(1)
		}
	} else {
		if portType == FilePort {
			port = ":8000"
		} else if portType == ToolkitPort {
			port = ":9000"
		}
	}

	return port
}
<commit_msg>update version to V0.3<commit_after>package main

import (
	"context"
	"os"

	"github.com/Orientsoft/kubekit/utils"

	"github.com/Orientsoft/kubekit/controllers"

	"github.com/fatih/color"

	cli "gopkg.in/urfave/cli.v1"
)

const (
	FilePort = iota
	ToolkitPort
	VERSION = "0.3"
)

func initialize() {
	//Remove the install log file
	os.Remove("install.log")
	utils.DisplayLogo(VERSION)
}

func main() {
	initialize()

	app := cli.NewApp()
	app.Name = "KubeKit"
	app.Usage = "A toolkit for Kubernetes & apps offline deployment."
	app.Version = VERSION
	app.Action = func(c *cli.Context) error {
		return nil
	}

	app.Commands = []cli.Command{
		{
			Name:      "init",
			Aliases:   []string{"i"},
			Usage:     "Initialize current server with Docker engine & Kubernetes master.",
			ArgsUsage: "K8S_MASTER_IP FILE_SERVER_PORT TOOLKIT_SERVER_PORT",
			Action: func(c *cli.Context) error {

				masterIP := c.Args().Get(0)

				if masterIP == "" {
					color.Red("Please run kubekit with master IP: kubekit i K8S_MASTER_IP")
					os.Exit(0)
				}

				color.Blue("Initialization process started, with kubernetes master IP: %s\r\n", masterIP)
				utils.SaveMasterIP(masterIP)

				filePort := getServerPort(c, FilePort, 1)
				toolkitPort := getServerPort(c, ToolkitPort, 2)

				masterAddr := masterIP + filePort

				srv := utils.StartServer(filePort)
				defer srv.Shutdown(context.Background())

				if !utils.SetupDocker(masterAddr) {
					color.Red("%sProgram terminated...", utils.CrossSymbol)
					os.Exit(1)
				}

				if utils.SetupMaster(masterAddr) {
					// Launch toolkit server
					controllers.StartToolkitServer(filePort, toolkitPort)
				}

				return nil
			},
		},
		{
			Name:      "server",
			Aliases:   []string{"s"},
			Usage:     "Start kubekit file server & toolkit server.",
			ArgsUsage: "FILE_SERVER_PORT TOOLKIT_SERVER_PORT",
			Action: func(c *cli.Context) error {
				filePort := getServerPort(c, FilePort, 0)
				toolkitPort := getServerPort(c, ToolkitPort, 1)

				srv := utils.StartServer(filePort)
				defer srv.Shutdown(context.Background())
				// Launch toolkit server only
				controllers.StartToolkitServer(filePort, toolkitPort)
				return nil
			},
		},
	}

	app.Run(os.Args)
}

// getServerPort parse and get server port by port type and argument position
// portType:  1.File server 2.Toolkit server
func getServerPort(c *cli.Context, portType, argPos int) string {
	var port string

	if len(c.Args().Get(argPos)) > 0 {
		if utils.IsValidPort(c.Args().Get(argPos)) {
			port = ":" + c.Args().Get(argPos)
		} else {
			if portType == FilePort {
				color.Red("%sPlease input a valid file server port!", utils.CrossSymbol)
			} else if portType == ToolkitPort {
				color.Red("%sPlease input a valid toolkit server port!", utils.CrossSymbol)
			}
			os.Exit(1)
		}
	} else {
		if portType == FilePort {
			port = ":8000"
		} else if portType == ToolkitPort {
			port = ":9000"
		}
	}

	return port
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2014 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

/*
Party is a tool to automate the use of third party packages.

Many existing Go package managers provide solutions to the problem of relying on
other projects' code. Go handles the use of other packages elegantly, but is
prone to problems when that code has breaking changes or is removed from the
Internet. As a solution, the FAQ suggests copying any third party repositories
into your own project and changing your import paths to match. This is a
recursive process, as you must then perform the same procedure for those copied
packages and so on, until all leaf packages have no external imports.

party is a tool that performs all of this work. It:

1. Rewrites all external imports in both your project and the third party
directory from "other.com/user/package" to
"host.com/user/your_project/_third_party/other.com/user/package".

2. Updates or copies files from "$GOPATH/src/other.com/user/package" to
"$GOPATH/src/host.com/user/your_project/_third_party/other.com/user/package".

3. Returns to step 1 until no more rewrites are performed.

Usage

From your project's root directory, for the first run:

	$ party -c

Subsequent invocations can omit -c.

The flags are:
	-c
		create the third party directory if it does not exist
	-u
		run "go get -d -u <pkg>" on packages imported by party for the
		current package
	-d="_third_party"
		set the third party directory
	-v
		enable verbose output
	-n
		dry run: actions are printed instead of run
	-r
		import third party packages as relative, local imports

App Engine and Relative Imports

When working on an App Engine project, use the -r flag. This will cause all
third party imports to be local packages. This is needed because otherwise App
Engine will run init() routines from your .go files twice if your app is also in
your $GOPATH.

Bugs

There is no guarantee that the source directory exists. For example, a package
may have a OS-specific file (file_windows.go) which imports another package. On
non-Windows, that other package will not have been fetched during "go get" and
will thus not exist. These packages must be fetched using "go get" or some other
method by hand. Or, running party on a Windows machine (in this case) will
correctly copy those files. These cases are reported when running party.

party performs the equivalent of gofmt on any files affected by the import
rewrite, thus the diff may be somewhat larger than just the import lines.

Third party imports that use local (relative) packages will fail, as those
packages are assumed to be part of the standard library. Those packages will not
be copied.

*/
package main

import (
	"bytes"
	"flag"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"code.google.com/p/go.tools/astutil"
)

const dirMode os.FileMode = 0755

var (
	dryrun     = flag.Bool("n", false, "don't perform any action, instead print them")
	create     = flag.Bool("c", false, "create the third party directory if needed")
	relative   = flag.Bool("r", false, "use a relative third party directory (needed on App Engine)")
	verbose    = flag.Bool("v", false, "print actions")
	thirdParty = flag.String("d", "_third_party", "name of third party directory")
	flagUpdate = flag.Bool("u", false, "update (go get -d -u) used packages")

	relpath, gopath, ThirdParty string
)

func main() {
	flag.Parse()
	if *dryrun {
		*verbose = true
	}

	ThirdParty = filepath.Base(*thirdParty)
	if ThirdParty != *thirdParty {
		log.Fatal("third party directory cannot contain path separators")
	}
	if *create {
		os.Mkdir(ThirdParty, dirMode)
	}
	if _, err := os.Open(ThirdParty); err != nil {
		log.Fatalf("could not open %s directory: run with -c or create it manually", ThirdParty)
	}

	pwd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	var importSites []string
	for _, s := range filepath.SplitList(os.Getenv("GOPATH")) {
		s, err := filepath.Abs(s)
		if err != nil {
			log.Fatal(err)
		}
		s, err = filepath.EvalSymlinks(s)
		if err != nil {
			log.Fatal(err)
		}
		if sf, err := os.Open(filepath.Join(s, "src")); err == nil {
			names, err := sf.Readdirnames(0)
			if err == nil {
				for _, name := range names {
					importSites = append(importSites, name)
				}
			}
		}
		if !strings.HasPrefix(pwd, s) {
			continue
		}
		gopath = s
		relpath = strings.TrimPrefix(pwd, filepath.Join(s, "src")+string(os.PathSeparator))
		break
	}
	if relpath == "" {
		log.Fatal("couldn't determine current directory relative to $GOPATH")
	}

	if *verbose {
		log.Println("using GOPATH:", gopath)
		log.Println("using relative path:", relpath)
	}
	for {
		update()
		if !rewriteImports(importSites) {
			break
		}
	}
}

func update() {
	paths := make(map[string]struct{})
	reltp := ThirdParty + "/"
	if !*relative {
		reltp = filepath.ToSlash(relpath) + "/" + reltp
	}
	filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return nil
		}
		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
		if err != nil {
			return nil
		}
		for _, im := range f.Imports {
			v := im.Path.Value[1 : len(im.Path.Value)-1]
			if !strings.HasPrefix(v, reltp) {
				continue
			}
			paths[strings.TrimPrefix(v, reltp)] = struct{}{}
		}
		return nil
	})
	for k := range paths {
		fpath := filepath.Join(gopath, "src", k)
		f, err := os.Open(fpath)
		if err != nil {
			log.Printf("%s required, but could not be found at %s", k, fpath)
			continue
		}
		if *flagUpdate {
			if *verbose {
				log.Printf("go get -d -u %v", k)
			}
			if !*dryrun {
				out, err := exec.Command("go", "get", "-d", "-u", k).CombinedOutput()
				if len(out) > 0 {
					log.Println(string(out))
				}
				if err != nil {
					log.Println("go get", k, "err:", err)
				}
			}
		}
		files, err := f.Readdir(0)
		if err != nil {
			log.Fatal(err)
		}
		f.Close()
		for _, source := range files {
			if source.IsDir() || strings.HasPrefix(source.Name(), ".") {
				continue
			}
			spath := filepath.Join(f.Name(), source.Name())
			destdir := filepath.Join(ThirdParty, k)
			dest := filepath.Join(destdir, source.Name())
			if err := os.MkdirAll(destdir, dirMode); err != nil {
				log.Fatal(err)
			}
			dst, err := os.Open(dest)
			created := false
			newer := false
			if os.IsNotExist(err) {
				if *verbose {
					log.Println("create", dest)
				}
				created = true
			} else if err != nil {
				log.Fatal(err)
			} else {
				destination, err := dst.Stat()
				if err != nil {
					log.Fatal(err)
				}
				newer = source.ModTime().After(destination.ModTime())
				dst.Close()
			}
			if created || newer {
				if *verbose {
					log.Printf("copy %s -> %s\n", spath, dest)
				}
				if !*dryrun {
					src, err := os.Open(spath)
					if err != nil {
						log.Fatal(err)
					}
					dst, err = os.OpenFile(dest, os.O_RDWR|os.O_CREATE|os.O_TRUNC, source.Mode())
					if err != nil {
						log.Fatal(err)
					}
					if _, err := io.Copy(dst, src); err != nil {
						log.Fatal(err)
					}
					src.Close()
					dst.Close()
				}
			}
		}
	}
}

func rewriteImports(importSites []string) (rewritten bool) {
	filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return nil
		}
		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
		if err != nil {
			return nil
		}
		for _, im := range f.Imports {
			for _, site := range importSites {
				v, err := strconv.Unquote(im.Path.Value)
				if err != nil {
					log.Fatal(err)
				}
				// don't replace ourself
				if v == relpath || strings.HasPrefix(v, relpath+"/") {
					if *verbose {
						log.Printf("skipping %s because it is in %s\n", v, relpath)
					}
					continue
				}
				if strings.HasPrefix(v, site+"/") {
					target := ThirdParty + "/" + v
					if !*relative {
						target = relpath + "/" + target
					}
					if *verbose {
						log.Printf("rewrite %s: %s -> %s\n", path, v, target)
					}
					if *dryrun {
						break
					}
					if !astutil.RewriteImport(fset, f, v, target) {
						log.Fatalf("%s: could not rewrite %s", path, im.Path.Value)
					}
					rewritten = true
					break
				}
			}
		}
		var buf bytes.Buffer
		if err := printer.Fprint(&buf, fset, f); err != nil {
			log.Fatal(err)
		}
		fb, err := format.Source(buf.Bytes())
		if err != nil {
			log.Fatal(err)
		}
		if err := ioutil.WriteFile(path, fb, info.Mode()); err != nil {
			log.Fatal(err)
		}
		return nil
	})
	return
}
<commit_msg>Fix import path checking for go 1.4<commit_after>/*
 * Copyright (c) 2014 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

/*
Party is a tool to automate the use of third party packages.

Many existing Go package managers provide solutions to the problem of relying on
other projects' code. Go handles the use of other packages elegantly, but is
prone to problems when that code has breaking changes or is removed from the
Internet. As a solution, the FAQ suggests copying any third party repositories
into your own project and changing your import paths to match. This is a
recursive process, as you must then perform the same procedure for those copied
packages and so on, until all leaf packages have no external imports.

party is a tool that performs all of this work. It:

1. Rewrites all external imports in both your project and the third party
directory from "other.com/user/package" to
"host.com/user/your_project/_third_party/other.com/user/package".

2. Updates or copies files from "$GOPATH/src/other.com/user/package" to
"$GOPATH/src/host.com/user/your_project/_third_party/other.com/user/package".

3. Returns to step 1 until no more rewrites are performed.

Usage

From your project's root directory, for the first run:

	$ party -c

Subsequent invocations can omit -c.

The flags are:
	-c
		create the third party directory if it does not exist
	-u
		run "go get -d -u <pkg>" on packages imported by party for the
		current package
	-d="_third_party"
		set the third party directory
	-v
		enable verbose output
	-n
		dry run: actions are printed instead of run
	-r
		import third party packages as relative, local imports

App Engine and Relative Imports

When working on an App Engine project, use the -r flag. This will cause all
third party imports to be local packages. This is needed because otherwise App
Engine will run init() routines from your .go files twice if your app is also in
your $GOPATH.

Bugs

There is no guarantee that the source directory exists. For example, a package
may have a OS-specific file (file_windows.go) which imports another package. On
non-Windows, that other package will not have been fetched during "go get" and
will thus not exist. These packages must be fetched using "go get" or some other
method by hand. Or, running party on a Windows machine (in this case) will
correctly copy those files. These cases are reported when running party.

party performs the equivalent of gofmt on any files affected by the import
rewrite, thus the diff may be somewhat larger than just the import lines.

Third party imports that use local (relative) packages will fail, as those
packages are assumed to be part of the standard library. Those packages will not
be copied.

*/
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"code.google.com/p/go.tools/astutil"
)

const dirMode os.FileMode = 0755

var (
	dryrun     = flag.Bool("n", false, "don't perform any action, instead print them")
	create     = flag.Bool("c", false, "create the third party directory if needed")
	relative   = flag.Bool("r", false, "use a relative third party directory (needed on App Engine)")
	verbose    = flag.Bool("v", false, "print actions")
	thirdParty = flag.String("d", "_third_party", "name of third party directory")
	flagUpdate = flag.Bool("u", false, "update (go get -d -u) used packages")

	relpath, gopath, ThirdParty string
)

func main() {
	flag.Parse()
	if *dryrun {
		*verbose = true
	}

	ThirdParty = filepath.Base(*thirdParty)
	if ThirdParty != *thirdParty {
		log.Fatal("third party directory cannot contain path separators")
	}
	if *create {
		os.Mkdir(ThirdParty, dirMode)
	}
	if _, err := os.Open(ThirdParty); err != nil {
		log.Fatalf("could not open %s directory: run with -c or create it manually", ThirdParty)
	}

	pwd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	var importSites []string
	for _, s := range filepath.SplitList(os.Getenv("GOPATH")) {
		s, err := filepath.Abs(s)
		if err != nil {
			log.Fatal(err)
		}
		s, err = filepath.EvalSymlinks(s)
		if err != nil {
			log.Fatal(err)
		}
		if sf, err := os.Open(filepath.Join(s, "src")); err == nil {
			names, err := sf.Readdirnames(0)
			if err == nil {
				for _, name := range names {
					importSites = append(importSites, name)
				}
			}
		}
		if !strings.HasPrefix(pwd, s) {
			continue
		}
		gopath = s
		relpath = strings.TrimPrefix(pwd, filepath.Join(s, "src")+string(os.PathSeparator))
		break
	}
	if relpath == "" {
		log.Fatal("couldn't determine current directory relative to $GOPATH")
	}

	if *verbose {
		log.Println("using GOPATH:", gopath)
		log.Println("using relative path:", relpath)
	}
	for {
		update()
		if !rewriteImports(importSites) {
			break
		}
	}
}

func update() {
	paths := make(map[string]struct{})
	reltp := ThirdParty + "/"
	if !*relative {
		reltp = filepath.ToSlash(relpath) + "/" + reltp
	}
	filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return nil
		}
		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
		if err != nil {
			return nil
		}
		for _, im := range f.Imports {
			v := im.Path.Value[1 : len(im.Path.Value)-1]
			if !strings.HasPrefix(v, reltp) {
				continue
			}
			paths[strings.TrimPrefix(v, reltp)] = struct{}{}
		}
		return nil
	})
	for k := range paths {
		fpath := filepath.Join(gopath, "src", k)
		f, err := os.Open(fpath)
		if err != nil {
			log.Printf("%s required, but could not be found at %s", k, fpath)
			continue
		}
		if *flagUpdate {
			if *verbose {
				log.Printf("go get -d -u %v", k)
			}
			if !*dryrun {
				out, err := exec.Command("go", "get", "-d", "-u", k).CombinedOutput()
				if len(out) > 0 {
					log.Println(string(out))
				}
				if err != nil {
					log.Println("go get", k, "err:", err)
				}
			}
		}
		files, err := f.Readdir(0)
		if err != nil {
			log.Fatal(err)
		}
		f.Close()
		for _, source := range files {
			if source.IsDir() || strings.HasPrefix(source.Name(), ".") {
				continue
			}
			spath := filepath.Join(f.Name(), source.Name())
			destdir := filepath.Join(ThirdParty, k)
			dest := filepath.Join(destdir, source.Name())
			if err := os.MkdirAll(destdir, dirMode); err != nil {
				log.Fatal(err)
			}
			dst, err := os.Open(dest)
			created := false
			newer := false
			if os.IsNotExist(err) {
				if *verbose {
					log.Println("create", dest)
				}
				created = true
			} else if err != nil {
				log.Fatal(err)
			} else {
				destination, err := dst.Stat()
				if err != nil {
					log.Fatal(err)
				}
				newer = source.ModTime().After(destination.ModTime())
				dst.Close()
			}
			if created || newer {
				if *verbose {
					log.Printf("copy %s -> %s\n", spath, dest)
				}
				if !*dryrun {
					b, err := ioutil.ReadFile(spath)
					if err != nil {
						log.Fatal(err)
					}
					if filepath.Ext(spath) == ".go" {
						b, err = fixImportCheck(b, path.Join(relpath, destdir))
						if err != nil {
							log.Fatal(err)
						}
					}
					if err := ioutil.WriteFile(dest, b, source.Mode()); err != nil {
						log.Fatal(err)
					}
				}
			}
		}
	}
}

func fixImportCheck(body []byte, importPath string) ([]byte, error) {
	fset := token.NewFileSet()
	// todo: see if we can restrict the mode some more
	f, err := parser.ParseFile(fset, "", body, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}
	var after *ast.CommentGroup
	var pos token.Pos = token.Pos(len(body))
	for _, v := range f.Comments {
		text := strings.TrimSpace(v.Text())
		if v.Pos() > f.Package && v.Pos() < pos && strings.HasPrefix(text, "import") {
			pos = v.Pos()
			after = v
		}
	}
	if bytes.IndexByte(body[f.Package:pos], '\n') == -1 {
		comment := fmt.Sprintf(`// import "%s"`, importPath)
		buf := new(bytes.Buffer)
		buf.Write(body[:after.Pos()-1])
		buf.WriteString(comment)
		buf.Write(body[after.End()-1:])
		body = buf.Bytes()
	}
	return body, nil
}

func rewriteImports(importSites []string) (rewritten bool) {
	filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return nil
		}
		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
		if err != nil {
			return nil
		}
		for _, im := range f.Imports {
			for _, site := range importSites {
				v, err := strconv.Unquote(im.Path.Value)
				if err != nil {
					log.Fatal(err)
				}
				// don't replace ourself
				if v == relpath || strings.HasPrefix(v, relpath+"/") {
					if *verbose {
						log.Printf("skipping %s because it is in %s\n", v, relpath)
					}
					continue
				}
				if strings.HasPrefix(v, site+"/") {
					target := ThirdParty + "/" + v
					if !*relative {
						target = relpath + "/" + target
					}
					if *verbose {
						log.Printf("rewrite %s: %s -> %s\n", path, v, target)
					}
					if *dryrun {
						break
					}
					if !astutil.RewriteImport(fset, f, v, target) {
						log.Fatalf("%s: could not rewrite %s", path, im.Path.Value)
					}
					rewritten = true
					break
				}
			}
		}
		var buf bytes.Buffer
		if err := printer.Fprint(&buf, fset, f); err != nil {
			log.Fatal(err)
		}
		fb, err := format.Source(buf.Bytes())
		if err != nil {
			log.Fatal(err)
		}
		if err := ioutil.WriteFile(path, fb, info.Mode()); err != nil {
			log.Fatal(err)
		}
		return nil
	})
	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log/syslog"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"strings"

	cors "gopkg.in/gin-contrib/cors.v1"

	"github.com/gin-gonic/gin"
)

const version = "0.0.4"

var (
	// Commands mapping to control OMXPlayer, these are piped via STDIN to omxplayer process
	Commands = map[string]string{
		"pause":             "p",            // Pause/continue playback
		"stop":              "q",            // Stop playback and exit
		"volume_up":         "+",            // Change volume by +3dB
		"volume_down":       "-",            // Change volume by -3dB
		"subtitles":         "s",            // Enable/disable subtitles
		"seek_back":         "\x1b\x5b\x44", // Seek -30 seconds
		"seek_back_fast":    "\x1b\x5b\x42", // Seek -600 second
		"seek_forward":      "\x1b\x5b\x43", // Seek +30 second
		"seek_forward_fast": "\x1b\x5b\x41", // Seek +600 seconds
	}

	// OmxPath is path to omxplayer executable
	OmxPath string

	// Omx is a child process for spawning omxplayer
	Omx *exec.Cmd

	// OmxIn is a child process STDIN pipe to send commands
	OmxIn io.WriteCloser

	// Command is a channel to pass along commands to the player routine
	Command chan string

	Status = make(chan Content)

	// CurrentURL represents currently played media
	CurrentContent *Content

	// Syslog logger
	syslogger *syslog.Writer
)

// Content describes model of currently playable video.
type Content struct {
	RawURL    string                 `json:"url,omitempty"`
	MediaInfo map[string]interface{} `json:"mediaInfo,omitempty"`
}

// APIErr is a generic structure for all errors returned from API
type APIErr struct {
	Message string `json:"message,omitempty"`
}

// Determine the full path to omxplayer executable. Returns error if not found.
func omxDetect() error {
	buff, err := exec.Command("which", "omxplayer").Output()
	if err != nil {
		return err
	}

	// Set path in global variable
	OmxPath = strings.TrimSpace(string(buff))

	return nil
}

// Start command listener. Commands are coming in through a channel.
func omxListen() {
	Command = make(chan string)

	for {
		command := <-Command

		// Skip command handling of omx player is not active
		if Omx == nil {
			continue
		}

		// Send command to the player
		omxWrite(command)

		// Attempt to kill the process if stop command is requested
		if command == "stop" {
			err := Omx.Process.Kill()
			if err != nil {
				syslogger.Err(err.Error())
			}
		}

		broadcastStatus()
	}
}

// Start omxplayer playback for a given video file. Returns error if start fails.
func omxPlay(c Content) error {
	contentURL, err := url.Parse(c.RawURL)
	if err != nil {
		return err
	}

	Omx = exec.Command(
		OmxPath,             // path to omxplayer executable
		"--blank",           // set background to black
		"--adev",            // audio out device
		"hdmi",              // using hdmi for audio/video
		contentURL.String(), // path to video file
	)

	// Grab child process STDIN
	stdin, err := Omx.StdinPipe()
	if err != nil {
		return err
	}

	defer stdin.Close()

	// Redirect output for debugging purposes
	// Omx.Stdout = os.Stdout

	// Start omxplayer execution.
	// If successful, something will appear on HDMI display.
	err = Omx.Start()
	if err != nil {
		return err
	}

	// Set current file
	CurrentContent = &c

	broadcastStatus()

	// Make child's STDIN globally available
	OmxIn = stdin

	// Wait until child process is finished
	err = Omx.Wait()
	if err != nil {
		syslogger.Err(fmt.Sprintln("Process exited with error:", err))
	}

	omxCleanup()

	return nil
}

// Write a command string to the omxplayer process's STDIN
func omxWrite(command string) {
	if OmxIn != nil {
		io.WriteString(OmxIn, Commands[command])
	}
}

// Terminate any running omxplayer processes. Fixes random hangs.
func omxKill() {
	exec.Command("killall", "omxplayer.bin").Output()
	exec.Command("killall", "omxplayer").Output()
}

// Reset internal state and stop any running processes
func omxCleanup() {
	Omx = nil
	OmxIn = nil
	CurrentContent = nil

	omxKill()
}

// Check if player is currently active
func omxIsActive() bool {
	return Omx != nil
}

func httpCommand(c *gin.Context) {
	val := c.Params.ByName("command")

	if _, ok := Commands[val]; !ok {
		c.JSON(http.StatusBadRequest, APIErr{"Invalid command"})
		return
	}

	// Handle requested commmand
	Command <- val

	c.Status(http.StatusAccepted)
}

func httpPlay(c *gin.Context) {
	if omxIsActive() {
		c.JSON(http.StatusBadRequest, APIErr{"Player is already running"})
		return
	}

	content := Content{}
	err := c.BindJSON(&content)
	if err != nil {
		c.JSON(http.StatusBadRequest, APIErr{err.Error()})
		return
	}

	go omxPlay(content)

	c.Status(http.StatusAccepted)
}

func httpStatus(c *gin.Context) {
	result := struct {
		Running        bool     `json:"running"`
		CurrentContent *Content `json:"currentContent,omitempty"`
	}{
		Running:        omxIsActive(),
		CurrentContent: CurrentContent,
	}

	c.JSON(http.StatusOK, result)
}

func streamStatus(c *gin.Context) {
	c.Stream(func(w io.Writer) bool {
		c.SSEvent("status", <-Status)
		return true
	})
}

func broadcastStatus() {
	if CurrentContent != nil {
		Status <- *CurrentContent
	}
}

func terminate(message string, code int) {
	fmt.Println(message)
	os.Exit(code)
}

func main() {
	fmt.Printf("omx-remote-api v%v\n", version)

	syslogger, _ = syslog.New(syslog.LOG_NOTICE, "omx-remote-api")

	// Check if player is installed
	if omxDetect() != nil {
		terminate("omxplayer is not installed", 1)
	}

	// Make sure nothing is running
	omxCleanup()

	// Start a remote command listener
	go omxListen()

	// Disable debugging mode
	gin.SetMode("release")
	gin.LoggerWithWriter(syslogger)

	// Setup HTTP server
	router := gin.Default()

	// CORS
	router.Use(cors.Default())

	router.GET("/status", httpStatus)
	router.GET("/status/stream", streamStatus)
	router.POST("/play", httpPlay)
	router.POST("/commands/:command", httpCommand)

	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}

	fmt.Println("Starting server on 0.0.0.0:" + port)
	router.Run(":" + port)
}
<commit_msg>broadcast status at the end<commit_after>package main

import (
	"fmt"
	"io"
	"log/syslog"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"strings"

	cors "gopkg.in/gin-contrib/cors.v1"

	"github.com/gin-gonic/gin"
)

const version = "0.0.4"

var (
	// Commands mapping to control OMXPlayer, these are piped via STDIN to omxplayer process
	Commands = map[string]string{
		"pause":             "p",            // Pause/continue playback
		"stop":              "q",            // Stop playback and exit
		"volume_up":         "+",            // Change volume by +3dB
		"volume_down":       "-",            // Change volume by -3dB
		"subtitles":         "s",            // Enable/disable subtitles
		"seek_back":         "\x1b\x5b\x44", // Seek -30 seconds
		"seek_back_fast":    "\x1b\x5b\x42", // Seek -600 second
		"seek_forward":      "\x1b\x5b\x43", // Seek +30 second
		"seek_forward_fast": "\x1b\x5b\x41", // Seek +600 seconds
	}

	// OmxPath is path to omxplayer executable
	OmxPath string

	// Omx is a child process for spawning omxplayer
	Omx *exec.Cmd

	// OmxIn is a child process STDIN pipe to send commands
	OmxIn io.WriteCloser

	// Command is a channel to pass along commands to the player routine
	Command chan string

	Status = make(chan Content)

	// CurrentURL represents currently played media
	CurrentContent *Content

	// Syslog logger
	syslogger *syslog.Writer
)

// Content describes model of currently playable video.
type Content struct {
	RawURL    string                 `json:"url,omitempty"`
	MediaInfo map[string]interface{} `json:"mediaInfo,omitempty"`
}

// APIErr is a generic structure for all errors returned from API
type APIErr struct {
	Message string `json:"message,omitempty"`
}

// Determine the full path to omxplayer executable. Returns error if not found.
func omxDetect() error {
	buff, err := exec.Command("which", "omxplayer").Output()
	if err != nil {
		return err
	}

	// Set path in global variable
	OmxPath = strings.TrimSpace(string(buff))

	return nil
}

// Start command listener. Commands are coming in through a channel.
func omxListen() {
	Command = make(chan string)

	for {
		command := <-Command

		// Skip command handling of omx player is not active
		if Omx == nil {
			continue
		}

		// Send command to the player
		omxWrite(command)

		// Attempt to kill the process if stop command is requested
		if command == "stop" {
			err := Omx.Process.Kill()
			if err != nil {
				syslogger.Err(err.Error())
			}
		}

		broadcastStatus()
	}
}

// Start omxplayer playback for a given video file. Returns error if start fails.
func omxPlay(c Content) error {
	contentURL, err := url.Parse(c.RawURL)
	if err != nil {
		return err
	}

	Omx = exec.Command(
		OmxPath,             // path to omxplayer executable
		"--blank",           // set background to black
		"--adev",            // audio out device
		"hdmi",              // using hdmi for audio/video
		contentURL.String(), // path to video file
	)

	// Grab child process STDIN
	stdin, err := Omx.StdinPipe()
	if err != nil {
		return err
	}

	defer stdin.Close()

	// Redirect output for debugging purposes
	// Omx.Stdout = os.Stdout

	// Start omxplayer execution.
	// If successful, something will appear on HDMI display.
	err = Omx.Start()
	if err != nil {
		return err
	}

	// Set current file
	CurrentContent = &c

	broadcastStatus()

	// Make child's STDIN globally available
	OmxIn = stdin

	// Wait until child process is finished
	err = Omx.Wait()
	if err != nil {
		syslogger.Err(fmt.Sprintln("Process exited with error:", err))
	}

	broadcastStatus()
	omxCleanup()

	return nil
}

// Write a command string to the omxplayer process's STDIN
func omxWrite(command string) {
	if OmxIn != nil {
		io.WriteString(OmxIn, Commands[command])
	}
}

// Terminate any running omxplayer processes. Fixes random hangs.
func omxKill() {
	exec.Command("killall", "omxplayer.bin").Output()
	exec.Command("killall", "omxplayer").Output()
}

// Reset internal state and stop any running processes
func omxCleanup() {
	Omx = nil
	OmxIn = nil
	CurrentContent = nil

	omxKill()
}

// Check if player is currently active
func omxIsActive() bool {
	return Omx != nil
}

func httpCommand(c *gin.Context) {
	val := c.Params.ByName("command")

	if _, ok := Commands[val]; !ok {
		c.JSON(http.StatusBadRequest, APIErr{"Invalid command"})
		return
	}

	// Handle requested commmand
	Command <- val

	c.Status(http.StatusAccepted)
}

func httpPlay(c *gin.Context) {
	if omxIsActive() {
		c.JSON(http.StatusBadRequest, APIErr{"Player is already running"})
		return
	}

	content := Content{}
	err := c.BindJSON(&content)
	if err != nil {
		c.JSON(http.StatusBadRequest, APIErr{err.Error()})
		return
	}

	go omxPlay(content)

	c.Status(http.StatusAccepted)
}

func httpStatus(c *gin.Context) {
	result := struct {
		Running        bool     `json:"running"`
		CurrentContent *Content `json:"currentContent,omitempty"`
	}{
		Running:        omxIsActive(),
		CurrentContent: CurrentContent,
	}

	c.JSON(http.StatusOK, result)
}

func streamStatus(c *gin.Context) {
	c.Stream(func(w io.Writer) bool {
		c.SSEvent("status", <-Status)
		return true
	})
}

func broadcastStatus() {
	if CurrentContent != nil {
		Status <- *CurrentContent
	}
}

func terminate(message string, code int) {
	fmt.Println(message)
	os.Exit(code)
}

func main() {
	fmt.Printf("omx-remote-api v%v\n", version)

	syslogger, _ = syslog.New(syslog.LOG_NOTICE, "omx-remote-api")

	// Check if player is installed
	if omxDetect() != nil {
		terminate("omxplayer is not installed", 1)
	}

	// Make sure nothing is running
	omxCleanup()

	// Start a remote command listener
	go omxListen()

	// Disable debugging mode
	gin.SetMode("release")
	gin.LoggerWithWriter(syslogger)

	// Setup HTTP server
	router := gin.Default()

	// CORS
	router.Use(cors.Default())

	router.GET("/status", httpStatus)
	router.GET("/status/stream", streamStatus)
	router.POST("/play", httpPlay)
	router.POST("/commands/:command", httpCommand)

	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}

	fmt.Println("Starting server on 0.0.0.0:" + port)
	router.Run(":" + port)
}
<|endoftext|>


<commit_before>package main // import "github.com/gaia-adm/pumba"

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gaia-adm/pumba/action"
	"github.com/gaia-adm/pumba/container"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
)

var (
	wg      sync.WaitGroup
	client  container.Client
	cleanup bool
)

const (
	defaultKillSignal = "SIGKILL"
	re2prefix         = "re2:"
)

type commandT struct {
	pattern string
	names   []string
	command string
	signal  string
}

func init() {
	log.SetLevel(log.InfoLevel)
}

func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

func main() {
	rootCertPath := "/etc/ssl/docker"

	if os.Getenv("DOCKER_CERT_PATH") != "" {
		rootCertPath = os.Getenv("DOCKER_CERT_PATH")
	}

	app := cli.NewApp()
	app.Name = "pumba"
	app.Version = "0.1.2"
	app.Usage = "Pumba is a resiliency tool that helps applications tolerate random Docker container failures."
	app.Before = before
	app.Action = start
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "host, H",
			Usage:  "daemon socket to connect to",
			Value:  "unix:///var/run/docker.sock",
			EnvVar: "DOCKER_HOST",
		},
		cli.BoolFlag{
			Name:  "tls",
			Usage: "use TLS; implied by --tlsverify",
		},
		cli.BoolFlag{
			Name:   "tlsverify",
			Usage:  "use TLS and verify the remote",
			EnvVar: "DOCKER_TLS_VERIFY",
		},
		cli.StringFlag{
			Name:  "tlscacert",
			Usage: "trust certs signed only by this CA",
			Value: fmt.Sprintf("%s/ca.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlscert",
			Usage: "client certificate for TLS authentication",
			Value: fmt.Sprintf("%s/cert.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlskey",
			Usage: "client key for TLS authentication",
			Value: fmt.Sprintf("%s/key.pem", rootCertPath),
		},
		cli.BoolFlag{
			Name:  "debug",
			Usage: "enable debug mode with verbose logging",
		},
		cli.StringSliceFlag{
			Name:  "chaos",
			Usage: "chaos command: `container(s,)/re2:regex|interval(s/m/h postfix)|STOP/KILL(:SIGNAL)/RM`",
		},
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func before(c *cli.Context) error {
	if c.GlobalBool("debug") {
		log.SetLevel(log.DebugLevel)
	}

	cleanup = c.GlobalBool("cleanup")

	// Set-up container client
	tls, err := tlsConfig(c)
	if err != nil {
		return err
	}

	client = container.NewClient(c.GlobalString("host"), tls, !c.GlobalBool("no-pull"))

	handleSignals()
	return nil
}

func start(c *cli.Context) {
	if err := actions.CheckPrereqs(client, cleanup); err != nil {
		log.Fatal(err)
	}
	if err := createChaos(actions.Pumba{}, c.GlobalStringSlice("chaos"), 1, false); err != nil {
		log.Fatal(err)
	}
}

func createChaos(chaos actions.Chaos, args []string, limit int, test bool) error {
	// docker channel to pass all "stop" commands to
	dc := make(chan commandT)
	glimit := limit * len(args)

	// range over all chaos arguments
	for _, chaosArg := range args {
		s := strings.Split(chaosArg, "|")
		if len(s) != 3 {
			return errors.New("Unexpected format for chaos_arg: use | separated triple")
		}
		// get container name pattern
		var pattern string
		var names []string
		if strings.HasPrefix(s[0], re2prefix) {
			pattern = strings.Trim(s[0], re2prefix)
			log.Debugf("Pattern: '%s'", pattern)
		} else {
			names = strings.Split(s[0], ",")
			log.Debugf("Names: '%s'", s[0])
		}
		// get interval duration
		interval, err := time.ParseDuration(s[1])
		if err != nil {
			return err
		}
		log.Debugf("Interval: '%s'", interval.String())
		// get command and signal (if specified); convert everything to upper case
		cs := strings.Split(strings.ToUpper(s[2]), ":")
		command := cs[0]
		if !stringInSlice(command, []string{"STOP", "KILL", "RM"}) {
			return errors.New("Unexpected command in chaos_arg: can be STOP, KILL or RM")
		}
		log.Debugf("Command: '%s'", command)
		signal := defaultKillSignal
		if len(cs) == 2 {
			signal = cs[1]
			log.Debugf("Signal: '%s'", signal)
		}

		ticker := time.NewTicker(interval)
		go func(cmd commandT, limit int, test bool) {
			for range ticker.C {
				if limit > 0 {
					log.Debugf("Tick: '%s'", cmd)
					dc <- cmd
					if test {
						limit--
					}
				}
			}
		}(commandT{pattern, names, command, signal}, limit, test)
	}
	for cmd := range dc {
		if test {
			glimit--
		}
		if glimit == 0 {
			break
		}
		wg.Add(1)
		go func(cmd commandT) {
			defer wg.Done()
			var err error
			switch cmd.command {
			case "STOP":
				if cmd.pattern == "" {
					err = chaos.StopByName(client, cmd.names)
				} else {
					err = chaos.StopByPattern(client, cmd.pattern)
				}
			case "KILL":
				if cmd.pattern == "" {
					err = chaos.KillByName(client, cmd.names, cmd.signal)
				} else {
					err = chaos.KillByPattern(client, cmd.pattern, cmd.signal)
				}
			case "RM":
				if cmd.pattern == "" {
					err = chaos.RemoveByName(client, cmd.names, true)
				} else {
					err = chaos.RemoveByPattern(client, cmd.pattern, true)
				}
			}
			if err != nil {
				log.Error(err)
			}
		}(cmd)

	}
	return nil
}

func handleSignals() {
	// Graceful shut-down on SIGINT/SIGTERM
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)

	go func() {
		<-c
		wg.Wait()
		os.Exit(1)
	}()
}

// tlsConfig translates the command-line options into a tls.Config struct
func tlsConfig(c *cli.Context) (*tls.Config, error) {
	var tlsConfig *tls.Config
	var err error
	caCertFlag := c.GlobalString("tlscacert")
	certFlag := c.GlobalString("tlscert")
	keyFlag := c.GlobalString("tlskey")

	if c.GlobalBool("tls") || c.GlobalBool("tlsverify") {
		tlsConfig = &tls.Config{
			InsecureSkipVerify: !c.GlobalBool("tlsverify"),
		}

		// Load CA cert
		if caCertFlag != "" {
			var caCert []byte

			if strings.HasPrefix(caCertFlag, "/") {
				caCert, err = ioutil.ReadFile(caCertFlag)
				if err != nil {
					return nil, err
				}
			} else {
				caCert = []byte(caCertFlag)
			}

			caCertPool := x509.NewCertPool()
			caCertPool.AppendCertsFromPEM(caCert)

			tlsConfig.RootCAs = caCertPool
		}

		// Load client certificate
		if certFlag != "" && keyFlag != "" {
			var cert tls.Certificate

			if strings.HasPrefix(certFlag, "/") && strings.HasPrefix(keyFlag, "/") {
				cert, err = tls.LoadX509KeyPair(certFlag, keyFlag)
				if err != nil {
					return nil, err
				}
			} else {
				cert, err = tls.X509KeyPair([]byte(certFlag), []byte(keyFlag))
				if err != nil {
					return nil, err
				}
			}
			tlsConfig.Certificates = []tls.Certificate{cert}
		}
	}

	return tlsConfig, nil
}
<commit_msg>version up<commit_after>package main // import "github.com/gaia-adm/pumba"

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gaia-adm/pumba/action"
	"github.com/gaia-adm/pumba/container"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
)

var (
	wg      sync.WaitGroup
	client  container.Client
	cleanup bool
)

const (
	defaultKillSignal = "SIGKILL"
	re2prefix         = "re2:"
)

type commandT struct {
	pattern string
	names   []string
	command string
	signal  string
}

func init() {
	log.SetLevel(log.InfoLevel)
}

func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

func main() {
	rootCertPath := "/etc/ssl/docker"

	if os.Getenv("DOCKER_CERT_PATH") != "" {
		rootCertPath = os.Getenv("DOCKER_CERT_PATH")
	}

	app := cli.NewApp()
	app.Name = "pumba"
	app.Version = "0.1.3"
	app.Usage = "Pumba is a resiliency tool that helps applications tolerate random Docker container failures."
	app.Before = before
	app.Action = start
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "host, H",
			Usage:  "daemon socket to connect to",
			Value:  "unix:///var/run/docker.sock",
			EnvVar: "DOCKER_HOST",
		},
		cli.BoolFlag{
			Name:  "tls",
			Usage: "use TLS; implied by --tlsverify",
		},
		cli.BoolFlag{
			Name:   "tlsverify",
			Usage:  "use TLS and verify the remote",
			EnvVar: "DOCKER_TLS_VERIFY",
		},
		cli.StringFlag{
			Name:  "tlscacert",
			Usage: "trust certs signed only by this CA",
			Value: fmt.Sprintf("%s/ca.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlscert",
			Usage: "client certificate for TLS authentication",
			Value: fmt.Sprintf("%s/cert.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlskey",
			Usage: "client key for TLS authentication",
			Value: fmt.Sprintf("%s/key.pem", rootCertPath),
		},
		cli.BoolFlag{
			Name:  "debug",
			Usage: "enable debug mode with verbose logging",
		},
		cli.StringSliceFlag{
			Name:  "chaos",
			Usage: "chaos command: `container(s,)/re2:regex|interval(s/m/h postfix)|STOP/KILL(:SIGNAL)/RM`",
		},
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func before(c *cli.Context) error {
	if c.GlobalBool("debug") {
		log.SetLevel(log.DebugLevel)
	}

	cleanup = c.GlobalBool("cleanup")

	// Set-up container client
	tls, err := tlsConfig(c)
	if err != nil {
		return err
	}

	client = container.NewClient(c.GlobalString("host"), tls, !c.GlobalBool("no-pull"))

	handleSignals()
	return nil
}

func start(c *cli.Context) {
	if err := actions.CheckPrereqs(client, cleanup); err != nil {
		log.Fatal(err)
	}
	if err := createChaos(actions.Pumba{}, c.GlobalStringSlice("chaos"), 1, false); err != nil {
		log.Fatal(err)
	}
}

func createChaos(chaos actions.Chaos, args []string, limit int, test bool) error {
	// docker channel to pass all "stop" commands to
	dc := make(chan commandT)
	glimit := limit * len(args)

	// range over all chaos arguments
	for _, chaosArg := range args {
		s := strings.Split(chaosArg, "|")
		if len(s) != 3 {
			return errors.New("Unexpected format for chaos_arg: use | separated triple")
		}
		// get container name pattern
		var pattern string
		var names []string
		if strings.HasPrefix(s[0], re2prefix) {
			pattern = strings.Trim(s[0], re2prefix)
			log.Debugf("Pattern: '%s'", pattern)
		} else {
			names = strings.Split(s[0], ",")
			log.Debugf("Names: '%s'", s[0])
		}
		// get interval duration
		interval, err := time.ParseDuration(s[1])
		if err != nil {
			return err
		}
		log.Debugf("Interval: '%s'", interval.String())
		// get command and signal (if specified); convert everything to upper case
		cs := strings.Split(strings.ToUpper(s[2]), ":")
		command := cs[0]
		if !stringInSlice(command, []string{"STOP", "KILL", "RM"}) {
			return errors.New("Unexpected command in chaos_arg: can be STOP, KILL or RM")
		}
		log.Debugf("Command: '%s'", command)
		signal := defaultKillSignal
		if len(cs) == 2 {
			signal = cs[1]
			log.Debugf("Signal: '%s'", signal)
		}

		ticker := time.NewTicker(interval)
		go func(cmd commandT, limit int, test bool) {
			for range ticker.C {
				if limit > 0 {
					log.Debugf("Tick: '%s'", cmd)
					dc <- cmd
					if test {
						limit--
					}
				}
			}
		}(commandT{pattern, names, command, signal}, limit, test)
	}
	for cmd := range dc {
		if test {
			glimit--
		}
		if glimit == 0 {
			break
		}
		wg.Add(1)
		go func(cmd commandT) {
			defer wg.Done()
			var err error
			switch cmd.command {
			case "STOP":
				if cmd.pattern == "" {
					err = chaos.StopByName(client, cmd.names)
				} else {
					err = chaos.StopByPattern(client, cmd.pattern)
				}
			case "KILL":
				if cmd.pattern == "" {
					err = chaos.KillByName(client, cmd.names, cmd.signal)
				} else {
					err = chaos.KillByPattern(client, cmd.pattern, cmd.signal)
				}
			case "RM":
				if cmd.pattern == "" {
					err = chaos.RemoveByName(client, cmd.names, true)
				} else {
					err = chaos.RemoveByPattern(client, cmd.pattern, true)
				}
			}
			if err != nil {
				log.Error(err)
			}
		}(cmd)

	}
	return nil
}

func handleSignals() {
	// Graceful shut-down on SIGINT/SIGTERM
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)

	go func() {
		<-c
		wg.Wait()
		os.Exit(1)
	}()
}

// tlsConfig translates the command-line options into a tls.Config struct
func tlsConfig(c *cli.Context) (*tls.Config, error) {
	var tlsConfig *tls.Config
	var err error
	caCertFlag := c.GlobalString("tlscacert")
	certFlag := c.GlobalString("tlscert")
	keyFlag := c.GlobalString("tlskey")

	if c.GlobalBool("tls") || c.GlobalBool("tlsverify") {
		tlsConfig = &tls.Config{
			InsecureSkipVerify: !c.GlobalBool("tlsverify"),
		}

		// Load CA cert
		if caCertFlag != "" {
			var caCert []byte

			if strings.HasPrefix(caCertFlag, "/") {
				caCert, err = ioutil.ReadFile(caCertFlag)
				if err != nil {
					return nil, err
				}
			} else {
				caCert = []byte(caCertFlag)
			}

			caCertPool := x509.NewCertPool()
			caCertPool.AppendCertsFromPEM(caCert)

			tlsConfig.RootCAs = caCertPool
		}

		// Load client certificate
		if certFlag != "" && keyFlag != "" {
			var cert tls.Certificate

			if strings.HasPrefix(certFlag, "/") && strings.HasPrefix(keyFlag, "/") {
				cert, err = tls.LoadX509KeyPair(certFlag, keyFlag)
				if err != nil {
					return nil, err
				}
			} else {
				cert, err = tls.X509KeyPair([]byte(certFlag), []byte(keyFlag))
				if err != nil {
					return nil, err
				}
			}
			tlsConfig.Certificates = []tls.Certificate{cert}
		}
	}

	return tlsConfig, nil
}
<|endoftext|>


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

package main

import (
	"fmt"
	flag "github.com/docker/docker/pkg/mflag"
	_ "github.com/trivago/gollum/consumer"
	_ "github.com/trivago/gollum/contrib"
	_ "github.com/trivago/gollum/distributor"
	_ "github.com/trivago/gollum/format"
	_ "github.com/trivago/gollum/producer"
	"github.com/trivago/gollum/shared"
	"io/ioutil"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"time"
)

const (
	gollumMajorVer = 0
	gollumMinorVer = 1
	gollumPatchVer = 0
)

func dumpMemoryProfile() {
	if file, err := os.Create(*memProfilePtr); err != nil {
		panic(err)
	} else {
		defer file.Close()
		pprof.WriteHeapProfile(file)
	}
}

func main() {
	flag.Parse()

	if *helpPtr || *configFilePtr == "" {
		flag.Usage()
		fmt.Println("Nothing to do. We must go.")
		return
	}

	if *versionPtr {
		fmt.Printf("Gollum v%d.%d.%d\n", gollumMajorVer, gollumMinorVer, gollumPatchVer)
		return
	}

	if *numCPU == 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		runtime.GOMAXPROCS(*numCPU)
	}

	if *pidFilePtr != "" {
		ioutil.WriteFile(*pidFilePtr, []byte(strconv.Itoa(os.Getpid())), 0644)
	}

	if *cpuProfilePtr != "" {
		if file, err := os.Create(*cpuProfilePtr); err != nil {
			panic(err)
		} else {
			defer file.Close()
			pprof.StartCPUProfile(file)
			defer pprof.StopCPUProfile()
		}
	}

	if *memProfilePtr != "" {
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		go func() {
			for {
				<-ticker.C
				dumpMemoryProfile()
			}
		}()
	}

	// Start the gollum multiplexer

	conf, err := shared.ReadConfig(*configFilePtr)
	if err != nil {
		fmt.Printf("Error: %s", err.Error())
		os.Exit(-1)
	}

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

package main

import (
	"fmt"
	flag "github.com/docker/docker/pkg/mflag"
	_ "github.com/trivago/gollum/consumer"
	_ "github.com/trivago/gollum/contrib"
	_ "github.com/trivago/gollum/distributor"
	_ "github.com/trivago/gollum/format"
	_ "github.com/trivago/gollum/producer"
	"github.com/trivago/gollum/shared"
	"io/ioutil"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"time"
)

const (
	gollumMajorVer = 0
	gollumMinorVer = 2
	gollumPatchVer = 0
)

func dumpMemoryProfile() {
	if file, err := os.Create(*memProfilePtr); err != nil {
		panic(err)
	} else {
		defer file.Close()
		pprof.WriteHeapProfile(file)
	}
}

func main() {
	flag.Parse()

	if *helpPtr || *configFilePtr == "" {
		flag.Usage()
		fmt.Println("Nothing to do. We must go.")
		return
	}

	if *versionPtr {
		fmt.Printf("Gollum v%d.%d.%d\n", gollumMajorVer, gollumMinorVer, gollumPatchVer)
		return
	}

	if *numCPU == 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		runtime.GOMAXPROCS(*numCPU)
	}

	if *pidFilePtr != "" {
		ioutil.WriteFile(*pidFilePtr, []byte(strconv.Itoa(os.Getpid())), 0644)
	}

	if *cpuProfilePtr != "" {
		if file, err := os.Create(*cpuProfilePtr); err != nil {
			panic(err)
		} else {
			defer file.Close()
			pprof.StartCPUProfile(file)
			defer pprof.StopCPUProfile()
		}
	}

	if *memProfilePtr != "" {
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		go func() {
			for {
				<-ticker.C
				dumpMemoryProfile()
			}
		}()
	}

	// Start the gollum multiplexer

	conf, err := shared.ReadConfig(*configFilePtr)
	if err != nil {
		fmt.Printf("Error: %s", err.Error())
		os.Exit(-1)
	}

	plex := newMultiplexer(conf, *msgProfilePtr)
	plex.run()
}
<|endoftext|>


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

// Program stackdriver-reverse-proxy provides a Stackdriver reverse
// proxy that creates traces for the incoming requests, logs request
// details, and reports errors.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"

	"cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/trace"
)

var (
	projectID string
	creds     string

	listen    string
	target    string
	tlsCert   string
	tlsKey    string
	traceFrac float64

	enableErrorReports bool
)

const usage = `stackdriver-reverse-proxy [opts...] -http=<host:port> -target=<target_url>

For example, to start at localhost:8080 to proxy requests to localhost:6060,
  $ stackdriver-reverse-proxy -http=:8080 -target=http://localhost:6060

Options:
  -project        Google Cloud Platform project ID if running outside of GCP.
  -creds          Google Cloud credentials file. If not given, ADC is used.
  -enable-errors  Set true to enable error reporting to Stackdriver.

Tracing options:
  -trace-fraction Tracing sampling fraction, between 0 and 1.0.

HTTPS options:
  -tls-cert TLS cert file to start an HTTPS proxy.
  -tls-key  TLS key file to start an HTTPS proxy.
`

func main() {
	ctx := context.Background()
	flag.Usage = func() {
		fmt.Println(usage)
	}

	flag.StringVar(&projectID, "project", "", "")
	flag.StringVar(&creds, "credentials", "", "")

	flag.StringVar(&listen, "http", ":6996", "host:port proxy listens")
	flag.StringVar(&target, "target", "", "target server")
	flag.Float64Var(&traceFrac, "trace-fraction", 1, "sampling fraction for tracing")
	flag.BoolVar(&enableErrorReports, "enable-errorreports", false, "set to enable error reporting to stackdriver")
	flag.StringVar(&tlsCert, "tls-cert", "", "TLS cert file to start an HTTPS proxy")
	flag.StringVar(&tlsKey, "tls-key", "", "TLS key file to start an HTTPS proxy")
	flag.Parse()

	if target == "" {
		usageExit()
	}
	if projectID == "" {
		// Try to retrieve it from metadata server.
		if metadata.OnGCE() {
			pid, err := metadata.ProjectID()
			if err != nil {
				log.Fatalf("Cannot get project ID from metadata server: %v", err)
			}
			projectID = pid
		}
	}
	if projectID == "" {
		usageExit()
	}

	url, err := url.Parse(target)
	if err != nil {
		log.Fatalf("Cannot URL parse -target: %v", err)
	}

	tc, err := trace.NewClient(ctx, projectID)
	if err != nil {
		log.Fatalf("Cannot initiate trace client: %v", err)
	}
	sp, _ := trace.NewLimitedSampler(traceFrac, 1<<32)
	tc.SetSamplingPolicy(sp)

	proxy := httputil.NewSingleHostReverseProxy(url)
	proxy.Transport = &transport{
		Trace: tc,
	}
	if tlsCert != "" && tlsKey != "" {
		log.Fatal(http.ListenAndServeTLS(listen, tlsCert, tlsKey, proxy))
	} else {
		log.Fatal(http.ListenAndServe(listen, proxy))
	}
}

type transport struct {
	Trace *trace.Client
}

func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) {
	s := t.Trace.SpanFromRequest(req)
	defer s.Finish()

	req.Header.Set("X-Cloud-Trace-Context", s.Header())
	resp, err := http.DefaultTransport.RoundTrip(req)
	if err != nil {
		s.SetLabel("error", err.Error())
	}
	return resp, err
}

func usageExit() {
	flag.Usage()
	os.Exit(1)
}
<commit_msg>remove unused flags, improve the usage text (#7)<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Program stackdriver-reverse-proxy provides a Stackdriver reverse
// proxy that creates traces for the incoming requests, logs request
// details, and reports errors.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"

	"cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/trace"
)

var (
	projectID string

	listen    string
	target    string
	tlsCert   string
	tlsKey    string
	traceFrac float64
)

const usage = `stackdriver-reverse-proxy [opts...] -target=<host:port>

For example, to start at localhost:6996 to proxy requests to localhost:6060,
  $ stackdriver-reverse-proxy -target=http://localhost:6060

Options:
  -http           hostname:port to start the proxy server, by default localhost:6996.
  -target         hostname:port where the app server is running.
  -project        Google Cloud Platform project ID if running outside of GCP.

Tracing options:
  -trace-fraction Tracing sampling fraction, between 0 and 1.0.

HTTPS options:
  -tls-cert TLS cert file to start an HTTPS proxy.
  -tls-key  TLS key file to start an HTTPS proxy.
`

func main() {
	ctx := context.Background()
	flag.Usage = func() {
		fmt.Println(usage)
	}

	flag.StringVar(&projectID, "project", "", "")
	flag.StringVar(&listen, "http", ":6996", "host:port proxy listens")
	flag.StringVar(&target, "target", "", "target server")
	flag.Float64Var(&traceFrac, "trace-fraction", 1, "sampling fraction for tracing")
	flag.StringVar(&tlsCert, "tls-cert", "", "TLS cert file to start an HTTPS proxy")
	flag.StringVar(&tlsKey, "tls-key", "", "TLS key file to start an HTTPS proxy")
	flag.Parse()

	if target == "" {
		usageExit()
	}
	if projectID == "" {
		// Try to retrieve it from metadata server.
		if metadata.OnGCE() {
			pid, err := metadata.ProjectID()
			if err != nil {
				log.Fatalf("Cannot get project ID from metadata server: %v", err)
			}
			projectID = pid
		}
	}
	if projectID == "" {
		usageExit()
	}

	url, err := url.Parse(target)
	if err != nil {
		log.Fatalf("Cannot URL parse -target: %v", err)
	}

	tc, err := trace.NewClient(ctx, projectID)
	if err != nil {
		log.Fatalf("Cannot initiate trace client: %v", err)
	}
	sp, _ := trace.NewLimitedSampler(traceFrac, 1<<32)
	tc.SetSamplingPolicy(sp)

	proxy := httputil.NewSingleHostReverseProxy(url)
	proxy.Transport = &transport{
		Trace: tc,
	}
	if tlsCert != "" && tlsKey != "" {
		log.Fatal(http.ListenAndServeTLS(listen, tlsCert, tlsKey, proxy))
	} else {
		log.Fatal(http.ListenAndServe(listen, proxy))
	}
}

type transport struct {
	Trace *trace.Client
}

func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) {
	s := t.Trace.SpanFromRequest(req)
	defer s.Finish()

	req.Header.Set("X-Cloud-Trace-Context", s.Header())
	resp, err := http.DefaultTransport.RoundTrip(req)
	if err != nil {
		s.SetLabel("error", err.Error())
	}
	return resp, err
}

func usageExit() {
	flag.Usage()
	os.Exit(1)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"

	"github.com/hidez8891/albero-server/module"
	"github.com/hidez8891/albero-server/utility/array"
	"github.com/hidez8891/albero-server/utility/json"

	_ "github.com/hidez8891/albero-server/module/image"
	_ "github.com/hidez8891/albero-server/module/zip"
)

const (
	defaultPort = 5358
)

func main() {
	var port uint
	flag.UintVar(&port, "port", defaultPort, "listen port")
	flag.UintVar(&port, "p", defaultPort, "listen port")
	flag.Parse()

	// path = url.encode(path1/path2/arch/image)

	// /support return support file type (json format)
	http.HandleFunc("/support", supportRouting)

	// /type?path=path_enc return file types (json format)
	http.HandleFunc("/type", typeRouting)

	// /files?path=path_enc return files (json format)
	http.HandleFunc("/files", filesRouting)

	// /image?path=path_enc return image (binary format)
	http.HandleFunc("/image", imageRouting)

	log.Printf("Listening on %d\n", port)
	err := http.ListenAndServe(fmt.Sprintf("localhost:%d", defaultPort), nil)
	if err != nil {
		log.Fatal("Listen And Serve:", err)
	}
}

func queryPath(r *http.Request) (string, error) {
	param := r.URL.Query()
	pathEnc := param.Get("path")
	if len(pathEnc) == 0 {
		return "", fmt.Errorf("need 'path' query")
	}

	path, err := url.QueryUnescape(pathEnc)
	if err != nil {
		return "", err
	}
	return path, nil
}

func splitRealVirtualPath(path string) (string, []string, error) {
	paths := strings.SplitAfter(path, "/")
	path = ""
	for len(paths) > 0 {
		pt := filepath.Join(path, paths[0])
		_, err := os.Stat(pt)
		if os.IsNotExist(err) {
			break
		}
		if err != nil {
			return "", nil, err
		}
		path = pt
		paths = paths[1:]
	}
	return path, paths, nil
}

func supportRouting(w http.ResponseWriter, r *http.Request) {
	exts := struct {
		Image []string `json:"image"`
		Arch  []string `json:"arch"`
	}{}

	types := module.SupportType()
	exts.Image = types[module.MODULE_IMAGE]
	exts.Arch = types[module.MODULE_ARCH]

	json.WriteResponse(w, exts)
}

func typeRouting(w http.ResponseWriter, r *http.Request) {
	path, err := queryPath(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// TODO path validation
	ext := filepath.Ext(path)
	supports := module.SupportType()

	ret := struct {
		Type string `json:"type"`
	}{}

	types := []struct {
		mtype module.ModuleType
		stype string
	}{
		{module.MODULE_IMAGE, "image"},
		{module.MODULE_ARCH, "arch"},
	}

	for _, tp := range types {
		exts := supports[tp.mtype]
		for _, ex := range exts {
			if ex == ext {
				// return file's type string
				ret.Type = tp.stype
				json.WriteResponse(w, ret)
				return
			}
		}
	}

	// not found
	http.NotFound(w, r)
}

func filesRouting(w http.ResponseWriter, r *http.Request) {
	path, err := queryPath(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	path, paths, err := splitRealVirtualPath(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	stat, _ := os.Stat(path)
	var files []string

	// 'paths' is not empty, 'path' needs archive file path
	// 'paths' is empty, 'path' needs archive file or directory path
	if len(paths) == 0 && stat.IsDir() {
		// directory
		files = getFilesFromDirectory(w, path)
	} else {
		files = getFilesFromArchive(w, path, paths)
	}
	if files == nil {
		return // still response error
	}

	// return json format
	ret := struct {
		Files []string `json:"files"`
	}{files}
	json.WriteResponse(w, ret)
}

func getFilesFromDirectory(w http.ResponseWriter, path string) []string {
	infos, err := ioutil.ReadDir(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return nil
	}

	files := make([]string, 0)
	for _, v := range infos {
		name := strings.Replace(v.Name(), "\\", "/", -1)
		if ind := strings.Index(name, "/"); ind >= 0 {
			name = name[ind+1:]
		}
		files = append(files, name)
	}
	return files
}

func getFilesFromArchive(w http.ResponseWriter, path string, paths []string) []string {
	mod := module.GetSupportModule(path)
	if mod == nil {
		http.Error(w, "no support type", http.StatusBadRequest)
		return nil
	}
	if mod.Type != module.MODULE_ARCH {
		http.Error(w, "wrong module operation", http.StatusBadRequest)
		return nil
	}

	file, err := module.NewReaderAt(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}
	defer file.Close()

	// TODO
	// support nested archive's path
	vpath := strings.Join(paths, "")
	files := mod.FuncArchFiles(file)
	if files == nil {
		http.Error(w, "fail read file", http.StatusInternalServerError)
		return nil
	}
	files = array.Filter(files, func(s string) bool {
		return strings.HasPrefix(s, vpath)
	})
	files = array.Map(files, func(s string) string {
		return s[len(vpath):] // remove prefix
	})
	files = array.Uniq(array.Map(files, func(s string) string {
		if index := strings.Index(s, "/"); index >= 0 {
			return s[:index] // remove suffix (when directory)
		}
		return s
	}))
	return files
}

func imageRouting(w http.ResponseWriter, r *http.Request) {
	path, err := queryPath(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	path, paths, err := splitRealVirtualPath(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	stat, _ := os.Stat(path)

	// 'paths' is not empty, 'path' needs archive file path
	// 'paths' is empty, 'path' needs image file
	if len(paths) == 0 {
		// reject directory path
		if stat.IsDir() {
			http.Error(w, "path is directory", http.StatusBadRequest)
			return
		}
		file := getImageDataFromFile(w, path)
		defer file.Close()
		imageRoutingResponse(w, path, file)
	} else {
		file := getImageDataFromArchive(w, path, paths)
		defer file.Close()
		imageRoutingResponse(w, paths[len(paths)-1], file)
	}
}

func getImageDataFromFile(w http.ResponseWriter, path string) module.Reader {
	r, err := module.NewReaderAt(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}
	return r
}

func getImageDataFromArchive(w http.ResponseWriter, path string, paths []string) module.Reader {
	mod := module.GetSupportModule(path)
	if mod == nil {
		http.Error(w, "no support type", http.StatusBadRequest)
		return nil
	}
	if mod.Type != module.MODULE_ARCH {
		http.Error(w, "wrong module operation", http.StatusBadRequest)
		return nil
	}

	arch, err := module.NewReaderAt(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}
	defer arch.Close()

	// TODO
	// support nested archive's path
	imgpath := strings.Join(paths, "")
	file := mod.FuncArchRead(arch, imgpath)
	if file == nil {
		http.Error(w, "fail read file", http.StatusInternalServerError)
		return nil
	}
	defer file.Data.Close()

	r, err := module.NewReader(file.Data, file.Size)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}
	return r
}

func imageRoutingResponse(w http.ResponseWriter, path string, r module.Reader) {
	mod := module.GetSupportModule(path)
	if mod == nil {
		http.Error(w, "no support type", http.StatusBadRequest)
		return
	}
	if mod.Type != module.MODULE_IMAGE {
		http.Error(w, "wrong module operation", http.StatusBadRequest)
		return
	}

	file := mod.FuncImageRead(r)
	if file == nil {
		http.Error(w, "fail read file", http.StatusInternalServerError)
		return
	}
	defer file.Data.Close()

	w.Header().Set("Content-Type", file.Mime)
	w.Header().Set("Content-Length", fmt.Sprintf("%d", file.Size))
	if _, err := io.Copy(w, file.Data); err != nil {
		log.Printf("ERR: WriteResponse: %v\n", err)
	}
}
<commit_msg>fix: forget file == nil case<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"

	"github.com/hidez8891/albero-server/module"
	"github.com/hidez8891/albero-server/utility/array"
	"github.com/hidez8891/albero-server/utility/json"

	_ "github.com/hidez8891/albero-server/module/image"
	_ "github.com/hidez8891/albero-server/module/zip"
)

const (
	defaultPort = 5358
)

func main() {
	var port uint
	flag.UintVar(&port, "port", defaultPort, "listen port")
	flag.UintVar(&port, "p", defaultPort, "listen port")
	flag.Parse()

	// path = url.encode(path1/path2/arch/image)

	// /support return support file type (json format)
	http.HandleFunc("/support", supportRouting)

	// /type?path=path_enc return file types (json format)
	http.HandleFunc("/type", typeRouting)

	// /files?path=path_enc return files (json format)
	http.HandleFunc("/files", filesRouting)

	// /image?path=path_enc return image (binary format)
	http.HandleFunc("/image", imageRouting)

	log.Printf("Listening on %d\n", port)
	err := http.ListenAndServe(fmt.Sprintf("localhost:%d", defaultPort), nil)
	if err != nil {
		log.Fatal("Listen And Serve:", err)
	}
}

func queryPath(r *http.Request) (string, error) {
	param := r.URL.Query()
	pathEnc := param.Get("path")
	if len(pathEnc) == 0 {
		return "", fmt.Errorf("need 'path' query")
	}

	path, err := url.QueryUnescape(pathEnc)
	if err != nil {
		return "", err
	}
	return path, nil
}

func splitRealVirtualPath(path string) (string, []string, error) {
	paths := strings.SplitAfter(path, "/")
	path = ""
	for len(paths) > 0 {
		pt := filepath.Join(path, paths[0])
		_, err := os.Stat(pt)
		if os.IsNotExist(err) {
			break
		}
		if err != nil {
			return "", nil, err
		}
		path = pt
		paths = paths[1:]
	}
	return path, paths, nil
}

func supportRouting(w http.ResponseWriter, r *http.Request) {
	exts := struct {
		Image []string `json:"image"`
		Arch  []string `json:"arch"`
	}{}

	types := module.SupportType()
	exts.Image = types[module.MODULE_IMAGE]
	exts.Arch = types[module.MODULE_ARCH]

	json.WriteResponse(w, exts)
}

func typeRouting(w http.ResponseWriter, r *http.Request) {
	path, err := queryPath(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// TODO path validation
	ext := filepath.Ext(path)
	supports := module.SupportType()

	ret := struct {
		Type string `json:"type"`
	}{}

	types := []struct {
		mtype module.ModuleType
		stype string
	}{
		{module.MODULE_IMAGE, "image"},
		{module.MODULE_ARCH, "arch"},
	}

	for _, tp := range types {
		exts := supports[tp.mtype]
		for _, ex := range exts {
			if ex == ext {
				// return file's type string
				ret.Type = tp.stype
				json.WriteResponse(w, ret)
				return
			}
		}
	}

	// not found
	http.NotFound(w, r)
}

func filesRouting(w http.ResponseWriter, r *http.Request) {
	path, err := queryPath(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	path, paths, err := splitRealVirtualPath(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	stat, _ := os.Stat(path)
	var files []string

	// 'paths' is not empty, 'path' needs archive file path
	// 'paths' is empty, 'path' needs archive file or directory path
	if len(paths) == 0 && stat.IsDir() {
		// directory
		files = getFilesFromDirectory(w, path)
	} else {
		files = getFilesFromArchive(w, path, paths)
	}
	if files == nil {
		return // still response error
	}

	// return json format
	ret := struct {
		Files []string `json:"files"`
	}{files}
	json.WriteResponse(w, ret)
}

func getFilesFromDirectory(w http.ResponseWriter, path string) []string {
	infos, err := ioutil.ReadDir(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return nil
	}

	files := make([]string, 0)
	for _, v := range infos {
		name := strings.Replace(v.Name(), "\\", "/", -1)
		if ind := strings.Index(name, "/"); ind >= 0 {
			name = name[ind+1:]
		}
		files = append(files, name)
	}
	return files
}

func getFilesFromArchive(w http.ResponseWriter, path string, paths []string) []string {
	mod := module.GetSupportModule(path)
	if mod == nil {
		http.Error(w, "no support type", http.StatusBadRequest)
		return nil
	}
	if mod.Type != module.MODULE_ARCH {
		http.Error(w, "wrong module operation", http.StatusBadRequest)
		return nil
	}

	file, err := module.NewReaderAt(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}
	defer file.Close()

	// TODO
	// support nested archive's path
	vpath := strings.Join(paths, "")
	files := mod.FuncArchFiles(file)
	if files == nil {
		http.Error(w, "fail read file", http.StatusInternalServerError)
		return nil
	}
	files = array.Filter(files, func(s string) bool {
		return strings.HasPrefix(s, vpath)
	})
	files = array.Map(files, func(s string) string {
		return s[len(vpath):] // remove prefix
	})
	files = array.Uniq(array.Map(files, func(s string) string {
		if index := strings.Index(s, "/"); index >= 0 {
			return s[:index] // remove suffix (when directory)
		}
		return s
	}))
	return files
}

func imageRouting(w http.ResponseWriter, r *http.Request) {
	path, err := queryPath(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	path, paths, err := splitRealVirtualPath(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	stat, _ := os.Stat(path)
	var file module.Reader

	// 'paths' is not empty, 'path' needs archive file path
	// 'paths' is empty, 'path' needs image file
	if len(paths) == 0 {
		// reject directory path
		if stat.IsDir() {
			http.Error(w, "path is directory", http.StatusBadRequest)
			return
		}
		file = getImageDataFromFile(w, path)
	} else {
		file = getImageDataFromArchive(w, path, paths)
		path = paths[len(paths)-1]
	}

	if file == nil {
		return // still response error
	}
	defer file.Close()

	// return binary format
	imageRoutingResponse(w, path, file)
}

func getImageDataFromFile(w http.ResponseWriter, path string) module.Reader {
	r, err := module.NewReaderAt(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}
	return r
}

func getImageDataFromArchive(w http.ResponseWriter, path string, paths []string) module.Reader {
	mod := module.GetSupportModule(path)
	if mod == nil {
		http.Error(w, "no support type", http.StatusBadRequest)
		return nil
	}
	if mod.Type != module.MODULE_ARCH {
		http.Error(w, "wrong module operation", http.StatusBadRequest)
		return nil
	}

	arch, err := module.NewReaderAt(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}
	defer arch.Close()

	// TODO
	// support nested archive's path
	imgpath := strings.Join(paths, "")
	file := mod.FuncArchRead(arch, imgpath)
	if file == nil {
		http.Error(w, "fail read file", http.StatusInternalServerError)
		return nil
	}
	defer file.Data.Close()

	r, err := module.NewReader(file.Data, file.Size)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}
	return r
}

func imageRoutingResponse(w http.ResponseWriter, path string, r module.Reader) {
	mod := module.GetSupportModule(path)
	if mod == nil {
		http.Error(w, "no support type", http.StatusBadRequest)
		return
	}
	if mod.Type != module.MODULE_IMAGE {
		http.Error(w, "wrong module operation", http.StatusBadRequest)
		return
	}

	file := mod.FuncImageRead(r)
	if file == nil {
		http.Error(w, "fail read file", http.StatusInternalServerError)
		return
	}
	defer file.Data.Close()

	w.Header().Set("Content-Type", file.Mime)
	w.Header().Set("Content-Length", fmt.Sprintf("%d", file.Size))
	if _, err := io.Copy(w, file.Data); err != nil {
		log.Printf("ERR: WriteResponse: %v\n", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"log"
	"os"
)

type funcDescription struct {
	Recv *ast.FieldList
	Type *ast.FuncType
	Body *ast.BlockStmt
}

var (
	receivers = flag.Bool("recv", false, "include receivers in scan")
	returns   = flag.Bool("ret", false, "include returns in scan")
	overwrite = flag.Bool("w", false, "overwrite file with changes")
)

func main() {

	flag.Parse()

	if flag.NArg() == 0 {
		fmt.Println("Usage:", os.Args[0], "[-w] [-recv] [-ret] <file[s]>")
		flag.PrintDefaults()
		os.Exit(1)
	}

	for _, file := range flag.Args() {

		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
		if err != nil {
			fmt.Println(file, err)
			continue
		}

		altered := blankId(f)

		if *overwrite {
			if altered {
				fmt.Println(file, "overwritten")
				outfile, err := os.OpenFile(file, os.O_TRUNC|os.O_WRONLY, 0)
				if err != nil {
					log.Fatal(err)
				}
				if err := printer.Fprint(outfile, fset, f); err != nil {
					log.Fatal(err)
				}
				if err := outfile.Close(); err != nil {
					log.Fatal(err)
				}
			} else {
				fmt.Println(file, "unaltered")
			}
		} else {
			printer.Fprint(os.Stdout, fset, f)
		}

	}

}

func blankId(f ast.Node) bool {

	var altered bool

	ast.Inspect(f, func(n ast.Node) bool {

		switch x := n.(type) {
		case *ast.FuncDecl:
			if changed := scan(&funcDescription{
				Recv: x.Recv,
				Type: x.Type,
				Body: x.Body,
			}); changed {
				altered = true
			}
		case *ast.FuncLit:
			if changed := scan(&funcDescription{
				Type: x.Type,
				Body: x.Body,
			}); changed {
				altered = true
			}
		}

		return true

	})

	return altered

}

func scan(x *funcDescription) bool {

	if x.Body == nil {
		return false
	}

	type v struct {
		ident *ast.Ident
		used  bool
	}

	vars := map[string]*v{}

	// store receiver
	if *receivers && x.Recv != nil {
		// only if a named receiver
		if len(x.Recv.List[0].Names) > 0 {
			vars[x.Recv.List[0].Names[0].Name] = &v{ident: x.Recv.List[0].Names[0]}
		}
	}

	// store function params
	for _, p := range x.Type.Params.List {
		for _, ident := range p.Names {
			if ident.Name != "_" {
				vars[ident.Name] = &v{ident: ident}
			}
		}
	}

	// store return params
	if *returns && x.Type.Results != nil {
		for _, r := range x.Type.Results.List {
			for _, ident := range r.Names {
				if ident.Name != "_" {
					vars[ident.Name] = &v{ident: ident}
				}
			}
		}
	}

	// scan for idents that are used in the function body
	ast.Inspect(x.Body, func(n ast.Node) bool {
		switch x := n.(type) {
		case *ast.Ident:
			if x.Name != "_" {
				if iv, exists := vars[x.Name]; exists {
					if x.Obj != nil {
						// matching ident declaration means same var
						if f, ok := x.Obj.Decl.(*ast.Field); ok && f != nil {
							for _, ident := range f.Names {
								if ident == iv.ident {
									iv.used = true
									return true
								}
							}
						}
					}
				}
			}
		}
		return true
	})

	var altered bool

	// set to _ where unused
	for _, v := range vars {
		if !v.used {
			v.ident.Name = "_"
			altered = true
		}
	}

	return altered

}
<commit_msg>minor refactor<commit_after>package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"log"
	"os"
)

type funcDescription struct {
	Recv *ast.FieldList
	Type *ast.FuncType
	Body *ast.BlockStmt
}

var (
	receivers = flag.Bool("recv", false, "include receivers in scan")
	returns   = flag.Bool("ret", false, "include returns in scan")
	overwrite = flag.Bool("w", false, "overwrite file with changes")
)

func main() {

	flag.Parse()

	if flag.NArg() == 0 {
		fmt.Println("Usage:", os.Args[0], "[-w] [-recv] [-ret] <file[s]>")
		flag.PrintDefaults()
		os.Exit(1)
	}

	for _, file := range flag.Args() {

		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
		if err != nil {
			fmt.Println(file, err)
			continue
		}

		altered := blankId(f)

		if *overwrite {
			if altered {
				fmt.Println(file, "overwritten")
				outfile, err := os.OpenFile(file, os.O_TRUNC|os.O_WRONLY, 0)
				if err != nil {
					log.Fatal(err)
				}
				if err := printer.Fprint(outfile, fset, f); err != nil {
					log.Fatal(err)
				}
				if err := outfile.Close(); err != nil {
					log.Fatal(err)
				}
			} else {
				fmt.Println(file, "unaltered")
			}
		} else {
			printer.Fprint(os.Stdout, fset, f)
		}

	}

}

func blankId(f ast.Node) bool {

	var altered bool

	ast.Inspect(f, func(n ast.Node) bool {

		switch x := n.(type) {
		case *ast.FuncDecl:
			if changed := scan(&funcDescription{
				Recv: x.Recv,
				Type: x.Type,
				Body: x.Body,
			}); changed {
				altered = true
			}
		case *ast.FuncLit:
			if changed := scan(&funcDescription{
				Type: x.Type,
				Body: x.Body,
			}); changed {
				altered = true
			}
		}

		return true

	})

	return altered

}

func scan(x *funcDescription) bool {

	if x.Body == nil {
		return false
	}

	vars := map[*ast.Ident]bool{}

	// store receiver
	if *receivers && x.Recv != nil {
		// only if a named receiver
		if len(x.Recv.List[0].Names) > 0 {
			vars[x.Recv.List[0].Names[0]] = false
		}
	}

	// store function params
	for _, p := range x.Type.Params.List {
		for _, ident := range p.Names {
			if ident.Name != "_" {
				vars[ident] = false
			}
		}
	}

	// store return params
	if *returns && x.Type.Results != nil {
		for _, r := range x.Type.Results.List {
			for _, ident := range r.Names {
				if ident.Name != "_" {
					vars[ident] = false
				}
			}
		}
	}

	// scan for idents that are used in the function body
	ast.Inspect(x.Body, func(n ast.Node) bool {
		if x, ok := n.(*ast.Ident); ok {
			if x.Obj != nil && x.Name != "_" {
				if f, ok := x.Obj.Decl.(*ast.Field); ok && f != nil {
					for _, ident := range f.Names {
						if used, exists := vars[ident]; exists && !used {
							vars[ident] = true
							return true
						}
					}
				}
			}
		}
		return true
	})

	var altered bool

	// set to _ where unused
	for ident, used := range vars {
		if !used {
			ident.Name = "_"
			altered = true
		}
	}

	return altered

}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"os/signal"
	"syscall"

	"github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/garden/server"
	"github.com/pivotal-cf-experimental/garden-dot-net/backend"
	"github.com/pivotal-golang/lager"
)

func main() {
	listenAddr := "0.0.0.0:3000"
	if os.Getenv("PORT") != "" {
		listenAddr = "0.0.0.0:" + os.Getenv("PORT")
	}

	logger := cf_lager.New("garden-dotnet")

	netBackend := backend.DotNetBackend{}

	gardenServer := server.New("tcp", listenAddr, 0, netBackend, logger)
	err := gardenServer.Start()
	if err != nil {
		logger.Fatal("Server Failed to Start", err)
		os.Exit(1)
	}

	logger.Info("started", lager.Data{
		"addr":    listenAddr,
	})

	signals := make(chan os.Signal, 1)

	go func() {
		<-signals
		gardenServer.Stop()
		os.Exit(0)
	}()

	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
	select {}
}
<commit_msg>Revert "Simplify listen port setup"<commit_after>package main

import (
	"flag"
	"os"
	"os/signal"
	"syscall"

	"github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/garden/server"
	"github.com/pivotal-cf-experimental/garden-dot-net/backend"
	"github.com/pivotal-golang/lager"
)

var containerGraceTime = flag.Duration(
	"containerGraceTime",
	0,
	"time after which to destroy idle containers",
)

func main() {
	defaultListNetwork := "unix"
	defaultListAddr := "/tmp/garden.sock"
	if os.Getenv("PORT") != "" {
		defaultListNetwork = "tcp"
		defaultListAddr = "0.0.0.0:" + os.Getenv("PORT")
	}
	var listenNetwork = flag.String(
		"listenNetwork",
		defaultListNetwork,
		"how to listen on the address (unix, tcp, etc.)",
	)
	var listenAddr = flag.String(
		"listenAddr",
		defaultListAddr,
		"address to listen on",
	)

	logger := cf_lager.New("garden-dotnet")

	netBackend := backend.DotNetBackend{}

	gardenServer := server.New(*listenNetwork, *listenAddr, *containerGraceTime, netBackend, logger)
	err := gardenServer.Start()
	if err != nil {
		logger.Fatal("Server Failed to Start", err)
		os.Exit(1)
	}

	logger.Info("started", lager.Data{
		"network": *listenNetwork,
		"addr":    *listenAddr,
	})

	signals := make(chan os.Signal, 1)

	go func() {
		<-signals
		gardenServer.Stop()
		os.Exit(0)
	}()

	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
	select {}
}
<|endoftext|>


<commit_before>// operators
package engine

import "fmt"

type Op int

const (
	EQ = Op(iota + 1)
	NE
	LT
	GT
)

func (o Op) String() string {

	var s string
	switch o {
	case EQ:
		s = "="
	case NE:
		s = "!="
	case LT:
		s = "<"
	case GT:
		s = ">"
	}
	return s
}
func parseOp(op string) (Op, error) {
	var finalOp Op
	switch op {
	case "<":
		finalOp = LT
	case ">":
		finalOp = GT
	case "=":
		finalOp = EQ
	case "!=":
		finalOp = NE
	default:
		return 0, fmt.Errorf("unknown operator %q", op)
	}
	return finalOp, nil
}
<commit_msg>Remove blank line<commit_after>// operators
package engine

import "fmt"

type Op int

const (
	EQ = Op(iota + 1)
	NE
	LT
	GT
)

func (o Op) String() string {
	var s string
	switch o {
	case EQ:
		s = "="
	case NE:
		s = "!="
	case LT:
		s = "<"
	case GT:
		s = ">"
	}
	return s
}
func parseOp(op string) (Op, error) {
	var finalOp Op
	switch op {
	case "<":
		finalOp = LT
	case ">":
		finalOp = GT
	case "=":
		finalOp = EQ
	case "!=":
		finalOp = NE
	default:
		return 0, fmt.Errorf("unknown operator %q", op)
	}
	return finalOp, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/user"
	"strings"
	"sync"
	"text/template"
	"unicode"

	ext "github.com/MattAitchison/remotectl/providers"
	sshutil "github.com/MattAitchison/remotectl/ssh"

	// Enabled Providers
	_ "github.com/MattAitchison/remotectl/digitalocean"
	_ "github.com/MattAitchison/remotectl/stdin"

	env "github.com/MattAitchison/envconfig"
)

var (
	// Version number that gets set at compile time.
	Version string

	curUser, _ = user.Current()

	sshPort = env.Int("remotectl_port", 22, "ssh port to use. Note: all ports must be the same on hosts for a run.")
	// Show default ident.
	ident         = env.String("remotectl_identity", "", "private key file")
	usr           = env.String("remotectl_user", strings.ToLower(curUser.Username), "user to connect as")
	provider      = env.String("remotectl_provider", "do", "name or comma-sep list of provider modules to use for selecting hosts")
	namespace     = env.String("remotectl_namespace", "", "")
	prefixTmplStr = env.String("remotectl_prefix", "{{.Name}}: ", "prefix template for host log output")
	prefixTmpl    = template.Must(template.New("prefix").Parse(prefixTmplStr))

	showVersion = flag.Bool("version", false, "show version")
	showHelp    = flag.Bool("help", false, "show this help message")
	verbose     = flag.Bool("verbose", false, "enable verbose status output")
	localMode   = flag.Bool("local", false, "runs a local shell with <cmd> instead of ssh")
	showList    = flag.Bool("list", false, "lists selected ips and names. /etc/hosts friendly output")

	// Use current working path.
	profile = flag.String("profile", "", "sources a bash profile to load a config") // Maybe a name will default to a file in ~/.remotectl
)

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

func fields(s string) []string {
	f := func(c rune) bool {
		return c == ',' || unicode.IsSpace(c)
	}
	return strings.FieldsFunc(s, f)
}

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

	switch {
	case *showHelp:
		helpCmd()
		os.Exit(0)
	case *showVersion:
		fmt.Printf("remotectl: %s", Version)
		os.Exit(0)
	}

	query := flag.Args()[0]
	cmd := strings.Join(flag.Args()[1:], " ")

	// Get all of the hosts from each provider.
	// This will have a way to filter soon.
	var hosts []ext.Host
	for _, p := range ext.Providers.Select(fields(provider)) {
		if p == nil {
			log.Fatal("unknown provider")
		}

		// Setup the provider
		p.Setup()

		// Query the provider for hosts
		extHosts, err := p.Query(namespace, query)
		fatalErr(err)
		hosts = append(hosts, extHosts...)
	}

	if len(hosts) == 0 {
		log.Fatal("no hosts")
	}

	if *showList {
		log.Println(hosts)
		return
	}

	cfg, err := sshutil.NewClientConfig(usr)
	fatalErr(err)

	log.Println("Connecting...")
	var sessions []*sshutil.Session
	for _, host := range hosts {

		addr := fmt.Sprint(host.Addr, ":", sshPort)
		sess, err := cfg.NewSession(addr)
		if err != nil {
			log.Fatalln(host.Name+":", addr, err)
		}

		outPipe, err := sess.StdoutPipe()
		fatalErr(err)

		errPipe, err := sess.StderrPipe()
		fatalErr(err)

		go hostStreamer(host, outPipe, os.Stdout)
		go hostStreamer(host, errPipe, os.Stderr)
		sessions = append(sessions, sess)
	}

	if len(sessions) == 0 {
		log.Fatal("no sessions")
	}

	group := &sync.WaitGroup{}
	for _, s := range sessions {
		group.Add(1)
		// TODO: Close s when done.
		go s.RunWaitGroup(group, cmd)
	}
	group.Wait()
}

func hostStreamer(host ext.Host, r io.Reader, w io.Writer) {
	// Locking? One writer?
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		prefixTmpl.Execute(w, host)
		fmt.Println(scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		log.Fatalf("err streaming logs %s", err)
	}

}

func helpCmd() {
	usage := `Usage: remotectl <flags> <query> [--] <cmd>

Providers:
%s

Environment Vars:
`
	fmt.Printf(usage, ext.Providers.Names())
	env.PrintDefaults()
	fmt.Println("\nFlags:")
	flag.PrintDefaults()

}
<commit_msg>Cleanup and prep for features (wait)<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/user"
	"strings"
	"sync"
	"text/template"

	ext "github.com/MattAitchison/remotectl/providers"
	sshutil "github.com/MattAitchison/remotectl/ssh"

	// Enabled Providers
	_ "github.com/MattAitchison/remotectl/digitalocean"
	_ "github.com/MattAitchison/remotectl/stdin"

	env "github.com/MattAitchison/envconfig"
)

var (
	// Version number that gets set at compile time.
	Version string

	curUser, _ = user.Current()

	sshPort = env.Int("remotectl_port", 22, "ssh port to use. Note: all ports must be the same on hosts for a run.")
	// Show default ident.
	ident         = env.String("remotectl_identity", "", "private key file")
	usr           = env.String("remotectl_user", strings.ToLower(curUser.Username), "user to connect as")
	provider      = env.String("remotectl_provider", "do", "name or comma-sep list of provider modules to use for selecting hosts")
	namespace     = env.String("remotectl_namespace", "", "")
	prefixTmplStr = env.String("remotectl_prefix", "{{.Name}}: ", "prefix template for host log output")
	prefixTmpl    = template.Must(template.New("prefix").Parse(prefixTmplStr))

	showVersion = flag.Bool("version", false, "show version")
	showHelp    = flag.Bool("help", false, "show this help message")
	verbose     = flag.Bool("verbose", false, "enable verbose status output")
	wait        = flag.Bool("wait", false, "wait for all hosts to be up")
	localMode   = flag.Bool("local", false, "runs a local shell with <cmd> instead of ssh")
	showList    = flag.Bool("list", false, "lists selected ips and names. /etc/hosts friendly output")

	// Use current working path.
	profile = flag.String("profile", "", "sources a bash profile to load a config") // Maybe a name will default to a file in ~/.remotectl
)

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

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

	switch {
	case *showHelp:
		helpCmd()
		os.Exit(0)
	case *showVersion:
		fmt.Printf("remotectl: %s", Version)
		os.Exit(0)
	}

	query := flag.Args()[0]
	cmd := strings.Join(flag.Args()[1:], " ")

	fmt.Println(flag.Args())
	fmt.Println(flag.NArg())
	fmt.Println(flag.NFlag())
	fmt.Println(os.Args)

	// Get all of the hosts from each provider.
	// This will have a way to filter soon.
	var hosts []ext.Host

	for _, p := range ext.Providers.Select(strings.Fields(provider)) {
		if p == nil {
			log.Fatal("unknown provider")
		}

		// Setup the provider
		p.Setup()

		// Query the provider for hosts
		extHosts, err := p.Query(namespace, query)
		fatalErr(err)
		hosts = append(hosts, extHosts...)
	}

	if len(hosts) == 0 {
		log.Fatal("no hosts")
	}

	if *showList {
		log.Println(hosts)
		return
	}

	cfg, err := sshutil.NewClientConfig(usr)
	fatalErr(err)

	sessions := make(chan *sshutil.Session, len(hosts))
	for _, host := range hosts {
		sessions <- setupSession(cfg, host)

	}

	// TODO: Implement wait

	group := &sync.WaitGroup{}
	for {
		s := <-sessions
		// TODO: Does this actually work?
		defer s.Close()

		group.Add(1)
		go s.RunWaitGroup(group, cmd)
	}

	// group := &sync.WaitGroup{}
	// for _, s := range sessions {
	// 	group.Add(1)
	//
	// 	go s.RunWaitGroup(group, cmd)
	// }
	// group.Wait()
}

func setupSession(cfg *sshutil.ClientConfig, host ext.Host) *sshutil.Session {
	addr := fmt.Sprint(host.Addr, ":", sshPort)
	s, err := cfg.NewSession(addr)
	if err != nil {
		log.Fatalln(host.Name+":", addr, err)
	}

	outPipe, err := s.StdoutPipe()
	fatalErr(err)

	errPipe, err := s.StderrPipe()
	fatalErr(err)

	go hostStreamer(host, outPipe, os.Stdout)
	go hostStreamer(host, errPipe, os.Stderr)

	return s
}

func hostStreamer(host ext.Host, r io.Reader, w io.Writer) {
	// Locking? One writer?
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		prefixTmpl.Execute(w, host)
		fmt.Println(scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		log.Fatalf("err streaming logs %s", err)
	}

}

func helpCmd() {
	usage := `Usage: remotectl <flags> <query> [--] <cmd>

Providers:
%s

Environment Vars:
`
	fmt.Printf(usage, ext.Providers.Names())
	env.PrintDefaults()
	fmt.Println("\nFlags:")
	flag.PrintDefaults()

}
<|endoftext|>