tellaro-query-language 1.3.8

//! TQL AST to OpenSearch Query DSL translator.
//!
//! This module translates TQL abstract syntax trees into OpenSearch Query DSL.

use super::error::{OpenSearchError, Result};
use super::field_mappings::FieldMappings;
use crate::parser::{AstNode, CollectionOpNode, NslookupExprNode, Value as AstValue};
use serde_json::{json, Value as JsonValue};

/// Query builder for translating TQL to OpenSearch DSL
pub struct QueryBuilder {
    field_mappings: Option<FieldMappings>,
}

impl QueryBuilder {
    /// Create a new query builder
    ///
    /// # Arguments
    ///
    /// * `field_mappings` - Optional field mappings for intelligent query generation
    pub fn new(field_mappings: Option<FieldMappings>) -> Self {
        Self { field_mappings }
    }

    /// Convert AstValue to JsonValue
    fn ast_value_to_json(value: &AstValue) -> JsonValue {
        match value {
            AstValue::String(s) => json!(s),
            AstValue::Integer(i) => json!(i),
            AstValue::Float(f) => json!(f),
            AstValue::Boolean(b) => json!(b),
            AstValue::List(list) => {
                json!(list.iter().map(Self::ast_value_to_json).collect::<Vec<_>>())
            }
            AstValue::Null => json!(null),
        }
    }

    /// Build an OpenSearch query from a TQL AST
    ///
    /// # Arguments
    ///
    /// * `ast` - The TQL abstract syntax tree
    ///
    /// # Returns
    ///
    /// OpenSearch Query DSL as JSON
    ///
    /// # Example
    ///
    /// ```ignore
    /// use tql::parser::TqlParser;
    /// use tql::opensearch::QueryBuilder;
    ///
    /// let parser = TqlParser::new();
    /// let ast = parser.parse("age > 25 AND status eq 'active'").unwrap();
    /// let builder = QueryBuilder::new(None);
    /// let query = builder.build_query(&ast).unwrap();
    /// ```
    pub fn build_query(&self, ast: &AstNode) -> Result<JsonValue> {
        let query_clause = self.build_query_clause(ast)?;
        Ok(json!({
            "query": query_clause
        }))
    }

    /// Check if mutators require post-processing (cannot be pushed to OpenSearch)
    fn has_post_processing_mutators(mutators: &Option<Vec<crate::parser::Mutator>>) -> bool {
        mutators.as_ref().is_some_and(|mutators| {
            mutators.iter().any(|m| {
                let name = m.name.to_lowercase();
                // These mutators require post-processing - they transform values
                // and cannot be evaluated by OpenSearch
                matches!(
                    name.as_str(),
                    "is_global"
                        | "is_private"
                        | "is_multicast"
                        | "is_loopback"
                        | "is_link_local"
                        | "nslookup"
                        | "geoip"
                        | "geoip_lookup"
                        | "geo"
                        | "lowercase"
                        | "uppercase"
                        | "trim"
                        | "length"
                        | "split"
                        | "replace"
                        | "b64encode"
                        | "b64decode"
                        | "urldecode"
                        | "hexencode"
                        | "hexdecode"
                        | "md5"
                        | "sha256"
                        | "refang"
                        | "defang"
                )
            })
        })
    }

