1	/*
     2	Copyright 2013 The Perkeep Authors.
     3	
     4	Licensed under the Apache License, Version 2.0 (the "License");
     5	you may not use this file except in compliance with the License.
     6	You may obtain a copy of the License at
     7	
     8	     http://www.apache.org/licenses/LICENSE-2.0
     9	
    10	Unless required by applicable law or agreed to in writing, software
    11	distributed under the License is distributed on an "AS IS" BASIS,
    12	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13	See the License for the specific language governing permissions and
    14	limitations under the License.
    15	*/
    16	
    17	// Package cmdmain contains the shared implementation for pk-get,
    18	// pk-put, pk, and other Perkeep command-line tools.
    19	package cmdmain // import "perkeep.org/pkg/cmdmain"
    20	
    21	import (
    22		"flag"
    23		"fmt"
    24		"io"
    25		"log"
    26		"os"
    27		"os/exec"
    28		"path/filepath"
    29		"sort"
    30		"strings"
    31		"sync"
    32	
    33		"perkeep.org/pkg/buildinfo"
    34	
    35		"go4.org/legal"
    36	)
    37	
    38	var (
    39		FlagVersion = flag.Bool("version", false, "show version")
    40		FlagHelp    = flag.Bool("help", false, "print usage")
    41		FlagVerbose = flag.Bool("verbose", false, "extra debug logging")
    42		FlagLegal   = flag.Bool("legal", false, "show licenses")
    43	)
    44	
    45	var (
    46		// ExtraFlagRegistration allows to add more flags from
    47		// other packages (with AddFlags) when Main starts.
    48		ExtraFlagRegistration = func() {}
    49		// PostFlag runs code that needs to happen after flags were parsed, but
    50		// before the subcommand is run.
    51		PostFlag = func() {}
    52		// PreExit runs after the subcommand, but before Main terminates
    53		// with either success or the error from the subcommand.
    54		PreExit = func() {}
    55		// ExitWithFailure determines whether the command exits
    56		// with a non-zero exit status.
    57		ExitWithFailure bool
    58	)
    59	
    60	var ErrUsage = UsageError("invalid command")
    61	
    62	type UsageError string
    63	
    64	func (ue UsageError) Error() string {
    65		return "Usage error: " + string(ue)
    66	}
    67	
    68	var (
    69		// mode name to actual subcommand mapping
    70		modeCommand = make(map[string]CommandRunner)
    71		modeFlags   = make(map[string]*flag.FlagSet)
    72		wantHelp    = make(map[string]*bool)
    73		// asNewCommand stores whether the mode should actually be run as a new
    74		// independent command.
    75		asNewCommand = make(map[string]bool)
    76	
    77		// Indirections for replacement by tests
    78		Stderr io.Writer = os.Stderr
    79		Stdout io.Writer = os.Stdout
    80		Stdin  io.Reader = os.Stdin
    81	
    82		Exit = realExit
    83		// TODO: abstract out vfs operation. should never call os.Stat, os.Open, os.Create, etc.
    84		// Only use fs.Stat, fs.Open, where vs is an interface type.
    85		// TODO: switch from using the global flag FlagSet and use our own. right now
    86		// running "go test -v" dumps the flag usage data to the global stderr.
    87	
    88		logger = log.New(Stderr, "", log.LstdFlags)
    89	)
    90	
    91	func realExit(code int) {
    92		os.Exit(code)
    93	}
    94	
    95	// CommandRunner is the type that a command mode should implement.
    96	type CommandRunner interface {
    97		Usage()
    98		RunCommand(args []string) error
    99	}
   100	
   101	// ExecRunner is the type that a command mode should implement when that mode
   102	// just calls a new executable that will run as a new command.
   103	type ExecRunner interface {
   104		CommandRunner
   105		LookPath() (string, error)
   106	}
   107	
   108	// Demoter is an interface that boring commands can implement to
   109	// demote themselves in the tool listing, for boring or low-level
   110	// subcommands. They only show up in --help mode.
   111	type Demoter interface {
   112		CommandRunner
   113		Demote() bool
   114	}
   115	
   116	type exampler interface {
   117		Examples() []string
   118	}
   119	
   120	type describer interface {
   121		Describe() string
   122	}
   123	
   124	func demote(c CommandRunner) bool {
   125		i, ok := c.(Demoter)
   126		return ok && i.Demote()
   127	}
   128	
   129	// RegisterMode adds a mode to the list of modes for the main command.
   130	// It is meant to be called in init() for each subcommand.
   131	func RegisterMode(mode string, makeCmd func(Flags *flag.FlagSet) CommandRunner) {
   132		if _, dup := modeCommand[mode]; dup {
   133			log.Fatalf("duplicate command %q registered", mode)
   134		}
   135		flags := flag.NewFlagSet(mode+" options", flag.ContinueOnError)
   136		flags.Usage = func() {}
   137	
   138		var cmdHelp bool
   139		flags.BoolVar(&cmdHelp, "help", false, "Help for this mode.")
   140		wantHelp[mode] = &cmdHelp
   141		modeFlags[mode] = flags
   142		modeCommand[mode] = makeCmd(flags)
   143	}
   144	
   145	// RegisterCommand adds a mode to the list of modes for the main command, and
   146	// also specifies that this mode is just another executable that runs as a new
   147	// cmdmain command. The executable to run is determined by the LookPath implementation
   148	// for this mode.
   149	func RegisterCommand(mode string, makeCmd func(Flags *flag.FlagSet) CommandRunner) {
   150		RegisterMode(mode, makeCmd)
   151		asNewCommand[mode] = true
   152	}
   153	
   154	func hasFlags(flags *flag.FlagSet) bool {
   155		any := false
   156		flags.VisitAll(func(*flag.Flag) {
   157			any = true
   158		})
   159		return any
   160	}
   161	
   162	func usage(msg string) {
   163		cmdName := filepath.Base(os.Args[0])
   164		if msg != "" {
   165			Errorf("Error: %v\n", msg)
   166		}
   167		var modesQualifer string
   168		if !*FlagHelp {
   169			modesQualifer = " (use --help to see all modes)"
   170		}
   171		Errorf(`
   172	Usage: `+cmdName+` [globalopts] <mode> [commandopts] [commandargs]
   173	
   174	Modes:%s
   175	
   176	`, modesQualifer)
   177		var modes []string
   178		for mode, cmd := range modeCommand {
   179			if des, ok := cmd.(describer); ok && (*FlagHelp || !demote(cmd)) {
   180				modes = append(modes, fmt.Sprintf("  %s: %s\n", mode, des.Describe()))
   181			}
   182		}
   183		sort.Strings(modes)
   184		for i := range modes {
   185			Errorf("%s", modes[i])
   186		}
   187	
   188		Errorf("\nExamples:\n")
   189		modes = nil
   190		for mode, cmd := range modeCommand {
   191			if ex, ok := cmd.(exampler); ok && (*FlagHelp || !demote(cmd)) {
   192				line := ""
   193				exs := ex.Examples()
   194				if len(exs) > 0 {
   195					line = "\n"
   196				}
   197				for _, example := range exs {
   198					line += fmt.Sprintf("  %s %s %s\n", cmdName, mode, example)
   199				}
   200				modes = append(modes, line)
   201			}
   202		}
   203		sort.Strings(modes)
   204		for i := range modes {
   205			Errorf("%s", modes[i])
   206		}
   207	
   208		Errorf(`
   209	For mode-specific help:
   210	
   211	  ` + cmdName + ` <mode> -help
   212	
   213	Global options:
   214	`)
   215		flag.PrintDefaults()
   216		Exit(1)
   217	}
   218	
   219	func help(mode string) {
   220		cmdName := os.Args[0]
   221		// We can skip all the checks as they're done in Main
   222		cmd := modeCommand[mode]
   223		cmdFlags := modeFlags[mode]
   224		cmdFlags.SetOutput(Stderr)
   225		if des, ok := cmd.(describer); ok {
   226			Errorf("%s\n", des.Describe())
   227		}
   228		Errorf("\n")
   229		cmd.Usage()
   230		if hasFlags(cmdFlags) {
   231			cmdFlags.PrintDefaults()
   232		}
   233		if ex, ok := cmd.(exampler); ok {
   234			Errorf("\nExamples:\n")
   235			for _, example := range ex.Examples() {
   236				Errorf("  %s %s %s\n", cmdName, mode, example)
   237			}
   238		}
   239	}
   240	
   241	// registerFlagOnce guards ExtraFlagRegistration. Tests may invoke
   242	// Main multiple times, but duplicate flag registration is fatal.
   243	var registerFlagOnce sync.Once
   244	
   245	var setCommandLineOutput func(io.Writer) // or nil if before Go 1.2
   246	
   247	// PrintLicenses prints all the licences registered by go4.org/legal for this program.
   248	func PrintLicenses() {
   249		for _, text := range legal.Licenses() {
   250			fmt.Fprintln(Stderr, text)
   251		}
   252	}
   253	
   254	// Main is meant to be the core of a command that has
   255	// subcommands (modes), such as pk-put or pk.
   256	func Main() {
   257		registerFlagOnce.Do(ExtraFlagRegistration)
   258		if setCommandLineOutput != nil {
   259			setCommandLineOutput(Stderr)
   260		}
   261		flag.Usage = func() {
   262			usage("")
   263		}
   264		flag.Parse()
   265		flag.CommandLine.SetOutput(Stderr)
   266		PostFlag()
   267	
   268		args := flag.Args()
   269		if *FlagVersion {
   270			fmt.Fprintf(Stderr, "%s version: %s\n", os.Args[0], buildinfo.Summary())
   271			return
   272		}
   273		if *FlagHelp {
   274			usage("")
   275		}
   276		if *FlagLegal {
   277			PrintLicenses()
   278			return
   279		}
   280		if len(args) == 0 {
   281			usage("No mode given.")
   282		}
   283	
   284		mode := args[0]
   285		cmd, ok := modeCommand[mode]
   286		if !ok {
   287			usage(fmt.Sprintf("Unknown mode %q", mode))
   288		}
   289	
   290		if _, ok := asNewCommand[mode]; ok {
   291			runAsNewCommand(cmd, mode)
   292			return
   293		}
   294	
   295		cmdFlags := modeFlags[mode]
   296		cmdFlags.SetOutput(Stderr)
   297		err := cmdFlags.Parse(args[1:])
   298		if err != nil {
   299			// We want -h to behave as -help, but without having to define another flag for
   300			// it, so we handle it here.
   301			// TODO(mpl): maybe even remove -help and just let them both be handled here?
   302			if err == flag.ErrHelp {
   303				help(mode)
   304				return
   305			}
   306			err = ErrUsage
   307		} else {
   308			if *wantHelp[mode] {
   309				help(mode)
   310				return
   311			}
   312			err = cmd.RunCommand(cmdFlags.Args())
   313		}
   314		if ue, isUsage := err.(UsageError); isUsage {
   315			if isUsage {
   316				Errorf("%s\n", ue)
   317			}
   318			cmd.Usage()
   319			Errorf("\nGlobal options:\n")
   320			flag.PrintDefaults()
   321	
   322			if hasFlags(cmdFlags) {
   323				Errorf("\nMode-specific options for mode %q:\n", mode)
   324				cmdFlags.PrintDefaults()
   325			}
   326			Exit(1)
   327		}
   328		PreExit()
   329		if err != nil {
   330			if !ExitWithFailure {
   331				// because it was already logged if ExitWithFailure
   332				Errorf("Error: %v\n", err)
   333			}
   334			Exit(2)
   335		}
   336	}
   337	
   338	// runAsNewCommand runs the executable specified by cmd's LookPath, which means
   339	// cmd must implement the ExecRunner interface. The executable must be a binary of
   340	// a program that runs Main.
   341	func runAsNewCommand(cmd CommandRunner, mode string) {
   342		execCmd, ok := cmd.(ExecRunner)
   343		if !ok {
   344			panic(fmt.Sprintf("%v does not implement ExecRunner", mode))
   345		}
   346		cmdPath, err := execCmd.LookPath()
   347		if err != nil {
   348			Errorf("Error: %v\n", err)
   349			Exit(2)
   350		}
   351		allArgs := shiftFlags(mode)
   352		if err := runExec(cmdPath, allArgs, newCopyEnv()); err != nil {
   353			panic(fmt.Sprintf("running %v should have ended with an os.Exit, and not leave us with that error: %v", cmdPath, err))
   354		}
   355	}
   356	
   357	// shiftFlags prepends all the arguments (global flags) passed before the given
   358	// mode to the list of arguments after that mode, and returns that list.
   359	func shiftFlags(mode string) []string {
   360		modePos := 0
   361		for k, v := range os.Args {
   362			if v == mode {
   363				modePos = k
   364				break
   365			}
   366		}
   367		globalFlags := os.Args[1:modePos]
   368		return append(globalFlags, os.Args[modePos+1:]...)
   369	}
   370	
   371	// Errorf prints to Stderr, regardless of FlagVerbose.
   372	func Errorf(format string, args ...interface{}) {
   373		fmt.Fprintf(Stderr, format, args...)
   374	}
   375	
   376	// Printf prints to Stderr if FlagVerbose, and is silent otherwise.
   377	func Printf(format string, args ...interface{}) {
   378		if *FlagVerbose {
   379			fmt.Fprintf(Stderr, format, args...)
   380		}
   381	}
   382	
   383	// Logf logs to Stderr if FlagVerbose, and is silent otherwise.
   384	func Logf(format string, v ...interface{}) {
   385		if !*FlagVerbose {
   386			return
   387		}
   388		logger.Printf(format, v...)
   389	}
   390	
   391	// sysExec is set to syscall.Exec on platforms that support it.
   392	var sysExec func(argv0 string, argv []string, envv []string) (err error)
   393	
   394	// runExec execs bin. If the platform doesn't support exec, it runs it and waits
   395	// for it to finish.
   396	func runExec(bin string, args []string, e *env) error {
   397		if sysExec != nil {
   398			sysExec(bin, append([]string{filepath.Base(bin)}, args...), e.flat())
   399		}
   400	
   401		cmd := exec.Command(bin, args...)
   402		cmd.Env = e.flat()
   403		cmd.Stdout = Stdout
   404		cmd.Stderr = Stderr
   405		return cmd.Run()
   406	}
   407	
   408	type env struct {
   409		m     map[string]string
   410		order []string
   411	}
   412	
   413	func (e *env) set(k, v string) {
   414		_, dup := e.m[k]
   415		e.m[k] = v
   416		if !dup {
   417			e.order = append(e.order, k)
   418		}
   419	}
   420	
   421	func (e *env) flat() []string {
   422		vv := make([]string, 0, len(e.order))
   423		for _, k := range e.order {
   424			if v, ok := e.m[k]; ok {
   425				vv = append(vv, k+"="+v)
   426			}
   427		}
   428		return vv
   429	}
   430	
   431	func newCopyEnv() *env {
   432		e := &env{make(map[string]string), nil}
   433		for _, kv := range os.Environ() {
   434			eq := strings.Index(kv, "=")
   435			if eq > 0 {
   436				e.set(kv[:eq], kv[eq+1:])
   437			}
   438		}
   439		return e
   440	}