diff --git a/cmd/zgrab2/main.go b/cmd/zgrab2/main.go index bd04a2c..f44db0e 100644 --- a/cmd/zgrab2/main.go +++ b/cmd/zgrab2/main.go @@ -2,6 +2,7 @@ package main import ( "encoding/json" + "os" "time" flags "github.com/ajholland/zflags" @@ -11,7 +12,8 @@ import ( ) 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. flagsErr, ok := err.(*flags.Error) if ok && flagsErr.Type == flags.ErrHelp { @@ -22,6 +24,24 @@ func main() { 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() m := zgrab2.MakeMonitor() start := time.Now() diff --git a/module.go b/module.go index bd779db..4c04abc 100644 --- a/module.go +++ b/module.go @@ -7,8 +7,8 @@ import ( ) type Scanner interface { - // Init runs once for this module at library init time. It is passed the parsed command-line flags - Init(name string, flags ScanFlags) error + // Init runs once for this module at library init time + Init(flags ScanFlags) error // InitPerSender runs once per Goroutine. A single Goroutine will scan some non-deterministics // 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 // the command-line. The framework will then call scanner.Init(name, flags). - NewScanner() interface{} + NewScanner() Scanner } type ScanFlags interface { @@ -50,14 +50,20 @@ func (b *BaseFlags) GetName() string { return b.Name } +func GetModule(name string) *ScanModule { + return modules[name] +} + +var modules map[string]*ScanModule var scanners map[string]*Scanner var orderedScanners []string func init() { 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 if scanners[name] != nil { log.Fatal("name already used") diff --git a/multiple.go b/multiple.go index e0e6783..ae1b2c9 100644 --- a/multiple.go +++ b/multiple.go @@ -1,11 +1,6 @@ package zgrab2 -import ( - "errors" - "os" - - "github.com/ajholland/zflags" -) +import "errors" type MultipleCommand struct { 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") } - var err error - parse := flags.NewIniParser(parser) - switch x.ConfigFileName { - case "-": - err = parse.Parse(os.Stdin) - default: - err = parse.ParseFile(x.ConfigFileName) - } - - return err + return nil +} + +func (x *MultipleCommand) Help() string { + return "" } diff --git a/utility.go b/utility.go index bbd0556..772543a 100644 --- a/utility.go +++ b/utility.go @@ -15,6 +15,10 @@ func init() { 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 // a flags.command object or an 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("name").Default = []string{command} + modules[command] = &m return cmd, nil } -// ParseFlags abstracts away the parser and validates the framework -// configuration (global options) immediately after parsing -func ParseFlags() ([]string, error) { - r, err := parser.Parse() +// ParseFlags parses the commands given on the command line +// and validates the framework configuration (global options) +// immediately after parsing +func ParseCommandLine(flags []string) ([]string, string, ScanFlags, error) { + posArgs, moduleType, f, err := parser.ParseCommandLine(flags) if err == nil { 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 diff --git a/zmodules/http.go b/zmodules/http.go index 311cdc7..ebeae2a 100644 --- a/zmodules/http.go +++ b/zmodules/http.go @@ -47,6 +47,7 @@ type HTTPModule struct { } type HTTPScanner struct { + config *HTTPFlags } func init() { @@ -57,11 +58,11 @@ func init() { } } -func (m *HTTPModule) NewFlags() ScanFlags { +func (m *HTTPModule) NewFlags() interface{} { return new(HTTPFlags) } -func (m *HTTPModule) NewScanner() Scanner { +func (m *HTTPModule) NewScanner() zgrab2.Scanner { return new(HTTPScanner) } @@ -73,10 +74,9 @@ func (f *HTTPFlags) Help() string { return "" } -func (s *HTTPScanner) Init(name string, flags zgrab2.ScanFlags) error { - //httpFlags := flags.(*HTTPFlags) - - zgrab2.RegisterScanner(name, s) +func (s *HTTPScanner) Init(flags zgrab2.ScanFlags) error { + fl, _ := flags.(*HTTPFlags) + s.config = fl return nil } diff --git a/zmodules/ssh.go b/zmodules/ssh.go index 5bfc27e..0209762 100644 --- a/zmodules/ssh.go +++ b/zmodules/ssh.go @@ -17,7 +17,7 @@ type SSHModule struct { } type SSHScanner struct { - SSHFlags + config *SSHFlags } func init() { @@ -28,11 +28,11 @@ func init() { } } -func (m *SSHModule) NewFlags() ScanFlags { +func (m *SSHModule) NewFlags() interface{} { return new(SSHFlags) } -func (m *SSHModule) NewScanner() Scanner { +func (m *SSHModule) NewScanner() zgrab2.Scanner { return new(SSHScanner) } @@ -44,10 +44,9 @@ func (f *SSHFlags) Help() string { return "" } -func (s *SSHScanner) Init(name string, flags interface{}) error { - sshFlags := flags.(*SSHFlags) - // set vars based on flags - zgrab2.RegisterScanner(name, s) +func (s *SSHScanner) Init(flags zgrab2.ScanFlags) error { + f, _ := flags.(*SSHFlags) + s.config = f return nil } diff --git a/zmodules/tls.go b/zmodules/tls.go index d4be410..9166093 100644 --- a/zmodules/tls.go +++ b/zmodules/tls.go @@ -22,7 +22,7 @@ type TLSModule struct { } type TLSScanner struct { - TLSFlags + config *TLSFlags } func init() { @@ -37,7 +37,7 @@ func (m *TLSModule) NewFlags() interface{} { return new(TLSFlags) } -func (m *TLSModule) NewScanner() interface{} { +func (m *TLSModule) NewScanner() zgrab2.Scanner { return new(TLSScanner) } @@ -49,10 +49,9 @@ func (f *TLSFlags) Help() string { return "" } -func (s *TLSScanner) Init(name string, flags zgrab2.ScanFlags) error { - //tlsFlags := flags.(*TLSFlags) - - zgrab2.RegisterScanner(name, s) +func (s *TLSScanner) Init(flags zgrab2.ScanFlags) error { + f, _ := flags.(*TLSFlags) + s.config = f return nil }