Add custom regexp banner grabber (#199)

https://github.com/zmap/zgrab2/pull/199

modules/banner.go

@@ -0,0 +1,9 @@
+package modules
+
+import (
+	"github.com/zmap/zgrab2/modules/banner"
+)
+
+func init() {
+	banner.RegisterModule()
+}

modules/banner/scanner.go

@@ -0,0 +1,158 @@
+// Package banner provides simple banner grab and matching implementation of the zgrab2.Module.
+// It sends a customizble probe (default to "\n") and filters the results based on custom regexp (--pattern)
+
+package banner
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"log"
+	"net"
+	"regexp"
+	"strconv"
+
+	"github.com/zmap/zgrab2"
+)
+
+// Flags give the command-line flags for the banner module.
+type Flags struct {
+	zgrab2.BaseFlags
+	Probe    string `long:"probe" default:"\\n" description:"Probe to send to the server. Use triple slashes to escape, for example \\\\\\n is literal \\n" `
+	Pattern  string `long:"pattern" description:"Pattern to match, must be valid regexp."`
+	MaxTries int    `long:"max-tries" default:"1" description:"Number of tries for timeouts and connection errors before giving up."`
+}
+
+// Module is the implementation of the zgrab2.Module interface.
+type Module struct {
+}
+
+// Scanner is the implementation of the zgrab2.Scanner interface.
+type Scanner struct {
+	config *Flags
+	regex  *regexp.Regexp
+	probe  []byte
+}
+
+type Results struct {
+	Banner string `json:"banner,omitempty"`
+	Length int    `json:"length,omitempty"`
+}
+
+// RegisterModule is called by modules/banner.go to register the scanner.
+func RegisterModule() {
+	var module Module
+	_, err := zgrab2.AddCommand("banner", "Banner", "Grab banner by sending probe and match with regexp", 80, &module)
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+// NewFlags returns a new default flags object.
+func (m *Module) NewFlags() interface{} {
+	return new(Flags)
+}
+
+// GetName returns the Scanner name defined in the Flags.
+func (scanner *Scanner) GetName() string {
+	return scanner.config.Name
+}
+
+// GetTrigger returns the Trigger defined in the Flags.
+func (scanner *Scanner) GetTrigger() string {
+	return scanner.config.Trigger
+}
+
+// Protocol returns the protocol identifier of the scan.
+func (scanner *Scanner) Protocol() string {
+	return "banner"
+}
+
+// GetPort returns the port being scanned.
+func (scanner *Scanner) GetPort() uint {
+	return scanner.config.Port
+}
+
+// InitPerSender initializes the scanner for a given sender.
+func (scanner *Scanner) InitPerSender(senderID int) error {
+	return nil
+}
+
+// NewScanner returns a new Scanner object.
+func (m *Module) NewScanner() zgrab2.Scanner {
+	return new(Scanner)
+}
+
+// Validate validates the flags and returns nil on success.
+func (f *Flags) Validate(args []string) error {
+	return nil
+}
+
+// Help returns the module's help string.
+func (f *Flags) Help() string {
+	return ""
+}
+
+// Init initializes the Scanner with the command-line flags.
+func (scanner *Scanner) Init(flags zgrab2.ScanFlags) error {
+	f, _ := flags.(*Flags)
+	scanner.config = f
+	scanner.regex = regexp.MustCompile(scanner.config.Pattern)
+	probe, err := strconv.Unquote(fmt.Sprintf(`"%s"`, scanner.config.Probe))
+	if err != nil {
+		panic("Probe error")
+	}
+	scanner.probe = []byte(probe)
+	return nil
+}
+
+var NoMatchError = errors.New("pattern did not match")
+
+func (scanner *Scanner) Scan(target zgrab2.ScanTarget) (zgrab2.ScanStatus, interface{}, error) {
+	try := 0
+	var (
+		conn    net.Conn
+		err     error
+		readerr error
+	)
+	for try < scanner.config.MaxTries {
+		try += 1
+		conn, err = target.Open(&scanner.config.BaseFlags)
+		if err != nil {
+			continue
+		}
+		break
+	}
+	if err != nil {
+		return zgrab2.TryGetScanStatus(err), nil, err
+	}
+	defer conn.Close()
+
+	var ret []byte
+	try = 0
+	for try < scanner.config.MaxTries {
+		try += 1
+		_, err = conn.Write(scanner.probe)
+		ret, readerr = zgrab2.ReadAvailable(conn)
+		if err != nil {
+			continue
+		}
+		if readerr != io.EOF && readerr != nil {
+			continue
+		}
+		break
+	}
+	if err != nil {
+		return zgrab2.TryGetScanStatus(err), nil, err
+	}
+	if readerr != io.EOF && readerr != nil {
+		return zgrab2.TryGetScanStatus(readerr), nil, readerr
+	}
+	results := Results{Banner: string(ret), Length: len(ret)}
+	if scanner.regex.Match(ret) {
+		return zgrab2.SCAN_SUCCESS, &results, nil
+	}
+
+	return zgrab2.SCAN_PROTOCOL_ERROR, &results, NoMatchError
+
+}

zgrab2_schemas/zgrab2/__init__.py

@@ -20,3 +20,4 @@ from . import smtp
 from . import ssh
 from . import telnet
 from . import ipp
+from . import banner

zgrab2_schemas/zgrab2/banner.py

@@ -0,0 +1,20 @@
+# zschema sub-schema for zgrab2's banner module
+# Registers zgrab2-banner globally, and banner with the main zgrab2 schema.
+from zschema.leaves import *
+from zschema.compounds import *
+import zschema.registry
+
+import zcrypto_schemas.zcrypto as zcrypto
+from . import zgrab2
+
+# modules/banner/scanner.go - Results
+banner_scan_response = SubRecord({
+    "result": SubRecord({
+        "banner": String(),
+        "length": Unsigned32BitInteger()
+    })
+}, extends=zgrab2.base_scan_response)
+
+zschema.registry.register_schema("zgrab2-banner", banner_scan_response)
+
+zgrab2.register_scan_response_type("banner", banner_scan_response)