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
This commit is contained in:
parent
2ec074a082
commit
075924400f
|
@ -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
|
||||
|
|
|
@ -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
|
||||
}
|
|
@ -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
|
||||
}
|
Loading…
Reference in New Issue