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

package cmd

import (
	"errors"
	"io"
	"os"
	"os/exec"
	"path"
	"syscall"

	"launchpad.net/gnuflag"
	"launchpad.net/gocheck"
)

var appflag = &gnuflag.Flag{
	Name:     "app",
	Usage:    "The name of the app.",
	Value:    nil,
	DefValue: "",
}

var appshortflag = &gnuflag.Flag{
	Name:     "a",
	Usage:    "The name of the app.",
	Value:    nil,
	DefValue: "",
}

func writeConfig(sourceFile string, c *gocheck.C) string {
	srcConfig, err := os.Open(sourceFile)
	c.Assert(err, gocheck.IsNil)
	defer srcConfig.Close()
	p := path.Join(os.TempDir(), "guesser-tests")
	err = os.MkdirAll(p, 0700)
	c.Assert(err, gocheck.IsNil)
	cmd := exec.Command("git", "init")
	cmd.Dir = p
	c.Assert(err, gocheck.IsNil)
	err = cmd.Run()
	c.Assert(err, gocheck.IsNil)
	dstConfig, err := os.OpenFile(path.Join(p, ".git", "config"), syscall.O_WRONLY|syscall.O_TRUNC|syscall.O_CREAT|syscall.O_CLOEXEC, 0644)
	c.Assert(err, gocheck.IsNil)
	defer dstConfig.Close()
	_, err = io.Copy(dstConfig, srcConfig)
	c.Assert(err, gocheck.IsNil)
	return p
}

func (s *S) TestGitGuesser(c *gocheck.C) {
	p := writeConfig("testdata/gitconfig-ok", c)
	defer os.RemoveAll(p)
	dirPath := path.Join(p, "somepath")
	err := os.MkdirAll(dirPath, 0700) // Will be removed when p is removed.
	c.Assert(err, gocheck.IsNil)
	g := GitGuesser{}
	name, err := g.GuessName(p) // repository root
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "gopher")
	name, err = g.GuessName(dirPath) // subdirectory
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "gopher")
}

func (s *S) TestGitGuesserNonGitUser(c *gocheck.C) {
	p := writeConfig("testdata/gitconfig-ok-nongit-user", c)
	defer os.RemoveAll(p)
	dirPath := path.Join(p, "somepath")
	err := os.MkdirAll(dirPath, 0700) // Will be removed when p is removed.
	c.Assert(err, gocheck.IsNil)
	g := GitGuesser{}
	name, err := g.GuessName(p) // repository root
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "gopher")
	name, err = g.GuessName(dirPath) // subdirectory
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "gopher")
}

// This test may fail if you have a git repository in /tmp. By the way, if you
// do have a repository in the temporary file hierarchy, please kill yourself.
func (s *S) TestGitGuesserWhenTheDirectoryIsNotAGitRepository(c *gocheck.C) {
	p := path.Join(os.TempDir(), "guesser-tests")
	err := os.MkdirAll(p, 0700)
	c.Assert(err, gocheck.IsNil)
	defer os.RemoveAll(p)
	name, err := GitGuesser{}.GuessName(p)
	c.Assert(name, gocheck.Equals, "")
	c.Assert(err, gocheck.NotNil)
	c.Assert(err, gocheck.ErrorMatches, "^Git repository not found:.*")
}

func (s *S) TestGitGuesserWithoutTsuruRemote(c *gocheck.C) {
	p := writeConfig("testdata/gitconfig-without-tsuru-remote", c)
	defer os.RemoveAll(p)
	name, err := GitGuesser{}.GuessName(p)
	c.Assert(name, gocheck.Equals, "")
	c.Assert(err, gocheck.NotNil)
	c.Assert(err.Error(), gocheck.Equals, "tsuru remote not declared.")
}

func (s *S) TestGitGuesserWithTsuruRemoteNotMatchingTsuruPattern(c *gocheck.C) {
	p := writeConfig("testdata/gitconfig-not-matching", c)
	defer os.RemoveAll(p)
	name, err := GitGuesser{}.GuessName(p)
	c.Assert(name, gocheck.Equals, "")
	c.Assert(err, gocheck.NotNil)
	c.Assert(err.Error(), gocheck.Equals, `"tsuru" remote did not match the pattern. Want something like <user>@<host>:<app-name>.git, got git://myhost.com/gopher.git`)
}

func (s *S) TestGuessingCommandGuesserNil(c *gocheck.C) {
	g := GuessingCommand{G: nil}
	c.Assert(g.guesser(), gocheck.FitsTypeOf, GitGuesser{})
}

func (s *S) TestGuessingCommandGuesserNonNil(c *gocheck.C) {
	fake := &FakeGuesser{}
	g := GuessingCommand{G: fake}
	c.Assert(g.guesser(), gocheck.DeepEquals, fake)
}

func (s *S) TestGuessingCommandWithFlagDefined(c *gocheck.C) {
	fake := &FakeGuesser{name: "other-app"}
	g := GuessingCommand{G: fake}
	g.Flags().Parse(true, []string{"--app", "myapp"})
	name, err := g.Guess()
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "myapp")
	pwd, err := os.Getwd()
	c.Assert(err, gocheck.IsNil)
	c.Assert(fake.HasGuess(pwd), gocheck.Equals, false)
}

func (s *S) TestGuessingCommandWithShortFlagDefined(c *gocheck.C) {
	fake := &FakeGuesser{name: "other-app"}
	g := GuessingCommand{G: fake}
	g.Flags().Parse(true, []string{"-a", "myapp"})
	name, err := g.Guess()
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "myapp")
	pwd, err := os.Getwd()
	c.Assert(err, gocheck.IsNil)
	c.Assert(fake.HasGuess(pwd), gocheck.Equals, false)
}

func (s *S) TestGuessingCommandWithoutFlagDefined(c *gocheck.C) {
	fake := &FakeGuesser{name: "other-app"}
	g := GuessingCommand{G: fake}
	name, err := g.Guess()
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "other-app")
	pwd, err := os.Getwd()
	c.Assert(err, gocheck.IsNil)
	c.Assert(fake.HasGuess(pwd), gocheck.Equals, true)
}

func (s *S) TestGuessingCommandFailToGuess(c *gocheck.C) {
	fake := &FailingFakeGuesser{}
	g := GuessingCommand{G: fake}
	name, err := g.Guess()
	c.Assert(name, gocheck.Equals, "")
	c.Assert(err, gocheck.NotNil)
	c.Assert(err.Error(), gocheck.Equals, `tsuru wasn't able to guess the name of the app.

Use the --app flag to specify it.`)
	pwd, err := os.Getwd()
	c.Assert(err, gocheck.IsNil)
	c.Assert(fake.HasGuess(pwd), gocheck.Equals, true)
}

func (s *S) TestGuessingCommandFlags(c *gocheck.C) {
	var flags []gnuflag.Flag
	expected := []gnuflag.Flag{*appshortflag, *appflag}
	command := GuessingCommand{}
	flagset := command.Flags()
	flagset.VisitAll(func(f *gnuflag.Flag) {
		f.Value = nil
		flags = append(flags, *f)
	})
	c.Assert(flags, gocheck.DeepEquals, expected)
}

type FakeGuesser struct {
	guesses []string
	name    string
}

func (f *FakeGuesser) HasGuess(path string) bool {
	for _, g := range f.guesses {
		if g == path {
			return true
		}
	}
	return false
}

func (f *FakeGuesser) GuessName(path string) (string, error) {
	f.guesses = append(f.guesses, path)
	return f.name, nil
}

type FailingFakeGuesser struct {
	FakeGuesser
	message string
}

func (f *FailingFakeGuesser) GuessName(path string) (string, error) {
	f.FakeGuesser.GuessName(path)
	return "", errors.New(f.message)
}
<commit_msg>cmd: use FakeGuesser from cmd/testing<commit_after>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cmd

import (
	"io"
	"os"
	"os/exec"
	"path"
	"syscall"

	"github.com/tsuru/tsuru/cmd/testing"
	"launchpad.net/gnuflag"
	"launchpad.net/gocheck"
)

var appflag = &gnuflag.Flag{
	Name:     "app",
	Usage:    "The name of the app.",
	Value:    nil,
	DefValue: "",
}

var appshortflag = &gnuflag.Flag{
	Name:     "a",
	Usage:    "The name of the app.",
	Value:    nil,
	DefValue: "",
}

func writeConfig(sourceFile string, c *gocheck.C) string {
	srcConfig, err := os.Open(sourceFile)
	c.Assert(err, gocheck.IsNil)
	defer srcConfig.Close()
	p := path.Join(os.TempDir(), "guesser-tests")
	err = os.MkdirAll(p, 0700)
	c.Assert(err, gocheck.IsNil)
	cmd := exec.Command("git", "init")
	cmd.Dir = p
	c.Assert(err, gocheck.IsNil)
	err = cmd.Run()
	c.Assert(err, gocheck.IsNil)
	dstConfig, err := os.OpenFile(path.Join(p, ".git", "config"), syscall.O_WRONLY|syscall.O_TRUNC|syscall.O_CREAT|syscall.O_CLOEXEC, 0644)
	c.Assert(err, gocheck.IsNil)
	defer dstConfig.Close()
	_, err = io.Copy(dstConfig, srcConfig)
	c.Assert(err, gocheck.IsNil)
	return p
}

func (s *S) TestGitGuesser(c *gocheck.C) {
	p := writeConfig("testdata/gitconfig-ok", c)
	defer os.RemoveAll(p)
	dirPath := path.Join(p, "somepath")
	err := os.MkdirAll(dirPath, 0700) // Will be removed when p is removed.
	c.Assert(err, gocheck.IsNil)
	g := GitGuesser{}
	name, err := g.GuessName(p) // repository root
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "gopher")
	name, err = g.GuessName(dirPath) // subdirectory
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "gopher")
}

func (s *S) TestGitGuesserNonGitUser(c *gocheck.C) {
	p := writeConfig("testdata/gitconfig-ok-nongit-user", c)
	defer os.RemoveAll(p)
	dirPath := path.Join(p, "somepath")
	err := os.MkdirAll(dirPath, 0700) // Will be removed when p is removed.
	c.Assert(err, gocheck.IsNil)
	g := GitGuesser{}
	name, err := g.GuessName(p) // repository root
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "gopher")
	name, err = g.GuessName(dirPath) // subdirectory
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "gopher")
}

// This test may fail if you have a git repository in /tmp. By the way, if you
// do have a repository in the temporary file hierarchy, please kill yourself.
func (s *S) TestGitGuesserWhenTheDirectoryIsNotAGitRepository(c *gocheck.C) {
	p := path.Join(os.TempDir(), "guesser-tests")
	err := os.MkdirAll(p, 0700)
	c.Assert(err, gocheck.IsNil)
	defer os.RemoveAll(p)
	name, err := GitGuesser{}.GuessName(p)
	c.Assert(name, gocheck.Equals, "")
	c.Assert(err, gocheck.NotNil)
	c.Assert(err, gocheck.ErrorMatches, "^Git repository not found:.*")
}

func (s *S) TestGitGuesserWithoutTsuruRemote(c *gocheck.C) {
	p := writeConfig("testdata/gitconfig-without-tsuru-remote", c)
	defer os.RemoveAll(p)
	name, err := GitGuesser{}.GuessName(p)
	c.Assert(name, gocheck.Equals, "")
	c.Assert(err, gocheck.NotNil)
	c.Assert(err.Error(), gocheck.Equals, "tsuru remote not declared.")
}

func (s *S) TestGitGuesserWithTsuruRemoteNotMatchingTsuruPattern(c *gocheck.C) {
	p := writeConfig("testdata/gitconfig-not-matching", c)
	defer os.RemoveAll(p)
	name, err := GitGuesser{}.GuessName(p)
	c.Assert(name, gocheck.Equals, "")
	c.Assert(err, gocheck.NotNil)
	c.Assert(err.Error(), gocheck.Equals, `"tsuru" remote did not match the pattern. Want something like <user>@<host>:<app-name>.git, got git://myhost.com/gopher.git`)
}

func (s *S) TestGuessingCommandGuesserNil(c *gocheck.C) {
	g := GuessingCommand{G: nil}
	c.Assert(g.guesser(), gocheck.FitsTypeOf, GitGuesser{})
}

func (s *S) TestGuessingCommandGuesserNonNil(c *gocheck.C) {
	fake := &testing.FakeGuesser{}
	g := GuessingCommand{G: fake}
	c.Assert(g.guesser(), gocheck.DeepEquals, fake)
}

func (s *S) TestGuessingCommandWithFlagDefined(c *gocheck.C) {
	fake := &testing.FakeGuesser{Name: "other-app"}
	g := GuessingCommand{G: fake}
	g.Flags().Parse(true, []string{"--app", "myapp"})
	name, err := g.Guess()
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "myapp")
	pwd, err := os.Getwd()
	c.Assert(err, gocheck.IsNil)
	c.Assert(fake.HasGuess(pwd), gocheck.Equals, false)
}

func (s *S) TestGuessingCommandWithShortFlagDefined(c *gocheck.C) {
	fake := &testing.FakeGuesser{Name: "other-app"}
	g := GuessingCommand{G: fake}
	g.Flags().Parse(true, []string{"-a", "myapp"})
	name, err := g.Guess()
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "myapp")
	pwd, err := os.Getwd()
	c.Assert(err, gocheck.IsNil)
	c.Assert(fake.HasGuess(pwd), gocheck.Equals, false)
}

func (s *S) TestGuessingCommandWithoutFlagDefined(c *gocheck.C) {
	fake := &testing.FakeGuesser{Name: "other-app"}
	g := GuessingCommand{G: fake}
	name, err := g.Guess()
	c.Assert(err, gocheck.IsNil)
	c.Assert(name, gocheck.Equals, "other-app")
	pwd, err := os.Getwd()
	c.Assert(err, gocheck.IsNil)
	c.Assert(fake.HasGuess(pwd), gocheck.Equals, true)
}

