tellaro-query-language 1.3.8

//! DNS lookup mutators for TQL.
//!
//! Provides DNS resolution functionality using hickory-resolver.
//! Supports both forward DNS lookups (hostname → IP) and reverse DNS lookups (IP → hostname).

use super::{Mutator, MutatorParams};
use crate::error::{Result, TqlError};
use hickory_resolver::config::{ResolverConfig, ResolverOpts};
use hickory_resolver::name_server::TokioConnectionProvider;
use hickory_resolver::Resolver;
use once_cell::sync::Lazy;
use serde_json::{json, Value as JsonValue};
use std::net::IpAddr;
use std::str::FromStr;
use std::sync::Arc;
use tokio::runtime::{Handle, Runtime};

/// Type alias for the tokio-based resolver
type TokioResolver = Resolver<TokioConnectionProvider>;

/// Shared Tokio runtime for DNS resolution, created once and reused.
static SHARED_RUNTIME: Lazy<Arc<Runtime>> =
    Lazy::new(|| Arc::new(Runtime::new().expect("Failed to create shared tokio runtime for DNS")));

/// Enrichment mutator that performs DNS lookups on hostnames or IP addresses.
///
/// This mutator can:
/// - Perform forward DNS lookups (hostname to IP)
/// - Perform reverse DNS lookups (IP to hostname)
/// - Return ECS-compliant DNS data structure
///
/// Field Storage (ECS-compliant):
/// - destination.ip | nslookup → returns domain name, enrichment stored at destination.domain/dns
/// - source.ip | nslookup → returns domain name, enrichment stored at source.domain/dns
/// - ip | nslookup → returns domain name, enrichment stored at domain/dns
///
/// The mutator returns the resolved domain name (for filtering/comparison) and stores
/// the full ECS-compliant DNS data in the record for enrichment.
pub struct NSLookupMutator {
    params: MutatorParams,
    resolver: Arc<TokioResolver>,
    runtime: Arc<Runtime>,
}

impl NSLookupMutator {
    pub fn new(params: MutatorParams) -> Self {
        // Reuse the shared static runtime instead of creating one per instantiation
        let runtime = SHARED_RUNTIME.clone();

        // Create resolver with default configuration using the new builder API
        let resolver = Resolver::builder_with_config(
            ResolverConfig::default(),
            TokioConnectionProvider::default(),
        )
        .with_options(ResolverOpts::default())
        .build();

        Self {
            params,
            resolver: Arc::new(resolver),
            runtime,
        }
    }

    /// Execute an async future, handling the case where we may already be inside
    /// a Tokio runtime (which would panic on nested `block_on`).
    fn run_async<F, T>(&self, future: F) -> T
    where
        F: std::future::Future<Output = T> + Send + 'static,
        T: Send + 'static,
    {
        if let Ok(handle) = Handle::try_current() {
            // Already inside a Tokio runtime — use spawn_blocking to avoid nested block_on
            std::thread::scope(|s| {
                s.spawn(move || handle.block_on(future))
                    .join()
                    .expect("DNS lookup thread panicked")
            })
        } else {
            // No active runtime — safe to block_on
            self.runtime.block_on(future)
        }
    }

    /// Perform forward DNS lookup (hostname → IP addresses)
    fn lookup_hostname(&self, hostname: &str) -> Result<DnsResult> {
        let resolver = self.resolver.clone();
        let hostname_owned = hostname.to_string();

        let result: std::result::Result<DnsResult, TqlError> = self.run_async(async move {
            match resolver.lookup_ip(&hostname_owned).await {
                Ok(lookup) => {
                    let addresses: Vec<String> = lookup.iter().map(|ip| ip.to_string()).collect();
                    Ok(DnsResult {
                        query: hostname_owned,
                        query_type: "A".to_string(), // Could be A or AAAA
                        answers: addresses,
                        response_code: "NOERROR".to_string(),
                    })
                }
                Err(_) => Ok(DnsResult {
                    query: hostname_owned,
                    query_type: "A".to_string(),
                    answers: vec![],
                    response_code: "NXDOMAIN".to_string(),
                }),
            }
        });

        result
    }

