// Package cli is adapted from https://github.com/rancher/wrangler-cli package cli import ( "fmt" "os" "reflect" "regexp" "strconv" "strings" "unsafe" "github.com/sirupsen/logrus" "github.com/spf13/cobra" "github.com/spf13/pflag" "go.linka.cloud/grpc-toolkit/signals" ) var ( caseRegexp = regexp.MustCompile("([a-z])([A-Z])") ) type PersistentPreRunnable interface { PersistentPre(cmd *cobra.Command, args []string) error } type PreRunnable interface { Pre(cmd *cobra.Command, args []string) error } type Runnable interface { Run(cmd *cobra.Command, args []string) error } type customizer interface { Customize(cmd *cobra.Command) } type fieldInfo struct { FieldType reflect.StructField FieldValue reflect.Value } func fields(obj interface{}) []fieldInfo { ptrValue := reflect.ValueOf(obj) objValue := ptrValue.Elem() var result []fieldInfo for i := 0; i < objValue.NumField(); i++ { fieldType := objValue.Type().Field(i) if !fieldType.IsExported() { continue } if fieldType.Anonymous && fieldType.Type.Kind() == reflect.Struct { result = append(result, fields(objValue.Field(i).Addr().Interface())...) } else if !fieldType.Anonymous { result = append(result, fieldInfo{ FieldValue: objValue.Field(i), FieldType: objValue.Type().Field(i), }) } } return result } func Name(obj interface{}) string { ptrValue := reflect.ValueOf(obj) objValue := ptrValue.Elem() commandName := strings.Replace(objValue.Type().Name(), "Command", "", 1) commandName, _ = name(commandName, "", "") return commandName } func Main(cmd *cobra.Command) { ctx := signals.SetupSignalHandler() if err := cmd.ExecuteContext(ctx); err != nil { logrus.Fatal(err) } } func makeEnvVar[T comparable](to []func(), name string, vars []string, defValue T, flags *pflag.FlagSet, fn func(flag string) (T, error)) []func() { for _, v := range vars { to = append(to, func() { v := os.Getenv(v) if v == "" { return } fv, err := fn(name) if err == nil && fv == defValue { flags.Set(name, v) } }) } return to } // Command populates a obj.Command() object by extracting args from struct tags of the // Runnable obj passed. Also the Run method is assigned to the RunE of the cli. // name = Override the struct field with func Command(obj Runnable, c *cobra.Command) *cobra.Command { var ( envs []func() arrays = map[string]reflect.Value{} slices = map[string]reflect.Value{} maps = map[string]reflect.Value{} ptrValue = reflect.ValueOf(obj) objValue = ptrValue.Elem() ) if c.Use == "" { c.Use = Name(obj) } for _, info := range fields(obj) { fieldType := info.FieldType v := info.FieldValue name, alias := name(fieldType.Name, fieldType.Tag.Get("name"), fieldType.Tag.Get("short")) usage := fieldType.Tag.Get("usage") envVars := strings.Split(fieldType.Tag.Get("env"), ",") defValue := fieldType.Tag.Get("default") if len(envVars) == 1 && envVars[0] == "" { envVars = nil } for _, v := range envVars { if v == "" { continue } usage += fmt.Sprintf(" [$%s]", v) } defInt, err := strconv.Atoi(defValue) if err != nil { defInt = 0 } defValueLower := strings.ToLower(defValue) defBool := defValueLower == "true" || defValueLower == "1" || defValueLower == "yes" || defValueLower == "y" flags := c.PersistentFlags() switch fieldType.Type.Kind() { case reflect.Int: flags.IntVarP((*int)(unsafe.Pointer(v.Addr().Pointer())), name, alias, defInt, usage) envs = append(envs, makeEnvVar(envs, name, envVars, defInt, flags, flags.GetInt)...) case reflect.String: flags.StringVarP((*string)(unsafe.Pointer(v.Addr().Pointer())), name, alias, defValue, usage) envs = append(envs, makeEnvVar(envs, name, envVars, defValue, flags, flags.GetString)...) case reflect.Slice: // env is not supported for slices switch fieldType.Tag.Get("split") { case "false": arrays[name] = v flags.StringArrayP(name, alias, nil, usage) default: slices[name] = v flags.StringSliceP(name, alias, nil, usage) } case reflect.Map: maps[name] = v flags.StringSliceP(name, alias, nil, usage) case reflect.Bool: flags.BoolVarP((*bool)(unsafe.Pointer(v.Addr().Pointer())), name, alias, defBool, usage) envs = append(envs, makeEnvVar(envs, name, envVars, defBool, flags, flags.GetBool)...) default: panic("Unknown kind on field " + fieldType.Name + " on " + objValue.Type().Name()) } if len(envVars) == 0 { continue } } if p, ok := obj.(PersistentPreRunnable); ok { c.PersistentPreRunE = p.PersistentPre } if p, ok := obj.(PreRunnable); ok { c.PreRunE = p.Pre } c.RunE = obj.Run c.PersistentPreRunE = bind(c.PersistentPreRunE, arrays, slices, maps, envs) c.PreRunE = bind(c.PreRunE, arrays, slices, maps, envs) c.RunE = bind(c.RunE, arrays, slices, maps, envs) cust, ok := obj.(customizer) if ok { cust.Customize(c) } return c } func assignMaps(app *cobra.Command, maps map[string]reflect.Value) error { for k, v := range maps { k = contextKey(k) s, err := app.Flags().GetStringSlice(k) if err != nil { return err } if s != nil { values := map[string]string{} for _, part := range s { parts := strings.SplitN(part, "=", 2) if len(parts) == 1 { values[parts[0]] = "" } else { values[parts[0]] = parts[1] } } v.Set(reflect.ValueOf(values)) } } return nil } func assignSlices(app *cobra.Command, slices map[string]reflect.Value) error { for k, v := range slices { k = contextKey(k) s, err := app.Flags().GetStringSlice(k) if err != nil { return err } if s != nil { v.Set(reflect.ValueOf(s[:])) } } return nil } func assignArrays(app *cobra.Command, arrays map[string]reflect.Value) error { for k, v := range arrays { k = contextKey(k) s, err := app.Flags().GetStringArray(k) if err != nil { return err } if s != nil { v.Set(reflect.ValueOf(s[:])) } } return nil } func contextKey(name string) string { parts := strings.Split(name, ",") return parts[len(parts)-1] } func name(name, setName, short string) (string, string) { if setName != "" { return setName, short } parts := strings.Split(name, "_") i := len(parts) - 1 name = caseRegexp.ReplaceAllString(parts[i], "$1-$2") name = strings.ToLower(name) result := append([]string{name}, parts[0:i]...) for i := 0; i < len(result); i++ { result[i] = strings.ToLower(result[i]) } if short == "" && len(result) > 1 { short = result[1] } return result[0], short } func bind(next func(*cobra.Command, []string) error, arrays map[string]reflect.Value, slices map[string]reflect.Value, maps map[string]reflect.Value, envs []func()) func(*cobra.Command, []string) error { if next == nil { return nil } return func(cmd *cobra.Command, args []string) error { for _, envCallback := range envs { envCallback() } if err := assignArrays(cmd, arrays); err != nil { return err } if err := assignSlices(cmd, slices); err != nil { return err } if err := assignMaps(cmd, maps); err != nil { return err } if next != nil { return next(cmd, args) } return nil } }