change structure of zgrab2

cmd/zgrab2/main.go

@@ -22,6 +22,7 @@ func main() {
 		log.Fatalf("could not parse flags: %s", err)
 	}
 
+	zgrab2.PrintScanners()
 	m := zgrab2.MakeMonitor()
 	start := time.Now()
 	log.Infof("started grab at %s", start.Format(time.RFC3339))

module.go

@@ -1,71 +1,89 @@
 package zgrab2
 
 import (
+	"fmt"
 	"log"
-	"net"
-	"strconv"
 	"time"
-
-	"github.com/ajholland/zflags"
 )
 
-type ScanModule interface {
+type Scanner interface {
-	Scan(ip net.IP) (interface{}, error)
+	// Init runs once for this module at library init time. It is passed the parsed command-line flags
-	PerRoutineInitialize()
+	Init(name string, flags ScanFlags) error
-	GetPort() uint
+
+	// InitPerSender runs once per Goroutine. A single Goroutine will scan some non-deterministics
+	// subset of the input scan targets
+	InitPerSender(senderID int) error
+
+	// Returns the name passed at init
 	GetName() string
-	New() interface{}
+
+	// Scan connects to a host. The result should be JSON-serializable
+	Scan(t ScanTarget, port uint) (interface{}, error)
+}
+
+type ScanModule interface {
+	// Called by the framework to pass to the argument parser. The parsed flags will be passed
+	// to the scanner created by NewScanner().
+	NewFlags() 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{}
+}
+
+type ScanFlags interface {
+	// Help optionally returns any additional help text, e.g. specifying what empty defaults
+	// are interpreted as.
+	Help() string
+
+	// Validate enforces all command-line flags and positional arguments have valid values.
 	Validate(args []string) error
 }
 
-type BaseScanModule struct {
+type BaseFlags struct {
 	Port    uint   `short:"p" long:"port" description:"Specify port to grab on"`
 	Name    string `short:"n" long:"name" description:"Specify name for output json, only necessary if scanning multiple modules"`
-	Timeout int    `short:"t" long:"timeout" description:"Set connection timeout in seconds"`
+	Timeout uint   `short:"t" long:"timeout" description:"Set connection timeout in seconds"`
 }
 
-func (b *BaseScanModule) GetPort() uint {
+func (b *BaseFlags) GetName() string {
-	return b.Port
-}
-
-func (b *BaseScanModule) GetName() string {
 	return b.Name
 }
 
-func (b *BaseScanModule) SetDefaultPortAndName(cmd *flags.Command, port uint, name string) {
+var scanners map[string]*Scanner
-	cmd.FindOptionByLongName("port").Default = []string{strconv.FormatUint(uint64(port), 10)}
+var orderedScanners []string
-	cmd.FindOptionByLongName("name").Default = []string{name}
-}
-
-var modules map[string]*ScanModule
-var orderedModules []string
 
 func init() {
-	modules = make(map[string]*ScanModule)
+	scanners = make(map[string]*Scanner)
 }
 
-func RegisterModule(name string, m ScanModule) {
+func RegisterScanner(name string, s Scanner) {
 	//add to list and map
-	if modules[name] != nil {
+	if scanners[name] != nil {
 		log.Fatal("name already used")
 	}
-	orderedModules = append(orderedModules, name)
+	orderedScanners = append(orderedScanners, name)
-	modules[name] = &m
+	scanners[name] = &s
+	fmt.Println("Registered: ", name, s)
 }
 
-// runHandler will call perRoutineInitialize, Scan, and respond with a protocol response, data unmarshalled, to the worker
+func PrintScanners() {
-func RunModule(module ScanModule, mon *Monitor, ip net.IP) (string, ModuleResponse) {
+	for k, v := range scanners {
+		fmt.Println(k, v)
+	}
+}
+
+func RunModule(s Scanner, mon *Monitor, target ScanTarget) (string, ScanResponse) {
 	t := time.Now()
-	module.PerRoutineInitialize()
+	res, e := s.Scan(target, uint(22))
-	res, e := module.Scan(ip)
 	var err *error //nil pointers are null in golang, which is not nil and not empty
 	if e == nil {
-		mon.statusesChan <- moduleStatus{name: module.GetName(), st: status_success}
+		mon.statusesChan <- moduleStatus{name: s.GetName(), st: status_success}
 		err = nil
 	} else {
-		mon.statusesChan <- moduleStatus{name: module.GetName(), st: status_failure}
+		mon.statusesChan <- moduleStatus{name: s.GetName(), st: status_failure}
 		err = &e
 	}
-	resp := ModuleResponse{Result: res, Error: err, Time: t.Format(time.RFC3339)}
+	resp := ScanResponse{Result: res, Error: err, Time: t.Format(time.RFC3339)}
-	return module.GetName(), resp
+	return s.GetName(), resp
 }

processing.go

@@ -11,17 +11,17 @@ import (
 )
 
 type Grab struct {
-	IP     string                    `json:"ip,omitempty"`
+	IP     string                  `json:"ip,omitempty"`
-	Domain string                    `json:"domain,omitempty"`
+	Domain string                  `json:"domain,omitempty"`
-	Data   map[string]ModuleResponse `json:"data,omitempty"`
+	Data   map[string]ScanResponse `json:"data,omitempty"`
 }
 
-type target struct {
+type ScanTarget struct {
 	IP     net.IP
 	Domain string
 }
 
-type ModuleResponse struct {
+type ScanResponse struct {
 	Result         interface{} `json:"result,omitempty"`
 	Time           string      `json:"time,omitempty"`
 	Error          *error      `json:"error,omitempty"`
@@ -29,12 +29,12 @@ type ModuleResponse struct {
 }
 
 // grabTarget calls handler for each action
-func grabTarget(input target, m *Monitor) []byte {
+func grabTarget(input ScanTarget, m *Monitor) []byte {
-	moduleResult := make(map[string]ModuleResponse)
+	moduleResult := make(map[string]ScanResponse)
 
-	for _, moduleName := range orderedModules {
+	for _, scannerName := range orderedScanners {
-		module := modules[moduleName]
+		scanner := scanners[scannerName]
-		name, res := RunModule(*module, m, input.IP)
+		name, res := RunModule(*scanner, m, input)
 		moduleResult[name] = res
 		if res.Error != nil && !config.Multiple.ContinueOnError {
 			break
@@ -61,7 +61,7 @@ func grabTarget(input target, m *Monitor) []byte {
 // Process sets up an output encoder, input reader, and starts grab workers
 func Process(mon *Monitor) {
 	workers := config.Senders
-	processQueue := make(chan target, workers*4)
+	processQueue := make(chan ScanTarget, workers*4)
 	outputQueue := make(chan []byte, workers*4)
 
 	//Create wait groups
@@ -116,14 +116,14 @@ func Process(mon *Monitor) {
 		if ipnet != nil {
 			if ipnet.Mask != nil {
 				for ip = ipnet.IP.Mask(ipnet.Mask); ipnet.Contains(ip); incrementIP(ip) {
-					processQueue <- target{IP: duplicateIP(ip), Domain: domain}
+					processQueue <- ScanTarget{IP: duplicateIP(ip), Domain: domain}
 				}
 				continue
 			} else {
 				ip = ipnet.IP
 			}
 		}
-		processQueue <- target{IP: ip, Domain: domain}
+		processQueue <- ScanTarget{IP: ip, Domain: domain}
 	}
 
 	close(processQueue)

utility.go

@@ -3,6 +3,7 @@ package zgrab2
 import (
 	"errors"
 	"net"
+	"strconv"
 	"strings"
 
 	"github.com/ajholland/zflags"
@@ -15,8 +16,14 @@ func init() {
 }
 
 // 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, m ScanModule) (*flags.Command, error) {
+func AddCommand(command string, shortDescription string, longDescription string, port int, m ScanModule) (*flags.Command, error) {
-	return parser.AddCommand(command, shortDescription, longDescription, m)
+	cmd, err := parser.AddCommand(command, shortDescription, longDescription, m)
+	if err != nil {
+		return nil, err
+	}
+	cmd.FindOptionByLongName("port").Default = []string{strconv.FormatUint(uint64(port), 10)}
+	cmd.FindOptionByLongName("name").Default = []string{command}
+	return cmd, nil
 }
 
 // ParseFlags abstracts away the parser and validates the framework configuration (global options) immediately after parsing

zmodules/http.go

@@ -1,14 +1,12 @@
 package zmodules
 
 import (
-	"net"
-
 	log "github.com/sirupsen/logrus"
 	"github.com/zmap/zgrab2"
 )
 
-type HTTPModule struct {
+type HTTPFlags struct {
-	zgrab2.BaseScanModule
+	zgrab2.BaseFlags
 	HTTP HTTPOptions `json:"http"`
 }
 
@@ -45,32 +43,52 @@ type HTTPResults struct {
 	//RedirectResponseChain []*http.Response `json:"redirect_response_chain,omitempty"`
 }
 
-// Per module initialization call
+type HTTPModule struct {
+}
+
+type HTTPScanner struct {
+}
+
 func init() {
 	var httpModule HTTPModule
-	cmd, err := zgrab2.AddCommand("http", "HTTP Banner Grab", "Grab a banner over HTTP", &httpModule)
+	_, err := zgrab2.AddCommand("http", "HTTP Banner Grab", "Grab a banner over HTTP", 80, &httpModule)
 	if err != nil {
 		log.Fatal(err)
 	}
-	httpModule.SetDefaultPortAndName(cmd, uint(80), "http")
 }
 
-func (x *HTTPModule) New() interface{} {
+func (m *HTTPModule) NewFlags() ScanFlags {
-	return new(HTTPModule)
+	return new(HTTPFlags)
 }
 
-// Per module per goroutine initialization call
+func (m *HTTPModule) NewScanner() Scanner {
-func (x *HTTPModule) PerRoutineInitialize() {
+	return new(HTTPScanner)
-
 }
 
-// Validates the options sent to HTTPConfig, registers the config module, and then passes operation back to main
+func (f *HTTPFlags) Validate(args []string) error {
-func (x *HTTPModule) Validate(args []string) error {
-	zgrab2.RegisterModule(x.Name, x)
 	return nil
 }
 
-func (x *HTTPModule) Scan(ip net.IP) (interface{}, error) {
+func (f *HTTPFlags) Help() string {
+	return ""
+}
+
+func (s *HTTPScanner) Init(name string, flags zgrab2.ScanFlags) error {
+	//httpFlags := flags.(*HTTPFlags)
+
+	zgrab2.RegisterScanner(name, s)
+	return nil
+}
+
+func (s *HTTPScanner) InitPerSender(senderID int) error {
+	return nil
+}
+
+func (s *HTTPScanner) GetName() string {
+	return ""
+}
+
+func (s *HTTPScanner) Scan(t zgrab2.ScanTarget, port uint) (interface{}, error) {
 	http := HTTPRequest{Method: "Get", Body: "testing"}
 	ret := HTTPResults{ProxyRequest: &http}
 	return ret, nil

zmodules/ssh.go

@@ -1,69 +1,63 @@
 package zmodules
 
 import (
-	"net"
-	"strconv"
-	"time"
-
 	log "github.com/sirupsen/logrus"
 	"github.com/zmap/zgrab2"
-	"github.com/zmap/zgrab2/zimports/ssh"
 )
 
-type SSHModule struct {
+type SSHFlags struct {
-	zgrab2.BaseModule
+	zgrab2.BaseFlags
 	ClientID          string `long:"client" description:"Specify the client ID string to use" default:"SSH-2.0-Go"`
 	KexAlgorithms     string `long:"kex-algorithms" description:"Set SSH Key Exchange Algorithms"`
 	HostKeyAlgorithms string `long:"host-key-algorithms" description:"Set SSH Host Key Algorithms"`
 	NegativeOne       bool   `long:"negative-one" description:"Set SSH DH kex value to -1 in the selected group"`
 }
 
+type SSHModule struct {
+}
+
+type SSHScanner struct {
+	SSHFlags
+}
+
 func init() {
 	var sshModule SSHModule
-	cmd, err := zgrab2.AddCommand("ssh", "SSH Banner Grab", "Grab a banner over SSH", &sshModule)
+	_, err := zgrab2.AddCommand("ssh", "SSH Banner Grab", "Grab a banner over SSH", 22, &sshModule)
 	if err != nil {
 		log.Fatal(err)
 	}
-	sshModule.SetDefaultPortAndName(cmd, uint(22), "ssh")
 }
 
-func (x *SSHModule) New() interface{} {
+func (m *SSHModule) NewFlags() ScanFlags {
-	return new(SSHModule)
+	return new(SSHFlags)
 }
 
-// per module per routine initialization call
+func (m *SSHModule) NewScanner() Scanner {
-func (x *SSHModule) PerRoutineInitialize() {
+	return new(SSHScanner)
-
 }
 
-// Execute validates the options sent to SSHModule and then passes operation back to main
+func (f *SSHFlags) Validate(args []string) error {
-func (x *SSHModule) Validate(args []string) error {
-	zgrab2.RegisterModule(x.Name, x)
 	return nil
 }
 
-func (x *SSHModule) makeSSHGrabber(hlog *ssh.HandshakeLog) func(string) error {
+func (f *SSHFlags) Help() string {
-	return func(netAddr string) error {
+	return ""
-		sshConfig := ssh.MakeSSHConfig()
-		sshConfig.Timeout = time.Duration(x.Timeout) * time.Second
-		sshConfig.ConnLog = hlog
-		_, err := ssh.Dial("tcp", netAddr, sshConfig)
-		if err != nil {
-			return err
-		}
-
-		return nil
-	}
 }
 
-func (x *SSHModule) Scan(ip net.IP) (interface{}, error) {
+func (s *SSHScanner) Init(name string, flags interface{}) error {
-	data := new(ssh.HandshakeLog)
+	sshFlags := flags.(*SSHFlags)
-	sshGrabber := x.makeSSHGrabber(data)
+	// set vars based on flags
-
+	zgrab2.RegisterScanner(name, s)
-	port := strconv.FormatUint(uint64(x.Port), 10)
+	return nil
-	rhost := net.JoinHostPort(ip.String(), port)
+}
-
+
-	err := sshGrabber(rhost)
+func (s *SSHScanner) InitPerSender(senderID int) error {
-
+	return nil
-	return data, err
+}
+
+func (s *SSHScanner) GetName() string {
+	return ""
+}
+func (s *SSHScanner) Scan(t zgrab2.ScanTarget, port uint) (interface{}, error) {
+	return nil, nil
 }

zmodules/tls.go

@@ -1,14 +1,12 @@
 package zmodules
 
 import (
-	"net"
-
 	log "github.com/sirupsen/logrus"
 	"github.com/zmap/zgrab2"
 )
 
-type TLSModule struct {
+type TLSFlags struct {
-	zgrab2.BaseScanModule
+	zgrab2.BaseFlags
 	Heartbleed           bool        `long:"heartbleed" description:"Check if server is vulnerable to Heartbleed"`
 	Version              int         `long:"version" description:"Max TLS version to use"`
 	Verbose              bool        `long:"verbose" description:"Add extra TLS information to JSON output (client hello, client KEX, key material, etc)" json:"verbose"`
@@ -20,29 +18,52 @@ type TLSModule struct {
 	HTTP                 HTTPOptions `json:"http"`
 }
 
+type TLSModule struct {
+}
+
+type TLSScanner struct {
+	TLSFlags
+}
+
 func init() {
 	var tlsModule TLSModule
-	cmd, err := zgrab2.AddCommand("tls", "TLS Banner Grab", "Grab banner over TLS", &tlsModule)
+	_, err := zgrab2.AddCommand("tls", "TLS Banner Grab", "Grab banner over TLS", 443, &tlsModule)
 	if err != nil {
 		log.Fatal(err)
 	}
-	tlsModule.SetDefaultPortAndName(cmd, uint(443), "ssh")
 }
 
-func (x *TLSModule) New() interface{} {
+func (m *TLSModule) NewFlags() interface{} {
-	return new(TLSModule)
+	return new(TLSFlags)
 }
 
-func (x *TLSModule) PerRoutineInitialize() {
+func (m *TLSModule) NewScanner() interface{} {
-
+	return new(TLSScanner)
 }
 
-// Execute validates the options sent to TLSModule and then passes operation back to main
+func (f *TLSFlags) Validate(args []string) error {
-func (x *TLSModule) Validate(args []string) error {
-	zgrab2.RegisterModule(x.Name, x)
 	return nil
 }
 
-func (x *TLSModule) Scan(ip net.IP) (interface{}, error) {
+func (f *TLSFlags) Help() string {
-	return x, nil
+	return ""
+}
+
+func (s *TLSScanner) Init(name string, flags zgrab2.ScanFlags) error {
+	//tlsFlags := flags.(*TLSFlags)
+
+	zgrab2.RegisterScanner(name, s)
+	return nil
+}
+
+func (s *TLSScanner) GetName() string {
+	return ""
+}
+
+func (s *TLSScanner) InitPerSender(senderID int) error {
+	return nil
+}
+
+func (s *TLSScanner) Scan(t zgrab2.ScanTarget, port uint) (interface{}, error) {
+	return s, nil
 }