    /// Perform reverse DNS lookup (IP → hostname)
    fn reverse_lookup(&self, ip_str: &str) -> Result<DnsResult> {
        let ip: IpAddr = IpAddr::from_str(ip_str).map_err(|e| {
            TqlError::MutatorError(format!("Invalid IP address '{}': {}", ip_str, e))
        })?;

        let resolver = self.resolver.clone();
        let ip_for_query = ip_str.to_string();

        let result: std::result::Result<DnsResult, TqlError> = self.run_async(async move {
            match resolver.reverse_lookup(ip).await {
                Ok(lookup) => {
                    let hostnames: Vec<String> = lookup
                        .iter()
                        .map(|name| name.to_string().trim_end_matches('.').to_string())
                        .collect();
                    Ok(DnsResult {
                        query: ip_for_query,
                        query_type: "PTR".to_string(),
                        answers: hostnames,
                        response_code: "NOERROR".to_string(),
                    })
                }
                Err(_) => Ok(DnsResult {
                    query: ip_for_query,
                    query_type: "PTR".to_string(),
                    answers: vec![],
                    response_code: "NXDOMAIN".to_string(),
                }),
            }
        });

        result
    }

    /// Check if a string is an IP address
    fn is_ip_address(value: &str) -> bool {
        IpAddr::from_str(value).is_ok()
    }

    /// Build ECS-compliant DNS data structure
    fn build_ecs_dns_data(result: &DnsResult) -> JsonValue {
        let mut ecs_data = json!({
            "question": {
                "name": result.query,
                "type": result.query_type
            },
            "answers": result.answers,
            "response_code": result.response_code
        });

        // Add resolved_ip for forward lookups (A/AAAA records)
        if (result.query_type == "A" || result.query_type == "AAAA") && !result.answers.is_empty() {
            ecs_data["resolved_ip"] = json!(result.answers);
        }

        ecs_data
    }

    /// Get the enrichment field paths based on the source field name
    fn get_enrichment_paths(field_name: &str) -> (String, String) {
        if field_name.contains('.') {
            // Nested field like destination.ip → destination.domain, destination.dns
            let parent_path = field_name.rsplit('.').skip(1).collect::<Vec<_>>();
            let parent = parent_path.into_iter().rev().collect::<Vec<_>>().join(".");
            (format!("{}.domain", parent), format!("{}.dns", parent))
        } else {
            // Top-level field → domain, dns
            ("domain".to_string(), "dns".to_string())
        }
    }
}

/// Result of a DNS lookup
struct DnsResult {
    query: String,
    query_type: String,
    answers: Vec<String>,
    response_code: String,
}

impl Mutator for NSLookupMutator {
    fn apply(&self, field_name: &str, _record: &JsonValue, value: &JsonValue) -> Result<JsonValue> {
        // Check if we should save enrichment (default: true)
        let save_enrichment = self
            .params
            .get("save")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);

