tellaro-query-language 0.1.0__py3-none-any.whl

tql/evaluator_components/field_access.py

@@ -0,0 +1,176 @@
+"""Field access utilities for TQL evaluator.
+
+This module provides utilities for accessing nested fields in records,
+handling field mappings, and type conversions.
+"""
+
+from typing import Any, Dict
+
+
+class FieldAccessor:
+    """Handles field access operations for TQL evaluation."""
+
+    # Sentinel value to distinguish missing fields from None values
+    _MISSING_FIELD = object()
+
+    def get_field_value(self, record: Dict[str, Any], field_path: str) -> Any:  # noqa: C901
+        """Get a field value from a record, supporting nested field access.
+
+        Args:
+            record: The record dictionary
+            field_path: Dot-separated field path (e.g., "user.name")
+
+        Returns:
+            The field value or _MISSING_FIELD if not found
+        """
+        # Split the field path into parts
+        parts = field_path.split(".")
+        current = record
+
+        for part in parts:
+            if isinstance(current, dict):
+                if part in current:
+                    current = current[part]
+                elif part.isdigit() and isinstance(current, (list, tuple)):
+                    # Support array indexing like "items.0"
+                    try:
+                        index = int(part)
+                        if 0 <= index < len(current):
+                            current = current[index]
+                        else:
+                            return self._MISSING_FIELD
+                    except (ValueError, IndexError):
+                        return self._MISSING_FIELD
+                else:
+                    return self._MISSING_FIELD
+            elif isinstance(current, (list, tuple)) and part.isdigit():
+                # Support direct array indexing
+                try:
+                    index = int(part)
+                    if 0 <= index < len(current):
+                        current = current[index]
+                    else:
+                        return self._MISSING_FIELD
+                except (ValueError, IndexError):
+                    return self._MISSING_FIELD
+            else:
+                return self._MISSING_FIELD
+
+        return current
+
+    def apply_field_mapping(self, field_name: str, field_mappings: Dict[str, Any]) -> str:
+        """Apply field mapping to get the actual field name.
+
+        Args:
+            field_name: Original field name from query
+            field_mappings: Field mapping configuration
+
+        Returns:
+            Mapped field name
+        """
+        if field_name not in field_mappings:
+            return field_name
+
+        mapping = field_mappings[field_name]
+
+        if isinstance(mapping, str):
+            # Simple string mapping
+            if mapping not in [
+                "keyword",
+                "text",
+                "long",
+                "integer",
+                "short",
+                "byte",
+                "double",
+                "float",
+                "boolean",
+                "date",
+                "ip",
+            ]:
+                # This is a field name mapping, not a type
+                return mapping
+            else:
+                # This is a type specification, use original field
+                return field_name
+        elif isinstance(mapping, dict) and mapping:
+            # Intelligent mapping - extract the base field
+            if "type" in mapping and len(mapping) == 1:
+                # Just a type specification, use original field
+                return field_name
+            else:
+                # Find the first field that's not a meta field
+                for key in mapping:
+                    if key not in ["analyzer", "type", "fields"]:
+                        return key
+
+        return field_name
+
+    def apply_type_hint(  # noqa: C901
+        self, value: Any, type_hint: str, field_name: str, operator: str, field_mappings: Dict[str, str]
+    ) -> Any:
+        """Apply type hint to convert value to the appropriate type.
+
+        Args:
+            value: Value to convert
+            type_hint: Type hint (e.g., 'ip', 'integer', 'boolean')
+            field_name: Field name for error messages
+            operator: Operator being used
+            field_mappings: Field mappings
+
+        Returns:
+            Converted value
+
+        Raises:
+            TQLError: If conversion fails
+        """
+        from ..exceptions import TQLError
+
+        if value is None or value is self._MISSING_FIELD:
+            return value
+
+        if type_hint == "ip":
+            # For IP type hint, validate the IP address format
+            import ipaddress
+
+            try:
+                # Try to parse as IP address to validate format
+                ipaddress.ip_address(str(value))
+                # Return as string for comparison
+                return str(value)
+            except ValueError:
+                # For CIDR operator, allow CIDR notation
+                if operator in ["cidr", "not_cidr"]:
+                    try:
+                        ipaddress.ip_network(str(value), strict=False)
+                        return str(value)
+                    except ValueError:
+                        pass
+                raise TQLError(f"Invalid IP address format for field '{field_name}': {value}")
+        elif type_hint == "integer":
+            try:
+                return int(value)
+            except (ValueError, TypeError):
+                raise TQLError(f"Cannot convert value to integer for field '{field_name}': {value}")
+        elif type_hint == "float":
+            try:
+                return float(value)
+            except (ValueError, TypeError):
+                raise TQLError(f"Cannot convert value to float for field '{field_name}': {value}")
+        elif type_hint == "boolean" or type_hint == "bool":
+            if isinstance(value, bool):
+                return value
+            elif isinstance(value, str):
+                if value.lower() == "true" or value == "1":
+                    return True
+                elif value.lower() == "false" or value == "0":
+                    return False
+                else:
+                    raise TQLError(f"Cannot convert value to boolean for field '{field_name}': {value}")
+            else:
+                raise TQLError(f"Cannot convert value to boolean for field '{field_name}': {value}")
+        elif type_hint == "string":
+            return str(value)
+        else:
+            # Unknown type hint, return value as-is
+            return value

