redo modules/make command line parsing sane and easy to understand

This commit is contained in:
Alex 2017-09-29 18:59:37 -04:00
parent 424ee8c071
commit b7f50acf5d
7 changed files with 66 additions and 45 deletions

@ -2,6 +2,7 @@ package main
import ( import (
"encoding/json" "encoding/json"
"os"
"time" "time"
flags "github.com/ajholland/zflags" flags "github.com/ajholland/zflags"
@ -11,7 +12,8 @@ import (
) )
func main() { func main() {
if _, err := zgrab2.ParseFlags(); err != nil { _, modType, fl, err := zgrab2.ParseCommandLine(os.Args[1:])
if err != nil { //blanked arg is positional arguments
// Outputting help is returned as an error. Exit successfuly on help output. // Outputting help is returned as an error. Exit successfuly on help output.
flagsErr, ok := err.(*flags.Error) flagsErr, ok := err.(*flags.Error)
if ok && flagsErr.Type == flags.ErrHelp { if ok && flagsErr.Type == flags.ErrHelp {
@ -22,6 +24,24 @@ func main() {
log.Fatalf("could not parse flags: %s", err) log.Fatalf("could not parse flags: %s", err)
} }
if modType == "multiple" {
iniParser := zgrab2.NewIniParser()
m, _ := fl.(*zgrab2.MultipleCommand)
if m.ConfigFileName == "-" {
err = iniParser.Parse(os.Stdin)
} else {
err = iniParser.ParseFile(m.ConfigFileName)
}
if err != nil {
log.Fatalf("could not parse multiple: %s", err)
}
} else {
f, _ := fl.(zgrab2.ScanFlags)
mod := *zgrab2.GetModule(modType)
s := mod.NewScanner()
s.Init(f)
zgrab2.RegisterScan(modType, s)
}
zgrab2.PrintScanners() zgrab2.PrintScanners()
m := zgrab2.MakeMonitor() m := zgrab2.MakeMonitor()
start := time.Now() start := time.Now()

@ -7,8 +7,8 @@ import (
) )
type Scanner interface { type Scanner interface {
// Init runs once for this module at library init time. It is passed the parsed command-line flags // Init runs once for this module at library init time
Init(name string, flags ScanFlags) error Init(flags ScanFlags) error
// InitPerSender runs once per Goroutine. A single Goroutine will scan some non-deterministics // InitPerSender runs once per Goroutine. A single Goroutine will scan some non-deterministics
// subset of the input scan targets // subset of the input scan targets
@ -28,7 +28,7 @@ type ScanModule interface {
// Called by the framework for each time an individual scan is specified in the config or on // Called by the framework for each time an individual scan is specified in the config or on
// the command-line. The framework will then call scanner.Init(name, flags). // the command-line. The framework will then call scanner.Init(name, flags).
NewScanner() interface{} NewScanner() Scanner
} }
type ScanFlags interface { type ScanFlags interface {
@ -50,14 +50,20 @@ func (b *BaseFlags) GetName() string {
return b.Name return b.Name
} }
func GetModule(name string) *ScanModule {
return modules[name]
}
var modules map[string]*ScanModule
var scanners map[string]*Scanner var scanners map[string]*Scanner
var orderedScanners []string var orderedScanners []string
func init() { func init() {
scanners = make(map[string]*Scanner) scanners = make(map[string]*Scanner)
modules = make(map[string]*ScanModule)
} }
func RegisterScanner(name string, s Scanner) { func RegisterScan(name string, s Scanner) {
//add to list and map //add to list and map
if scanners[name] != nil { if scanners[name] != nil {
log.Fatal("name already used") log.Fatal("name already used")

@ -1,11 +1,6 @@
package zgrab2 package zgrab2
import ( import "errors"
"errors"
"os"
"github.com/ajholland/zflags"
)
type MultipleCommand struct { type MultipleCommand struct {
ConfigFileName string `short:"c" long:"config-file" default:"-" description:"Config filename, use - for stdin"` ConfigFileName string `short:"c" long:"config-file" default:"-" description:"Config filename, use - for stdin"`
@ -18,14 +13,9 @@ func (x *MultipleCommand) Validate(args []string) error {
return errors.New("cannot receive config file and input file from same source") return errors.New("cannot receive config file and input file from same source")
} }
var err error return nil
parse := flags.NewIniParser(parser) }
switch x.ConfigFileName {
case "-": func (x *MultipleCommand) Help() string {
err = parse.Parse(os.Stdin) return ""
default:
err = parse.ParseFile(x.ConfigFileName)
}
return err
} }

@ -15,6 +15,10 @@ func init() {
parser = flags.NewParser(&config, flags.Default) parser = flags.NewParser(&config, flags.Default)
} }
func NewIniParser() *flags.IniParser {
return flags.NewIniParser(parser)
}
// AddCommand adds a module to the parser and returns a pointer to // AddCommand adds a module to the parser and returns a pointer to
// a flags.command object or an error // a flags.command object or an error
func AddCommand(command string, shortDescription string, longDescription string, port int, m ScanModule) (*flags.Command, error) { func AddCommand(command string, shortDescription string, longDescription string, port int, m ScanModule) (*flags.Command, error) {
@ -24,17 +28,20 @@ func AddCommand(command string, shortDescription string, longDescription string,
} }
cmd.FindOptionByLongName("port").Default = []string{strconv.FormatUint(uint64(port), 10)} cmd.FindOptionByLongName("port").Default = []string{strconv.FormatUint(uint64(port), 10)}
cmd.FindOptionByLongName("name").Default = []string{command} cmd.FindOptionByLongName("name").Default = []string{command}
modules[command] = &m
return cmd, nil return cmd, nil
} }
// ParseFlags abstracts away the parser and validates the framework // ParseFlags parses the commands given on the command line
// configuration (global options) immediately after parsing // and validates the framework configuration (global options)
func ParseFlags() ([]string, error) { // immediately after parsing
r, err := parser.Parse() func ParseCommandLine(flags []string) ([]string, string, ScanFlags, error) {
posArgs, moduleType, f, err := parser.ParseCommandLine(flags)
if err == nil { if err == nil {
validateFrameworkConfiguration() validateFrameworkConfiguration()
} }
return r, err sf, _ := f.(ScanFlags)
return posArgs, moduleType, sf, err
} }
// ParseTarget takes input as a string and parses it into either an IPNet // ParseTarget takes input as a string and parses it into either an IPNet

@ -47,6 +47,7 @@ type HTTPModule struct {
} }
type HTTPScanner struct { type HTTPScanner struct {
config *HTTPFlags
} }
func init() { func init() {
@ -57,11 +58,11 @@ func init() {
} }
} }
func (m *HTTPModule) NewFlags() ScanFlags { func (m *HTTPModule) NewFlags() interface{} {
return new(HTTPFlags) return new(HTTPFlags)
} }
func (m *HTTPModule) NewScanner() Scanner { func (m *HTTPModule) NewScanner() zgrab2.Scanner {
return new(HTTPScanner) return new(HTTPScanner)
} }
@ -73,10 +74,9 @@ func (f *HTTPFlags) Help() string {
return "" return ""
} }
func (s *HTTPScanner) Init(name string, flags zgrab2.ScanFlags) error { func (s *HTTPScanner) Init(flags zgrab2.ScanFlags) error {
//httpFlags := flags.(*HTTPFlags) fl, _ := flags.(*HTTPFlags)
s.config = fl
zgrab2.RegisterScanner(name, s)
return nil return nil
} }

@ -17,7 +17,7 @@ type SSHModule struct {
} }
type SSHScanner struct { type SSHScanner struct {
SSHFlags config *SSHFlags
} }
func init() { func init() {
@ -28,11 +28,11 @@ func init() {
} }
} }
func (m *SSHModule) NewFlags() ScanFlags { func (m *SSHModule) NewFlags() interface{} {
return new(SSHFlags) return new(SSHFlags)
} }
func (m *SSHModule) NewScanner() Scanner { func (m *SSHModule) NewScanner() zgrab2.Scanner {
return new(SSHScanner) return new(SSHScanner)
} }
@ -44,10 +44,9 @@ func (f *SSHFlags) Help() string {
return "" return ""
} }
func (s *SSHScanner) Init(name string, flags interface{}) error { func (s *SSHScanner) Init(flags zgrab2.ScanFlags) error {
sshFlags := flags.(*SSHFlags) f, _ := flags.(*SSHFlags)
// set vars based on flags s.config = f
zgrab2.RegisterScanner(name, s)
return nil return nil
} }

@ -22,7 +22,7 @@ type TLSModule struct {
} }
type TLSScanner struct { type TLSScanner struct {
TLSFlags config *TLSFlags
} }
func init() { func init() {
@ -37,7 +37,7 @@ func (m *TLSModule) NewFlags() interface{} {
return new(TLSFlags) return new(TLSFlags)
} }
func (m *TLSModule) NewScanner() interface{} { func (m *TLSModule) NewScanner() zgrab2.Scanner {
return new(TLSScanner) return new(TLSScanner)
} }
@ -49,10 +49,9 @@ func (f *TLSFlags) Help() string {
return "" return ""
} }
func (s *TLSScanner) Init(name string, flags zgrab2.ScanFlags) error { func (s *TLSScanner) Init(flags zgrab2.ScanFlags) error {
//tlsFlags := flags.(*TLSFlags) f, _ := flags.(*TLSFlags)
s.config = f
zgrab2.RegisterScanner(name, s)
return nil return nil
} }