        match value {
            JsonValue::String(s) => {
                // Determine if this is an IP (reverse lookup) or hostname (forward lookup)
                let dns_result = if Self::is_ip_address(s) {
                    self.reverse_lookup(s)?
                } else {
                    self.lookup_hostname(s)?
                };

                // Build ECS DNS data
                let ecs_data = Self::build_ecs_dns_data(&dns_result);

                // Get the first answer (domain name) for the return value
                let domain = dns_result.answers.first().cloned();

                if save_enrichment {
                    // Build enrichment response with domain, dns, and the original value
                    let (domain_field, dns_field) = Self::get_enrichment_paths(field_name);

                    // Return a special structure that the post-processor can use for enrichment
                    // The actual enrichment is handled by the evaluator/post-processor
                    //
                    // _tql_return_value: The domain name for filtering/comparison purposes
                    //   (e.g., `destination.ip | nslookup contains 'google'` should match)
                    // _tql_preserve_original: Tells post-processor NOT to overwrite the original
                    //   field (destination.ip should stay as an IP, not become a hostname)
                    // _tql_original_value: The original value to preserve in the field
                    let return_value = match &domain {
                        Some(d) => json!(d),
                        None => value.clone(), // If no domain resolved, use original for comparison
                    };

                    let enrichment = json!({
                        "_tql_enrichment": {
                            "type": "dns",
                            "domain_field": domain_field,
                            "dns_field": dns_field,
                            "domain": domain,
                            "dns": ecs_data
                        },
                        "_tql_return_value": return_value,
                        "_tql_preserve_original": true,
                        "_tql_original_value": value
                    });

                    Ok(enrichment)
                } else {
                    // Without save_enrichment, just return the original value
                    // The DNS lookup was performed but no enrichment is stored
                    Ok(value.clone())
                }
            }
            JsonValue::Array(arr) => {
                let mut results = Vec::new();
                let mut all_domains = Vec::new();
                let mut all_dns = Vec::new();

                for item in arr {
                    if let JsonValue::String(s) = item {
                        let dns_result = if Self::is_ip_address(s) {
                            self.reverse_lookup(s)?
                        } else {
                            self.lookup_hostname(s)?
                        };

                        let ecs_data = Self::build_ecs_dns_data(&dns_result);
                        all_dns.push(ecs_data);

                        if let Some(domain) = dns_result.answers.first() {
                            all_domains.push(domain.clone());
                            results.push(JsonValue::String(domain.clone()));
                        } else {
                            results.push(item.clone());
                        }
                    } else {
                        results.push(item.clone());
                    }
                }

                if save_enrichment {
                    let (domain_field, dns_field) = Self::get_enrichment_paths(field_name);

                    Ok(json!({
                        "_tql_enrichment": {
                            "type": "dns",
                            "domain_field": domain_field,
                            "dns_field": dns_field,
                            "domain": all_domains,
                            "dns": all_dns
                        },
                        "_tql_return_value": results,
                        "_tql_preserve_original": true,
                        "_tql_original_value": value
                    }))
                } else {
                    Ok(JsonValue::Array(results))
                }
            }
            _ => Ok(value.clone()),
        }
    }

    fn name(&self) -> &str {
        "nslookup"
    }

    fn is_enrichment(&self) -> bool {
        true
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use std::collections::HashMap;

    #[test]
    fn test_nslookup_reverse_lookup_public_ip() {
        let mutator = NSLookupMutator::new(HashMap::new());
        let record = json!({});

        // Test reverse DNS lookup on Google's public DNS
        let value = json!("8.8.8.8");
        let result = mutator.apply("destination.ip", &record, &value).unwrap();

        // Should return enrichment structure with dns.google
        assert!(result.is_object());
        let enrichment = result
            .get("_tql_enrichment")
            .expect("Should have enrichment");
        assert_eq!(enrichment.get("type").unwrap(), "dns");
        assert_eq!(
            enrichment.get("domain_field").unwrap(),
            "destination.domain"
        );
        assert_eq!(enrichment.get("dns_field").unwrap(), "destination.dns");

        // Check that we got a domain (dns.google)
        let domain = enrichment.get("domain");
        assert!(domain.is_some());
        if let Some(JsonValue::String(d)) = domain {
            assert!(
                d.contains("dns.google") || d.contains("google"),
                "Expected google domain, got: {}",
                d
            );
        }

        // Check ECS DNS structure
        let dns = enrichment.get("dns").expect("Should have dns data");
        assert!(dns.get("question").is_some());
        assert!(dns.get("answers").is_some());
        assert_eq!(dns.get("response_code").unwrap(), "NOERROR");
    }

    #[test]
    fn test_nslookup_forward_lookup() {
        let mutator = NSLookupMutator::new(HashMap::new());
        let record = json!({});

        // Test forward DNS lookup on localhost
        let value = json!("localhost");
        let result = mutator.apply("hostname", &record, &value).unwrap();

        // Should return enrichment structure
        assert!(result.is_object());
        let enrichment = result
            .get("_tql_enrichment")
            .expect("Should have enrichment");
        assert_eq!(enrichment.get("type").unwrap(), "dns");

        // Check that we got resolved IPs
        let dns = enrichment.get("dns").expect("Should have dns data");
        let answers = dns.get("answers").expect("Should have answers");
        assert!(answers.is_array());
    }

    #[test]
    fn test_nslookup_private_ip_no_reverse() {
        let mutator = NSLookupMutator::new(HashMap::new());
        let record = json!({});

        // Private IPs typically don't have reverse DNS
        let value = json!("192.168.1.1");
        let result = mutator.apply("source.ip", &record, &value).unwrap();

        // Should still return enrichment structure, but with empty answers
        assert!(result.is_object());
        let enrichment = result
            .get("_tql_enrichment")
            .expect("Should have enrichment");
        let dns = enrichment.get("dns").expect("Should have dns data");

        // Response code should be NXDOMAIN for private IPs without reverse DNS
        let response_code = dns.get("response_code").unwrap().as_str().unwrap();
        assert!(response_code == "NXDOMAIN" || response_code == "NOERROR");
    }

    #[test]
    fn test_nslookup_without_save_enrichment() {
        let mut params = HashMap::new();
        params.insert("save".to_string(), json!(false));
        let mutator = NSLookupMutator::new(params);
        let record = json!({});

        // Test with save=false - should return the original value (not the DNS answer)
        // to prevent overwriting the original field with incompatible types
        let value = json!("8.8.8.8");
        let result = mutator.apply("ip", &record, &value).unwrap();

        // Should return the original value unchanged
        assert_eq!(
            result, value,
            "With save=false, should return original value"
        );
    }

    #[test]
    fn test_nslookup_array_of_ips() {
        let mutator = NSLookupMutator::new(HashMap::new());
        let record = json!({});

        // Test array of IPs
        let value = json!(["8.8.8.8", "8.8.4.4"]);
        let result = mutator.apply("destination.ip", &record, &value).unwrap();

        // Should return enrichment structure with arrays
        assert!(result.is_object());
        let enrichment = result
            .get("_tql_enrichment")
            .expect("Should have enrichment");

        let domain = enrichment.get("domain");
        assert!(domain.is_some());

        let dns = enrichment.get("dns");
        assert!(dns.is_some());
        assert!(dns.unwrap().is_array());
    }

    #[test]
    fn test_nslookup_non_string() {
        let mutator = NSLookupMutator::new(HashMap::new());
        let record = json!({});

        // Numbers should pass through unchanged
        let value = json!(42);
        assert_eq!(mutator.apply("field", &record, &value).unwrap(), json!(42));

        // Booleans should pass through unchanged
        let value = json!(true);
        assert_eq!(
            mutator.apply("field", &record, &value).unwrap(),
            json!(true)
        );

        // Null should pass through unchanged
        let value = json!(null);
        assert_eq!(
            mutator.apply("field", &record, &value).unwrap(),
            json!(null)
        );
    }

    #[test]
    fn test_enrichment_field_paths() {
        // Test nested field
        let (domain, dns) = NSLookupMutator::get_enrichment_paths("destination.ip");
        assert_eq!(domain, "destination.domain");
        assert_eq!(dns, "destination.dns");

        // Test deeply nested field
        let (domain, dns) = NSLookupMutator::get_enrichment_paths("network.outer.ip");
        assert_eq!(domain, "network.outer.domain");
        assert_eq!(dns, "network.outer.dns");

        // Test top-level field
        let (domain, dns) = NSLookupMutator::get_enrichment_paths("ip");
        assert_eq!(domain, "domain");
        assert_eq!(dns, "dns");
    }

    #[test]
    fn test_is_ip_address() {
        // Valid IPv4
        assert!(NSLookupMutator::is_ip_address("192.168.1.1"));
        assert!(NSLookupMutator::is_ip_address("8.8.8.8"));
        assert!(NSLookupMutator::is_ip_address("0.0.0.0"));

        // Valid IPv6
        assert!(NSLookupMutator::is_ip_address("::1"));
        assert!(NSLookupMutator::is_ip_address("2001:db8::1"));
        assert!(NSLookupMutator::is_ip_address("fe80::1"));

        // Invalid - hostnames
        assert!(!NSLookupMutator::is_ip_address("localhost"));
        assert!(!NSLookupMutator::is_ip_address("google.com"));
        assert!(!NSLookupMutator::is_ip_address("dns.google"));
    }
}