doggo/cmd/parse.go

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")
	}
}