package main import ( "os" "strings" "github.com/miekg/dns" "github.com/sirupsen/logrus" ) func (hub *Hub) loadQueryArgs() error { // set log level if k.Bool("debug") { // Set logger level hub.Logger.SetLevel(logrus.DebugLevel) } else { hub.Logger.SetLevel(logrus.InfoLevel) } err := hub.loadNamedArgs() err = hub.loadFreeArgs() if err != nil { hub.Logger.WithError(err).Error("Error parsing arguments") hub.Logger.Exit(2) } err = hub.initResolver() if err != nil { hub.Logger.WithError(err).Error("Error parsing nameservers") hub.Logger.Exit(2) } hub.loadFallbacks() return err } // loadFreeArgs tries to parse all the arguments // given to the CLI. These arguments don't have any specific // order so we have to deduce based on the pattern of argument. // For eg, a nameserver must always begin with `@`. In this // pattern we deduce the arguments and map it to internal query // options. In case an argument isn't able to fit in any of the existing // pattern it is considered to be a "query name". func (hub *Hub) loadFreeArgs() error { args := os.Args[1:] for _, arg := range args { if strings.HasPrefix(arg, "--") || strings.HasPrefix(arg, "-") { continue } if strings.HasPrefix(arg, "@") { hub.QueryFlags.Nameservers = append(hub.QueryFlags.Nameservers, strings.Trim(arg, "@")) } else if _, ok := dns.StringToType[strings.ToUpper(arg)]; ok { hub.QueryFlags.QTypes = append(hub.QueryFlags.QTypes, arg) } else if _, ok := dns.StringToClass[strings.ToUpper(arg)]; ok { hub.QueryFlags.QClasses = append(hub.QueryFlags.QClasses, arg) } else { // if nothing matches, consider it's a query name. hub.QueryFlags.QNames = append(hub.QueryFlags.QNames, arg) } } return nil } // loadNamedArgs checks for all flags and loads their // values inside the Hub. func (hub *Hub) loadNamedArgs() error { // Unmarshall flags to the struct. err := k.Unmarshal("", &hub.QueryFlags) if err != nil { return err } return nil } // loadFallbacks sets fallbacks for options // that are not specified by the user. func (hub *Hub) loadFallbacks() { if len(hub.QueryFlags.QTypes) == 0 { hub.QueryFlags.QTypes = append(hub.QueryFlags.QTypes, "A") } if len(hub.QueryFlags.QClasses) == 0 { hub.QueryFlags.QClasses = append(hub.QueryFlags.QClasses, "IN") } }