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/hex"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"math"
|
|
||||||
"net"
|
"net"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
log "github.com/sirupsen/logrus"
|
log "github.com/sirupsen/logrus"
|
||||||
|
"github.com/zmap/zgrab2"
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
|
@ -92,20 +92,6 @@ type Config struct {
|
||||||
ReservedData []byte
|
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
|
// GetServerStatusFlags returns a map[string]bool representation of the
|
||||||
// given flags. The keys are the constant names defined in the MySQL
|
// 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
|
// 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_STATUS_IN_TRANS_READONLY",
|
||||||
"SERVER_SESSION_STATE_CHANGED",
|
"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
|
// GetClientCapabilityFlags returns a map[string]bool representation of
|
||||||
|
@ -162,7 +149,8 @@ func GetClientCapabilityFlags(flags uint32) map[string]bool {
|
||||||
"CLIENT_SESSION_TRACK",
|
"CLIENT_SESSION_TRACK",
|
||||||
"CLIENT_DEPRECATED_EOF",
|
"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
|
// 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