tql/evaluator_components/special_expressions.py

@@ -0,0 +1,296 @@
+"""Special expression evaluators for TQL.
+
+This module handles evaluation of special expressions like geo() and nslookup()
+that require external lookups or enrichment.
+"""
+
+from typing import Any, Dict, Optional
+
+from ..mutators import apply_mutators
+
+
+class SpecialExpressionEvaluator:
+    """Evaluates special TQL expressions like geo() and nslookup()."""
+
+    # Sentinel value to distinguish missing fields from None values
+    _MISSING_FIELD = object()
+
+    def __init__(self, get_field_value_func, evaluate_node_func):
+        """Initialize the special expression evaluator.
+
+        Args:
+            get_field_value_func: Function to get field values from records
+            evaluate_node_func: Function to evaluate AST nodes
+        """
+        self._get_field_value = get_field_value_func
+        self._evaluate_node = evaluate_node_func
+
+    def evaluate_geo_expr(  # noqa: C901
+        self, node: Dict[str, Any], record: Dict[str, Any], field_mappings: Dict[str, str]
+    ) -> bool:
+        """Evaluate a geo() expression.
+
+        Args:
+            node: Geo expression AST node
+            record: Record to evaluate against
+            field_mappings: Field name mappings
+
+        Returns:
+            Boolean result of the geo expression
+        """
+        field_name = node["field"]
+        field_mutators = node.get("field_mutators", [])
+        conditions = node["conditions"]
+        geo_params = node.get("geo_params", {})
+
+        # Apply field mapping if available
+        actual_field = field_name
+        if field_name in field_mappings:
+            mapping = field_mappings[field_name]
+            if isinstance(mapping, str):
+                # Simple string mapping
+                if mapping not in [
+                    "keyword",
+                    "text",
+                    "long",
+                    "integer",
+                    "short",
+                    "byte",
+                    "double",
+                    "float",
+                    "boolean",
+                    "date",
+                    "ip",
+                ]:
+                    # This is a field name mapping, not a type
+                    actual_field = mapping
+            elif isinstance(mapping, dict) and mapping:
+                # Intelligent mapping - extract the base field
+                if "type" in mapping and len(mapping) == 1:
+                    # Just a type specification, use original field
+                    actual_field = field_name
+                else:
+                    # Find the first field that's not a meta field
+                    for key in mapping:
+                        if key != "analyzer" and key != "type":
+                            actual_field = key
+                            break
+
+        # Get the field value (IP address)
+        field_value = self._get_field_value(record, actual_field)
+
+        # If field is missing or None, return False
+        if field_value is self._MISSING_FIELD or field_value is None:
+            return False
+
+        # Check if the record already has geo data (from post-processing)
+        # Geo data would be nested under the parent of the IP field
+        geo_data = None
+
+        # Check if a custom field location was specified
+        custom_field = geo_params.get("field")
+
+        if custom_field:
+            # Check the custom field location
+            custom_data = self._get_field_value(record, custom_field)
+            if custom_data is not self._MISSING_FIELD and isinstance(custom_data, dict):
+                # Check if this looks like geo data
+                if any(key in custom_data for key in ["country_iso_code", "city_name", "location"]):
+                    geo_data = {"geo": custom_data}
+                    # Also check for AS data as sibling
+                    if "." in custom_field:
+                        as_parent_path = custom_field.rsplit(".", 1)[0]
+                        as_parent = self._get_field_value(record, as_parent_path)
+                        if isinstance(as_parent, dict) and "as" in as_parent:
+                            geo_data["as"] = as_parent["as"]
+                    elif "as" in record:
+                        geo_data["as"] = record["as"]
+        else:
+            # Default locations
+            if "." in actual_field:
+                # For nested fields like destination.ip, check destination.geo
+                parent_path = actual_field.rsplit(".", 1)[0]
+                parent = self._get_field_value(record, parent_path)
+                if isinstance(parent, dict) and "geo" in parent:
+                    # Found geo data under parent
+                    geo_data = parent
+            else:
+                # For top-level fields, check enrichment.geo
+                if "enrichment" in record and isinstance(record["enrichment"], dict):
+                    if "geo" in record["enrichment"]:
+                        geo_data = record["enrichment"]
+
+        # Check if we should use existing geo data or force a new lookup
+        force_lookup = geo_params.get("force", False)
+
+        if geo_data and not force_lookup:
+            # Use existing geo data
+            pass
+        else:
+            # Apply the geo mutator to get geo data
+            # Build mutator params from geo_params
+            mutator_params = []
+            for param_name, param_value in geo_params.items():
+                mutator_params.append([param_name, param_value])
+
+            # If no force parameter was specified, add the default
+            if "force" not in geo_params:
+                mutator_params.append(["force", force_lookup])
+
+            geo_mutator: Dict[str, Any] = {"name": "geoip_lookup"}
+            if mutator_params:
+                geo_mutator["params"] = mutator_params
+
+            # Apply any field mutators before the geo lookup
+            if field_mutators:
+                field_value = apply_mutators(field_value, field_mutators, actual_field, record)
+
+            # Apply geo lookup
+            geo_data = apply_mutators(field_value, [geo_mutator], actual_field, record)
+
+        # Now evaluate the conditions against the geo data
+        if conditions:
+            # Handle None geo_data (e.g., private IPs or lookup failures)
+            if geo_data is None:
+                geo_data = {}
+
+            # Create a temporary record with the geo data
+            # The conditions are evaluated against the geo fields directly
+            temp_record = geo_data.get("geo", {})
+            # Also include AS data if present
+            if "as" in geo_data:
+                temp_record["as"] = geo_data["as"]
+            return self._evaluate_node(conditions, temp_record, {})
+        else:
+            # No conditions, just checking if geo lookup succeeded
+            return bool(geo_data and "geo" in geo_data)
+
+    def evaluate_nslookup_expr(  # noqa: C901
+        self, node: Dict[str, Any], record: Dict[str, Any], field_mappings: Dict[str, str]
+    ) -> bool:
+        """Evaluate a nslookup() expression.
+
+        Args:
+            node: NSLookup expression AST node
+            record: Record to evaluate against
+            field_mappings: Field name mappings
+
+        Returns:
+            Boolean result of the nslookup expression
+        """
+        field_name = node["field"]
+        field_mutators = node.get("field_mutators", [])
+        conditions = node["conditions"]
+        nslookup_params = node.get("nslookup_params", {})
+
+        # Apply field mapping if available
+        actual_field = field_name
+        if field_name in field_mappings:
+            mapping = field_mappings[field_name]
+            if isinstance(mapping, str):
+                # Simple string mapping
+                if mapping not in [
+                    "keyword",
+                    "text",
+                    "long",
+                    "integer",
+                    "short",
+                    "byte",
+                    "double",
+                    "float",
+                    "boolean",
+                    "date",
+                    "ip",
+                ]:
+                    # This is a field name mapping, not a type
+                    actual_field = mapping
+            elif isinstance(mapping, dict) and mapping:
+                # Intelligent mapping - extract the base field
+                if "type" in mapping and len(mapping) == 1:
+                    # Just a type specification, use original field
+                    actual_field = field_name
+                else:
+                    # Find the first field that's not a meta field
+                    for key in mapping:
+                        if key != "analyzer" and key != "type":
+                            actual_field = key
+                            break
+
+        # Get the field value (hostname or IP)
+        field_value = self._get_field_value(record, actual_field)
+
+        # If field is missing or None, return False
+        if field_value is self._MISSING_FIELD or field_value is None:
+            return False
+
+        # Check if the record already has DNS data (from post-processing)
+        dns_data = None
+
+        # Check if a custom field location was specified
+        custom_field = nslookup_params.get("field")
+
+        if custom_field:
+            # Check the custom field location
+            custom_data = self._get_field_value(record, custom_field)
+            if custom_data is not self._MISSING_FIELD and isinstance(custom_data, dict):
+                # Check if this looks like DNS data
+                if any(key in custom_data for key in ["question", "answers", "resolved_ip"]):
+                    dns_data = custom_data
+        else:
+            # Default locations
+            # If field is like "destination.ip", DNS data should be in "destination.domain"
+            if "." in field_name:
+                # Nested field like destination.ip or source.hostname
+                parent_path = field_name.rsplit(".", 1)[0]
+                parent: Optional[Dict[str, Any]] = record
+                for part in parent_path.split("."):
+                    if isinstance(parent, dict) and part in parent:
+                        parent = parent[part]
+                    else:
+                        parent = None
+                        break
+
+                if parent and isinstance(parent, dict) and "domain" in parent:
+                    dns_data = parent["domain"]
+            else:
+                # Top-level field - check enrichment.domain
+                enrichment = record.get("enrichment", {})
+                if "domain" in enrichment:
+                    dns_data = enrichment["domain"]
+
+        # Check if we should use existing DNS data or force a new lookup
+        force_lookup = nslookup_params.get("force", False)
+
+        if dns_data and not force_lookup:
+            # Use existing DNS data
+            pass
+        else:
+            # Apply the nslookup mutator to get DNS data
+            # Build mutator params from nslookup_params
+            mutator_params = []
+            for param_name, param_value in nslookup_params.items():
+                mutator_params.append([param_name, param_value])
+
+            # If no force parameter was specified, add the default
+            if "force" not in nslookup_params:
+                mutator_params.append(["force", force_lookup])
+
+            nslookup_mutator: Dict[str, Any] = {"name": "nslookup"}
+            if mutator_params:
+                nslookup_mutator["params"] = mutator_params
+
+            # Apply any field mutators before the nslookup
+            if field_mutators:
+                field_value = apply_mutators(field_value, field_mutators, field_name, record)
+
+            # Apply nslookup
+            dns_data = apply_mutators(field_value, [nslookup_mutator], field_name, record)
+
+        # Now evaluate the conditions against the DNS data
+        if conditions:
+            # Create a temporary record with the DNS data at root level
+            temp_record = dns_data if dns_data else {}
+            return self._evaluate_node(conditions, temp_record, {})
+        else:
+            # No conditions, just checking if nslookup succeeded
+            return bool(dns_data)