Add operator to SigmaStatus and SigmaLevel

.github/workflows/test.yml

@@ -18,7 +18,7 @@ jobs:
       - name: Install Poetry
         run: pipx install poetry
       - name: Set up Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
           cache: poetry

sigma/rule.py

@@ -37,11 +37,47 @@ def __str__(self) -> str:
 
 
 class SigmaStatus(EnumLowercaseStringMixin, Enum):
-    STABLE = auto()
+    UNSUPPORTED = auto()
+    DEPRECATED = auto()
     EXPERIMENTAL = auto()
     TEST = auto()
-    DEPRECATED = auto()
-    UNSUPPORTED = auto()
+    STABLE = auto()
+
+    def __eq__(self, other):
+        if isinstance(other, SigmaStatus):
+            return self.value == other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaStatus", source=other)
+
+    def __ge__(self, other):
+        if isinstance(other, SigmaStatus):
+            return self.value >= other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaStatus", source=other)
+
+    def __gt__(self, other):
+        if isinstance(other, SigmaStatus):
+            return self.value > other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaStatus", source=other)
+
+    def __ne__(self, other):
+        if isinstance(other, SigmaStatus):
+            return self.value != other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaStatus", source=other)
+
+    def __le__(self, other):
+        if isinstance(other, SigmaStatus):
+            return self.value <= other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaStatus", source=other)
+
+    def __lt__(self, other):
+        if isinstance(other, SigmaStatus):
+            return self.value < other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaStatus", source=other)
 
 
 class SigmaLevel(EnumLowercaseStringMixin, Enum):
@@ -51,6 +87,42 @@ class SigmaLevel(EnumLowercaseStringMixin, Enum):
     HIGH = auto()
     CRITICAL = auto()
 
+    def __eq__(self, other):
+        if isinstance(other, SigmaLevel):
+            return self.value == other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaLevel", source=other)
+
+    def __ge__(self, other):
+        if isinstance(other, SigmaLevel):
+            return self.value >= other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaLevel", source=other)
+
+    def __gt__(self, other):
+        if isinstance(other, SigmaLevel):
+            return self.value > other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaLevel", source=other)
+
+    def __ne__(self, other):
+        if isinstance(other, SigmaLevel):
+            return self.value != other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaLevel", source=other)
+
+    def __le__(self, other):
+        if isinstance(other, SigmaLevel):
+            return self.value <= other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaLevel", source=other)
+
+    def __lt__(self, other):
+        if isinstance(other, SigmaLevel):
+            return self.value < other.value
+
+        raise sigma_exceptions.SigmaTypeError("Must be a SigmaLevel", source=other)
+
 
 class SigmaRelatedType(EnumLowercaseStringMixin, Enum):
     OBSOLETES = auto()

tests/test_rule.py

@@ -45,10 +45,48 @@ def test_sigmalevel_str():
     assert str(SigmaLevel.MEDIUM) == "medium"
 
 
+def test_sigmalevel_comparation():
+    assert SigmaLevel.HIGH == SigmaLevel.HIGH
+    assert SigmaLevel.HIGH >= SigmaLevel.LOW
+    assert SigmaLevel.HIGH > SigmaLevel.LOW
+    assert SigmaLevel.HIGH != SigmaLevel.LOW
+    assert SigmaLevel.LOW <= SigmaLevel.HIGH
+    assert SigmaLevel.LOW < SigmaLevel.HIGH
+
+
+def test_sigmalevel_comparation_invalid():
+    with pytest.raises(sigma_exceptions.SigmaTypeError, match="Must be a SigmaLevel .*"):
+        assert SigmaLevel.HIGH == "HIGH"
+        assert SigmaLevel.HIGH >= "LOW"
+        assert SigmaLevel.HIGH > "LOW"
+        assert SigmaLevel.HIGH != "LOW"
+        assert SigmaLevel.LOW <= "HIGH"
+        assert SigmaLevel.LOW < "HIGH"
+
+
 def test_sigmastatus_str():
     assert str(SigmaStatus.STABLE) == "stable"
 
 
+def test_sigmastatus_comparation():
+    assert SigmaStatus.STABLE == SigmaStatus.STABLE
+    assert SigmaStatus.STABLE >= SigmaStatus.EXPERIMENTAL
+    assert SigmaStatus.STABLE > SigmaStatus.EXPERIMENTAL
+    assert SigmaStatus.STABLE != SigmaStatus.EXPERIMENTAL
+    assert SigmaStatus.EXPERIMENTAL <= SigmaStatus.STABLE
+    assert SigmaStatus.EXPERIMENTAL < SigmaStatus.STABLE
+
+
+def test_sigmastatus_comparation_invalid():
+    with pytest.raises(sigma_exceptions.SigmaTypeError, match="Must be a SigmaStatus .*"):
+        assert SigmaStatus.STABLE == "STABLE"
+        assert SigmaStatus.STABLE >= "EXPERIMENTAL"
+        assert SigmaStatus.STABLE > "EXPERIMENTAL"
+        assert SigmaStatus.STABLE != "EXPERIMENTAL"
+        assert SigmaStatus.EXPERIMENTAL <= "STABLE"
+        assert SigmaStatus.EXPERIMENTAL < "STABLE"
+
+
 ### SigmaRuleTag tests ###
 def test_sigmaruletag_fromstr():
     assert SigmaRuleTag.from_str("namespace.name") == SigmaRuleTag("namespace", "name")