Skip to content

Commit

Permalink
adding more to our CLI
Browse files Browse the repository at this point in the history
  • Loading branch information
phillip-stephens committed Apr 8, 2024
1 parent 6682adc commit cbf1020
Show file tree
Hide file tree
Showing 9 changed files with 504 additions and 266 deletions.
55 changes: 36 additions & 19 deletions pkg/cmd/cli.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,29 +15,44 @@ package cmd

import (
"fmt"
"github.com/zmap/dns"
"net"
"os"
"strings"
"time"

"github.com/spf13/cobra"
"github.com/spf13/viper"
"github.com/zmap/dns"
"github.com/zmap/zdns/internal/util"
"github.com/zmap/zdns/pkg/zdns"
"net"
"os"
"strings"
"sync"
)

type InputHandler interface {
FeedChannel(in chan<- interface{}, wg *sync.WaitGroup) error
}
type OutputHandler interface {
WriteResults(results <-chan string, wg *sync.WaitGroup) error
}

type CLIConf struct {
Threads int
Timeout time.Duration
IterationTimeout time.Duration
Timeout int
IterationTimeout int
Retries int
AlexaFormat bool
MetadataFormat bool
NameServerInputFormat bool
IterativeResolution bool
LookupAllNameServers bool

NameServersString string
LocalAddrString string
LocalIfaceString string
ConfigFilePath string
ClassString string
UseNanoseconds bool
ClientSubnetString string
RequestNSID bool

ResultVerbosity string
IncludeInOutput string
OutputGroups []string
Expand All @@ -62,6 +77,8 @@ type CLIConf struct {

InputFilePath string
OutputFilePath string
InputHandler InputHandler
OutputHandler OutputHandler
LogFilePath string
MetadataFilePath string

Expand Down Expand Up @@ -182,17 +199,17 @@ func init() {
rootCmd.PersistentFlags().BoolVar(&GC.RecycleSockets, "recycle-sockets", true, "Create long-lived unbound UDP socket for each thread at launch and reuse for all (UDP) queries")
rootCmd.PersistentFlags().BoolVar(&GC.NameServerMode, "name-server-mode", false, "Treats input as nameservers to query with a static query rather than queries to send to a static name server")

//rootCmd.PersistentFlags().StringVar(&Servers_string, "name-servers", "", "List of DNS servers to use. Can be passed as comma-delimited string or via @/path/to/file. If no port is specified, defaults to 53.")
//rootCmd.PersistentFlags().StringVar(&Localaddr_string, "local-addr", "", "comma-delimited list of local addresses to use")
//rootCmd.PersistentFlags().StringVar(&Localif_string, "local-interface", "", "local interface to use")
//rootCmd.PersistentFlags().StringVar(&Config_file, "conf-file", "/etc/resolv.conf", "config file for DNS servers")
//rootCmd.PersistentFlags().IntVar(&Timeout, "timeout", 15, "timeout for resolving an individual name")
//rootCmd.PersistentFlags().IntVar(&IterationTimeout, "iteration-timeout", 4, "timeout for resolving a single iteration in an iterative query")
//rootCmd.PersistentFlags().StringVar(&Class_string, "class", "INET", "DNS class to query. Options: INET, CSNET, CHAOS, HESIOD, NONE, ANY. Default: INET.")
//rootCmd.PersistentFlags().BoolVar(&NanoSeconds, "nanoseconds", false, "Use nanosecond resolution timestamps")
//rootCmd.PersistentFlags().StringVar(&ClientSubnet_string, "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(&NSID, "nsid", false, "Request NSID.")
rootCmd.PersistentFlags().StringVar(&GC.NameServersString, "name-servers", "", "List of DNS servers to use. Can be passed as comma-delimited string or via @/path/to/file. If no port is specified, defaults to 53.")
rootCmd.PersistentFlags().StringVar(&GC.LocalAddrString, "local-addr", "", "comma-delimited list of local addresses to use")
rootCmd.PersistentFlags().StringVar(&GC.LocalIfaceString, "local-interface", "", "local interface to use")
rootCmd.PersistentFlags().StringVar(&GC.ConfigFilePath, "conf-file", "/etc/resolv.conf", "config file for DNS servers")
rootCmd.PersistentFlags().IntVar(&GC.Timeout, "timeout", 15, "timeout for resolving an individual name")
rootCmd.PersistentFlags().IntVar(&GC.IterationTimeout, "iteration-timeout", 4, "timeout for resolving a single iteration in an iterative query")
rootCmd.PersistentFlags().StringVar(&GC.ClassString, "class", "INET", "DNS class to query. Options: INET, CSNET, CHAOS, HESIOD, NONE, ANY. Default: INET.")
rootCmd.PersistentFlags().BoolVar(&GC.UseNanoseconds, "nanoseconds", false, "Use nanosecond resolution timestamps")
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.RequestNSID, "nsid", false, "Request NSID.")

rootCmd.PersistentFlags().Bool("ipv4-lookup", false, "Perform an IPv4 Lookup in modules")
rootCmd.PersistentFlags().Bool("ipv6-lookup", false, "Perform an IPv6 Lookup in modules")
Expand Down
87 changes: 87 additions & 0 deletions pkg/cmd/file_handlers.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
/*
* ZDNS Copyright 2022 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 cmd

import (
"bufio"
"os"
"sync"

log "github.com/sirupsen/logrus"
)

type FileInputHandler struct {
filepath string
}

func NewFileInputHandler(filepath string) *FileInputHandler {
return &FileInputHandler{
filepath: filepath,
}
}

func (h *FileInputHandler) FeedChannel(in chan<- interface{}, wg *sync.WaitGroup) error {
defer close(in)
defer (*wg).Done()

var f *os.File
if h.filepath == "" || h.filepath == "-" {
f = os.Stdin
} else {
var err error
f, err = os.Open(h.filepath)
if err != nil {
log.Fatalf("unable to open input file: %v", err)
}
}
s := bufio.NewScanner(f)
for s.Scan() {
in <- s.Text()
}
if err := s.Err(); err != nil {
log.Fatalf("input unable to read file: %v", err)
}
return nil
}

type FileOutputHandler struct {
filepath string
}

func NewFileOutputHandler(filepath string) *FileOutputHandler {
return &FileOutputHandler{
filepath: filepath,
}
}

func (h *FileOutputHandler) WriteResults(results <-chan string, wg *sync.WaitGroup) error {
defer (*wg).Done()

var f *os.File
if h.filepath == "" || h.filepath == "-" {
f = os.Stdout
} else {
var err error
f, err = os.OpenFile(h.filepath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
if err != nil {
log.Fatalf("unable to open output file: %v", err)
}
defer f.Close()
}
for n := range results {
f.WriteString(n + "\n")
}
return nil
}
65 changes: 65 additions & 0 deletions pkg/cmd/stream_handlers.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
/*
* ZDNS Copyright 2022 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 cmd

import (
"bufio"
"io"
"sync"

log "github.com/sirupsen/logrus"
)

type StreamInputHandler struct {
reader io.Reader
}

func NewStreamInputHandler(r io.Reader) *StreamInputHandler {
return &StreamInputHandler{
reader: r,
}
}

func (h *StreamInputHandler) FeedChannel(in chan<- interface{}, wg *sync.WaitGroup) error {
defer close(in)
defer (*wg).Done()

s := bufio.NewScanner(h.reader)
for s.Scan() {
in <- s.Text()
}
if err := s.Err(); err != nil {
log.Fatalf("unable to read input stream: %v", err)
}
return nil
}

type StreamOutputHandler struct {
writer io.Writer
}

func NewStreamOutputHandler(w io.Writer) *StreamOutputHandler {
return &StreamOutputHandler{
writer: w,
}
}

func (h *StreamOutputHandler) WriteResults(results <-chan string, wg *sync.WaitGroup) error {
defer (*wg).Done()
for n := range results {
io.WriteString(h.writer, n+"\n")
}
return nil
}
43 changes: 43 additions & 0 deletions pkg/cmd/ulimit_check.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
//go:build linux || darwin
// +build linux darwin

/*
* ZDNS Copyright 2020 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 cmd

import (
"syscall"

log "github.com/sirupsen/logrus"
)

func ulimitCheck(maxOpenFiles uint64) {
var rLimit syscall.Rlimit
err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
if err != nil {
log.Fatal("Failed to fetch ulimit ", err)
}

if maxOpenFiles > rLimit.Cur {
log.Warn("Current nofile limit (", rLimit.Cur, ") lower than maximum connection count (", maxOpenFiles, "), trying to update.")

rLimit.Max = maxOpenFiles
rLimit.Cur = maxOpenFiles
err = syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit)
if err != nil {
log.Fatal("Error setting nofile limit to ", rLimit.Cur, ": ", err)
}
log.Info("Updated nofile limit to ", rLimit.Cur)
}
}
21 changes: 21 additions & 0 deletions pkg/cmd/ulimit_check_unknown.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
//go:build !linux && !darwin
// +build !linux,!darwin

/*
* ZDNS Copyright 2020 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 cmd

func ulimitCheck(maxOpenFiles uint64) {
// fallback for ulimit check on unsupported platform
}
Loading

0 comments on commit cbf1020

Please sign in to comment.