    fn build_query_clause(&self, node: &AstNode) -> Result<JsonValue> {
        match node {
            AstNode::Comparison(comp) => {
                // Check if there are post-processing mutators on the field
                // If so, we can only check that the field exists - the actual
                // filtering will be done in post-processing
                if Self::has_post_processing_mutators(&comp.field_mutators) {
                    // For mutators that require post-processing, return exists query
                    // The actual filtering (is_global eq true, etc.) happens after
                    // results are fetched from OpenSearch
                    return Ok(json!({
                        "exists": {
                            "field": &comp.field
                        }
                    }));
                }

                // Handle "exists" operator with no value (field-only expressions like `field | nslookup`)
                if comp.operator == "exists" && comp.value.is_none() {
                    // Check if there are enrichment mutators (nslookup, geoip, etc.)
                    // These require post-processing, so we return match_all to fetch all docs
                    let has_enrichment_mutator =
                        comp.field_mutators.as_ref().is_some_and(|mutators| {
                            mutators.iter().any(|m| {
                                let name = m.name.to_lowercase();
                                name == "nslookup" || name == "geoip" || name == "geoip_lookup"
                            })
                        });

                    if has_enrichment_mutator {
                        // Enrichment-only expression - return match_all for post-processing
                        return Ok(json!({
                            "match_all": {}
                        }));
                    } else {
                        // Regular exists check
                        return Ok(json!({
                            "exists": {
                                "field": &comp.field
                            }
                        }));
                    }
                }

                let value = comp.value.as_ref().ok_or_else(|| {
                    OpenSearchError::TranslationError("Comparison requires a value".to_string())
                })?;
                self.build_comparison(&comp.field, &comp.operator, value)
            }
            AstNode::LogicalOp(logical) => {
                self.build_logical(&logical.operator, &logical.left, &logical.right)
            }
            AstNode::UnaryOp(unary) => {
                let inner = self.build_query_clause(&unary.operand)?;
                Ok(json!({
                    "bool": {
                        "must_not": inner
                    }
                }))
            }
            AstNode::MatchAll => Ok(json!({
                "match_all": {}
            })),
            AstNode::CollectionOp(collection) => self.build_collection_op(collection),
            AstNode::NslookupExpr(nslookup) => self.build_nslookup_expr(nslookup),
            AstNode::GeoExpr(geo) => {
                // Geo expressions require post-processing, similar to nslookup
                // If there are conditions, use exists query on the field
                // If no conditions, return match_all
                if geo.conditions.is_some() {
                    Ok(json!({
                        "exists": {
                            "field": &geo.field
                        }
                    }))
                } else {
                    Ok(json!({
                        "match_all": {}
                    }))
                }
            }
            _ => Err(OpenSearchError::TranslationError(format!(
                "Unsupported AST node type: {:?}",
                node
            ))),
        }
    }

    fn build_comparison(&self, field: &str, operator: &str, value: &AstValue) -> Result<JsonValue> {
        // Convert AstValue to JsonValue
        let json_value = Self::ast_value_to_json(value);

        // Determine the actual field name to use (may include .keyword suffix)
        let query_field = self
            .field_mappings
            .as_ref()
            .map(|m| m.get_query_field(field, operator))
            .unwrap_or_else(|| field.to_string());

        match operator {
            "eq" => {
                // Use term for exact match
                Ok(json!({
                    "term": {
                        query_field: json_value
                    }
                }))
            }
            "ne" => {
                // Use bool must_not with term
                Ok(json!({
                    "bool": {
                        "must_not": {
                            "term": {
                                query_field: json_value
                            }
                        }
                    }
                }))
            }
            "gt" => Ok(json!({
                "range": {
                    query_field: {
                        "gt": json_value
                    }
                }
            })),
            "gte" => Ok(json!({
                "range": {
                    query_field: {
                        "gte": json_value
                    }
                }
            })),
            "lt" => Ok(json!({
                "range": {
                    query_field: {
                        "lt": json_value
                    }
                }
            })),
            "lte" => Ok(json!({
                "range": {
                    query_field: {
                        "lte": json_value
                    }
                }
            })),
            "contains" => {
                // Use wildcard or match_phrase depending on field type
                Ok(json!({
                    "wildcard": {
                        query_field: format!("*{}*", json_value.as_str().unwrap_or(""))
                    }
                }))
            }
            "startswith" => {
                // Use prefix query
                Ok(json!({
                    "prefix": {
                        query_field: json_value
                    }
                }))
            }
            "endswith" => {
                // Use wildcard with *value pattern
                Ok(json!({
                    "wildcard": {
                        query_field: format!("*{}", json_value.as_str().unwrap_or(""))
                    }
                }))
            }
            "matches" => {
                // Use regexp query
                Ok(json!({
                    "regexp": {
                        query_field: json_value
                    }
                }))
            }
            "in" => {
                // Use terms query
                let values = if let JsonValue::Array(arr) = &json_value {
                    arr.clone()
                } else {
                    vec![json_value.clone()]
                };
                Ok(json!({
                    "terms": {
                        query_field: values
                    }
                }))
            }
            "between" => {
                // Use range query with gte and lte
                if let JsonValue::Array(arr) = &json_value {
                    if arr.len() == 2 {
                        Ok(json!({
                            "range": {
                                query_field: {
                                    "gte": arr[0],
                                    "lte": arr[1]
                                }
                            }
                        }))
                    } else {
                        Err(OpenSearchError::TranslationError(
                            "between operator requires array of 2 values".to_string(),
                        ))
                    }
                } else {
                    Err(OpenSearchError::TranslationError(
                        "between operator requires array value".to_string(),
                    ))
                }
            }
            _ => Err(OpenSearchError::TranslationError(format!(
                "Unsupported operator: {}",
                operator
            ))),
        }
    }

