Pass domains as arguments similar to dig

src/cli/cli.go

@@ -17,7 +17,7 @@ import (
 	"fmt"
 	"net"
 	"os"
-	"strings"
+	"sort"
 	"sync"
 
 	"github.com/spf13/cobra"
@@ -29,9 +29,18 @@ import (
 )
 
 const (
-	zdnsCLIVersion = "1.1.0"
+	zdnsCLIVersion  = "1.1.0"
+	ModulesColWidth = 14
+	ModulesPerRow   = 6
 )
 
+// cmds is a set of commands, special "modules" with their own flags. They should not be printed as "modules" in --help or used with --module
+// They must be used as their own command, e.g. zdns MXLOOKUP
+var cmds = map[string]struct{}{
+	"MXLOOKUP": {},
+	"NSLOOKUP": {},
+}
+
 type InputHandler interface {
 	FeedChannel(in chan<- string, wg *sync.WaitGroup) error
 }
@@ -59,6 +68,7 @@ type CLIConf struct {
 	ResultVerbosity string
 	IncludeInOutput string
 	OutputGroups    []string
+	ModuleString    string
 
 	MaxDepth             int
 	CacheSize            int
@@ -105,23 +115,21 @@ var GC CLIConf
 var rootCmd = &cobra.Command{
 	Use:   "zdns",
 	Short: "High-speed, low-drag DNS lookups",
-	Long: `ZDNS is a library and CLI tool for making very fast DNS requests. It's built upon
-https://github.com/zmap/dns (and in turn https://github.com/miekg/dns) for constructing
-and parsing raw DNS packets.
-
-ZDNS also includes its own recursive resolution and a cache to further optimize performance.`,
-	ValidArgs: GetValidLookups(),
-	Args:      cobra.MatchAll(cobra.ExactArgs(1), cobra.OnlyValidArgs),
+	// Cannot register Long description here, we'll wait for all modules to register themselves in their init() functions
+	// and then generate the long description during Execute.
+	Args: cobra.MatchAll(),
 	Run: func(cmd *cobra.Command, args []string) {
-		GC.Module = strings.ToUpper(args[0])
-		Run(GC, cmd.Flags())
+		Run(GC, cmd.Flags(), args)
 	},
 	Version: zdnsCLIVersion,
 }
 
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
+	// This must be run after all modules are registered in the various init() functions
+	rootCmd.Long = getRootCmdLongText()
+
 	err := rootCmd.Execute()
 	if err != nil {
 		os.Exit(1)
@@ -182,6 +190,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVar(&GC.ClientSubnetString, "client-subnet", "", "Client subnet in CIDR format for EDNS0.")
 	rootCmd.PersistentFlags().BoolVar(&GC.Dnssec, "dnssec", false, "Requests DNSSEC records by setting the DNSSEC OK (DO) bit")
 	rootCmd.PersistentFlags().BoolVar(&GC.UseNSID, "nsid", false, "Request NSID.")
+	rootCmd.PersistentFlags().StringVar(&GC.ModuleString, "module", "", "DNS module to use, can be any DNS query type or one of the modules that provides extra processing (see zdns --help for a complete list). This is required if passing domains as arguments: zdns example.com google.com --module='A'.")
 
 	rootCmd.PersistentFlags().Bool("ipv4-lookup", false, "Perform an IPv4 Lookup (requests A records) in modules")
 	rootCmd.PersistentFlags().Bool("ipv6-lookup", false, "Perform an IPv6 Lookup (requests AAAA recoreds) in modules")
@@ -214,3 +223,43 @@ func initConfig() {
 	// Bind the current command's flags to viper
 	util.BindFlags(rootCmd, viper.GetViper(), util.EnvPrefix)
 }
+
+// getRootCmdLongText dynamically generates the list of available modules for the long --help text
+func getRootCmdLongText() string {
+	rootCmdLongText := `ZDNS is a library and CLI tool for making very fast DNS requests. It's built upon
+https://github.com/zmap/dns (and in turn https://github.com/miekg/dns) for constructing
+and parsing raw DNS packets.
+
+ZDNS also includes its own recursive resolution and a cache to further optimize performance.
+
+ZDNS can take input (usually domains and a module name) in the following ways:
+- file (./zdns A --input-file=domains.txt)
+- stream (echo "example.com" | ./zdns A)
+- as arguments (./zdns --module=A example.com google.com).`
+
+	modules := make([]string, 0, len(moduleToLookupModule))
+	for module := range moduleToLookupModule {
+		if _, ok := cmds[module]; ok {
+			// these are their own command, do not print them as a module
+			continue
+		}
+		modules = append(modules, module)
+	}
+	// sort modules alphabetically
+	sort.Strings(modules)
+	rootCmdLongText += "\n\nAvailable modules:\n"
+	// print in grid format for readability
+	for i, module := range modules {
+		rootCmdLongText += fmt.Sprintf("%-*s", ModulesColWidth, module)
+		if (i+1)%ModulesPerRow == 0 {
+			rootCmdLongText += "\n"
+		}
+	}
+
+	// If the number of modules isn't a multiple of ModulesPerRow, ensure a final newline
+	if len(modules)%ModulesPerRow != 0 {
+		fmt.Println()
+		rootCmdLongText += "\n"
+	}
+	return rootCmdLongText
+}

src/cli/iohandlers/string_slice_input_handler.go

@@ -0,0 +1,42 @@
+/*
+ * ZDNS Copyright 2024 Regents of the University of Michigan
+ *
+ * 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 iohandlers
+
+import (
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// StringSliceInputHandler Feeds a channel with the strings in the slice.
+type StringSliceInputHandler struct {
+	Names []string
+}
+
+func NewStringSliceInputHandler(domains []string) *StringSliceInputHandler {
+	if len(domains) == 0 {
+		log.Fatal("No domains provided, cannot create a string slice input handler")
+	}
+	return &StringSliceInputHandler{Names: domains}
+}
+
+func (h *StringSliceInputHandler) FeedChannel(in chan<- string, wg *sync.WaitGroup) error {
+	defer close(in)
+	defer wg.Done()
+	for _, name := range h.Names {
+		in <- name
+	}
+	return nil
+}

src/cli/modules.go

@@ -154,10 +154,10 @@ func GetLookupModule(name string) (LookupModule, error) {
 	return module, nil
 }
 
-func GetValidLookups() []string {
-	lookups := make([]string, 0, len(moduleToLookupModule))
+func GetValidLookups() map[string]struct{} {
+	lookups := make(map[string]struct{}, len(moduleToLookupModule))
 	for lookup := range moduleToLookupModule {
-		lookups = append(lookups, lookup)
+		lookups[lookup] = struct{}{}
 	}
 	return lookups
 }

src/cli/mxlookup.go

@@ -31,7 +31,7 @@ var mxlookupCmd = &cobra.Command{
 correspond with an exchange record.`,
 	Run: func(cmd *cobra.Command, args []string) {
 		GC.Module = strings.ToUpper("mxlookup")
-		Run(GC, cmd.Flags())
+		Run(GC, cmd.Flags(), args)
 	},
 }
 

src/cli/nslookup.go

@@ -29,7 +29,7 @@ var nslookupCmd = &cobra.Command{
 	Long:  `nslookup will additionally do an A/AAAA lookup for the IP addresses that correspond with name server records.`,
 	Run: func(cmd *cobra.Command, args []string) {
 		GC.Module = strings.ToUpper("nslookup")
-		Run(GC, cmd.Flags())
+		Run(GC, cmd.Flags(), args)
 	},
 }
 

src/cli/worker_manager.go

@@ -61,7 +61,8 @@ type Metadata struct {
 	ZDNSVersion string         `json:"zdns_version"`
 }
 
-func populateCLIConfig(gc *CLIConf, flags *pflag.FlagSet) *CLIConf {
+// populateCLIConfig populates the CLIConf struct with the values from the command line arguments or the defaults.
+func populateCLIConfig(gc *CLIConf, domains []string) *CLIConf {
 	if gc.LogFilePath != "" && gc.LogFilePath != "-" {
 		f, err := os.OpenFile(gc.LogFilePath, os.O_WRONLY|os.O_CREATE, util.DefaultFilePermissions)
 		if err != nil {
@@ -153,7 +154,10 @@ func populateCLIConfig(gc *CLIConf, flags *pflag.FlagSet) *CLIConf {
 	gc.OutputGroups = append(gc.OutputGroups, groups...)
 
 	// setup i/o if not specified
-	if gc.InputHandler == nil {
+	if len(domains) > 0 {
+		// using domains from command line
+		gc.InputHandler = iohandlers.NewStringSliceInputHandler(domains)
+	} else if gc.InputHandler == nil {
 		gc.InputHandler = iohandlers.NewFileInputHandler(gc.InputFilePath)
 	}
 	if gc.OutputHandler == nil {
@@ -164,9 +168,7 @@ func populateCLIConfig(gc *CLIConf, flags *pflag.FlagSet) *CLIConf {
 
 func populateResolverConfig(gc *CLIConf) *zdns.ResolverConfig {
 	config := zdns.NewResolverConfig()
-
 	config.TransportMode = zdns.GetTransportMode(gc.UDPOnly, gc.TCPOnly)
-
 	config.Timeout = time.Second * time.Duration(gc.Timeout)
 	config.IterativeTimeout = time.Second * time.Duration(gc.IterationTimeout)
 	config.LookupAllNameServers = gc.LookupAllNameServers
@@ -321,7 +323,6 @@ func populateNameServers(gc *CLIConf, config *zdns.ResolverConfig) (*zdns.Resolv
 
 	// Additionally, both Root and External nameservers must be populated, since the Resolver doesn't know we'll only
 	// be performing either iterative or recursive lookups, not both.
-
 	// IPv4 Name Servers/Local Address only needs to be populated if we're doing IPv4 lookups, same for IPv6
 	if len(gc.NameServers) != 0 {
 		// User provided name servers, use them.
@@ -452,18 +453,39 @@ func populateLocalAddresses(gc *CLIConf, config *zdns.ResolverConfig) (*zdns.Res
 	return config, nil
 }
 
-func Run(gc CLIConf, flags *pflag.FlagSet) {
-	gc = *populateCLIConfig(&gc, flags)
+func Run(gc CLIConf, flags *pflag.FlagSet, args []string) {
+	// if this is one of the special commands registered with Cobra, it'll have it's module already set
+	var module string
+	var domains []string
+	var err error
+	if _, ok := cmds[strings.ToUpper(gc.Module)]; ok {
+		domains, err = parseArgsWithCobraCmd(args, gc.ModuleString, gc.Module)
+		module = strings.ToUpper(gc.Module)
+	} else if len(gc.ModuleString) != 0 {
+		module, domains, err = parseArgsWithModuleFlag(args, gc.ModuleString)
+	} else {
+		module, domains, err = parseArgsWithoutModuleFlag(args)
+	}
+	if err != nil {
+		log.Fatal("could not parse arguments: ", err)
+	}
+	if len(module) == 0 {
+		// no module specified, we cannot continue
+		log.Fatal("No valid DNS lookup module specified. Please provide a module to run.")
+	}
+	gc.Module = module
+
+	gc = *populateCLIConfig(&gc, domains)
 	resolverConfig := populateResolverConfig(&gc)
 	// Log any information about the resolver configuration, according to log level
 	resolverConfig.PrintInfo()
-	err := resolverConfig.Validate()
+	err = resolverConfig.Validate()
 	if err != nil {
 		log.Fatalf("resolver config did not pass validation: %v", err)
 	}
 	lookupModule, err := GetLookupModule(gc.Module)
 	if err != nil {
-		log.Fatal("could not get lookup module: ", err)
+		log.Fatalf("could not get lookup module %s: %v", gc.Module, err)
 	}
 	err = lookupModule.CLIInit(&gc, resolverConfig, flags)
 	if err != nil {
@@ -696,3 +718,46 @@ func aggregateMetadata(c <-chan routineMetadata) Metadata {
 	}
 	return meta
 }
+
+func parseArgsWithModuleFlag(args []string, moduleString string) (module string, domains []string, err error) {
+	// --module flag is set, so we must have a module name as the first arg
+	module = strings.ToUpper(moduleString)
+	if _, ok := cmds[module]; ok {
+		return "", nil, fmt.Errorf("the module specified (--module=%s) has its own arguements and must be called with 'zdns %s'. See 'zdns %s --help' for more", module, module, module)
+	}
+	if _, ok := GetValidLookups()[module]; !ok {
+		return "", nil, fmt.Errorf("invalid lookup module specified - %s. ex: zdns A or zdns --module=A. See 'zdns --help' for applicable modules", moduleString)
+	}
+	// treating all args as domains
+	return module, args, nil
+}
+
+// parseArgsWithCobraCmd parses the arguments for a ZDNS command like mxlookup or nslookup, which are registered as Cobra
+// commands and have their own special flags.
+// Other `--modules` are not compatible with these.
+func parseArgsWithCobraCmd(args []string, moduleString, cmd string) (domains []string, err error) {
+	if len(moduleString) != 0 {
+		return nil, fmt.Errorf("%s command is not compatabile with --modules and should be invoked alone", cmd)
+	}
+	if len(args) == 0 {
+		// no args, we'll read from stdin
+		return nil, nil
+	}
+	// treat args as domains
+	return args, nil
+}
+
+func parseArgsWithoutModuleFlag(args []string) (module string, domains []string, err error) {
+	// no --module flag, so we must have a module name as the first arg
+	if len(args) == 0 {
+		return "", nil, errors.New("no module specified. Valid usages are 1) zdns <module> (where domains come from std. in) or 2) zdns --module=<module> <domain1> <domain2> ... See zdns --help for applicable modules")
+	}
+	module = strings.ToUpper(args[0])
+	if _, ok := GetValidLookups()[module]; !ok {
+		return "", nil, fmt.Errorf("invalid lookup module specified - %s. ex: zdns A or zdns --module=A. See 'zdns --help' for applicable modules", args[0])
+	}
+	if len(args) > 1 {
+		return "", nil, errors.New("invalid args, cannot mix domains with module in args. Valid usages are 1) zdns <module> (where domains come from std. in) or 2) zdns --module=<module> <domain1> <domain2> ...")
+	}
+	return module, nil, nil
+}

src/zdns/conf.go

@@ -14,17 +14,6 @@
 
 package zdns
 
-//type Metadata struct {
-//	Names       int            `json:"names"`
-//	Status      map[string]int `json:"statuses"`
-//	StartTime   string         `json:"start_time"`
-//	EndTime     string         `json:"end_time"`
-//	NameServers []string       `json:"name_servers"`
-//	Timeout     int            `json:"timeout"`
-//	Retries     int            `json:"retries"`
-//	Conf        *GlobalConf    `json:"conf"`
-//}
-
 type TargetedDomain struct {
 	Domain      string   `json:"domain"`
 	Nameservers []string `json:"nameservers"`

src/zdns/lookup.go

@@ -11,6 +11,7 @@
  * implied. See the License for the specific language governing
  * permissions and limitations under the License.
  */
+
 package zdns
 
 import (

testing/integration_tests.py

@@ -629,6 +629,13 @@ def test_a(self):
         self.assertSuccess(res, cmd)
         self.assertEqualAnswers(res, self.ROOT_A_ANSWERS, cmd)
 
+    def test_a_dig_style_args(self):
+        c = "--module=A zdns-testing.com"
+        name = ""
+        cmd, res = self.run_zdns(c, name)
+        self.assertSuccess(res, cmd)
+        self.assertEqualAnswers(res, self.ROOT_A_ANSWERS, cmd)
+
     def test_cname(self):
         c = "CNAME"
         name = "www.zdns-testing.com"
@@ -934,6 +941,20 @@ def test_server_mode_a_lookup_ipv4(self):
         self.assertEqual(res["data"]["resolver"], "8.8.8.8:53")
         self.assertEqualAnswers(res, self.ROOT_A_ANSWERS, cmd)
 
+    def test_type_option_server_mode_a_lookup_ipv4(self):
+        c = "--module=A --override-name=www.zdns-testing.com --name-server-mode"
+        name = "8.8.8.8"
+        cmd, res = self.run_zdns(c, name)
+        self.assertEqual(res["data"]["resolver"], "8.8.8.8:53")
+        self.assertEqualAnswers(res, self.WWW_CNAME_AND_A_ANSWERS, cmd)
+
+    def test_dig_style_type_option_server_mode_a_lookup_ipv4(self):
+        c = "--module=A 8.8.8.8 --override-name=www.zdns-testing.com --name-server-mode"
+        name = ""
+        cmd, res = self.run_zdns(c, name)
+        self.assertEqual(res["data"]["resolver"], "8.8.8.8:53")
+        self.assertEqualAnswers(res, self.WWW_CNAME_AND_A_ANSWERS, cmd)
+
     def test_mixed_mode_a_lookup_ipv4(self):
         c = "A --name-servers=0.0.0.0"
         name = "zdns-testing.com,8.8.8.8:53"