Add unique field support

Author: eric-forte-elastic

detection_rules/remote_validation.py

@@ -19,6 +19,7 @@
 from .config import load_current_package_version
 from .misc import ClientError, get_elasticsearch_client, get_kibana_client, getdefault
 from .rule import TOMLRule, TOMLRuleContents
+from .rule_validators import ESQLValidator
 from .schemas import definitions
 
 
@@ -180,18 +181,14 @@ def request(c: TOMLRuleContents) -> None:
     def validate_esql(self, contents: TOMLRuleContents) -> dict[str, Any]:
         query = contents.data.query  # type: ignore[reportAttributeAccessIssue]
         rule_id = contents.data.rule_id
-        headers = {"accept": "application/json", "content-type": "application/json"}
-        body = {"query": f"{query} | LIMIT 0"}
         if not self.es_client:
             raise ValueError("No ES client found")
+
+        if not self.kibana_client:
+            raise ValueError("No Kibana client found")
         try:
-            response = self.es_client.perform_request(
-                "POST",
-                "/_query",
-                headers=headers,
-                params={"pretty": True},
-                body=body,
-            )
+            validator = ESQLValidator(contents.data.query)  # type: ignore[reportIncompatibleMethodOverride]
+            response = validator.remote_validate_rule_contents(self.kibana_client, self.es_client, contents)
         except Exception as exc:
             if isinstance(exc, elasticsearch.BadRequestError):
                 raise ValidationError(f"ES|QL query failed: {exc} for rule: {rule_id}, query: \n{query}") from exc

detection_rules/rule.py

@@ -647,6 +647,10 @@ def unique_fields(self) -> Any:
     def validate(self, _: "QueryRuleData", __: RuleMeta) -> None:
         raise NotImplementedError
 
+    def get_unique_field_type(self, __: str) -> None:
+        """Used to get unique field types when schema is not used"""
+        raise NotImplementedError
+
     @cached
     def get_required_fields(self, index: str) -> list[dict[str, Any]]:
         """Retrieves fields needed for the query along with type information from the schema."""
@@ -663,7 +667,9 @@ def get_required_fields(self, index: str) -> list[dict[str, Any]]:
         # construct integration schemas
         packages_manifest = load_integrations_manifests()
         integrations_schemas = load_integrations_schemas()
-        datasets, _ = beats.get_datasets_and_modules(self.ast)
+        datasets: set[str] = set()
+        if self.ast:
+            datasets, _ = beats.get_datasets_and_modules(self.ast)
         package_integrations = parse_datasets(list(datasets), packages_manifest)
         int_schema: dict[str, Any] = {}
         data = {"notify": False}
@@ -691,6 +697,9 @@ def get_required_fields(self, index: str) -> list[dict[str, Any]]:
                 elif endgame_schema:
                     field_type = endgame_schema.endgame_schema.get(fld, None)
 
+            if not field_type and isinstance(self, ESQLValidator):
+                field_type = self.get_unique_field_type(fld)
+
             required.append({"name": fld, "type": field_type or "unknown", "ecs": is_ecs})
 
         return sorted(required, key=lambda f: f["name"])

detection_rules/rule_validators.py

@@ -16,6 +16,7 @@
 import click
 import eql  # type: ignore[reportMissingTypeStubs]
 import kql  # type: ignore[reportMissingTypeStubs]
+from elastic_transport import ObjectApiResponse
 from elasticsearch import Elasticsearch  # type: ignore[reportMissingTypeStubs]
 from eql import ast  # type: ignore[reportMissingTypeStubs]
 from eql.parser import KvTree, LarkToEQL, NodeInfo, TypeHint  # type: ignore[reportMissingTypeStubs]
@@ -618,18 +619,29 @@ def validate_rule_type_configurations(self, data: EQLRuleData, meta: RuleMeta) -
 class ESQLValidator(QueryValidator):
     """Validate specific fields for ESQL query event types."""
 
-    esql_unique_fields: list[str]
+    def __init__(self, query: str) -> None:
+        """Initialize the ESQLValidator with the given query."""
+        super().__init__(query)
+        self.esql_unique_fields: list[dict[str, str]] = []
 
     @cached_property
     def ast(self) -> None:  # type: ignore[reportIncompatibleMethodOverride]