func (s *S) TestGuessingCommandFailToGuess(c *gocheck.C) {
	fake := &testing.FailingFakeGuesser{}
	g := GuessingCommand{G: fake}
	name, err := g.Guess()
	c.Assert(name, gocheck.Equals, "")
	c.Assert(err, gocheck.NotNil)
	c.Assert(err.Error(), gocheck.Equals, `tsuru wasn't able to guess the name of the app.

Use the --app flag to specify it.`)
	pwd, err := os.Getwd()
	c.Assert(err, gocheck.IsNil)
	c.Assert(fake.HasGuess(pwd), gocheck.Equals, true)
}

func (s *S) TestGuessingCommandFlags(c *gocheck.C) {
	var flags []gnuflag.Flag
	expected := []gnuflag.Flag{*appshortflag, *appflag}
	command := GuessingCommand{}
	flagset := command.Flags()
	flagset.VisitAll(func(f *gnuflag.Flag) {
		f.Value = nil
		flags = append(flags, *f)
	})
	c.Assert(flags, gocheck.DeepEquals, expected)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"runtime/pprof"

	logging "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-logging"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	manet "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr/net"

	cmds "github.com/jbenet/go-ipfs/commands"
	cmdsCli "github.com/jbenet/go-ipfs/commands/cli"
	cmdsHttp "github.com/jbenet/go-ipfs/commands/http"
	"github.com/jbenet/go-ipfs/config"
	"github.com/jbenet/go-ipfs/core"
	daemon "github.com/jbenet/go-ipfs/daemon2"
	u "github.com/jbenet/go-ipfs/util"
)

// log is the command logger
var log = u.Logger("cmd/ipfs")

// signal to output help
var errHelpRequested = errors.New("Help Requested")

const (
	cpuProfile  = "ipfs.cpuprof"
	heapProfile = "ipfs.memprof"
	errorFormat = "ERROR: %v\n\n"
)

type cmdInvocation struct {
	path []string
	cmd  *cmds.Command
	req  cmds.Request
}

// main roadmap:
// - parse the commandline to get a cmdInvocation
// - if user requests, help, print it and exit.
// - run the command invocation
// - output the response
// - if anything fails, print error, maybe with help
func main() {
	var invoc cmdInvocation
	var err error

	// we'll call this local helper to output errors.
	// this is so we control how to print errors in one place.
	printErr := func(err error) {
		fmt.Fprintf(os.Stderr, "Error: %s\n", err.Error())
	}

	// this is a local helper to print out help text.
	// there's some considerations that this makes easier.
	printHelp := func(long bool) {
		helpFunc := cmdsCli.ShortHelp
		if long {
			helpFunc = cmdsCli.LongHelp
		}

		helpFunc("ipfs", Root, invoc.path, os.Stderr)
	}

	// parse the commandline into a command invocation
	parseErr := invoc.Parse(os.Args[1:])

	// BEFORE handling the parse error, if we have enough information
	// AND the user requested help, print it out and exit
	if invoc.req != nil {
		longH, shortH, err := invoc.requestedHelp()
		if err != nil {
			printErr(err)
			os.Exit(1)
		}
		if longH || shortH {
			printHelp(longH)
			os.Exit(0)
		}
	}

	// here we handle the cases where
	// - commands with no Run func are invoked directly.
	// - the main command is invoked.
	if invoc.cmd == nil || invoc.cmd.Run == nil {
		printHelp(false)
		os.Exit(0)
	}

	// ok now handle parse error (which means cli input was wrong,
	// e.g. incorrect number of args, or nonexistent subcommand)
	if parseErr != nil {
		printErr(parseErr)

		// this was a user error, print help.
		if invoc.cmd != nil {
			// we need a newline space.
			fmt.Fprintf(os.Stderr, "\n")
			printHelp(false)
		}
		os.Exit(1)
	}

	// ok, finally, run the command invocation.
	output, err := invoc.Run()
	if err != nil {
		printErr(err)

		// if this error was a client error, print short help too.
		if isClientError(err) {
			printHelp(false)
		}
		os.Exit(1)
	}

	// everything went better than expected :)
	io.Copy(os.Stdout, output)
}

func (i *cmdInvocation) Run() (output io.Reader, err error) {
	handleInterrupt()

	// check if user wants to debug. option OR env var.
	debug, _, err := i.req.Option("debug").Bool()
	if err != nil {
		return nil, err
	}
	if debug || u.GetenvBool("DEBUG") {
		u.Debug = true
		u.SetAllLoggers(logging.DEBUG)
	}

	// if debugging, let's profile.
	// TODO maybe change this to its own option... profiling makes it slower.
	if u.Debug {
		stopProfilingFunc, err := startProfiling()
		if err != nil {
			return nil, err
		}
		defer stopProfilingFunc() // to be executed as late as possible
	}

	res, err := callCommand(i.req, Root)
	if err != nil {
		return nil, err
	}

	if err := res.Error(); err != nil {
		return nil, err
	}

	return res.Reader()
}

func (i *cmdInvocation) Parse(args []string) error {
	var err error

	i.req, i.cmd, i.path, err = cmdsCli.Parse(args, Root)
	if err != nil {
		return err
	}

	configPath, err := getConfigRoot(i.req)
	if err != nil {
		return err
	}

	// this sets up the function that will initialize the config lazily.
	ctx := i.req.Context()
	ctx.ConfigRoot = configPath
	ctx.LoadConfig = loadConfig

	// if no encoding was specified by user, default to plaintext encoding
	// (if command doesn't support plaintext, use JSON instead)
	if !i.req.Option("encoding").Found() {
		if i.req.Command().Marshalers != nil && i.req.Command().Marshalers[cmds.Text] != nil {
			i.req.SetOption("encoding", cmds.Text)
		} else {
			i.req.SetOption("encoding", cmds.JSON)
		}
	}

	return nil
}

func (i *cmdInvocation) requestedHelp() (short bool, long bool, err error) {
	longHelp, _, err := i.req.Option("help").Bool()
	if err != nil {
		return false, false, err
	}
	shortHelp, _, err := i.req.Option("h").Bool()
	if err != nil {
		return false, false, err
	}
	return longHelp, shortHelp, nil
}

func callCommand(req cmds.Request, root *cmds.Command) (cmds.Response, error) {
	var res cmds.Response

	local, found, err := req.Option("local").Bool()
	if err != nil {
		return nil, err
	}

	remote := !isLocal(req.Command()) && (!found || !local)

	log.Info("Checking if daemon is running...")
	if remote && daemon.Locked(req.Context().ConfigRoot) {

		cfg, err := req.Context().GetConfig()
		if err != nil {
			return nil, err
		}

		addr, err := ma.NewMultiaddr(cfg.Addresses.API)
		if err != nil {
			return nil, err
		}

		_, host, err := manet.DialArgs(addr)
		if err != nil {
			return nil, err
		}

		client := cmdsHttp.NewClient(host)

		res, err = client.Send(req)
		if err != nil {
			return nil, err
		}

	} else {
		log.Info("Executing command locally: daemon not running")

		// this sets up the function that will initialize the node
		// this is so that we can construct the node lazily.
		ctx := req.Context()
		ctx.ConstructNode = func() (*core.IpfsNode, error) {
			cfg, err := ctx.GetConfig()
			if err != nil {
				return nil, err
			}
			return core.NewIpfsNode(cfg, false)
		}

		// Okay!!!!! NOW we can call the command.
		res = root.Call(req)

		// let's not forget teardown. If a node was initialized, we must close it.
		// Note that this means the underlying req.Context().Node variable is exposed.
		// this is gross, and should be changed when we extract out the exec Context.
		node := req.Context().NodeWithoutConstructing()
		if node != nil {
			node.Close()
		}
	}
	return res, nil
}

func isClientError(err error) bool {

	// Somewhat suprisingly, the pointer cast fails to recognize commands.Error
	// passed as values, so we check both.

	// cast to cmds.Error
	switch e := err.(type) {
	case *cmds.Error:
		return e.Code == cmds.ErrClient
	case cmds.Error:
		return e.Code == cmds.ErrClient
	}
	return false
}

func getConfigRoot(req cmds.Request) (string, error) {
	configOpt, found, err := req.Option("config").String()
	if err != nil {
		return "", err
	}
	if found && configOpt != "" {
		return configOpt, nil
	}

	configPath, err := config.PathRoot()
	if err != nil {
		return "", err
	}
	return configPath, nil
}

func loadConfig(path string) (*config.Config, error) {
	configFile, err := config.Filename(path)
	if err != nil {
		return nil, err
	}

	return config.Load(configFile)
}

// startProfiling begins CPU profiling and returns a `stop` function to be
// executed as late as possible. The stop function captures the memprofile.
func startProfiling() (func(), error) {

	// start CPU profiling as early as possible
	ofi, err := os.Create(cpuProfile)
	if err != nil {
		return nil, err
	}
	pprof.StartCPUProfile(ofi)

	stopProfiling := func() {
		pprof.StopCPUProfile()
		defer ofi.Close() // captured by the closure
		err := writeHeapProfileToFile()
		if err != nil {
			log.Critical(err)
		}
	}
	return stopProfiling, nil
}

func writeHeapProfileToFile() error {
	mprof, err := os.Create(heapProfile)
	if err != nil {
		return err
	}
	defer mprof.Close() // _after_ writing the heap profile
	return pprof.WriteHeapProfile(mprof)
}

// listen for and handle SIGTERM
func handleInterrupt() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		for _ = range c {
			log.Info("Received interrupt signal, terminating...")
			os.Exit(0)
		}
	}()
}
<commit_msg>cmds2: commandShouldRunOnDaemon<commit_after>package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"runtime/pprof"

	logging "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-logging"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	manet "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr/net"

	cmds "github.com/jbenet/go-ipfs/commands"
	cmdsCli "github.com/jbenet/go-ipfs/commands/cli"
	cmdsHttp "github.com/jbenet/go-ipfs/commands/http"
	"github.com/jbenet/go-ipfs/config"
	"github.com/jbenet/go-ipfs/core"
	daemon "github.com/jbenet/go-ipfs/daemon2"
	u "github.com/jbenet/go-ipfs/util"
)

// log is the command logger
var log = u.Logger("cmd/ipfs")

// signal to output help
var errHelpRequested = errors.New("Help Requested")

const (
	cpuProfile  = "ipfs.cpuprof"
	heapProfile = "ipfs.memprof"
	errorFormat = "ERROR: %v\n\n"
)

type cmdInvocation struct {
	path []string
	cmd  *cmds.Command
	req  cmds.Request
}

// main roadmap:
// - parse the commandline to get a cmdInvocation
// - if user requests, help, print it and exit.
// - run the command invocation
// - output the response
// - if anything fails, print error, maybe with help
func main() {
	var invoc cmdInvocation
	var err error

	// we'll call this local helper to output errors.
	// this is so we control how to print errors in one place.
	printErr := func(err error) {
		fmt.Fprintf(os.Stderr, "Error: %s\n", err.Error())
	}

	// this is a local helper to print out help text.
	// there's some considerations that this makes easier.
	printHelp := func(long bool) {
		helpFunc := cmdsCli.ShortHelp
		if long {
			helpFunc = cmdsCli.LongHelp
		}

		helpFunc("ipfs", Root, invoc.path, os.Stderr)
	}

	// parse the commandline into a command invocation
	parseErr := invoc.Parse(os.Args[1:])

	// BEFORE handling the parse error, if we have enough information
	// AND the user requested help, print it out and exit
	if invoc.req != nil {
		longH, shortH, err := invoc.requestedHelp()
		if err != nil {
			printErr(err)
			os.Exit(1)
		}
		if longH || shortH {
			printHelp(longH)
			os.Exit(0)
		}
	}

	// here we handle the cases where
	// - commands with no Run func are invoked directly.
	// - the main command is invoked.
	if invoc.cmd == nil || invoc.cmd.Run == nil {
		printHelp(false)
		os.Exit(0)
	}

	// ok now handle parse error (which means cli input was wrong,
	// e.g. incorrect number of args, or nonexistent subcommand)
	if parseErr != nil {
		printErr(parseErr)

		// this was a user error, print help.
		if invoc.cmd != nil {
			// we need a newline space.
			fmt.Fprintf(os.Stderr, "\n")
			printHelp(false)
		}
		os.Exit(1)
	}

	// ok, finally, run the command invocation.
	output, err := invoc.Run()
	if err != nil {
		printErr(err)

		// if this error was a client error, print short help too.
		if isClientError(err) {
			printHelp(false)
		}
		os.Exit(1)
	}

	// everything went better than expected :)
	io.Copy(os.Stdout, output)
}

func (i *cmdInvocation) Run() (output io.Reader, err error) {
	handleInterrupt()

	// check if user wants to debug. option OR env var.
	debug, _, err := i.req.Option("debug").Bool()
	if err != nil {
		return nil, err
	}
	if debug || u.GetenvBool("DEBUG") {
		u.Debug = true
		u.SetAllLoggers(logging.DEBUG)
	}

	// if debugging, let's profile.
	// TODO maybe change this to its own option... profiling makes it slower.
	if u.Debug {
		stopProfilingFunc, err := startProfiling()
		if err != nil {
			return nil, err
		}
		defer stopProfilingFunc() // to be executed as late as possible
	}

	res, err := callCommand(i.req, Root)
	if err != nil {
		return nil, err
	}

	if err := res.Error(); err != nil {
		return nil, err
	}

	return res.Reader()
}

func (i *cmdInvocation) Parse(args []string) error {
	var err error

	i.req, i.cmd, i.path, err = cmdsCli.Parse(args, Root)
	if err != nil {
		return err
	}

	configPath, err := getConfigRoot(i.req)
	if err != nil {
		return err
	}

	// this sets up the function that will initialize the config lazily.
	ctx := i.req.Context()
	ctx.ConfigRoot = configPath
	ctx.LoadConfig = loadConfig

	// if no encoding was specified by user, default to plaintext encoding
	// (if command doesn't support plaintext, use JSON instead)
	if !i.req.Option("encoding").Found() {
		if i.req.Command().Marshalers != nil && i.req.Command().Marshalers[cmds.Text] != nil {
			i.req.SetOption("encoding", cmds.Text)
		} else {
			i.req.SetOption("encoding", cmds.JSON)
		}
	}

	return nil
}

