74 lines
2.2 KiB
Go
74 lines
2.2 KiB
Go
package main
|
|
|
|
import (
|
|
"strings"
|
|
|
|
"github.com/miekg/dns"
|
|
flag "github.com/spf13/pflag"
|
|
)
|
|
|
|
func (hub *Hub) loadQueryArgs() error {
|
|
// Appends a list of unparsed args to
|
|
// internal query flags.
|
|
err := hub.loadUnparsedArgs()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// check if ndots is set
|
|
hub.QueryFlags.isNdotsSet = isFlagPassed("ndots", hub.flag)
|
|
|
|
// Load all fallbacks in internal query flags.
|
|
hub.loadFallbacks()
|
|
return nil
|
|
}
|
|
|
|
// loadUnparsedArgs tries to parse all the arguments
|
|
// which are unparsed by `flag` library. 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 append it to the
|
|
// list of internal query flags.
|
|
// In case an argument isn't able to fit in any of the existing
|
|
// pattern it is considered to be a "hostname".
|
|
// Eg of unparsed argument: `dig mrkaran.dev @1.1.1.1 AAAA`
|
|
// where `@1.1.1.1` and `AAAA` are "unparsed" args.
|
|
func (hub *Hub) loadUnparsedArgs() error {
|
|
for _, arg := range hub.UnparsedArgs {
|
|
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
|
|
}
|
|
|
|
// loadFallbacks sets fallbacks for options
|
|
// that are not specified by the user but necessary
|
|
// for the resolver.
|
|
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")
|
|
}
|
|
}
|
|
|
|
// isFlagPassed checks if the flag is supplied by
|
|
//user or not.
|
|
func isFlagPassed(name string, f *flag.FlagSet) bool {
|
|
found := false
|
|
f.Visit(func(f *flag.Flag) {
|
|
if f.Name == name {
|
|
found = true
|
|
}
|
|
})
|
|
return found
|
|
}
|