From 66dcb9c21e72633c877c76e1da6cb3fc559f34b9 Mon Sep 17 00:00:00 2001
From: frack113 <62423083+frack113@users.noreply.github.com>
Date: Tue, 13 Aug 2024 09:21:59 +0200
Subject: [PATCH] Update SigmaLogSource class

---
 pyproject.toml                      |  3 +++
 sigma/rule.py                       | 15 +++++++++++++++
 sigma/validators/core/logsources.py |  1 -
 tests/test_rule.py                  | 22 ++++++++++++++++++++++
 4 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index c1d6c54..816129a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,3 +45,6 @@ line-length = 100
 [build-system]
 requires = ["poetry-core>=1.8.1"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.pytest.ini_options]
+filterwarnings = ["ignore::DeprecationWarning"]
diff --git a/sigma/rule.py b/sigma/rule.py
index a91b116..00ababd 100644
--- a/sigma/rule.py
+++ b/sigma/rule.py
@@ -216,6 +216,7 @@ class SigmaLogSource:
     category: Optional[str] = field(default=None)
     product: Optional[str] = field(default=None)
     service: Optional[str] = field(default=None)
+    definition: Optional[str] = field(default=None)
     source: Optional[SigmaRuleLocation] = field(default=None, compare=False)
     custom_attributes: Optional[Dict[str, Any]] = field(default=None, compare=False)
 
@@ -225,6 +226,18 @@ def __post_init__(self):
             raise sigma_exceptions.SigmaLogsourceError(
                 "Sigma log source can't be empty", source=self.source
             )
+        if self.category and not isinstance(self.category, str):
+            raise sigma_exceptions.SigmaLogsourceError(
+                "Sigma log source category must be string", source=self.source
+            )
+        if self.product and not isinstance(self.product, str):
+            raise sigma_exceptions.SigmaLogsourceError(
+                "Sigma log source product must be string", source=self.source
+            )
+        if self.service and not isinstance(self.service, str):
+            raise sigma_exceptions.SigmaLogsourceError(
+                "Sigma log source service must be string", source=self.source
+            )
 
     @classmethod
     def from_dict(
@@ -234,10 +247,12 @@ def from_dict(
         custom_attributes = {
             k: v for k, v in logsource.items() if k not in set(cls.__dataclass_fields__.keys())
         }
+
         return cls(
             logsource.get("category"),
             logsource.get("product"),
             logsource.get("service"),
+            logsource.get("definition"),
             source,
             custom_attributes if len(custom_attributes) > 0 else None,
         )
diff --git a/sigma/validators/core/logsources.py b/sigma/validators/core/logsources.py
index e076d08..798c80d 100644
--- a/sigma/validators/core/logsources.py
+++ b/sigma/validators/core/logsources.py
@@ -115,7 +115,6 @@ def validate(self, rule: SigmaRule) -> List[SigmaValidationIssue]:
             return [
                 FieldnameLogsourceIssue(rules=[rule], fieldname=name)
                 for name in rule.logsource.custom_attributes
-                if not name == "definition"
             ]
         else:
             return []
diff --git a/tests/test_rule.py b/tests/test_rule.py
index ee994cf..28662f2 100644
--- a/tests/test_rule.py
+++ b/tests/test_rule.py
@@ -147,6 +147,28 @@ def test_sigmalogsource_fromdict_no_service():
     assert logsource == SigmaLogSource("category-id", "product-id", None)
 
 
+def test_sigmalogsource_fromdict_definition():
+    logsource = SigmaLogSource.from_dict(
+        {"category": "category-id", "product": "product-id", "definition": "use it"}
+    )
+    assert logsource == SigmaLogSource("category-id", "product-id", None, "use it")
+
+
+def test_sigmalogsource_fromdict_category_not_str():
+    with pytest.raises(sigma_exceptions.SigmaLogsourceError):
+        SigmaLogSource.from_dict({"category": 1234, "product": "product-id"})
+
+
+def test_sigmalogsource_fromdict_product_not_str():
+    with pytest.raises(sigma_exceptions.SigmaLogsourceError):
+        SigmaLogSource.from_dict({"category": "category-id", "product": {"a": "b"}})
+
+
+def test_sigmalogsource_fromdict_service_not_str():
+    with pytest.raises(sigma_exceptions.SigmaLogsourceError):
+        SigmaLogSource.from_dict({"category": "category-id", "service": ["1", "2", "3"]})
+
+
 def test_sigmalogsource_empty():
     with pytest.raises(sigma_exceptions.SigmaLogsourceError, match="can't be empty.*test.yml"):
         SigmaLogSource(None, None, None, source=sigma_exceptions.SigmaRuleLocation("test.yml"))