    fn build_logical(&self, operator: &str, left: &AstNode, right: &AstNode) -> Result<JsonValue> {
        let left_clause = self.build_query_clause(left)?;
        let right_clause = self.build_query_clause(right)?;

        match operator.to_lowercase().as_str() {
            "and" => Ok(json!({
                "bool": {
                    "must": [left_clause, right_clause]
                }
            })),
            "or" => Ok(json!({
                "bool": {
                    "should": [left_clause, right_clause],
                    "minimum_should_match": 1
                }
            })),
            _ => Err(OpenSearchError::TranslationError(format!(
                "Unsupported logical operator: {}",
                operator
            ))),
        }
    }

    fn build_collection_op(&self, collection: &CollectionOpNode) -> Result<JsonValue> {
        let json_value = Self::ast_value_to_json(&collection.value);
        let operator = collection.operator.to_lowercase();
        let comparison_op = &collection.comparison_operator;

        // Build the inner comparison query
        let inner_query = match comparison_op.as_str() {
            "eq" => json!({
                "term": {
                    collection.field.clone(): json_value
                }
            }),
            "ne" => json!({
                "bool": {
                    "must_not": {
                        "term": {
                            collection.field.clone(): json_value
                        }
                    }
                }
            }),
            "gt" => json!({
                "range": {
                    collection.field.clone(): {
                        "gt": json_value
                    }
                }
            }),
            "gte" => json!({
                "range": {
                    collection.field.clone(): {
                        "gte": json_value
                    }
                }
            }),
            "lt" => json!({
                "range": {
                    collection.field.clone(): {
                        "lt": json_value
                    }
                }
            }),
            "lte" => json!({
                "range": {
                    collection.field.clone(): {
                        "lte": json_value
                    }
                }
            }),
            "contains" => json!({
                "wildcard": {
                    collection.field.clone(): format!("*{}*", json_value.as_str().unwrap_or(""))
                }
            }),
            _ => {
                return Err(OpenSearchError::TranslationError(format!(
                    "Unsupported collection comparison operator: {}",
                    comparison_op
                )));
            }
        };

        // Wrap in appropriate collection operator
        match operator.as_str() {
            "any" => {
                // ANY: At least one element matches (OpenSearch handles arrays automatically)
                Ok(inner_query)
            }
            "all" => {
                // ALL: All elements match - requires script query or post-processing
                // For now, we'll return a placeholder that indicates post-processing is needed
                Err(OpenSearchError::TranslationError(
                    "ALL collection operator requires post-processing".to_string(),
                ))
            }
            "none" => {
                // NONE: No elements match
                Ok(json!({
                    "bool": {
                        "must_not": inner_query
                    }
                }))
            }
            _ => Err(OpenSearchError::TranslationError(format!(
                "Unsupported collection operator: {}",
                operator
            ))),
        }
    }

