From 784dd1090b7cb41a62bce17fff28bcc8421fafa3 Mon Sep 17 00:00:00 2001 From: forgedhallpass <13679401+forgedhallpass@users.noreply.github.com> Date: Thu, 8 Jul 2021 13:28:44 +0300 Subject: [PATCH] RES-84 # Improve Nuclei CLI interface * "enum" for safely working with severities --- v2/internal/severity/severity.go | 132 ++++++++++++++++++++++++++ v2/internal/severity/severity_test.go | 91 ++++++++++++++++++ 2 files changed, 223 insertions(+) create mode 100644 v2/internal/severity/severity.go create mode 100644 v2/internal/severity/severity_test.go diff --git a/v2/internal/severity/severity.go b/v2/internal/severity/severity.go new file mode 100644 index 000000000..b947f1a90 --- /dev/null +++ b/v2/internal/severity/severity.go @@ -0,0 +1,132 @@ +package severity + +import ( + "encoding/json" + "github.com/pkg/errors" + "strings" +) + +type Severity int + +const ( + Info Severity = iota + Low + Medium + High + Critical +) + +var severityMappings = map[Severity]string{ + Info: "info", + Low: "low", + Medium: "medium", + High: "high", + Critical: "critical", +} + +type SeverityStruct struct { + Key Severity +} + +func toSeverity(valueToMap string) (Severity, error) { + for key, currentValue := range severityMappings { + if normalizeValue(valueToMap) == currentValue { + return key, nil + } + } + return -1, errors.New("Invalid severity: " + valueToMap) +} + +func GetSupportedSeverities() []Severity { + var result []Severity + for key := range severityMappings { + result = append(result, key) + } + return result +} + +func (severity SeverityStruct) MarshalYAML() (interface{}, error) { + if value, found := severityMappings[severity.Key]; found { + return &struct{ Key string }{value}, nil + } else { + panic("Invalid field to marshall") + } +} + +func (severity SeverityStruct) MarshalJSON() ([]byte, error) { + if value, found := severityMappings[severity.Key]; found { + return json.Marshal(&struct{ Key string }{value}) + } else { + panic("Invalid field to marshall") + } +} + +func (severity *SeverityStruct) UnmarshalYAML(unmarshal func(interface{}) error) error { + var objMap map[string]string + if err := unmarshal(&objMap); err != nil { + return err + } + + return mapToSeverity(objMap, severity) +} + +func mapToSeverity(objMap map[string]string, severity *SeverityStruct) error { + stringSeverity := getFirstElement(objMap) + if readableSeverity, err := toSeverity(stringSeverity); err == nil { + severity = &SeverityStruct{readableSeverity} + return nil + } else { + return err + } +} + +func (severity *SeverityStruct) UnmarshalJSON(data []byte) error { + var objMap map[string]string + if err := json.Unmarshal(data, &objMap); err != nil { + return err + } + + return mapToSeverity(objMap, severity) +} + +func normalizeValue(value string) string { + return strings.TrimSpace(strings.ToLower(value)) +} + +func getFirstElement(stringMap map[string]string) string { + var result string + for _, value := range stringMap { + result = value + break + } + return result +} + +/* Alternative implementation +func (severity *SeverityStruct) UnmarshalJSON(data []byte) error { + var objMap map[string]*json.RawMessage + if err := json.Unmarshal(data, &objMap); err != nil { + return err + } + severityStructFirstFieldName := reflect.Indirect(reflect.ValueOf(severity)).Type().Field(0).Name + + var stringSeverity string + if err := json.Unmarshal(*objMap[severityStructFirstFieldName], &stringSeverity); err != nil { + return err + } + + if readableSeverity, err := toSeverity(stringSeverity); err == nil { + severity = &SeverityStruct{readableSeverity} + return nil + } else { + return err + } +}*/ + +func (severity Severity) normalize() string { + return strings.TrimSpace(strings.ToLower(severity.String())) +} + +func (severity Severity) String() string { + return severityMappings[severity] +} diff --git a/v2/internal/severity/severity_test.go b/v2/internal/severity/severity_test.go new file mode 100644 index 000000000..ed9bca12a --- /dev/null +++ b/v2/internal/severity/severity_test.go @@ -0,0 +1,91 @@ +package severity + +import ( + "encoding/json" + "fmt" + "github.com/stretchr/testify/assert" + "gopkg.in/yaml.v2" + "testing" +) + +func TestJsonUnmarshal(t *testing.T) { + testUnmarshal(t, json.Unmarshal, createJson) +} + +func TestYamlUnmarshal(t *testing.T) { + testUnmarshal(t, yaml.Unmarshal, createYaml) +} + +func TestJsonUnmarshalFail(t *testing.T) { + testUnmarshalFail(t, json.Unmarshal, createJson) +} + +func TestYamlUnmarshalFail(t *testing.T) { + testUnmarshalFail(t, yaml.Unmarshal, createYaml) +} + +func TestJsonMarshalFails(t *testing.T) { + testMarshallerFails(t, json.Marshal) +} + +func TestYamlMarshalFails(t *testing.T) { + testMarshallerFails(t, yaml.Marshal) +} + +func TestJsonMarshalSucceed(t *testing.T) { + testMarshal(t, json.Marshal, createJson) +} + +func TestYamlMarshal(t *testing.T) { + testMarshal(t, yaml.Marshal, createYaml) +} + +func testUnmarshal(t *testing.T, unmarshaler func(data []byte, v interface{}) error, payloadCreator func(value string) string) { + payloads := [...]string{ + payloadCreator("Info"), + payloadCreator("info"), + payloadCreator("inFo "), + payloadCreator("infO "), + payloadCreator(" INFO "), + } + + for _, payload := range payloads { + t.Run(payload, func(t *testing.T) { + result := unmarshal(payload, unmarshaler) + assert.Equal(t, result.Key, Info) + assert.Equal(t, result.Key.String(), "info") + }) + } +} + +func testMarshal(t *testing.T, marshaller func(v interface{}) ([]byte, error), payloadCreator func(value string) string) { + for _, severity := range GetSupportedSeverities() { + result, _ := marshaller(&SeverityStruct{Key: severity}) + assert.Equal(t, string(result), payloadCreator(severity.String())) + } +} + +func testUnmarshalFail(t *testing.T, unmarshaler func(data []byte, v interface{}) error, payloadCreator func(value string) string) bool { + return assert.Panics(t, func() { unmarshal(payloadCreator("invalid"), unmarshaler) }) +} + +func testMarshallerFails(t *testing.T, marshaller func(v interface{}) ([]byte, error)) { + assert.Panics(t, func() { marshaller(&SeverityStruct{Key: 13}) }) +} + +func unmarshal(value string, unmarshaller func(data []byte, v interface{}) error) SeverityStruct { + severityStruct := SeverityStruct{} + var err = unmarshaller([]byte(value), &severityStruct) + if err != nil { + panic(err) + } + return severityStruct +} + +func createJson(severityString string) string { + return fmt.Sprintf(`{"Key":"%s"}`, severityString) +} + +func createYaml(value string) string { + return "key: " + value + "\n" +}