func (i *cmdInvocation) requestedHelp() (short bool, long bool, err error) {
	longHelp, _, err := i.req.Option("help").Bool()
	if err != nil {
		return false, false, err
	}
	shortHelp, _, err := i.req.Option("h").Bool()
	if err != nil {
		return false, false, err
	}
	return longHelp, shortHelp, nil
}

func callCommand(req cmds.Request, root *cmds.Command) (cmds.Response, error) {
	var res cmds.Response

	useDaemon, err := commandShouldRunOnDaemon(req, root)
	if err != nil {
		return nil, err
	}

	if useDaemon {

		cfg, err := req.Context().GetConfig()
		if err != nil {
			return nil, err
		}

		addr, err := ma.NewMultiaddr(cfg.Addresses.API)
		if err != nil {
			return nil, err
		}

		log.Infof("Executing command on daemon running at %s", addr)
		_, host, err := manet.DialArgs(addr)
		if err != nil {
			return nil, err
		}

		client := cmdsHttp.NewClient(host)

		res, err = client.Send(req)
		if err != nil {
			return nil, err
		}

	} else {
		log.Info("Executing command locally")

		// this sets up the function that will initialize the node
		// this is so that we can construct the node lazily.
		ctx := req.Context()
		ctx.ConstructNode = func() (*core.IpfsNode, error) {
			cfg, err := ctx.GetConfig()
			if err != nil {
				return nil, err
			}
			return core.NewIpfsNode(cfg, false)
		}

		// Okay!!!!! NOW we can call the command.
		res = root.Call(req)

		// let's not forget teardown. If a node was initialized, we must close it.
		// Note that this means the underlying req.Context().Node variable is exposed.
		// this is gross, and should be changed when we extract out the exec Context.
		node := req.Context().NodeWithoutConstructing()
		if node != nil {
			node.Close()
		}
	}
	return res, nil
}

func commandShouldRunOnDaemon(req cmds.Request, root *cmds.Command) (bool, error) {
	path := req.Path()
	// root command.
	if len(path) < 1 {
		return false, nil
	}

	cmd, found := root.Subcommands[path[0]]
	if !found {
		return false, fmt.Errorf("subcommand %s should be in root", path[0])
	}

	details, found := cmdDetailsMap[cmd]
	if !found {
		details = cmdDetails{} // defaults
	}

	if details.cannotRunOnClient && details.cannotRunOnDaemon {
		return false, fmt.Errorf("command disabled: %s", path[0])
	}

	if details.doesNotUseRepo && !details.cannotRunOnClient {
		return false, nil
	}

	// at this point need to know whether daemon is running. we defer
	// to this point so that some commands dont open files unnecessarily.
	daemonLocked := daemon.Locked(req.Context().ConfigRoot)
	log.Info("Daemon is running.")

	if daemonLocked {

		if details.cannotRunOnDaemon {
			e := "ipfs daemon is running. please stop it to run this command"
			return false, cmds.ClientError(e)
		}

		return true, nil
	}

	if details.cannotRunOnClient {
		return false, cmds.ClientError("must run on the ipfs daemon")
	}

	return false, nil
}

func isClientError(err error) bool {

	// Somewhat suprisingly, the pointer cast fails to recognize commands.Error
	// passed as values, so we check both.

	// cast to cmds.Error
	switch e := err.(type) {
	case *cmds.Error:
		return e.Code == cmds.ErrClient
	case cmds.Error:
		return e.Code == cmds.ErrClient
	}
	return false
}

func getConfigRoot(req cmds.Request) (string, error) {
	configOpt, found, err := req.Option("config").String()
	if err != nil {
		return "", err
	}
	if found && configOpt != "" {
		return configOpt, nil
	}

	configPath, err := config.PathRoot()
	if err != nil {
		return "", err
	}
	return configPath, nil
}

func loadConfig(path string) (*config.Config, error) {
	configFile, err := config.Filename(path)
	if err != nil {
		return nil, err
	}

	return config.Load(configFile)
}

// startProfiling begins CPU profiling and returns a `stop` function to be
// executed as late as possible. The stop function captures the memprofile.
func startProfiling() (func(), error) {

	// start CPU profiling as early as possible
	ofi, err := os.Create(cpuProfile)
	if err != nil {
		return nil, err
	}
	pprof.StartCPUProfile(ofi)

	stopProfiling := func() {
		pprof.StopCPUProfile()
		defer ofi.Close() // captured by the closure
		err := writeHeapProfileToFile()
		if err != nil {
			log.Critical(err)
		}
	}
	return stopProfiling, nil
}

func writeHeapProfileToFile() error {
	mprof, err := os.Create(heapProfile)
	if err != nil {
		return err
	}
	defer mprof.Close() // _after_ writing the heap profile
	return pprof.WriteHeapProfile(mprof)
}

// listen for and handle SIGTERM
func handleInterrupt() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		for _ = range c {
			log.Info("Received interrupt signal, terminating...")
			os.Exit(0)
		}
	}()
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"crypto/x509"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"text/tabwriter"

	"github.com/go-kit/kit/log"
	"github.com/micromdm/micromdm/pkg/crypto"
	"github.com/micromdm/micromdm/platform/blueprint"
	"github.com/micromdm/micromdm/platform/device"
	"github.com/micromdm/micromdm/platform/profile"
	"github.com/pkg/errors"
)

type getCommand struct {
	config *ServerConfig
	*remoteServices
}

func (cmd *getCommand) setup() error {
	cfg, err := LoadServerConfig()
	if err != nil {
		return err
	}
	cmd.config = cfg
	logger := log.NewLogfmtLogger(os.Stderr)

	remote, err := setupClient(logger)
	if err != nil {
		return err
	}
	cmd.remoteServices = remote
	return nil
}

func (cmd *getCommand) Run(args []string) error {
	if len(args) < 1 {
		cmd.Usage()
		os.Exit(1)
	}

	if err := cmd.setup(); err != nil {
		return err
	}

	var run func([]string) error
	switch strings.ToLower(args[0]) {
	case "dev", "device", "devices":
		run = cmd.getDevices
	case "dep-devices":
		run = cmd.getDEPDevices
	case "dep-account":
		run = cmd.getDEPAccount
	case "dep-profiles":
		run = cmd.getDEPProfiles
	case "dep-tokens":
		run = cmd.getDepTokens
	case "blueprints":
		run = cmd.getBlueprints
	case "profiles":
		run = cmd.getProfiles
	case "users":
		run = cmd.getUsers
	case "apps":
		run = cmd.getApps
	case "dep-autoassigners":
		run = cmd.getDEPAutoAssigners
	default:
		cmd.Usage()
		os.Exit(1)
	}

	return run(args[1:])
}

func (cmd *getCommand) Usage() error {
	const getUsage = `
Display one or many resources.

Valid resource types:

  * devices
  * blueprints
  * dep-tokens
  * dep-devices
  * dep-account
  * dep-profiles
  * dep-autoassigners
  * users
  * profiles
  * apps

Examples:
  # Get a list of devices
  mdmctl get devices

  # Get a device by serial (TODO implement filtering)
  mdmctl get devices -serials=C02ABCDEF
`
	fmt.Println(getUsage)
	return nil
}

type devicesTableOutput struct{ w *tabwriter.Writer }

func (out *devicesTableOutput) BasicHeader() {
	fmt.Fprintf(out.w, "UDID\tSerialNumber\tEnrollmentStatus\tLastSeen\n")
}

func (out *devicesTableOutput) BasicFooter() {
	out.w.Flush()
}

