From 075924400f0ff10e351851919c568dfb248f44ca Mon Sep 17 00:00:00 2001
From: justinbastress <33579608+justinbastress@users.noreply.github.com>
Date: Thu, 1 Mar 2018 14:32:12 -0500
Subject: [PATCH] Add shared FlagsToSet function in output.go (#62)

* Add shared FlagsToSet function in output.go, use it in mysql. Add examples / tests.

* Add utility functions to widen map keys
---
 lib/mysql/mysql.go |  22 ++------
 output.go          | 124 +++++++++++++++++++++++++++++++++++++++++++++
 output_test.go     |  77 ++++++++++++++++++++++++++++
 3 files changed, 206 insertions(+), 17 deletions(-)
 create mode 100644 output.go
 create mode 100644 output_test.go

diff --git a/lib/mysql/mysql.go b/lib/mysql/mysql.go
index d6fbdc7..af3fef8 100644
--- a/lib/mysql/mysql.go
+++ b/lib/mysql/mysql.go
@@ -17,11 +17,11 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
-	"math"
 	"net"
 	"strings"
 
 	log "github.com/sirupsen/logrus"
+	"github.com/zmap/zgrab2"
 )
 
 const (
@@ -92,20 +92,6 @@ type Config struct {
 	ReservedData       []byte
 }
 
-// flagsToSet() converts an integer flags variable to a set of consts corresponding to each bit.
-// The result is a map from the labels to bool (true).
-// Example: flagsToSet(0x12, { "a", "b", "c", "d", "e" }) returns { "b": true, "e": true }.
-func flagsToSet(flags uint64, consts []string) (ret map[string]bool) {
-	ret = make(map[string]bool)
-	for i, label := range consts {
-		v := uint64(math.Pow(2, float64(i)))
-		if uint64(flags)&v == v {
-			ret[label] = true
-		}
-	}
-	return ret
-}
-
 // GetServerStatusFlags returns a map[string]bool representation of the
 // given flags. The keys are the constant names defined in the MySQL
 // docs, and the values are true (flags that are not set have no
@@ -127,7 +113,8 @@ func GetServerStatusFlags(flags uint16) map[string]bool {
 		"SERVER_STATUS_IN_TRANS_READONLY",
 		"SERVER_SESSION_STATE_CHANGED",
 	}
-	return flagsToSet(uint64(flags), consts)
+	ret, _ := zgrab2.ListFlagsToSet(uint64(flags), consts)
+	return ret
 }
 
 // GetClientCapabilityFlags returns a map[string]bool representation of
@@ -162,7 +149,8 @@ func GetClientCapabilityFlags(flags uint32) map[string]bool {
 		"CLIENT_SESSION_TRACK",
 		"CLIENT_DEPRECATED_EOF",
 	}
-	return flagsToSet(uint64(flags), consts)
+	ret, _ := zgrab2.ListFlagsToSet(uint64(flags), consts)
+	return ret
 }
 
 // InitConfig fills in a (possibly newly-created) Config instance with