+        """There is no AST for ESQL until we have an ESQL parser."""
         return None
 
     @cached_property
     def unique_fields(self) -> list[str]:  # type: ignore[reportIncompatibleMethodOverride]
         """Return a list of unique fields in the query. Requires remote validation to have occurred."""
-        if not self.esql_unique_fields:
-            return []
-        return self.esql_unique_fields
+        if self.esql_unique_fields:
+            return [field["name"] for field in self.esql_unique_fields]
+        return []
+
+    def get_unique_field_type(self, field_name: str) -> str | None:  # type: ignore[reportIncompatibleMethodOverride]
+        """Get the type of the unique field. Requires remote validation to have occurred."""
+        for field in self.esql_unique_fields:
+            if field["name"] == field_name:
+                return field["type"]
+        return None
 
     def validate(self, rule_data: "QueryRuleData", rule_meta: RuleMeta) -> None:  # type: ignore[reportIncompatibleMethodOverride]
         """Validate an ESQL query while checking TOMLRule."""
@@ -648,7 +660,7 @@ def validate(self, rule_data: "QueryRuleData", rule_meta: RuleMeta) -> None:  #
                 elasticsearch_url=misc.getdefault("elasticsearch_url")(),
                 ignore_ssl_errors=misc.getdefault("ignore_ssl_errors")(),
             )
-            self.remote_validate_rule(
+            _ = self.remote_validate_rule(
                 kibana_client,
                 elastic_client,
                 rule_data.query,
@@ -774,7 +786,7 @@ def execute_query_against_indices(
         test_index_str: str,
         log: Callable[[str], None],
         delete_indices: bool = True,
-    ) -> list[Any]:
+    ) -> tuple[list[Any], ObjectApiResponse[Any]]:
         """Execute the ESQL query against the test indices on a remote Stack and return the columns."""
         try:
             log(f"Executing a query against `{test_index_str}`")
@@ -789,7 +801,7 @@ def execute_query_against_indices(
 
         query_column_names = [c["name"] for c in query_columns]
         log(f"Got query columns: {', '.join(query_column_names)}")
-        return query_columns
+        return query_columns, response
 
     def find_nested_multifields(self, mapping: dict[str, Any], path: str = "") -> list[Any]:
         """Recursively search for nested multi-fields in Elasticsearch mappings."""
@@ -886,9 +898,9 @@ def prepare_mappings(
 
     def remote_validate_rule_contents(
         self, kibana_client: Kibana, elastic_client: Elasticsearch, contents: TOMLRuleContents, verbosity: int = 0
-    ) -> None:
+    ) -> ObjectApiResponse[Any]:
         """Remote validate a rule's ES|QL query using an Elastic Stack."""
-        self.remote_validate_rule(
+        return self.remote_validate_rule(
             kibana_client=kibana_client,
             elastic_client=elastic_client,
             query=contents.data.query,  # type: ignore[reportUnknownVariableType]
@@ -905,7 +917,7 @@ def remote_validate_rule(  # noqa: PLR0913
         metadata: RuleMeta,
         rule_id: str = "",
         verbosity: int = 0,
-    ) -> None:
+    ) -> ObjectApiResponse[Any]:
         """Uses remote validation from an Elastic Stack to validate ES|QL a given rule"""
 
         def log(val: str) -> None:
@@ -939,7 +951,7 @@ def log(val: str) -> None:
         # Replace all sources with the test indices
         query = query.replace(indices_str, full_index_str)  # type: ignore[reportUnknownVariableType]
 
-        query_columns = self.execute_query_against_indices(elastic_client, query, full_index_str, log)  # type: ignore[reportUnknownVariableType]
+        query_columns, response = self.execute_query_against_indices(elastic_client, query, full_index_str, log)  # type: ignore[reportUnknownVariableType]
         self.esql_unique_fields = query_columns
 
         # Validate that all fields (columns) are either dynamic fields or correctly mapped
@@ -949,6 +961,8 @@ def log(val: str) -> None:
         else:
             log("Dynamic column(s) have improper formatting.")
 
+        return response
+
 
 def extract_error_field(source: str, exc: eql.EqlParseError | kql.KqlParseError) -> str | None:
     """Extract the field name from an EQL or KQL parse error."""