refactor: improve structure

cmd/root.go

@@ -6,24 +6,16 @@ import (
 	"os"
 
 	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 
-	"github.com/rs/zerolog/log"
+	"github.com/jpts/coredns-enum/internal/types"
+	"github.com/jpts/coredns-enum/internal/util"
+	"github.com/jpts/coredns-enum/pkg/dnsclient"
+	"github.com/jpts/coredns-enum/pkg/scanners"
 )
 
-type cliOpts struct {
-	loglevel   int
-	maxWorkers int
-	cidrRange  string
-	nameserver string
-	nameport   int
-	timeout    float32
-	mode       string
-	zone       string
-	proto      string
-}
-
-var opts cliOpts
+var opts types.CliOpts
 
 var rootCmd = &cobra.Command{
 	Use:   "coredns-enum",
@@ -34,46 +26,46 @@ var rootCmd = &cobra.Command{
 		var err error
 
 		log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
-		lvl, err := zerolog.ParseLevel(fmt.Sprint(opts.loglevel))
+		lvl, err := zerolog.ParseLevel(fmt.Sprint(opts.LogLevel))
 		if err != nil {
 			return errors.New("Error setting up logging")
 		}
 		zerolog.SetGlobalLevel(lvl)
 
-		if opts.proto != "udp" && opts.proto != "tcp" && opts.proto != "auto" {
+		if opts.Proto != "udp" && opts.Proto != "tcp" && opts.Proto != "auto" {
 			log.Error().Msg("Invalid protocol")
 		}
 
-		initDNS()
+		dclient := dnsclient.InitDNS(&opts)
 
-		if opts.nameserver == "" {
-			opts.nameserver, opts.nameport, err = getNSFromSystem()
+		if opts.Nameserver == "" {
+			opts.Nameserver, opts.Nameport, err = dclient.GetNSFromSystem()
 			if err != nil {
 				return err
 			}
-			log.Info().Msgf("Detected nameserver as %s:%d", opts.nameserver, opts.nameport)
+			log.Info().Msgf("Detected nameserver as %s:%d", opts.Nameserver, opts.Nameport)
 		}
 
-		if opts.mode == MODE_AUTO {
-			opts.mode = detectMode()
+		if opts.Mode == dnsclient.MODE_AUTO {
+			opts.Mode = dclient.DetectMode()
 		}
 
-		var res []*svcResult
-		switch opts.mode {
-		case MODE_BRUTEFORCE:
-			res, err = brute(&opts)
-		case MODE_WILDCARD:
-			res, err = wildcard(&opts)
-		case MODE_FAILED:
+		var res []*types.SvcResult
+		switch opts.Mode {
+		case dnsclient.MODE_BRUTEFORCE:
+			res, err = scanners.BruteScan(&opts, dclient)
+		case dnsclient.MODE_WILDCARD:
+			res, err = scanners.WildcardScan(&opts, dclient)
+		case dnsclient.MODE_FAILED:
 			err = fmt.Errorf("Failed to detect mode automatically")
 		default:
-			err = fmt.Errorf("Unsupported mode: %s", opts.mode)
+			err = fmt.Errorf("Unsupported mode: %s", opts.Mode)
 		}
 		if err != nil {
 			return err
 		}
 
-		renderResults(res)
+		util.RenderResults(res)
 
 		return nil
 	},
@@ -90,17 +82,17 @@ func Execute() {
 
 func init() {
 	// global flags
-	rootCmd.PersistentFlags().IntVarP(&opts.loglevel, "loglevel", "v", 1, "Set loglevel (-1 => 5)")
-	rootCmd.PersistentFlags().StringVarP(&opts.mode, "mode", "m", "auto", "Select mode: wildcard|bruteforce|auto")
-	rootCmd.PersistentFlags().StringVar(&opts.zone, "zone", "cluster.local", "DNS zone")
+	rootCmd.PersistentFlags().IntVarP(&opts.LogLevel, "loglevel", "v", 1, "Set loglevel (-1 => 5)")
+	rootCmd.PersistentFlags().StringVarP(&opts.Mode, "mode", "m", "auto", "Select mode: wildcard|bruteforce|auto")
+	rootCmd.PersistentFlags().StringVar(&opts.Zone, "zone", "cluster.local", "DNS zone")
 
 	// bruteforce
-	rootCmd.Flags().IntVarP(&opts.maxWorkers, "max-workers", "t", 50, "Number of 'workers' to use for concurrency")
-	rootCmd.Flags().StringVar(&opts.cidrRange, "cidr", "", "Range to scan in bruteforce mode")
+	rootCmd.Flags().IntVarP(&opts.MaxWorkers, "max-workers", "t", 50, "Number of 'workers' to use for concurrency")
+	rootCmd.Flags().StringVar(&opts.CidrRange, "cidr", "", "Range to scan in bruteforce mode")
 
 	// nameserver
-	rootCmd.Flags().StringVarP(&opts.nameserver, "nsip", "n", "", "Nameserver to use (detected by default)")
-	rootCmd.Flags().IntVar(&opts.nameport, "nsport", 53, "Nameserver port to use (detected by default)")
-	rootCmd.Flags().Float32Var(&opts.timeout, "timeout", 0.5, "DNS query timeout (seconds)")
-	rootCmd.Flags().StringVar(&opts.proto, "protocol", "auto", "DNS protocol: udp|tcp|auto")
+	rootCmd.Flags().StringVarP(&opts.Nameserver, "nsip", "n", "", "Nameserver to use (detected by default)")
+	rootCmd.Flags().IntVar(&opts.Nameport, "nsport", 53, "Nameserver port to use (detected by default)")
+	rootCmd.Flags().Float32Var(&opts.Timeout, "timeout", 0.5, "DNS query timeout (seconds)")
+	rootCmd.Flags().StringVar(&opts.Proto, "protocol", "auto", "DNS protocol: udp|tcp|auto")
 }

internal/types/cli.go

@@ -0,0 +1,13 @@
+package types
+
+type CliOpts struct {
+	LogLevel   int
+	MaxWorkers int
+	CidrRange  string
+	Nameserver string
+	Nameport   int
+	Timeout    float32
+	Mode       string
+	Zone       string
+	Proto      string
+}

internal/types/results.go

@@ -0,0 +1,42 @@
+package types
+
+import (
+	"fmt"
+	"net"
+	"time"
+
+	"github.com/miekg/dns"
+)
+
+type QueryResult struct {
+	Answers    []dns.RR
+	Additional []dns.RR
+	Raw        *dns.Msg
+	IP         *net.IP
+	RTT        *time.Duration
+}
+
+type SvcResult struct {
+	Name      string
+	Namespace string
+	IP        *net.IP
+	Ports     []*PortResult
+	Endpoints []*PodResult
+}
+
+func (s *SvcResult) String() string { return fmt.Sprintf("%s/%s", s.Namespace, s.Name) }
+
+type PodResult struct {
+	Name      string
+	Namespace string
+	IP        *net.IP
+	Ports     []*PortResult
+}
+
+type PortResult struct {
+	Proto    string
+	PortNo   int
+	PortName string
+}
+
+func (p *PortResult) String() string { return fmt.Sprintf("%d/%s", p.PortNo, p.Proto) }

internal/util/array.go

@@ -0,0 +1,18 @@
+package util
+
+func Reverse(numbers []string) []string {
+	newNumbers := make([]string, len(numbers))
+	for i, j := 0, len(numbers)-1; i <= j; i, j = i+1, j-1 {
+		newNumbers[i], newNumbers[j] = numbers[j], numbers[i]
+	}
+	return newNumbers
+}
+
+func IsElement(s []string, str string) bool {
+	for _, v := range s {
+		if v == str {
+			return true
+		}
+	}
+	return false
+}

internal/util/cidr.go

@@ -0,0 +1,33 @@
+package util
+
+import (
+	"fmt"
+
+	"github.com/seancfoley/ipaddress-go/ipaddr"
+	"github.com/seancfoley/ipaddress-go/ipaddr/addrstrparam"
+)
+
+func ParseIPv4CIDR(cidr string) (*ipaddr.IPAddress, error) {
+	pb := addrstrparam.IPAddressStringParamsBuilder{}
+	pb.AllowWildcardedSeparator(true)
+	pb.AllowIPv4(true)
+	pb.AllowIPv6(false)
+	pb.AllowMask(true)
+	pb.AllowPrefix(true)
+	pb.AllowEmpty(false)
+	pb.AllowSingleSegment(false)
+	params := pb.ToParams()
+
+	ipastr := ipaddr.NewIPAddressStringParams(cidr, params)
+
+	if !ipastr.IsPrefixed() {
+		return nil, fmt.Errorf("CIDR %s requires prefix, use /32 for a single host", cidr)
+	}
+
+	subnet, err := ipastr.ToAddress()
+	if err != nil {
+		return nil, err
+	}
+
+	return subnet.ToPrefixBlock(), nil
+}

cmd/printer.go → internal/util/printer.go

@@ -1,63 +1,15 @@
-package cmd
+package util
 
 import (
 	"fmt"
-	"net"
 	"os"
 	"sort"
-	"time"
 
-	"github.com/miekg/dns"
+	"github.com/jpts/coredns-enum/internal/types"
 	"github.com/olekukonko/tablewriter"
 )
 
-type queryResult struct {
-	answers    []dns.RR
-	additional []dns.RR
-	raw        *dns.Msg
-	ip         *net.IP
-	rtt        *time.Duration
-}
-
-type svcResult struct {
-	Name      string
-	Namespace string
-	IP        *net.IP
-	Ports     []*portResult
-	Endpoints []*podResult
-}
-
-func (s *svcResult) String() string { return fmt.Sprintf("%s/%s", s.Namespace, s.Name) }
-
-type podResult struct {
-	Name      string
-	Namespace string
-	IP        *net.IP
-	Ports     []*portResult
-}
-
-type portResult struct {
-	Proto    string
-	PortNo   int
-	PortName string
-}
-
-func (p *portResult) String() string { return fmt.Sprintf("%d/%s", p.PortNo, p.Proto) }
-
-type SortByNsName []*svcResult
-
-func (a SortByNsName) Len() int { return len(a) }
-
-func (a SortByNsName) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
-
-func (a SortByNsName) Less(i, j int) bool {
-	if a[i].Namespace == a[j].Namespace {
-		return a[i].Name < a[j].Name
-	}
-	return a[i].Namespace < a[j].Namespace
-}
-
-func renderResults(res []*svcResult) {
+func RenderResults(res []*types.SvcResult) {
 	var output [][]string
 
 	sort.Sort(SortByNsName(res))

internal/util/sortsvc.go

@@ -0,0 +1,16 @@
+package util
+
+import "github.com/jpts/coredns-enum/internal/types"
+
+type SortByNsName []*types.SvcResult
+
+func (a SortByNsName) Len() int { return len(a) }
+
+func (a SortByNsName) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
+
+func (a SortByNsName) Less(i, j int) bool {
+	if a[i].Namespace == a[j].Namespace {
+		return a[i].Name < a[j].Name
+	}
+	return a[i].Namespace < a[j].Namespace
+}