diff --git a/pkg/models/severity.go b/pkg/models/severity.go index a66a39f..863480d 100644 --- a/pkg/models/severity.go +++ b/pkg/models/severity.go @@ -7,3 +7,32 @@ const ( SeverityLow = "LOW" SeverityUnknown = "UNKNOWN" ) + +var severityToInt = map[string]int{ + SeverityUnknown: 0, + SeverityLow: 1, + SeverityMedium: 2, + SeverityHigh: 3, + SeverityCritical: 4, +} + +func CompareSeverity(severity1, severity2 string) int { + val1, ok1 := severityToInt[severity1] + val2, ok2 := severityToInt[severity2] + + if !ok1 && !ok2 { + return 0 + } else if !ok1 { + return -1 + } else if !ok2 { + return 1 + } + + if val1 < val2 { + return -1 + } else if val1 > val2 { + return 1 + } else { + return 0 + } +} diff --git a/pkg/models/severity_test.go b/pkg/models/severity_test.go new file mode 100644 index 0000000..9f1dd23 --- /dev/null +++ b/pkg/models/severity_test.go @@ -0,0 +1,35 @@ +package models + +import "testing" + +func TestCompareSeverity(t *testing.T) { + tests := []struct { + severity1 string + severity2 string + expected int + }{ + {"INVALID", "HIGH", -1}, + {"UNKNOWN", "CRITICAL", -1}, + {"LOW", "CRITICAL", -1}, + {"MEDIUM", "CRITICAL", -1}, + {"HIGH", "CRITICAL", -1}, + {"CRITICAL", "HIGH", 1}, + {"HIGH", "MEDIUM", 1}, + {"MEDIUM", "LOW", 1}, + {"LOW", "UNKNOWN", 1}, + {"CRITICAL", "UNKNOWN", 1}, + {"HIGH", "INVALID", 1}, + {"UNKNOWN", "UNKNOWN", 0}, + {"LOW", "LOW", 0}, + {"MEDIUM", "MEDIUM", 0}, + {"HIGH", "HIGH", 0}, + {"CRITICAL", "CRITICAL", 0}, + } + + for _, test := range tests { + result := CompareSeverity(test.severity1, test.severity2) + if result != test.expected { + t.Errorf("CompareSeverity(%q, %q) = %d; expected %d", test.severity1, test.severity2, result, test.expected) + } + } +}