func (cmd *getCommand) getDevices(args []string) error {
	flagset := flag.NewFlagSet("devices", flag.ExitOnError)
	var (
		flFilterSerials = flagset.String("serials", "", "device serial, optionally comma-separated")
	)
	flagset.Usage = usageFor(flagset, "mdmctl get devices [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}
	w := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
	out := &devicesTableOutput{w}
	out.BasicHeader()
	defer out.BasicFooter()
	ctx := context.Background()

	// convert string to []string in case a filter serial is defined
	// nil serials will return all serials from the database
	var flFilterSerialsSlice []string
	if *flFilterSerials != "" {
		flFilterSerialsSlice = strings.Split(*flFilterSerials, ",")
	}

	devices, err := cmd.devicesvc.ListDevices(ctx, device.ListDevicesOption{FilterSerial: flFilterSerialsSlice})
	if err != nil {
		return err
	}
	for _, d := range devices {
		fmt.Fprintf(out.w, "%s\t%s\t%v\t%s\n", d.UDID, d.SerialNumber, d.EnrollmentStatus, d.LastSeen)
	}
	return nil
}

const defaultmdmctlFilesPath = "mdm-files"

func (cmd *getCommand) getDepTokens(args []string) error {
	flagset := flag.NewFlagSet("dep-tokens", flag.ExitOnError)
	var (
		flFullCK        = flagset.Bool("v", false, "Display full ConsumerKey in summary list")
		flPublicKeyPath = flagset.String(
			"export-public-key",
			filepath.Join(defaultmdmctlFilesPath, "DEPPublicKey"),
			"Filename of public key to write (to be uploaded to deploy.apple.com)",
		)
		flTokenPath = flagset.String(
			"export-token",
			filepath.Join(defaultmdmctlFilesPath, "DEPOAuthToken.json"),
			"Filename to save decrypted oauth token (JSON)")
	)
	flagset.Usage = usageFor(flagset, "mdmctl get dep-tokens [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}

	w := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
	fmt.Fprintf(w, "ConsumerKey\tAccessTokenExpiry\n")
	ctx := context.Background()
	tokens, certBytes, err := cmd.configsvc.GetDEPTokens(ctx)
	if err != nil {
		return err
	}
	var ckTrimmed string
	for _, t := range tokens {
		if len(t.ConsumerKey) > 40 && !*flFullCK {
			ckTrimmed = t.ConsumerKey[0:39] + "…"
		} else {
			ckTrimmed = t.ConsumerKey
		}
		fmt.Fprintf(w, "%s\t%s\n", ckTrimmed, t.AccessTokenExpiry.String())
	}
	w.Flush()

	if *flPublicKeyPath != "" && certBytes != nil {

		if err := os.MkdirAll(filepath.Dir(*flPublicKeyPath), 0755); err != nil {
			return errors.Wrapf(err, "create directory %s", filepath.Dir(*flPublicKeyPath))
		}

		cert, err := x509.ParseCertificate(certBytes)
		if err != nil {
			return err
		}
		err = crypto.WritePEMCertificateFile(cert, *flPublicKeyPath)
		if err != nil {
			return err
		}
		fmt.Printf("\nWrote DEP public key to: %s\n", *flPublicKeyPath)
	}

	if *flTokenPath != "" && len(tokens) > 0 {
		t := tokens[0]

		if err := os.MkdirAll(filepath.Dir(*flTokenPath), 0755); err != nil {
			return errors.Wrapf(err, "create directory %s", filepath.Dir(*flTokenPath))
		}

		tokenFile, err := os.Create(*flTokenPath)
		if err != nil {
			return err
		}
		defer tokenFile.Close()

		err = json.NewEncoder(tokenFile).Encode(t)
		if err != nil {
			return err
		}

		fmt.Printf("\nWrote DEP token JSON to: %s\n", *flTokenPath)
		if len(tokens) > 1 {
			fmt.Println("WARNING: more than one DEP token returned; only saved first")
		}
	}

	return nil
}

func (cmd *getCommand) getBlueprints(args []string) error {
	flagset := flag.NewFlagSet("blueprints", flag.ExitOnError)
	var (
		flBlueprintName = flagset.String("name", "", "name of blueprint")
		flJSONName      = flagset.String("f", "-", "filename of JSON to save to")
	)
	flagset.Usage = usageFor(flagset, "mdmctl get blueprints [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}

	ctx := context.Background()
	blueprints, err := cmd.blueprintsvc.GetBlueprints(ctx, blueprint.GetBlueprintsOption{FilterName: *flBlueprintName})
	if err != nil {
		return err
	}

	if *flBlueprintName == "" || len(blueprints) < 1 {
		w := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
		fmt.Fprintf(w, "Name\tUUID\tManifests\tProfiles\tApply At\n")
		for _, bp := range blueprints {
			var applyAtStr string
			if len(bp.ApplyAt) > 0 {
				applyAtStr = strings.Join(bp.ApplyAt, ",")
			} else {
				applyAtStr = "(None)"
			}
			fmt.Fprintf(
				w,
				"%s\t%s\t%d\t%d\t%s\n",
				bp.Name,
				bp.UUID,
				len(bp.ApplicationURLs),
				len(bp.ProfileIdentifiers),
				applyAtStr,
			)
		}
		w.Flush()
	} else if *flJSONName != "" {
		bp := blueprints[0]

		var output *os.File
		{
			if *flJSONName == "-" {
				output = os.Stdout
			} else {
				var err error
				output, err = os.Create(*flJSONName)
				if err != nil {
					return err
				}
				defer output.Close()
			}
		}

		enc := json.NewEncoder(output)
		enc.SetIndent("", "  ")
		err = enc.Encode(bp)
		if err != nil {
			return err
		}

		if *flJSONName != "-" {
			fmt.Printf("wrote blueprint %s to: %s\n", *flBlueprintName, *flJSONName)
			if len(blueprints) > 1 {
				fmt.Println("WARNING: more than one Blueprint returned; only saved first")
			}
		}
	}

	return nil
}

func (cmd *getCommand) getProfiles(args []string) error {
	flagset := flag.NewFlagSet("profiles", flag.ExitOnError)
	var (
		flProfilePath = flagset.String("f", "-", "filename of profile to write")
		flIdentifier  = flagset.String("id", "", "profile Identifier")
	)
	flagset.Usage = usageFor(flagset, "mdmctl get blueprints [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}

	ctx := context.Background()
	profiles, err := cmd.profilesvc.GetProfiles(ctx, profile.GetProfilesOption{Identifier: *flIdentifier})
	if err != nil {
		return err
	}

	if *flIdentifier == "" || len(profiles) < 1 {
		w := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
		fmt.Fprintf(w, "Identifier\tLength\n")
		for _, p := range profiles {
			fmt.Fprintf(
				w,
				"%s\t%d\n",
				p.Identifier,
				len(p.Mobileconfig),
			)
		}
		w.Flush()
	} else if *flProfilePath != "" {
		p := profiles[0]

		var output *os.File
		{
			if *flProfilePath == "-" {
				output = os.Stdout
			} else {
				var err error
				output, err = os.Create(*flProfilePath)
				if err != nil {
					return err
				}
				defer output.Close()
			}
		}
		_, err = output.Write([]byte(p.Mobileconfig))
		if err != nil {
			return err
		}
		if *flProfilePath != "-" {
			fmt.Printf("wrote profile id %s to: %s\n", p.Identifier, *flProfilePath)
			if len(profiles) > 1 {
				fmt.Println("WARNING: more than one Profile returned; only saved first")
			}
		}
	}
	return nil
}
<commit_msg>fix typo (#799)<commit_after>package main

import (
	"context"
	"crypto/x509"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"text/tabwriter"

	"github.com/go-kit/kit/log"
	"github.com/micromdm/micromdm/pkg/crypto"
	"github.com/micromdm/micromdm/platform/blueprint"
	"github.com/micromdm/micromdm/platform/device"
	"github.com/micromdm/micromdm/platform/profile"
	"github.com/pkg/errors"
)

type getCommand struct {
	config *ServerConfig
	*remoteServices
}

func (cmd *getCommand) setup() error {
	cfg, err := LoadServerConfig()
	if err != nil {
		return err
	}
	cmd.config = cfg
	logger := log.NewLogfmtLogger(os.Stderr)

	remote, err := setupClient(logger)
	if err != nil {
		return err
	}
	cmd.remoteServices = remote
	return nil
}

func (cmd *getCommand) Run(args []string) error {
	if len(args) < 1 {
		cmd.Usage()
		os.Exit(1)
	}

	if err := cmd.setup(); err != nil {
		return err
	}

	var run func([]string) error
	switch strings.ToLower(args[0]) {
	case "dev", "device", "devices":
		run = cmd.getDevices
	case "dep-devices":
		run = cmd.getDEPDevices
	case "dep-account":
		run = cmd.getDEPAccount
	case "dep-profiles":
		run = cmd.getDEPProfiles
	case "dep-tokens":
		run = cmd.getDepTokens
	case "blueprints":
		run = cmd.getBlueprints
	case "profiles":
		run = cmd.getProfiles
	case "users":
		run = cmd.getUsers
	case "apps":
		run = cmd.getApps
	case "dep-autoassigners":
		run = cmd.getDEPAutoAssigners
	default:
		cmd.Usage()
		os.Exit(1)
	}

	return run(args[1:])
}

func (cmd *getCommand) Usage() error {
	const getUsage = `
Display one or many resources.

Valid resource types:

  * devices
  * blueprints
  * dep-tokens
  * dep-devices
  * dep-account
  * dep-profiles
  * dep-autoassigners
  * users
  * profiles
  * apps

Examples:
  # Get a list of devices
  mdmctl get devices

  # Get a device by serial (TODO implement filtering)
  mdmctl get devices -serials=C02ABCDEF
`
	fmt.Println(getUsage)
	return nil
}

type devicesTableOutput struct{ w *tabwriter.Writer }

func (out *devicesTableOutput) BasicHeader() {
	fmt.Fprintf(out.w, "UDID\tSerialNumber\tEnrollmentStatus\tLastSeen\n")
}

func (out *devicesTableOutput) BasicFooter() {
	out.w.Flush()
}

func (cmd *getCommand) getDevices(args []string) error {
	flagset := flag.NewFlagSet("devices", flag.ExitOnError)
	var (
		flFilterSerials = flagset.String("serials", "", "device serial, optionally comma-separated")
	)
	flagset.Usage = usageFor(flagset, "mdmctl get devices [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}
	w := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
	out := &devicesTableOutput{w}
	out.BasicHeader()
	defer out.BasicFooter()
	ctx := context.Background()

	// convert string to []string in case a filter serial is defined
	// nil serials will return all serials from the database
	var flFilterSerialsSlice []string
	if *flFilterSerials != "" {
		flFilterSerialsSlice = strings.Split(*flFilterSerials, ",")
	}

	devices, err := cmd.devicesvc.ListDevices(ctx, device.ListDevicesOption{FilterSerial: flFilterSerialsSlice})
	if err != nil {
		return err
	}
	for _, d := range devices {
		fmt.Fprintf(out.w, "%s\t%s\t%v\t%s\n", d.UDID, d.SerialNumber, d.EnrollmentStatus, d.LastSeen)
	}
	return nil
}

const defaultmdmctlFilesPath = "mdm-files"

func (cmd *getCommand) getDepTokens(args []string) error {
	flagset := flag.NewFlagSet("dep-tokens", flag.ExitOnError)
	var (
		flFullCK        = flagset.Bool("v", false, "Display full ConsumerKey in summary list")
		flPublicKeyPath = flagset.String(
			"export-public-key",
			filepath.Join(defaultmdmctlFilesPath, "DEPPublicKey"),
			"Filename of public key to write (to be uploaded to deploy.apple.com)",
		)
		flTokenPath = flagset.String(
			"export-token",
			filepath.Join(defaultmdmctlFilesPath, "DEPOAuthToken.json"),
			"Filename to save decrypted oauth token (JSON)")
	)
	flagset.Usage = usageFor(flagset, "mdmctl get dep-tokens [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}

	w := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
	fmt.Fprintf(w, "ConsumerKey\tAccessTokenExpiry\n")
	ctx := context.Background()
	tokens, certBytes, err := cmd.configsvc.GetDEPTokens(ctx)
	if err != nil {
		return err
	}
	var ckTrimmed string
	for _, t := range tokens {
		if len(t.ConsumerKey) > 40 && !*flFullCK {
			ckTrimmed = t.ConsumerKey[0:39] + "…"
		} else {
			ckTrimmed = t.ConsumerKey
		}
		fmt.Fprintf(w, "%s\t%s\n", ckTrimmed, t.AccessTokenExpiry.String())
	}
	w.Flush()

	if *flPublicKeyPath != "" && certBytes != nil {

		if err := os.MkdirAll(filepath.Dir(*flPublicKeyPath), 0755); err != nil {
			return errors.Wrapf(err, "create directory %s", filepath.Dir(*flPublicKeyPath))
		}

		cert, err := x509.ParseCertificate(certBytes)
		if err != nil {
			return err
		}
		err = crypto.WritePEMCertificateFile(cert, *flPublicKeyPath)
		if err != nil {
			return err
		}
		fmt.Printf("\nWrote DEP public key to: %s\n", *flPublicKeyPath)
	}

	if *flTokenPath != "" && len(tokens) > 0 {
		t := tokens[0]

		if err := os.MkdirAll(filepath.Dir(*flTokenPath), 0755); err != nil {
			return errors.Wrapf(err, "create directory %s", filepath.Dir(*flTokenPath))
		}

		tokenFile, err := os.Create(*flTokenPath)
		if err != nil {
			return err
		}
		defer tokenFile.Close()

		err = json.NewEncoder(tokenFile).Encode(t)
		if err != nil {
			return err
		}

		fmt.Printf("\nWrote DEP token JSON to: %s\n", *flTokenPath)
		if len(tokens) > 1 {
			fmt.Println("WARNING: more than one DEP token returned; only saved first")
		}
	}

	return nil
}

func (cmd *getCommand) getBlueprints(args []string) error {
	flagset := flag.NewFlagSet("blueprints", flag.ExitOnError)
	var (
		flBlueprintName = flagset.String("name", "", "name of blueprint")
		flJSONName      = flagset.String("f", "-", "filename of JSON to save to")
	)
	flagset.Usage = usageFor(flagset, "mdmctl get blueprints [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}

	ctx := context.Background()
	blueprints, err := cmd.blueprintsvc.GetBlueprints(ctx, blueprint.GetBlueprintsOption{FilterName: *flBlueprintName})
	if err != nil {
		return err
	}

	if *flBlueprintName == "" || len(blueprints) < 1 {
		w := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
		fmt.Fprintf(w, "Name\tUUID\tManifests\tProfiles\tApply At\n")
		for _, bp := range blueprints {
			var applyAtStr string
			if len(bp.ApplyAt) > 0 {
				applyAtStr = strings.Join(bp.ApplyAt, ",")
			} else {
				applyAtStr = "(None)"
			}
			fmt.Fprintf(
				w,
				"%s\t%s\t%d\t%d\t%s\n",
				bp.Name,
				bp.UUID,
				len(bp.ApplicationURLs),
				len(bp.ProfileIdentifiers),
				applyAtStr,
			)
		}
		w.Flush()
	} else if *flJSONName != "" {
		bp := blueprints[0]

		var output *os.File
		{
			if *flJSONName == "-" {
				output = os.Stdout
			} else {
				var err error
				output, err = os.Create(*flJSONName)
				if err != nil {
					return err
				}
				defer output.Close()
			}
		}

		enc := json.NewEncoder(output)
		enc.SetIndent("", "  ")
		err = enc.Encode(bp)
		if err != nil {
			return err
		}

		if *flJSONName != "-" {
			fmt.Printf("wrote blueprint %s to: %s\n", *flBlueprintName, *flJSONName)
			if len(blueprints) > 1 {
				fmt.Println("WARNING: more than one Blueprint returned; only saved first")
			}
		}
	}

	return nil
}

func (cmd *getCommand) getProfiles(args []string) error {
	flagset := flag.NewFlagSet("profiles", flag.ExitOnError)
	var (
		flProfilePath = flagset.String("f", "-", "filename of profile to write")
		flIdentifier  = flagset.String("id", "", "profile Identifier")
	)
	flagset.Usage = usageFor(flagset, "mdmctl get profiles [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}

	ctx := context.Background()
	profiles, err := cmd.profilesvc.GetProfiles(ctx, profile.GetProfilesOption{Identifier: *flIdentifier})
	if err != nil {
		return err
	}

	if *flIdentifier == "" || len(profiles) < 1 {
		w := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
		fmt.Fprintf(w, "Identifier\tLength\n")
		for _, p := range profiles {
			fmt.Fprintf(
				w,
				"%s\t%d\n",
				p.Identifier,
				len(p.Mobileconfig),
			)
		}
		w.Flush()
	} else if *flProfilePath != "" {
		p := profiles[0]

		var output *os.File
		{
			if *flProfilePath == "-" {
				output = os.Stdout
			} else {
				var err error
				output, err = os.Create(*flProfilePath)
				if err != nil {
					return err
				}
				defer output.Close()
			}
		}
		_, err = output.Write([]byte(p.Mobileconfig))
		if err != nil {
			return err
		}
		if *flProfilePath != "-" {
			fmt.Printf("wrote profile id %s to: %s\n", p.Identifier, *flProfilePath)
			if len(profiles) > 1 {
				fmt.Println("WARNING: more than one Profile returned; only saved first")
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package mtp

// These tests require a single Android MTP device that is connected
// and unlocked.

import (
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"strings"
	"testing"
	"time"
)

// VerboseTest returns true if the testing framework is run with -v.
func VerboseTest() bool {
	flag := flag.Lookup("test.v")
	return flag != nil && flag.Value.String() == "true"
}

func TestAndroid(t *testing.T) {
	dev, err := SelectDevice("")
	if err != nil {
		t.Fatal(err)
	}
	defer dev.Close()

	info := DeviceInfo{}
	err = dev.GetDeviceInfo(&info)
	if err != nil {
		t.Fatal("GetDeviceInfo failed:", err)
	}

	if !strings.Contains(info.MTPExtension, "android.com:") {
		t.Log("no android extensions", info.MTPExtension)
		return
	}

	if err = dev.Configure(); err != nil {
		t.Fatal("Configure failed:", err)
	}

	sids := Uint32Array{}
	err = dev.GetStorageIDs(&sids)
	if err != nil {
		t.Fatalf("GetStorageIDs failed: %v", err)
	}

	if len(sids.Values) == 0 {
		t.Fatalf("No storages")
	}

	id := sids.Values[0]

	// 500 + 512 triggers the null read case on both sides.
	const testSize = 500 + 512
	name := fmt.Sprintf("mtp-doodle-test%x", rand.Int31())

	send := ObjectInfo{
		StorageID:        id,
		ObjectFormat:     OFC_Undefined,
		ParentObject:     0xFFFFFFFF,
		Filename:         name,
		CompressedSize:   uint32(testSize),
		ModificationDate: time.Now(),
		Keywords:         "bla",
	}
	data := make([]byte, testSize)
	for i := range data {
		data[i] = byte('0' + i%10)
	}

	_, _, handle, err := dev.SendObjectInfo(id, 0xFFFFFFFF, &send)
	if err != nil {
		t.Fatal("SendObjectInfo failed:", err)
	} else {
		buf := bytes.NewBuffer(data)
		t.Logf("Sent objectinfo handle: 0x%x\n", handle)
		err = dev.SendObject(buf, int64(len(data)))
		if err != nil {
			t.Log("SendObject failed:", err)
		}
	}

	magicStr := "life universe"
	magicOff := 21
	magicSize := 42

	err = dev.AndroidBeginEditObject(handle)
	if err != nil {
		t.Errorf("AndroidBeginEditObject: %v", err)
		return
	} else {
		err = dev.AndroidTruncate(handle, int64(magicSize))
		if err != nil {
			t.Errorf("AndroidTruncate: %v", err)
		}
		buf := bytes.NewBufferString(magicStr)
		err = dev.AndroidSendPartialObject(handle, int64(magicOff), uint32(buf.Len()), buf)
		if err != nil {
			t.Errorf("AndroidSendPartialObject: %v", err)
		}
		if buf.Len() > 0 {
			t.Errorf("buffer not consumed")
		}
		err = dev.AndroidEndEditObject(handle)
		if err != nil {
			t.Errorf("AndroidEndEditObject: %v", err)
		}
	}

	buf := &bytes.Buffer{}
	err = dev.GetObject(handle, buf)
	if err != nil {
		t.Errorf("GetObject: %v", err)
	}

	if buf.Len() != magicSize {
		t.Errorf("truncate failed:: %v", err)
	}
	for i := 0; i < len(magicStr); i++ {
		data[i+magicOff] = magicStr[i]
	}
	want := string(data[:magicSize])
	if buf.String() != want {
		t.Errorf("read result was %q, want %q", buf.String(), want)
	}

	buf = &bytes.Buffer{}
	err = dev.AndroidGetPartialObject64(handle, buf, int64(magicOff), 5)
	if err != nil {
		t.Errorf("AndroidGetPartialObject64: %v", err)
	}
	want = magicStr[:5]
	got := buf.String()
	if got != want {
		t.Errorf("AndroidGetPartialObject64: got %q want %q", got, want)
	}

	// Try write at end of file.
	err = dev.AndroidBeginEditObject(handle)
	if err != nil {
		t.Errorf("AndroidBeginEditObject: %v", err)
		return
	} else {
		buf := bytes.NewBufferString(magicStr)
		err = dev.AndroidSendPartialObject(handle, int64(magicSize), uint32(buf.Len()), buf)
		if err != nil {
			t.Errorf("AndroidSendPartialObject: %v", err)
		}
		if buf.Len() > 0 {
			t.Errorf("buffer not consumed")
		}
		err = dev.AndroidEndEditObject(handle)
		if err != nil {
			t.Errorf("AndroidEndEditObject: %v", err)
		}
	}
	buf = &bytes.Buffer{}
	err = dev.GetObject(handle, buf)
	if err != nil {
		t.Errorf("GetObject: %v", err)
	}
	want = string(data[:magicSize]) + magicStr
	got = buf.String()
	if got != want {
		t.Errorf("GetObject: got %q want %q", got, want)
	}

	err = dev.DeleteObject(handle)
	if err != nil {
		t.Fatalf("DeleteObject failed: %v", err)
	}
}

func TestDeviceProperties(t *testing.T) {
	dev, err := SelectDevice("")
	if err != nil {
		t.Fatal(err)
	}
	defer dev.Close()

	dev.DataDebug = VerboseTest()
	dev.MTPDebug = VerboseTest()
	dev.USBDebug = VerboseTest()
	err = dev.Configure()
	if err != nil {
		t.Log("Configure failed:", err)
	}

	// Test non-supported device property first.
	var battery DevicePropDesc
	err = dev.GetDevicePropDesc(DPC_BatteryLevel, &battery)
	if err != nil {
		// Not an error; not supported on Android.
		t.Log("battery failed:", err)
	} else {
		t.Logf("%#v\n", battery)
	}

	var friendly DevicePropDesc
	err = dev.GetDevicePropDesc(DPC_MTP_DeviceFriendlyName, &friendly)
	if err != nil {
		t.Fatal("GetDevicePropDesc FriendlyName failed:", err)
	} else {
		t.Logf("%s: %#v\n", DPC_names[DPC_MTP_DeviceFriendlyName], friendly)
	}
	before := friendly.CurrentValue.(string)

	newVal := fmt.Sprintf("gomtp device_test %x", rand.Int31())
	str := StringValue{newVal}
	err = dev.SetDevicePropValue(DPC_MTP_DeviceFriendlyName, &str)
	if err != nil {
		t.Error("SetDevicePropValue failed:", err)
	}

	str.Value = ""
	err = dev.GetDevicePropValue(DPC_MTP_DeviceFriendlyName, &str)
	if err != nil {
		t.Error("GetDevicePropValue failed:", err)
	}
	if str.Value != newVal {
		t.Logf("got %q for property value, want %q\n", str.Value, newVal)
	}

	err = dev.ResetDevicePropValue(DPC_MTP_DeviceFriendlyName)
	if err != nil {
		// For some reason, this is not supported? Returns
		// unknown error code 0xffff
		t.Log("ResetDevicePropValue failed:", err)
	}

	str = StringValue{before}
	err = dev.SetDevicePropValue(DPC_MTP_DeviceFriendlyName, &str)
	if err != nil {
		t.Error("SetDevicePropValue failed:", err)
	}

	// Test object properties.
	props := Uint16Array{}
	err = dev.GetObjectPropsSupported(OFC_Undefined, &props)
	if err != nil {
		t.Errorf("GetObjectPropsSupported failed: %v\n", err)
	} else {
		t.Logf("GetObjectPropsSupported (OFC_Undefined) value: %s\n", getNames(OPC_names, props.Values))
	}

	for _, p := range props.Values {
		var objPropDesc ObjectPropDesc
		if p == OPC_PersistantUniqueObjectIdentifier {
			// can't deal with int128.
			continue
		}

		err = dev.GetObjectPropDesc(p, OFC_Undefined, &objPropDesc)
		name := OPC_names[int(p)]
		if err != nil {
			t.Errorf("GetObjectPropDesc(%s) failed: %v\n", name, err)
		} else {
			t.Logf("GetObjectPropDesc(%s) value: %#v %T\n", name, objPropDesc,
				InstantiateType(objPropDesc.DataType).Interface())
		}
	}

}

func TestDeviceInfo(t *testing.T) {
	dev, err := SelectDevice("")
	if err != nil {
		t.Fatal(err)
	}
	defer dev.Close()

	i, _ := dev.ID()
	t.Log("device:", i)
	info := DeviceInfo{}
	err = dev.GetDeviceInfo(&info)
	if err != nil {
		t.Error("GetDeviceInfo failed:", err)
	} else {
		t.Logf("%v\n", &info)
	}
}

func TestDeviceStorage(t *testing.T) {
	dev, err := SelectDevice("")
	if err != nil {
		t.Fatal(err)
	}
	defer dev.Close()

	i, _ := dev.ID()
	t.Log("device:", i)

	info := DeviceInfo{}
	err = dev.GetDeviceInfo(&info)
	if err != nil {
		t.Log("GetDeviceInfo failed:", err)
	} else {
		t.Logf("device info %v\n", &info)
	}

	err = dev.Configure()
	if err != nil {
		t.Log("Configure failed:", err)
	}

	sids := Uint32Array{}
	err = dev.GetStorageIDs(&sids)
	if err != nil {
		t.Fatalf("GetStorageIDs failed: %v", err)
	} else {
		t.Logf("%#v\n", sids)
	}

	if len(sids.Values) == 0 {
		t.Fatalf("No storages")
	}

	id := sids.Values[0]
	var storageInfo StorageInfo
	dev.GetStorageInfo(id, &storageInfo)
	if err != nil {
		t.Fatalf("GetStorageInfo failed:", err)
	} else {
		t.Logf("%#v\n", storageInfo)
	}

	resp, err := dev.GetNumObjects(id, 0x0, 0x0)
	if err != nil {
		t.Fatalf("GenericRPC failed:", err)
	} else {
		t.Logf("num objects %#v\n", resp)
	}

	// 500 + 512 triggers the null read case on both sides.
	const testSize = 500 + 512
	data := make([]byte, testSize)
	for i := range data {
		data[i] = byte(i % 256)
	}

	name := fmt.Sprintf("go-mtp-test%x", rand.Int31())
	buf := bytes.NewBuffer(data)
	send := ObjectInfo{
		StorageID:        id,
		ObjectFormat:     OFC_Undefined,
		ParentObject:     0xFFFFFFFF,
		Filename:         name,
		CompressedSize:   uint32(len(data)),
		ModificationDate: time.Now(),
		Keywords:         "bla",
	}

	_, _, handle, err := dev.SendObjectInfo(id, 0xFFFFFFFF, &send)
	if err != nil {
		t.Fatalf("SendObjectInfo failed:", err)
	} else {
		t.Logf("Sent objectinfo handle: 0x%x\n", handle)
		err = dev.SendObject(buf, int64(len(data)))
		if err != nil {
			t.Log("SendObject failed:", err)
		}
	}

	hs := Uint32Array{}
	err = dev.GetObjectHandles(id,
		OFC_Undefined,
		//OFC_Association,
		0xFFFFFFFF, &hs)

	if err != nil {
		t.Fatalf("GetObjectHandles failed: %v", err)
	} else {
		found := false
		for _, h := range hs.Values {
			if h == handle {
				found = true
				break
			}
		}

		if !found {
			t.Fatalf("uploaded file is not there")
		}
	}

	var backInfo ObjectInfo
	err = dev.GetObjectInfo(handle, &backInfo)
	if err != nil {
		t.Fatalf("GetObjectInfo failed: %v", err)
	} else {
		t.Logf("info %#v\n", backInfo)
	}

	var objSize Uint64Value
	err = dev.GetObjectPropValue(handle, OPC_ObjectSize, &objSize)
	if err != nil {
		t.Fatalf("GetObjectPropValue failed: %v", err)
	} else {
		t.Logf("info %#v\n", objSize)
	}

	if objSize.Value != testSize {
		t.Errorf("object size error: got %v, want %v", objSize.Value, testSize)
	}

	backBuf := &bytes.Buffer{}
	err = dev.GetObject(handle, backBuf)
	if err != nil {
		t.Fatalf("GetObject failed: %v", err)
	} else {
		if bytes.Compare(backBuf.Bytes(), data) != 0 {
			t.Fatalf("back comparison failed.")
		}
	}

	newName := fmt.Sprintf("mtp-doodle-test%x", rand.Int31())
	err = dev.SetObjectPropValue(handle, OPC_ObjectFileName, &StringValue{newName})
	if err != nil {
		t.Errorf("error renaming object: %v", err)
	}

	err = dev.DeleteObject(handle)
	if err != nil {
		t.Fatalf("DeleteObject failed: %v", err)
	}
}
<commit_msg>Make Configure failure fatal in tests.<commit_after>package mtp

// These tests require a single Android MTP device that is connected
// and unlocked.

import (
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"strings"
	"testing"
	"time"
)

// VerboseTest returns true if the testing framework is run with -v.
func VerboseTest() bool {
	flag := flag.Lookup("test.v")
	return flag != nil && flag.Value.String() == "true"
}

func TestAndroid(t *testing.T) {
	dev, err := SelectDevice("")
	if err != nil {
		t.Fatal(err)
	}
	defer dev.Close()

	info := DeviceInfo{}
	err = dev.GetDeviceInfo(&info)
	if err != nil {
		t.Fatal("GetDeviceInfo failed:", err)
	}

	if !strings.Contains(info.MTPExtension, "android.com:") {
		t.Log("no android extensions", info.MTPExtension)
		return
	}

	if err = dev.Configure(); err != nil {
		t.Fatal("Configure failed:", err)
	}

	sids := Uint32Array{}
	err = dev.GetStorageIDs(&sids)
	if err != nil {
		t.Fatalf("GetStorageIDs failed: %v", err)
	}

	if len(sids.Values) == 0 {
		t.Fatalf("No storages")
	}

	id := sids.Values[0]

	// 500 + 512 triggers the null read case on both sides.
	const testSize = 500 + 512
	name := fmt.Sprintf("mtp-doodle-test%x", rand.Int31())

	send := ObjectInfo{
		StorageID:        id,
		ObjectFormat:     OFC_Undefined,
		ParentObject:     0xFFFFFFFF,
		Filename:         name,
		CompressedSize:   uint32(testSize),
		ModificationDate: time.Now(),
		Keywords:         "bla",
	}
	data := make([]byte, testSize)
	for i := range data {
		data[i] = byte('0' + i%10)
	}

	_, _, handle, err := dev.SendObjectInfo(id, 0xFFFFFFFF, &send)
	if err != nil {
		t.Fatal("SendObjectInfo failed:", err)
	} else {
		buf := bytes.NewBuffer(data)
		t.Logf("Sent objectinfo handle: 0x%x\n", handle)
		err = dev.SendObject(buf, int64(len(data)))
		if err != nil {
			t.Log("SendObject failed:", err)
		}
	}

	magicStr := "life universe"
	magicOff := 21
	magicSize := 42

	err = dev.AndroidBeginEditObject(handle)
	if err != nil {
		t.Errorf("AndroidBeginEditObject: %v", err)
		return
	} else {
		err = dev.AndroidTruncate(handle, int64(magicSize))
		if err != nil {
			t.Errorf("AndroidTruncate: %v", err)
		}
		buf := bytes.NewBufferString(magicStr)
		err = dev.AndroidSendPartialObject(handle, int64(magicOff), uint32(buf.Len()), buf)
		if err != nil {
			t.Errorf("AndroidSendPartialObject: %v", err)
		}
		if buf.Len() > 0 {
			t.Errorf("buffer not consumed")
		}
		err = dev.AndroidEndEditObject(handle)
		if err != nil {
			t.Errorf("AndroidEndEditObject: %v", err)
		}
	}

	buf := &bytes.Buffer{}
	err = dev.GetObject(handle, buf)
	if err != nil {
		t.Errorf("GetObject: %v", err)
	}

	if buf.Len() != magicSize {
		t.Errorf("truncate failed:: %v", err)
	}
	for i := 0; i < len(magicStr); i++ {
		data[i+magicOff] = magicStr[i]
	}
	want := string(data[:magicSize])
	if buf.String() != want {
		t.Errorf("read result was %q, want %q", buf.String(), want)
	}

	buf = &bytes.Buffer{}
	err = dev.AndroidGetPartialObject64(handle, buf, int64(magicOff), 5)
	if err != nil {
		t.Errorf("AndroidGetPartialObject64: %v", err)
	}
	want = magicStr[:5]
	got := buf.String()
	if got != want {
		t.Errorf("AndroidGetPartialObject64: got %q want %q", got, want)
	}

	// Try write at end of file.
	err = dev.AndroidBeginEditObject(handle)
	if err != nil {
		t.Errorf("AndroidBeginEditObject: %v", err)
		return
	} else {
		buf := bytes.NewBufferString(magicStr)
		err = dev.AndroidSendPartialObject(handle, int64(magicSize), uint32(buf.Len()), buf)
		if err != nil {
			t.Errorf("AndroidSendPartialObject: %v", err)
		}
		if buf.Len() > 0 {
			t.Errorf("buffer not consumed")
		}
		err = dev.AndroidEndEditObject(handle)
		if err != nil {
			t.Errorf("AndroidEndEditObject: %v", err)
		}
	}
	buf = &bytes.Buffer{}
	err = dev.GetObject(handle, buf)
	if err != nil {
		t.Errorf("GetObject: %v", err)
	}
	want = string(data[:magicSize]) + magicStr
	got = buf.String()
	if got != want {
		t.Errorf("GetObject: got %q want %q", got, want)
	}

	err = dev.DeleteObject(handle)
	if err != nil {
		t.Fatalf("DeleteObject failed: %v", err)
	}
}

func TestDeviceProperties(t *testing.T) {
	dev, err := SelectDevice("")
	if err != nil {
		t.Fatal(err)
	}
	defer dev.Close()

	dev.DataDebug = VerboseTest()
	dev.MTPDebug = VerboseTest()
	dev.USBDebug = VerboseTest()
	err = dev.Configure()
	if err != nil {
		t.Fatal("Configure failed:", err)
	}

	// Test non-supported device property first.
	var battery DevicePropDesc
	err = dev.GetDevicePropDesc(DPC_BatteryLevel, &battery)
	if err != nil {
		// Not an error; not supported on Android.
		t.Log("battery failed:", err)
	} else {
		t.Logf("%#v\n", battery)
	}

	var friendly DevicePropDesc
	err = dev.GetDevicePropDesc(DPC_MTP_DeviceFriendlyName, &friendly)
	if err != nil {
		t.Fatal("GetDevicePropDesc FriendlyName failed:", err)
	} else {
		t.Logf("%s: %#v\n", DPC_names[DPC_MTP_DeviceFriendlyName], friendly)
	}
	before := friendly.CurrentValue.(string)

	newVal := fmt.Sprintf("gomtp device_test %x", rand.Int31())
	str := StringValue{newVal}
	err = dev.SetDevicePropValue(DPC_MTP_DeviceFriendlyName, &str)
	if err != nil {
		t.Error("SetDevicePropValue failed:", err)
	}

	str.Value = ""
	err = dev.GetDevicePropValue(DPC_MTP_DeviceFriendlyName, &str)
	if err != nil {
		t.Error("GetDevicePropValue failed:", err)
	}
	if str.Value != newVal {
		t.Logf("got %q for property value, want %q\n", str.Value, newVal)
	}

	err = dev.ResetDevicePropValue(DPC_MTP_DeviceFriendlyName)
	if err != nil {
		// For some reason, this is not supported? Returns
		// unknown error code 0xffff
		t.Log("ResetDevicePropValue failed:", err)
	}

	str = StringValue{before}
	err = dev.SetDevicePropValue(DPC_MTP_DeviceFriendlyName, &str)
	if err != nil {
		t.Error("SetDevicePropValue failed:", err)
	}

	// Test object properties.
	props := Uint16Array{}
	err = dev.GetObjectPropsSupported(OFC_Undefined, &props)
	if err != nil {
		t.Errorf("GetObjectPropsSupported failed: %v\n", err)
	} else {
		t.Logf("GetObjectPropsSupported (OFC_Undefined) value: %s\n", getNames(OPC_names, props.Values))
	}

	for _, p := range props.Values {
		var objPropDesc ObjectPropDesc
		if p == OPC_PersistantUniqueObjectIdentifier {
			// can't deal with int128.
			continue
		}

		err = dev.GetObjectPropDesc(p, OFC_Undefined, &objPropDesc)
		name := OPC_names[int(p)]
		if err != nil {
			t.Errorf("GetObjectPropDesc(%s) failed: %v\n", name, err)
		} else {
			t.Logf("GetObjectPropDesc(%s) value: %#v %T\n", name, objPropDesc,
				InstantiateType(objPropDesc.DataType).Interface())
		}
	}

}

func TestDeviceInfo(t *testing.T) {
	dev, err := SelectDevice("")
	if err != nil {
		t.Fatal(err)
	}
	defer dev.Close()

	i, _ := dev.ID()
	t.Log("device:", i)
	info := DeviceInfo{}
	err = dev.GetDeviceInfo(&info)
	if err != nil {
		t.Error("GetDeviceInfo failed:", err)
	} else {
		t.Logf("%v\n", &info)
	}
}

func TestDeviceStorage(t *testing.T) {
	dev, err := SelectDevice("")
	if err != nil {
		t.Fatal(err)
	}
	defer dev.Close()

	i, _ := dev.ID()
	t.Log("device:", i)

	info := DeviceInfo{}
	err = dev.GetDeviceInfo(&info)
	if err != nil {
		t.Log("GetDeviceInfo failed:", err)
	} else {
		t.Logf("device info %v\n", &info)
	}

	err = dev.Configure()
	if err != nil {
		t.Fatal("Configure failed:", err)
	}

	sids := Uint32Array{}
	err = dev.GetStorageIDs(&sids)
	if err != nil {
		t.Fatalf("GetStorageIDs failed: %v", err)
	} else {
		t.Logf("%#v\n", sids)
	}

	if len(sids.Values) == 0 {
		t.Fatalf("No storages")
	}

	id := sids.Values[0]
	var storageInfo StorageInfo
	dev.GetStorageInfo(id, &storageInfo)
	if err != nil {
		t.Fatalf("GetStorageInfo failed:", err)
	} else {
		t.Logf("%#v\n", storageInfo)
	}

	resp, err := dev.GetNumObjects(id, 0x0, 0x0)
	if err != nil {
		t.Fatalf("GenericRPC failed:", err)
	} else {
		t.Logf("num objects %#v\n", resp)
	}

	// 500 + 512 triggers the null read case on both sides.
	const testSize = 500 + 512
	data := make([]byte, testSize)
	for i := range data {
		data[i] = byte(i % 256)
	}

	name := fmt.Sprintf("go-mtp-test%x", rand.Int31())
	buf := bytes.NewBuffer(data)
	send := ObjectInfo{
		StorageID:        id,
		ObjectFormat:     OFC_Undefined,
		ParentObject:     0xFFFFFFFF,
		Filename:         name,
		CompressedSize:   uint32(len(data)),
		ModificationDate: time.Now(),
		Keywords:         "bla",
	}

	_, _, handle, err := dev.SendObjectInfo(id, 0xFFFFFFFF, &send)
	if err != nil {
		t.Fatalf("SendObjectInfo failed:", err)
	} else {
		t.Logf("Sent objectinfo handle: 0x%x\n", handle)
		err = dev.SendObject(buf, int64(len(data)))
		if err != nil {
			t.Log("SendObject failed:", err)
		}
	}

	hs := Uint32Array{}
	err = dev.GetObjectHandles(id,
		OFC_Undefined,
		//OFC_Association,
		0xFFFFFFFF, &hs)

	if err != nil {
		t.Fatalf("GetObjectHandles failed: %v", err)
	} else {
		found := false
		for _, h := range hs.Values {
			if h == handle {
				found = true
				break
			}
		}

		if !found {
			t.Fatalf("uploaded file is not there")
		}
	}

	var backInfo ObjectInfo
	err = dev.GetObjectInfo(handle, &backInfo)
	if err != nil {
		t.Fatalf("GetObjectInfo failed: %v", err)
	} else {
		t.Logf("info %#v\n", backInfo)
	}

	var objSize Uint64Value
	err = dev.GetObjectPropValue(handle, OPC_ObjectSize, &objSize)
	if err != nil {
		t.Fatalf("GetObjectPropValue failed: %v", err)
	} else {
		t.Logf("info %#v\n", objSize)
	}

	if objSize.Value != testSize {
		t.Errorf("object size error: got %v, want %v", objSize.Value, testSize)
	}

	backBuf := &bytes.Buffer{}
	err = dev.GetObject(handle, backBuf)
	if err != nil {
		t.Fatalf("GetObject failed: %v", err)
	} else {
		if bytes.Compare(backBuf.Bytes(), data) != 0 {
			t.Fatalf("back comparison failed.")
		}
	}

	newName := fmt.Sprintf("mtp-doodle-test%x", rand.Int31())
	err = dev.SetObjectPropValue(handle, OPC_ObjectFileName, &StringValue{newName})
	if err != nil {
		t.Errorf("error renaming object: %v", err)
	}

	err = dev.DeleteObject(handle)
	if err != nil {
		t.Fatalf("DeleteObject failed: %v", err)
	}
}
<|endoftext|>


<commit_before>package daemon

import (
	"encoding/json"
	"net"
	"os"
	"os/signal"
	"time"

	log "github.com/socketplane/socketplane/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/socketplane/socketplane/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/socketplane/socketplane/datastore"
	"github.com/socketplane/socketplane/ovs"
)

type Daemon struct {
	Configuration *Configuration
	Connections   map[string]*Connection
}

func NewDaemon() *Daemon {
	return &Daemon{
		&Configuration{},
		map[string]*Connection{},
	}
}

func (d *Daemon) Run(ctx *cli.Context) {
	if ctx.Bool("debug") {
		log.SetLevel(log.DebugLevel)
	}

	var bindInterface string
	if ctx.String("iface") != "auto" {
		bindInterface = ctx.String("iface")
	} else {
		intf := identifyInterfaceToBind()
		if intf != nil {
			bindInterface = intf.Name
		}
	}
	if bindInterface != "" {
		log.Printf("Binding to %s", bindInterface)
	} else {
		log.Errorf("Unable to identify any Interface to Bind to. Going with Defaults")
	}
	go ServeAPI(d)
	go func() {
		err := ovs.CreateBridge()
		if err != nil {
			log.Error(err.Error)
		}
		d.populateConnections()
		datastore.Init(bindInterface, ctx.Bool("bootstrap"))
		_, err = ovs.CreateDefaultNetwork()
		if err != nil {
			log.Error(err.Error)
		}
		Bonjour(bindInterface)
	}()
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		for _ = range c {
			os.Exit(0)
		}
	}()
	select {}
}

func identifyInterfaceToBind() *net.Interface {
	// If the user isnt binding an interface using --iface option and let the daemon to
	// identify the interface, the daemon will try its best to identify the best interface
	// for the job.
	// In a few auto-install / zerotouch config scenarios, eligible interfaces may
	// be identified after the socketplane daemon is up and running.

	for {
		intf := InterfaceToBind()
		if intf != nil {
			return intf
		}
		time.Sleep(time.Second * 5)
		log.Infof("Identifying interface to bind ... Use --iface option for static binding")
	}
	return nil
}

func (d *Daemon) populateConnections() {
	for key, val := range ovs.ContextCache {
		connection := &Connection{}
		err := json.Unmarshal([]byte(val), connection)
		if err == nil {
			d.Connections[key] = connection
		}
	}
}
<commit_msg>Fixes #47 with the help of ecc caching support<commit_after>package daemon

import (
	"encoding/json"
	"net"
	"os"
	"os/signal"
	"time"

	log "github.com/socketplane/socketplane/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/socketplane/socketplane/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/socketplane/socketplane/datastore"
	"github.com/socketplane/socketplane/ovs"
)

type Daemon struct {
	Configuration *Configuration
	Connections   map[string]*Connection
}

func NewDaemon() *Daemon {
	return &Daemon{
		&Configuration{},
		map[string]*Connection{},
	}
}

func (d *Daemon) Run(ctx *cli.Context) {
	if ctx.Bool("debug") {
		log.SetLevel(log.DebugLevel)
	}
	bootstrapNode := ctx.Bool("bootstrap")
	serialChan := make(chan bool)

	go ServeAPI(d)
	go func() {
		var bindInterface string
		if ctx.String("iface") != "auto" {
			bindInterface = ctx.String("iface")
		} else {
			intf := identifyInterfaceToBind()
			if intf != nil {
				bindInterface = intf.Name
			}
		}
		if bindInterface != "" {
			log.Printf("Binding to %s", bindInterface)
		} else {
			log.Errorf("Unable to identify any Interface to Bind to. Going with Defaults")
		}
		datastore.Init(bindInterface, bootstrapNode)
		Bonjour(bindInterface)
		if !bootstrapNode {
			serialChan <- true
		}
	}()

	go func() {
		if !bootstrapNode {
			log.Printf("Non-Bootstrap node waiting on peer discovery")
			<-serialChan
			log.Printf("Non-Bootstrap node admitted into cluster")
		}
		err := ovs.CreateBridge()
		if err != nil {
			log.Error(err.Error)
		}
		d.populateConnections()
		_, err = ovs.CreateDefaultNetwork()
		if err != nil {
			log.Error(err.Error)
		}
	}()
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		for _ = range c {
			os.Exit(0)
		}
	}()
	select {}
}

func identifyInterfaceToBind() *net.Interface {
	// If the user isnt binding an interface using --iface option and let the daemon to
	// identify the interface, the daemon will try its best to identify the best interface
	// for the job.
	// In a few auto-install / zerotouch config scenarios, eligible interfaces may
	// be identified after the socketplane daemon is up and running.

	for {
		intf := InterfaceToBind()
		if intf != nil {
			return intf
		}
		time.Sleep(time.Second * 5)
		log.Infof("Identifying interface to bind ... Use --iface option for static binding")
	}
	return nil
}

func (d *Daemon) populateConnections() {
	for key, val := range ovs.ContextCache {
		connection := &Connection{}
		err := json.Unmarshal([]byte(val), connection)
		if err == nil {
			d.Connections[key] = connection
		}
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"io"
	"sort"
	"strings"
	"text/template"

	cmds "github.com/jbenet/go-ipfs/commands"
)

const (
	requiredArg = "<%v>"
	optionalArg = "[<%v>]"
	variadicArg = "%v..."
	optionFlag  = "-%v"
	optionType  = "(%v)"

	whitespace = "\r\n\t "

	indentStr = "    "
)

type helpFields struct {
	Indent      string
	Usage       string
	Path        string
	ArgUsage    string
	Tagline     string
	Arguments   string
	Options     string
	Synopsis    string
	Subcommands string
	Description string
	MoreHelp    bool
}

// TrimNewlines removes extra newlines from fields. This makes aligning
// commands easier. Below, the leading + tralining newlines are removed:
//	Synopsis: `
//	    ipfs config <key>          - Get value of <key>
//	    ipfs config <key> <value>  - Set value of <key> to <value>
//	    ipfs config --show         - Show config file
//	    ipfs config --edit         - Edit config file in $EDITOR
//	`
func (f *helpFields) TrimNewlines() {
	f.Path = strings.Trim(f.Path, "\n")
	f.ArgUsage = strings.Trim(f.ArgUsage, "\n")
	f.Tagline = strings.Trim(f.Tagline, "\n")
	f.Arguments = strings.Trim(f.Arguments, "\n")
	f.Options = strings.Trim(f.Options, "\n")
	f.Synopsis = strings.Trim(f.Synopsis, "\n")
	f.Subcommands = strings.Trim(f.Subcommands, "\n")
	f.Description = strings.Trim(f.Description, "\n")
}

// Indent adds whitespace the lines of fields.
func (f *helpFields) IndentAll() {
	indent := func(s string) string {
		if s == "" {
			return s
		}
		return indentString(s, indentStr)
	}

	f.Arguments = indent(f.Arguments)
	f.Options = indent(f.Options)
	f.Synopsis = indent(f.Synopsis)
	f.Subcommands = indent(f.Subcommands)
	f.Description = indent(f.Description)
}

const usageFormat = "{{if .Usage}}{{.Usage}}{{else}}{{.Path}}{{if .ArgUsage}} {{.ArgUsage}}{{end}} - {{.Tagline}}{{end}}"

const longHelpFormat = `
{{.Indent}}{{template "usage" .}}

{{if .Arguments}}ARGUMENTS:

{{.Arguments}}

{{end}}{{if .Options}}OPTIONS:

{{.Options}}

{{end}}{{if .Subcommands}}SUBCOMMANDS:

{{.Subcommands}}

{{.Indent}}Use '{{.Path}} <subcmd> --help' for more information about each command.

{{end}}{{if .Description}}DESCRIPTION:

{{.Description}}

{{end}}
`
const shortHelpFormat = `USAGE:

{{.Indent}}{{template "usage" .}}
{{if .Synopsis}}
SYNOPSIS

{{.Synopsis}}
{{end}}{{if .Description}}
{{.Description}}
{{end}}
{{if .MoreHelp}}Use '{{.Path}} --help' for more information about this command.
{{end}}
`

var usageTemplate *template.Template
var longHelpTemplate *template.Template
var shortHelpTemplate *template.Template

func init() {
	usageTemplate = template.Must(template.New("usage").Parse(usageFormat))
	longHelpTemplate = template.Must(usageTemplate.New("longHelp").Parse(longHelpFormat))
	shortHelpTemplate = template.Must(usageTemplate.New("shortHelp").Parse(shortHelpFormat))
}

// LongHelp returns a formatted CLI helptext string, generated for the given command
func LongHelp(rootName string, root *cmds.Command, path []string, out io.Writer) error {
	cmd, err := root.Get(path)
	if err != nil {
		return err
	}

	pathStr := rootName
	if len(path) > 0 {
		pathStr += " " + strings.Join(path, " ")
	}

	fields := helpFields{
		Indent:      indentStr,
		Path:        pathStr,
		ArgUsage:    usageText(cmd),
		Tagline:     cmd.Helptext.Tagline,
		Arguments:   cmd.Helptext.Arguments,
		Options:     cmd.Helptext.Options,
		Synopsis:    cmd.Helptext.Synopsis,
		Subcommands: cmd.Helptext.Subcommands,
		Description: cmd.Helptext.ShortDescription,
		Usage:       cmd.Helptext.Usage,
		MoreHelp:    (cmd != root),
	}

	if len(cmd.Helptext.LongDescription) > 0 {
		fields.Description = cmd.Helptext.LongDescription
	}

	// autogen fields that are empty
	if len(fields.Arguments) == 0 {
		fields.Arguments = strings.Join(argumentText(cmd), "\n")
	}
	if len(fields.Options) == 0 {
		fields.Options = strings.Join(optionText(cmd), "\n")
	}
	if len(fields.Subcommands) == 0 {
		fields.Subcommands = strings.Join(subcommandText(cmd, rootName, path), "\n")
	}

	// trim the extra newlines (see TrimNewlines doc)
	fields.TrimNewlines()

	// indent all fields that have been set
	fields.IndentAll()

	return longHelpTemplate.Execute(out, fields)
}

// ShortHelp returns a formatted CLI helptext string, generated for the given command
func ShortHelp(rootName string, root *cmds.Command, path []string, out io.Writer) error {
	cmd, err := root.Get(path)
	if err != nil {
		return err
	}

	// default cmd to root if there is no path
	if path == nil && cmd == nil {
		cmd = root
	}

	pathStr := rootName
	if len(path) > 0 {
		pathStr += " " + strings.Join(path, " ")
	}

	fields := helpFields{
		Indent:      indentStr,
		Path:        pathStr,
		ArgUsage:    usageText(cmd),
		Tagline:     cmd.Helptext.Tagline,
		Synopsis:    cmd.Helptext.Synopsis,
		Description: cmd.Helptext.ShortDescription,
		Usage:       cmd.Helptext.Usage,
		MoreHelp:    (cmd != root),
	}

	// trim the extra newlines (see TrimNewlines doc)
	fields.TrimNewlines()

	// indent all fields that have been set
	fields.IndentAll()

	return shortHelpTemplate.Execute(out, fields)
}

func argumentText(cmd *cmds.Command) []string {
	lines := make([]string, len(cmd.Arguments))

	for i, arg := range cmd.Arguments {
		lines[i] = argUsageText(arg)
	}
	lines = align(lines)
	for i, arg := range cmd.Arguments {
		lines[i] += " - " + arg.Description
	}

	return lines
}

func optionText(cmd ...*cmds.Command) []string {
	// get a slice of the options we want to list out
	options := make([]cmds.Option, 0)
	for _, c := range cmd {
		for _, opt := range c.Options {
			options = append(options, opt)
		}
	}

	// add option names to output (with each name aligned)
	lines := make([]string, 0)
	j := 0
	for {
		done := true
		i := 0
		for _, opt := range options {
			if len(lines) < i+1 {
				lines = append(lines, "")
			}

			names := sortByLength(opt.Names)
			if len(names) >= j+1 {
				lines[i] += fmt.Sprintf(optionFlag, names[j])
			}
			if len(names) > j+1 {
				lines[i] += ", "
				done = false
			}

			i++
		}

		if done {
			break
		}

		lines = align(lines)
		j++
	}
	lines = align(lines)

	// add option types to output
	for i, opt := range options {
		lines[i] += " " + fmt.Sprintf("%v", opt.Type)
	}
	lines = align(lines)

	// add option descriptions to output
	for i, opt := range options {
		lines[i] += " - " + opt.Description
	}

	return lines
}

func subcommandText(cmd *cmds.Command, rootName string, path []string) []string {
	prefix := fmt.Sprintf("%v %v", rootName, strings.Join(path, " "))
	if len(path) > 0 {
		prefix += " "
	}
	subcmds := make([]*cmds.Command, len(cmd.Subcommands))
	lines := make([]string, len(cmd.Subcommands))

	i := 0
	for name, sub := range cmd.Subcommands {
		usage := usageText(sub)
		if len(usage) > 0 {
			usage = " " + usage
		}
		lines[i] = prefix + name + usage
		subcmds[i] = sub
		i++
	}

	lines = align(lines)
	for i, sub := range subcmds {
		lines[i] += " - " + sub.Helptext.Tagline
	}

	return lines
}

func usageText(cmd *cmds.Command) string {
	s := ""
	for i, arg := range cmd.Arguments {
		if i != 0 {
			s += " "
		}
		s += argUsageText(arg)
	}

	return s
}

func argUsageText(arg cmds.Argument) string {
	s := arg.Name

	if arg.Required {
		s = fmt.Sprintf(requiredArg, s)
	} else {
		s = fmt.Sprintf(optionalArg, s)
	}

	if arg.Variadic {
		s = fmt.Sprintf(variadicArg, s)
	}

	return s
}

func align(lines []string) []string {
	longest := 0
	for _, line := range lines {
		length := len(line)
		if length > longest {
			longest = length
		}
	}

	for i, line := range lines {
		length := len(line)
		if length > 0 {
			lines[i] += strings.Repeat(" ", longest-length)
		}
	}

	return lines
}

func indent(lines []string, prefix string) []string {
	for i, line := range lines {
		lines[i] = prefix + indentString(line, prefix)
	}
	return lines
}

func indentString(line string, prefix string) string {
	return prefix + strings.Replace(line, "\n", "\n"+prefix, -1)
}

type lengthSlice []string

func (ls lengthSlice) Len() int {
	return len(ls)
}
func (ls lengthSlice) Swap(a, b int) {
	ls[a], ls[b] = ls[b], ls[a]
}
func (ls lengthSlice) Less(a, b int) bool {
	return len(ls[a]) < len(ls[b])
}

func sortByLength(slice []string) []string {
	output := make(lengthSlice, len(slice))
	for i, val := range slice {
		output[i] = val
	}
	sort.Sort(output)
	return []string(output)
}
<commit_msg>Remove 'SYNOPSIS' from short help text<commit_after>package cli

import (
	"fmt"
	"io"
	"sort"
	"strings"
	"text/template"

	cmds "github.com/jbenet/go-ipfs/commands"
)

const (
	requiredArg = "<%v>"
	optionalArg = "[<%v>]"
	variadicArg = "%v..."
	optionFlag  = "-%v"
	optionType  = "(%v)"

	whitespace = "\r\n\t "

	indentStr = "    "
)

type helpFields struct {
	Indent      string
	Usage       string
	Path        string
	ArgUsage    string
	Tagline     string
	Arguments   string
	Options     string
	Synopsis    string
	Subcommands string
	Description string
	MoreHelp    bool
}

// TrimNewlines removes extra newlines from fields. This makes aligning
// commands easier. Below, the leading + tralining newlines are removed:
//	Synopsis: `
//	    ipfs config <key>          - Get value of <key>
//	    ipfs config <key> <value>  - Set value of <key> to <value>
//	    ipfs config --show         - Show config file
//	    ipfs config --edit         - Edit config file in $EDITOR
//	`
func (f *helpFields) TrimNewlines() {
	f.Path = strings.Trim(f.Path, "\n")
	f.ArgUsage = strings.Trim(f.ArgUsage, "\n")
	f.Tagline = strings.Trim(f.Tagline, "\n")
	f.Arguments = strings.Trim(f.Arguments, "\n")
	f.Options = strings.Trim(f.Options, "\n")
	f.Synopsis = strings.Trim(f.Synopsis, "\n")
	f.Subcommands = strings.Trim(f.Subcommands, "\n")
	f.Description = strings.Trim(f.Description, "\n")
}

// Indent adds whitespace the lines of fields.
func (f *helpFields) IndentAll() {
	indent := func(s string) string {
		if s == "" {
			return s
		}
		return indentString(s, indentStr)
	}

	f.Arguments = indent(f.Arguments)
	f.Options = indent(f.Options)
	f.Synopsis = indent(f.Synopsis)
	f.Subcommands = indent(f.Subcommands)
	f.Description = indent(f.Description)
}

const usageFormat = "{{if .Usage}}{{.Usage}}{{else}}{{.Path}}{{if .ArgUsage}} {{.ArgUsage}}{{end}} - {{.Tagline}}{{end}}"

const longHelpFormat = `
{{.Indent}}{{template "usage" .}}

{{if .Arguments}}ARGUMENTS:

{{.Arguments}}

{{end}}{{if .Options}}OPTIONS:

{{.Options}}

{{end}}{{if .Subcommands}}SUBCOMMANDS:

{{.Subcommands}}

{{.Indent}}Use '{{.Path}} <subcmd> --help' for more information about each command.

{{end}}{{if .Description}}DESCRIPTION:

{{.Description}}

{{end}}
`
const shortHelpFormat = `USAGE:

{{.Indent}}{{template "usage" .}}
{{if .Synopsis}}
{{.Synopsis}}
{{end}}{{if .Description}}
{{.Description}}
{{end}}
{{if .MoreHelp}}Use '{{.Path}} --help' for more information about this command.
{{end}}
`

var usageTemplate *template.Template
var longHelpTemplate *template.Template
var shortHelpTemplate *template.Template

func init() {
	usageTemplate = template.Must(template.New("usage").Parse(usageFormat))
	longHelpTemplate = template.Must(usageTemplate.New("longHelp").Parse(longHelpFormat))
	shortHelpTemplate = template.Must(usageTemplate.New("shortHelp").Parse(shortHelpFormat))
}

// LongHelp returns a formatted CLI helptext string, generated for the given command
func LongHelp(rootName string, root *cmds.Command, path []string, out io.Writer) error {
	cmd, err := root.Get(path)
	if err != nil {
		return err
	}

	pathStr := rootName
	if len(path) > 0 {
		pathStr += " " + strings.Join(path, " ")
	}

	fields := helpFields{
		Indent:      indentStr,
		Path:        pathStr,
		ArgUsage:    usageText(cmd),
		Tagline:     cmd.Helptext.Tagline,
		Arguments:   cmd.Helptext.Arguments,
		Options:     cmd.Helptext.Options,
		Synopsis:    cmd.Helptext.Synopsis,
		Subcommands: cmd.Helptext.Subcommands,
		Description: cmd.Helptext.ShortDescription,
		Usage:       cmd.Helptext.Usage,
		MoreHelp:    (cmd != root),
	}

	if len(cmd.Helptext.LongDescription) > 0 {
		fields.Description = cmd.Helptext.LongDescription
	}

	// autogen fields that are empty
	if len(fields.Arguments) == 0 {
		fields.Arguments = strings.Join(argumentText(cmd), "\n")
	}
	if len(fields.Options) == 0 {
		fields.Options = strings.Join(optionText(cmd), "\n")
	}
	if len(fields.Subcommands) == 0 {
		fields.Subcommands = strings.Join(subcommandText(cmd, rootName, path), "\n")
	}

	// trim the extra newlines (see TrimNewlines doc)
	fields.TrimNewlines()

	// indent all fields that have been set
	fields.IndentAll()

	return longHelpTemplate.Execute(out, fields)
}

// ShortHelp returns a formatted CLI helptext string, generated for the given command
func ShortHelp(rootName string, root *cmds.Command, path []string, out io.Writer) error {
	cmd, err := root.Get(path)
	if err != nil {
		return err
	}

	// default cmd to root if there is no path
	if path == nil && cmd == nil {
		cmd = root
	}

	pathStr := rootName
	if len(path) > 0 {
		pathStr += " " + strings.Join(path, " ")
	}

	fields := helpFields{
		Indent:      indentStr,
		Path:        pathStr,
		ArgUsage:    usageText(cmd),
		Tagline:     cmd.Helptext.Tagline,
		Synopsis:    cmd.Helptext.Synopsis,
		Description: cmd.Helptext.ShortDescription,
		Usage:       cmd.Helptext.Usage,
		MoreHelp:    (cmd != root),
	}

	// trim the extra newlines (see TrimNewlines doc)
	fields.TrimNewlines()

	// indent all fields that have been set
	fields.IndentAll()

	return shortHelpTemplate.Execute(out, fields)
}

func argumentText(cmd *cmds.Command) []string {
	lines := make([]string, len(cmd.Arguments))

	for i, arg := range cmd.Arguments {
		lines[i] = argUsageText(arg)
	}
	lines = align(lines)
	for i, arg := range cmd.Arguments {
		lines[i] += " - " + arg.Description
	}

	return lines
}

func optionText(cmd ...*cmds.Command) []string {
	// get a slice of the options we want to list out
	options := make([]cmds.Option, 0)
	for _, c := range cmd {
		for _, opt := range c.Options {
			options = append(options, opt)
		}
	}

	// add option names to output (with each name aligned)
	lines := make([]string, 0)
	j := 0
	for {
		done := true
		i := 0
		for _, opt := range options {
			if len(lines) < i+1 {
				lines = append(lines, "")
			}

			names := sortByLength(opt.Names)
			if len(names) >= j+1 {
				lines[i] += fmt.Sprintf(optionFlag, names[j])
			}
			if len(names) > j+1 {
				lines[i] += ", "
				done = false
			}

			i++
		}

		if done {
			break
		}

		lines = align(lines)
		j++
	}
	lines = align(lines)

	// add option types to output
	for i, opt := range options {
		lines[i] += " " + fmt.Sprintf("%v", opt.Type)
	}
	lines = align(lines)

	// add option descriptions to output
	for i, opt := range options {
		lines[i] += " - " + opt.Description
	}

	return lines
}

func subcommandText(cmd *cmds.Command, rootName string, path []string) []string {
	prefix := fmt.Sprintf("%v %v", rootName, strings.Join(path, " "))
	if len(path) > 0 {
		prefix += " "
	}
	subcmds := make([]*cmds.Command, len(cmd.Subcommands))
	lines := make([]string, len(cmd.Subcommands))

	i := 0
	for name, sub := range cmd.Subcommands {
		usage := usageText(sub)
		if len(usage) > 0 {
			usage = " " + usage
		}
		lines[i] = prefix + name + usage
		subcmds[i] = sub
		i++
	}

	lines = align(lines)
	for i, sub := range subcmds {
		lines[i] += " - " + sub.Helptext.Tagline
	}

	return lines
}

func usageText(cmd *cmds.Command) string {
	s := ""
	for i, arg := range cmd.Arguments {
		if i != 0 {
			s += " "
		}
		s += argUsageText(arg)
	}

	return s
}

func argUsageText(arg cmds.Argument) string {
	s := arg.Name

	if arg.Required {
		s = fmt.Sprintf(requiredArg, s)
	} else {
		s = fmt.Sprintf(optionalArg, s)
	}

	if arg.Variadic {
		s = fmt.Sprintf(variadicArg, s)
	}

	return s
}

func align(lines []string) []string {
	longest := 0
	for _, line := range lines {
		length := len(line)
		if length > longest {
			longest = length
		}
	}

	for i, line := range lines {
		length := len(line)
		if length > 0 {
			lines[i] += strings.Repeat(" ", longest-length)
		}
	}

	return lines
}

func indent(lines []string, prefix string) []string {
	for i, line := range lines {
		lines[i] = prefix + indentString(line, prefix)
	}
	return lines
}

func indentString(line string, prefix string) string {
	return prefix + strings.Replace(line, "\n", "\n"+prefix, -1)
}

type lengthSlice []string

func (ls lengthSlice) Len() int {
	return len(ls)
}
func (ls lengthSlice) Swap(a, b int) {
	ls[a], ls[b] = ls[b], ls[a]
}
func (ls lengthSlice) Less(a, b int) bool {
	return len(ls[a]) < len(ls[b])
}

func sortByLength(slice []string) []string {
	output := make(lengthSlice, len(slice))
	for i, val := range slice {
		output[i] = val
	}
	sort.Sort(output)
	return []string(output)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// 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"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"github.com/gorilla/mux"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	jaegerClientConfig "github.com/uber/jaeger-client-go/config"
	"github.com/uber/jaeger-lib/metrics/go-kit"
	"github.com/uber/jaeger-lib/metrics/go-kit/expvar"
	"go.uber.org/zap"

	basicB "github.com/jaegertracing/jaeger/cmd/builder"
	"github.com/jaegertracing/jaeger/cmd/flags"
	casFlags "github.com/jaegertracing/jaeger/cmd/flags/cassandra"
	esFlags "github.com/jaegertracing/jaeger/cmd/flags/es"
	"github.com/jaegertracing/jaeger/cmd/query/app"
	"github.com/jaegertracing/jaeger/cmd/query/app/builder"
	"github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/pkg/healthcheck"
	"github.com/jaegertracing/jaeger/pkg/recoveryhandler"
	"github.com/jaegertracing/jaeger/pkg/version"
)

func main() {
	var serverChannel = make(chan os.Signal, 0)
	signal.Notify(serverChannel, os.Interrupt, syscall.SIGTERM)

	casOptions := casFlags.NewOptions("cassandra", "cassandra.archive")
	esOptions := esFlags.NewOptions("es", "es.archive")
	v := viper.New()

	var command = &cobra.Command{
		Use:   "jaeger-query",
		Short: "Jaeger query is a service to access tracing data",
		Long:  `Jaeger query is a service to access tracing data and host UI.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			err := flags.TryLoadConfigFile(v)
			if err != nil {
				return err
			}

			sFlags := new(flags.SharedFlags).InitFromViper(v)
			logger, err := sFlags.NewLogger(zap.NewProductionConfig())
			if err != nil {
				return err
			}

			casOptions.InitFromViper(v)
			esOptions.InitFromViper(v)
			queryOpts := new(builder.QueryOptions).InitFromViper(v)

			hc, err := healthcheck.Serve(http.StatusServiceUnavailable, queryOpts.HealthCheckHTTPPort, logger)
			if err != nil {
				logger.Fatal("Could not start the health check server.", zap.Error(err))
			}

			metricsFactory := xkit.Wrap("jaeger-query", expvar.NewFactory(10))

			tracer, closer, err := jaegerClientConfig.Configuration{
				Sampler: &jaegerClientConfig.SamplerConfig{
					Type:  "probabilistic",
					Param: 1.0,
				},
				RPCMetrics: true,
			}.New("jaeger-query", jaegerClientConfig.Metrics(metricsFactory))
			if err != nil {
				logger.Fatal("Failed to initialize tracer", zap.Error(err))
			}
			defer closer.Close()

			storageBuild, err := builder.NewStorageBuilder(
				sFlags.SpanStorage.Type,
				sFlags.DependencyStorage.DataFrequency,
				basicB.Options.LoggerOption(logger),
				basicB.Options.MetricsFactoryOption(metricsFactory),
				basicB.Options.CassandraSessionOption(casOptions.GetPrimary()),
				basicB.Options.ElasticClientOption(esOptions.GetPrimary()),
			)
			if err != nil {
				logger.Fatal("Failed to init storage builder", zap.Error(err))
			}

			apiHandler := app.NewAPIHandler(
				storageBuild.SpanReader,
				storageBuild.DependencyReader,
				app.HandlerOptions.Prefix(queryOpts.Prefix),
				app.HandlerOptions.Logger(logger),
				app.HandlerOptions.Tracer(tracer))
			r := mux.NewRouter()
			apiHandler.RegisterRoutes(r)
			registerStaticHandler(r, logger, queryOpts)
			portStr := ":" + strconv.Itoa(queryOpts.Port)
			recoveryHandler := recoveryhandler.NewRecoveryHandler(logger, true)

			go func() {
				logger.Info("Starting jaeger-query HTTP server", zap.Int("port", queryOpts.Port))
				if err := http.ListenAndServe(portStr, recoveryHandler(r)); err != nil {
					logger.Fatal("Could not launch service", zap.Error(err))
				}
				hc.Set(http.StatusInternalServerError)
			}()

			hc.Ready()

			select {
			case <-serverChannel:
				logger.Info("Jaeger Query is finishing")
			}
			return nil
		},
	}

	command.AddCommand(version.Command())

	config.AddFlags(
		v,
		command,
		flags.AddConfigFileFlag,
		flags.AddFlags,
		casOptions.AddFlags,
		esOptions.AddFlags,
		builder.AddFlags,
	)

	if error := command.Execute(); error != nil {
		fmt.Println(error.Error())
		os.Exit(1)
	}
}

func registerStaticHandler(r *mux.Router, logger *zap.Logger, qOpts *builder.QueryOptions) {
	staticHandler, err := app.NewStaticAssetsHandler(qOpts.StaticAssets, qOpts.UIConfig)
	if err != nil {
		logger.Fatal("Could not create static assets handler", zap.Error(err))
	}
	if staticHandler != nil {
		staticHandler.RegisterRoutes(r)
	} else {
		logger.Info("Static handler is not registered")
	}
}
<commit_msg>Add HTTP compression handler (#545)<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// 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"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	jaegerClientConfig "github.com/uber/jaeger-client-go/config"
	"github.com/uber/jaeger-lib/metrics/go-kit"
	"github.com/uber/jaeger-lib/metrics/go-kit/expvar"
	"go.uber.org/zap"

	basicB "github.com/jaegertracing/jaeger/cmd/builder"
	"github.com/jaegertracing/jaeger/cmd/flags"
	casFlags "github.com/jaegertracing/jaeger/cmd/flags/cassandra"
	esFlags "github.com/jaegertracing/jaeger/cmd/flags/es"
	"github.com/jaegertracing/jaeger/cmd/query/app"
	"github.com/jaegertracing/jaeger/cmd/query/app/builder"
	"github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/pkg/healthcheck"
	"github.com/jaegertracing/jaeger/pkg/recoveryhandler"
	"github.com/jaegertracing/jaeger/pkg/version"
)

func main() {
	var serverChannel = make(chan os.Signal, 0)
	signal.Notify(serverChannel, os.Interrupt, syscall.SIGTERM)

	casOptions := casFlags.NewOptions("cassandra", "cassandra.archive")
	esOptions := esFlags.NewOptions("es", "es.archive")
	v := viper.New()

	var command = &cobra.Command{
		Use:   "jaeger-query",
		Short: "Jaeger query is a service to access tracing data",
		Long:  `Jaeger query is a service to access tracing data and host UI.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			err := flags.TryLoadConfigFile(v)
			if err != nil {
				return err
			}

			sFlags := new(flags.SharedFlags).InitFromViper(v)
			logger, err := sFlags.NewLogger(zap.NewProductionConfig())
			if err != nil {
				return err
			}

			casOptions.InitFromViper(v)
			esOptions.InitFromViper(v)
			queryOpts := new(builder.QueryOptions).InitFromViper(v)

			hc, err := healthcheck.Serve(http.StatusServiceUnavailable, queryOpts.HealthCheckHTTPPort, logger)
			if err != nil {
				logger.Fatal("Could not start the health check server.", zap.Error(err))
			}

			metricsFactory := xkit.Wrap("jaeger-query", expvar.NewFactory(10))

			tracer, closer, err := jaegerClientConfig.Configuration{
				Sampler: &jaegerClientConfig.SamplerConfig{
					Type:  "probabilistic",
					Param: 1.0,
				},
				RPCMetrics: true,
			}.New("jaeger-query", jaegerClientConfig.Metrics(metricsFactory))
			if err != nil {
				logger.Fatal("Failed to initialize tracer", zap.Error(err))
			}
			defer closer.Close()

			storageBuild, err := builder.NewStorageBuilder(
				sFlags.SpanStorage.Type,
				sFlags.DependencyStorage.DataFrequency,
				basicB.Options.LoggerOption(logger),
				basicB.Options.MetricsFactoryOption(metricsFactory),
				basicB.Options.CassandraSessionOption(casOptions.GetPrimary()),
				basicB.Options.ElasticClientOption(esOptions.GetPrimary()),
			)
			if err != nil {
				logger.Fatal("Failed to init storage builder", zap.Error(err))
			}

			apiHandler := app.NewAPIHandler(
				storageBuild.SpanReader,
				storageBuild.DependencyReader,
				app.HandlerOptions.Prefix(queryOpts.Prefix),
				app.HandlerOptions.Logger(logger),
				app.HandlerOptions.Tracer(tracer))
			r := mux.NewRouter()
			apiHandler.RegisterRoutes(r)
			registerStaticHandler(r, logger, queryOpts)
			portStr := ":" + strconv.Itoa(queryOpts.Port)
			compressHandler := handlers.CompressHandler(r)
			recoveryHandler := recoveryhandler.NewRecoveryHandler(logger, true)

			go func() {
				logger.Info("Starting jaeger-query HTTP server", zap.Int("port", queryOpts.Port))
				if err := http.ListenAndServe(portStr, recoveryHandler(compressHandler)); err != nil {
					logger.Fatal("Could not launch service", zap.Error(err))
				}
				hc.Set(http.StatusInternalServerError)
			}()

			hc.Ready()

			select {
			case <-serverChannel:
				logger.Info("Jaeger Query is finishing")
			}
			return nil
		},
	}

	command.AddCommand(version.Command())

	config.AddFlags(
		v,
		command,
		flags.AddConfigFileFlag,
		flags.AddFlags,
		casOptions.AddFlags,
		esOptions.AddFlags,
		builder.AddFlags,
	)

	if error := command.Execute(); error != nil {
		fmt.Println(error.Error())
		os.Exit(1)
	}
}

func registerStaticHandler(r *mux.Router, logger *zap.Logger, qOpts *builder.QueryOptions) {
	staticHandler, err := app.NewStaticAssetsHandler(qOpts.StaticAssets, qOpts.UIConfig)
	if err != nil {
		logger.Fatal("Could not create static assets handler", zap.Error(err))
	}
	if staticHandler != nil {
		staticHandler.RegisterRoutes(r)
	} else {
		logger.Info("Static handler is not registered")
	}
}
<|endoftext|>