diff --git a/output.go b/output.go
new file mode 100644
index 0000000..85c11e8
--- /dev/null
+++ b/output.go
@@ -0,0 +1,124 @@
+package zgrab2
+
+import "fmt"
+
+// FlagMap is a function that maps a single-bit bitmask (i.e. a number of the
+// form (1 << x)) to a string representing that bit.
+// If the input is not valid / recognized, it should return a non-nil error,
+// which will cause the flag to be added to the "unknowns" list.
+type FlagMap func(uint64) (string, error)
+
+// MapFlagsToSet gets the "set" (map of strings to true) of values corresponding
+// to the bits in flags. For each bit i set in flags, the result will have
+// result[mapping(i << i)] = true.
+// Any bits for which the mapping returns a non-nil error are instead appended
+// to the unknowns list.
+func MapFlagsToSet(flags uint64, mapping FlagMap) (map[string]bool, []uint64) {
+	ret := make(map[string]bool)
+	unknowns := []uint64{}
+	for i := uint8(0); i < 64; i++ {
+		if flags == 0 {
+			break
+		}
+		bit := (flags & 1) << i
+		if bit > 0 {
+			str, err := mapping(bit)
+			if err != nil {
+				unknowns = append(unknowns, bit)
+			} else {
+				ret[str] = true
+			}
+		}
+		flags >>= 1
+	}
+	return ret, unknowns
+}
+
+// GetFlagMapFromMap returns a FlagMap function that uses mapping to do the
+// mapping. Values not present in the map are treated as unknown, and a non-nil
+// error is returned in those cases.
+func GetFlagMapFromMap(mapping map[uint64]string) FlagMap {
+	return func(bit uint64) (string, error) {
+		ret, ok := mapping[bit]
+		if ok {
+			return ret, nil
+		}
+		return "", fmt.Errorf("Unknown flag 0x%x", bit)
+	}
+}
+
+// GetFlagMapFromList returns a FlagMap function mapping the ith bit to the
+// ith entry of bits.
+// bits is a list of labels for the corresponding bits; any empty strings (and
+// bits beyond the end of the list) are treated as unknown.
+func GetFlagMapFromList(bits []string) FlagMap {
+	mapping := make(map[uint64]string)
+	for i, v := range bits {
+		if v != "" {
+			mapping[uint64(1)<<uint8(i)] = v
+		}
+	}
+	return GetFlagMapFromMap(mapping)
+}
+
+// FlagsToSet converts an integer flags variable to a set of string labels
+// corresponding to each bit, in the format described by the wiki (see
+// https://github.com/zmap/zgrab2/wiki/Scanner-details).
+// The mapping maps the bit mask value (i.e. a number of the form (1 << x)) to
+// the label for that bit.
+// Flags not present in mapping are appended to the unknown list.
+func FlagsToSet(flags uint64, mapping map[uint64]string) (map[string]bool, []uint64) {
+	mapper := GetFlagMapFromMap(mapping)
+	return MapFlagsToSet(flags, mapper)
+}
+
+// ListFlagsToSet converts an integer flags variable to a set of string labels
+// corresponding to each bit, in the format described by the wiki (see
+// https://github.com/zmap/zgrab2/wiki/Scanner-details).
+// The ith entry of labels gives the label for the ith bit (i.e. flags & (1<<i)).
+// Empty strings in labels are treated as unknown, as are bits beyond the end
+// of the list. Unknown flags are appended to the unknown list.
+func ListFlagsToSet(flags uint64, labels []string) (map[string]bool, []uint64) {
+	mapper := GetFlagMapFromList(labels)
+	return MapFlagsToSet(flags, mapper)
+}
+
+// WidenMapKeys8 copies a map with uint8 keys into an equivalent map with uint64
+// keys for use in the FlagsToSet function.
+func WidenMapKeys8(input map[uint8]string) map[uint64]string {
+	ret := make(map[uint64]string, len(input))
+	for k, v := range input {
+		ret[uint64(k)] = v
+	}
+	return ret
+}
+
+// WidenMapKeys16 copies a map with uint8 keys into an equivalent map with
+// uint64 keys for use in the FlagsToSet function.
+func WidenMapKeys16(input map[uint16]string) map[uint64]string {
+	ret := make(map[uint64]string, len(input))
+	for k, v := range input {
+		ret[uint64(k)] = v
+	}
+	return ret
+}
+
+// WidenMapKeys32 copies a map with uint8 keys into an equivalent map with
+// uint64 keys for use in the FlagsToSet function.
+func WidenMapKeys32(input map[uint32]string) map[uint64]string {
+	ret := make(map[uint64]string, len(input))
+	for k, v := range input {
+		ret[uint64(k)] = v
+	}
+	return ret
+}
+
+// WidenMapKeys copies a map with int keys into an equivalent map with uint64
+// keys for use in the FlagsToSet function.
+func WidenMapKeys(input map[int]string) map[uint64]string {
+	ret := make(map[uint64]string, len(input))
+	for k, v := range input {
+		ret[uint64(k)] = v
+	}
+	return ret
+}
diff --git a/output_test.go b/output_test.go
new file mode 100644
index 0000000..f8fd240
--- /dev/null
+++ b/output_test.go
@@ -0,0 +1,77 @@
+package zgrab2
+
+import (
+	"fmt"
+)
+
+func ExampleMapFlagsToSet_success() {
+	output, unknowns := MapFlagsToSet(0xb, func(bit uint64) (string, error) {
+		return fmt.Sprintf("bit0x%01x", bit), nil
+	})
+	for k, v := range output {
+		fmt.Printf("%s: %v\n", k, v)
+	}
+	for _, v := range unknowns {
+		fmt.Printf("Unknown: 0x%01x", v)
+	}
+	// Unordered Output:
+	// bit0x1: true
+	// bit0x2: true
+	// bit0x8: true
+}
+
+func ExampleMapFlagsToSet_error() {
+	output, unknowns := MapFlagsToSet(0x1b, func(bit uint64) (string, error) {
+		if bit < 0x10 {
+			return fmt.Sprintf("bit0x%01x", bit), nil
+		} else {
+			return "", fmt.Errorf("Unrecognized flag 0x%02x", bit)
+		}
+	})
+	for k, v := range output {
+		fmt.Printf("%s: %v\n", k, v)
+	}
+	for _, v := range unknowns {
+		fmt.Printf("Unknown: 0x%02x", v)
+	}
+	// Unordered Output:
+	// bit0x1: true
+	// bit0x2: true
+	// bit0x8: true
+	// Unknown: 0x10
+}
+
+func ExampleFlagsToSet() {
+	output, unknowns := FlagsToSet(0x5, WidenMapKeys(map[int]string{
+		0x1: "bit0",
+		0x2: "bit1",
+		0x8: "bit3",
+	}))
+	for k, v := range output {
+		fmt.Printf("%s: %v\n", k, v)
+	}
+	for _, v := range unknowns {
+		fmt.Printf("Unknown: 0x%01x", v)
+	}
+	// Unordered Output:
+	// bit0: true
+	// Unknown: 0x4
+}
+
+func ExampleListFlagsToSet() {
+	output, unknowns := ListFlagsToSet(0x5, []string{
+		"bit0",
+		"bit1",
+		"",
+		"bit3",
+	})
+	for k, v := range output {
+		fmt.Printf("%s: %v\n", k, v)
+	}
+	for _, v := range unknowns {
+		fmt.Printf("Unknown: 0x%01x", v)
+	}
+	// Unordered Output:
+	// bit0: true
+	// Unknown: 0x4
+}