RES-84 # Improve Nuclei CLI interface

* "enum" for safely working with severities

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]
+}

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"
+}