    fn build_nslookup_expr(&self, nslookup: &NslookupExprNode) -> Result<JsonValue> {
        // Nslookup expressions are enrichment operations that require post-processing
        // If there are conditions (filters on the nslookup results), use exists query
        // If no conditions (just enrichment), return match_all to fetch all documents
        if nslookup.conditions.is_some() {
            Ok(json!({
                "exists": {
                    "field": &nslookup.field
                }
            }))
        } else {
            Ok(json!({
                "match_all": {}
            }))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::TqlParser;

    #[test]
    fn test_simple_equality() {
        let parser = TqlParser::new();
        let ast = parser.parse("status eq 'active'").unwrap();
        let builder = QueryBuilder::new(None);
        let query = builder.build_query(&ast).unwrap();

        assert_eq!(
            query,
            json!({
                "query": {
                    "term": {
                        "status": "active"
                    }
                }
            })
        );
    }

    #[test]
    fn test_range_query() {
        let parser = TqlParser::new();
        let ast = parser.parse("age > 25").unwrap();
        let builder = QueryBuilder::new(None);
        let query = builder.build_query(&ast).unwrap();

        assert_eq!(
            query,
            json!({
                "query": {
                    "range": {
                        "age": {
                            "gt": 25
                        }
                    }
                }
            })
        );
    }

    #[test]
    fn test_and_query() {
        let parser = TqlParser::new();
        let ast = parser.parse("age > 25 AND status eq 'active'").unwrap();
        let builder = QueryBuilder::new(None);
        let query = builder.build_query(&ast).unwrap();

        // Verify it's a bool/must query
        assert!(query["query"]["bool"]["must"].is_array());
        assert_eq!(query["query"]["bool"]["must"].as_array().unwrap().len(), 2);
    }

    #[test]
    fn test_or_query() {
        let parser = TqlParser::new();
        let ast = parser
            .parse("status eq 'active' OR status eq 'pending'")
            .unwrap();
        let builder = QueryBuilder::new(None);
        let query = builder.build_query(&ast).unwrap();

        // Verify it's a bool/should query
        assert!(query["query"]["bool"]["should"].is_array());
        assert_eq!(
            query["query"]["bool"]["should"].as_array().unwrap().len(),
            2
        );
        assert_eq!(query["query"]["bool"]["minimum_should_match"], 1);
    }

    #[test]
    fn test_not_query() {
        let parser = TqlParser::new();
        let ast = parser.parse("NOT (age < 18)").unwrap();
        let builder = QueryBuilder::new(None);
        let query = builder.build_query(&ast).unwrap();

        // Verify it's a bool/must_not query
        assert!(query["query"]["bool"]["must_not"]["range"].is_object());
    }

    #[test]
    fn test_detection_rule_with_is_global_and_nslookup() {
        let parser = TqlParser::new();
        let query_str =
            "event.code = 3 AND destination.ip | is_global eq true AND destination.ip | nslookup";
        let ast = parser.parse(query_str).unwrap();

        // Print the AST for debugging
        eprintln!("AST for detection rule query:\n{:#?}", ast);

        let builder = QueryBuilder::new(None);
        let query = builder.build_query(&ast).unwrap();

        // Print the generated DSL
        eprintln!(
            "Generated OpenSearch DSL:\n{}",
            serde_json::to_string_pretty(&query).unwrap()
        );

        // The query should have a bool with must clauses
        assert!(
            query["query"]["bool"]["must"].is_array(),
            "Expected bool/must query, got: {}",
            serde_json::to_string_pretty(&query).unwrap()
        );

        let must_clauses = query["query"]["bool"]["must"].as_array().unwrap();

        // Should have 3 clauses: event.code=3, is_global, nslookup
        // Actually, the structure depends on how AND is parsed
        eprintln!("Number of must clauses: {}", must_clauses.len());
    }
}