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,28 @@ import (
 )
 
 const (
-	zdnsCLIVersion = "1.1.0"
+	zdnsCLIVersion  = "1.1.0"
+	ModulesColWidth = 14
+	ModulesPerRow   = 6
 )
 
+// Unfortunately, we can't use the moduleToLookupModule map here, as it's not available at runtime yet.
+// Variables get initialized and then init() gets run (where modules are registered), so we'll have to hardcode the list.
+var allModules = []string{
+	"A", "AAAA", "AFSDB", "ANY", "ATMA", "AVC", "AXFR", "BINDVERSION", "CAA", "CDNSKEY", "CDS", "CERT",
+	"CNAME", "CSYNC", "DHCID", "DMARC", "DNAME", "DNSKEY", "DS", "EID", "EUI48", "EUI64", "GID", "GPOS",
+	"HINFO", "HIP", "HTTPS", "ISDN", "KEY", "KX", "L32", "L64", "LOC", "LP", "MB", "MD", "MF", "MG",
+	"MR", "MX", "MXLOOKUP", "NAPTR", "NID", "NIMLOC", "NINFO", "NS", "NSAPPTR", "NSEC", "NSEC3", "NSEC3PARAM",
+	"NSLOOKUP", "NULL", "NXT", "OPENPGPKEY", "PTR", "PX", "RP", "RRSIG", "RT", "SMIMEA", "SOA", "SPF", "SRV",
+	"SSHFP", "SVCB", "TALINK", "TKEY", "TLSA", "TXT", "UID", "UINFO", "UNSPEC", "URI",
+}
+
+// cmds is a set of commands, special "modules" with their own flags. They should not be printed as "modules" or used with --module
+var cmds = map[string]struct{}{
+	"MXLOOKUP": {},
+	"NSLOOKUP": {},
+}
+
 type InputHandler interface {
 	FeedChannel(in chan<- string, wg *sync.WaitGroup) error
 }
@@ -59,6 +78,7 @@ type CLIConf struct {
 	ResultVerbosity string
 	IncludeInOutput string
 	OutputGroups    []string
+	ModuleString    string
 
 	MaxDepth             int
 	CacheSize            int
@@ -105,16 +125,10 @@ 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),
+	Long:  getRootCmdLongText(),
+	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,
 }
@@ -182,6 +196,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 +229,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(allModules))
+	for _, module := range allModules {
+		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,31 @@ 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) {
+	// User can provide both a module and a list of domains to query as inputs, similar to dig
+	module, domains, err := parseArgs(args, gc.ModuleString)
+	if err != nil {
+		log.Fatal("could not parse arguments: ", err)
+	}
+	if len(gc.Module) == 0 && len(module) == 0 {
+		// no module specified by either user or command, we cannot continue
+		log.Fatal("No valid DNS lookup module specified. Please provide a module to run.")
+	} else if len(gc.Module) == 0 {
+		// Some commands set gc.Module, but most don't. If it's not set, set with module from parsing args
+		gc.Module = strings.ToUpper(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 +710,47 @@ func aggregateMetadata(c <-chan routineMetadata) Metadata {
 	}
 	return meta
 }
+
+// parseArgs parses and validates the command line arguments to ZDNS
+// Valid usages of ZDNS are:
+// Single arg as module/query type, input is taken from std. in: zdns <module>
+// 1+ args as domains, module/query type must be passed in with --module: zdns --module=<module> <domain1> <domain2> ...
+func parseArgs(args []string, moduleString string) (module string, domains []string, err error) {
+	if len(args) == 0 && len(moduleString) == 0 {
+		// some commands (nslookup) don't require a module, let the caller error check
+		return "", nil, nil
+	}
+	if len(args) > 1 {
+		// pre-alloc the domains slice, we know it will be at most the length of args - 1 for the mandatory module name
+		domains = make([]string, 0, len(args)-1)
+	}
+
+	// --module takes precedence
+	validLookupModulesMap := GetValidLookups()
+	if len(moduleString) != 0 {
+		module = strings.ToUpper(moduleString)
+		_, ok := validLookupModulesMap[module]
+		if !ok {
+			return "", nil, fmt.Errorf("invalid lookup module specified - %s. ex: zdns A or zdns --module=A. See 'zdns --help' for applicable modules", moduleString)
+		}
+		// check if --module is one of the special commands which should be called directly.
+		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)
+		}
+		// alright, found the module, all args are domains
+		domains = append(domains, args...)
+		return module, domains, nil
+	}
+
+	// no --module, so we must have a module name as the first arg
+	if len(args) > 1 {
+		return "", nil, errors.New("invalid args. Valid usages are 1) zdns <module> (where domains come from std. in) or 2) zdns --module=<module> <domain1> <domain2> ...")
+	}
+
+	// only one arg, must be a module name
+	module = strings.ToUpper(args[0])
+	if _, ok := validLookupModulesMap[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])
+	}
+	return module, nil, nil
+}

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"