chore: Refactor WazuhRuleExclusionRequest validator in ai.py

socfortress · May 15, 2024 · 312223b · 312223b
1 parent da58e8f
commit 312223b
Showing 1 changed file with 23 additions and 11 deletions.
diff --git a/app/schema/ai.py b/app/schema/ai.py
@@ -77,17 +77,29 @@ def check_integration(cls, v):
         return v
 
     @validator("prompt")
-    def check_rule_group1(cls, v):
-        if "rule_group1" not in v:
-            raise HTTPException(
-                status_code=400,
-                detail="Missing 'rule_group1' in prompt.",
-            )
-        if v["rule_group1"] != "windows":
-            raise HTTPException(
-                status_code=400,
-                detail="Invalid 'rule_group1'. Only 'windows' is supported.",
-            )
+    def check_rule_group(cls, v):
+        if "rule_group3" in v:
+            if "rule_group1" not in v and "rule_group3" not in v:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Missing 'rule_group1' or 'rule_group3' in prompt.",
+                )
+            if ("rule_group1" in v and v["rule_group1"] != "windows") and ("rule_group3" in v and v["rule_group3"] != "windows"):
+                raise HTTPException(
+                    status_code=400,
+                    detail="Invalid 'rule_group1' or 'rule_group3'. At least one must be 'windows'.",
+                )
+        else:
+            if "rule_group1" not in v:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Missing 'rule_group1' in prompt.",
+                )
+            if v["rule_group1"] != "windows":
+                raise HTTPException(
+                    status_code=400,
+                    detail="Invalid 'rule_group1'. Only 'windows' is supported.",
+                )
         return v