tellaro-query-language 3.0.1

//! 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, GOOGLE};
use hickory_resolver::net::runtime::TokioRuntimeProvider;
use hickory_resolver::proto::rr::{Name, RData};
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<TokioRuntimeProvider>;

/// 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")));

/// The two DNS operations `nslookup` needs.
///
/// [`NSLookupMutator::new`] uses [`SystemDnsResolver`], which queries real DNS.
/// [`NSLookupMutator::with_resolver`] accepts any implementation, which is how the
/// unit tests answer deterministically -- success and failure -- without the
/// network. Everything the mutator derives from an answer (the ECS structure,
/// response codes, trailing-dot stripping, the enrichment envelope) happens after
/// this boundary, so it is exercised identically either way.
pub trait DnsResolver: Send + Sync {
    /// Forward lookup: hostname to addresses. `Err` means resolution failed for
    /// any reason (NXDOMAIN, timeout, no reachable resolver); the mutator reports
    /// every failure as `NXDOMAIN` with no answers.
    fn lookup_ip(&self, hostname: &str) -> std::result::Result<Vec<IpAddr>, String>;

    /// Reverse (PTR) lookup: address to host names, as the resolver returns
    /// them. A trailing root dot is allowed; the mutator strips it. `Err` has the
    /// same meaning as for [`DnsResolver::lookup_ip`].
    fn reverse_lookup(&self, ip: IpAddr) -> std::result::Result<Vec<String>, String>;
}

/// The production resolver: the system DNS configuration, falling back to
/// Google's public resolvers when it cannot be read.
pub struct SystemDnsResolver {
    resolver: Arc<TokioResolver>,
    runtime: Arc<Runtime>,
}

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

        // Build the resolver: prefer the system DNS configuration (matches what
        // hickory-resolver 0.25 effectively gave us), fall back to Google's public
        // resolvers if the system config can't be read. In 0.26, the derived
        // `ResolverConfig::default()` is empty, so an explicit config is required.
        let resolver = Resolver::builder(TokioRuntimeProvider::default())
            .unwrap_or_else(|_| {
                Resolver::builder_with_config(
                    ResolverConfig::udp_and_tcp(&GOOGLE),
                    TokioRuntimeProvider::default(),
                )
            })
            .with_options(ResolverOpts::default())
            .build()
            .expect("Failed to build DNS resolver");

        Self {
            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)
        }
    }
}

impl Default for SystemDnsResolver {
    fn default() -> Self {
        Self::new()
    }
}

impl DnsResolver for SystemDnsResolver {
    fn lookup_ip(&self, hostname: &str) -> std::result::Result<Vec<IpAddr>, String> {
        let resolver = self.resolver.clone();
        let hostname = hostname.to_string();
        self.run_async(async move {
            resolver
                .lookup_ip(&hostname)
                .await
                .map(|lookup| lookup.iter().collect())
                .map_err(|e| e.to_string())
        })
    }

    fn reverse_lookup(&self, ip: IpAddr) -> std::result::Result<Vec<String>, String> {
        let resolver = self.resolver.clone();
        // hickory-resolver 0.26 dropped the IpAddr overload of reverse_lookup; convert
        // the IP to its in-addr.arpa / ip6.arpa name explicitly.
        let reverse_name = Name::from(ip).to_string();
        self.run_async(async move {
            resolver
                .reverse_lookup(reverse_name)
                .await
                .map(|lookup| {
                    lookup
                        .answers()
                        .iter()
                        .filter_map(|r| match &r.data {
                            RData::PTR(ptr) => Some(ptr.0.to_string()),
                            _ => None,
                        })
                        .collect()
                })
                .map_err(|e| e.to_string())
        })
    }
}

/// The resolver [`NSLookupMutator::new`] uses: real DNS.
#[cfg(not(all(test, not(feature = "integration-tests"))))]
fn default_resolver() -> Arc<dyn DnsResolver> {
    Arc::new(SystemDnsResolver::new())
}

/// In this crate's own unit-test build, the default resolver answers nothing.
///
/// Tests that exercise `nslookup` directly inject a stub through
/// [`NSLookupMutator::with_resolver`]. Tests that reach it through a QUERY --
/// `create_mutator`, the evaluator -- cannot, and two of them
/// (`test_compound_query_with_nslookup_expr`,
/// `test_nslookup_expr_evaluation_enrichment_only`) were sending a real PTR
/// query for 8.8.8.8 on every run, incidentally to what they assert. Here every
/// lookup fails, which the mutator reports as `NXDOMAIN`: the same answer a host
/// without DNS gave, now on every host.
///
/// Scope is exactly `cfg(test)` of this library: `tests/*.rs`, downstream crates
/// and release builds link the real resolver, and so do the `integration-tests`
/// feature's live evaluator tests. `tests/dns_live.rs` checks the real path.
#[cfg(all(test, not(feature = "integration-tests")))]
fn default_resolver() -> Arc<dyn DnsResolver> {
    struct NoDns;
    impl DnsResolver for NoDns {
        fn lookup_ip(&self, hostname: &str) -> std::result::Result<Vec<IpAddr>, String> {
            Err(format!("unit tests do not resolve DNS ({hostname})"))
        }
        fn reverse_lookup(&self, ip: IpAddr) -> std::result::Result<Vec<String>, String> {
            Err(format!("unit tests do not resolve DNS ({ip})"))
        }
    }
    Arc::new(NoDns)
}

/// 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<dyn DnsResolver>,
}

impl NSLookupMutator {
    /// A mutator that resolves through real DNS ([`SystemDnsResolver`]).
    ///
    /// This is what `create_mutator("nslookup", ..)` builds, so it is the path
    /// every query takes.
    pub fn new(params: MutatorParams) -> Self {
        Self::with_resolver(params, default_resolver())
    }

    /// A mutator that resolves through `resolver` instead of real DNS.
    pub fn with_resolver(params: MutatorParams, resolver: Arc<dyn DnsResolver>) -> Self {
        Self { params, resolver }
    }

    /// Perform forward DNS lookup (hostname → IP addresses)
    fn lookup_hostname(&self, hostname: &str) -> Result<DnsResult> {
        Ok(match self.resolver.lookup_ip(hostname) {
            Ok(addresses) => DnsResult {
                query: hostname.to_string(),
                query_type: "A".to_string(), // Could be A or AAAA
                answers: addresses.iter().map(|ip| ip.to_string()).collect(),
                response_code: "NOERROR".to_string(),
            },
            Err(_) => DnsResult {
                query: hostname.to_string(),
                query_type: "A".to_string(),
                answers: vec![],
                response_code: "NXDOMAIN".to_string(),
            },
        })
    }

    /// 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))
        })?;

        Ok(match self.resolver.reverse_lookup(ip) {
            Ok(hostnames) => DnsResult {
                query: ip_str.to_string(),
                query_type: "PTR".to_string(),
                answers: hostnames
                    .iter()
                    .map(|name| name.trim_end_matches('.').to_string())
                    .collect(),
                response_code: "NOERROR".to_string(),
            },
            Err(_) => DnsResult {
                query: ip_str.to_string(),
                query_type: "PTR".to_string(),
                answers: vec![],
                response_code: "NXDOMAIN".to_string(),
            },
        })
    }

    /// 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;
    use std::sync::Mutex;

    /// A resolver that answers from fixed tables and records every query, so
    /// these tests never touch the network and can assert what was asked.
    /// A name or address absent from its table fails, like NXDOMAIN.
    #[derive(Default)]
    struct StubResolver {
        forward: HashMap<String, Vec<IpAddr>>,
        reverse: HashMap<IpAddr, Vec<String>>,
        queries: Mutex<Vec<String>>,
    }

    impl StubResolver {
        fn forward(mut self, hostname: &str, addresses: &[&str]) -> Self {
            let parsed = addresses.iter().map(|a| a.parse().unwrap()).collect();
            self.forward.insert(hostname.to_string(), parsed);
            self
        }

        fn reverse(mut self, ip: &str, names: &[&str]) -> Self {
            let names = names.iter().map(|n| n.to_string()).collect();
            self.reverse.insert(ip.parse().unwrap(), names);
            self
        }

        fn queries(&self) -> Vec<String> {
            self.queries.lock().unwrap().clone()
        }
    }

    impl DnsResolver for StubResolver {
        fn lookup_ip(&self, hostname: &str) -> std::result::Result<Vec<IpAddr>, String> {
            self.queries.lock().unwrap().push(format!("A {hostname}"));
            self.forward
                .get(hostname)
                .cloned()
                .ok_or_else(|| format!("no record for {hostname}"))
        }

        fn reverse_lookup(&self, ip: IpAddr) -> std::result::Result<Vec<String>, String> {
            self.queries.lock().unwrap().push(format!("PTR {ip}"));
            self.reverse
                .get(&ip)
                .cloned()
                .ok_or_else(|| format!("no PTR for {ip}"))
        }
    }

    fn mutator(params: MutatorParams, stub: StubResolver) -> (NSLookupMutator, Arc<StubResolver>) {
        let stub = Arc::new(stub);
        let mutator = NSLookupMutator::with_resolver(params, stub.clone());
        (mutator, stub)
    }

    #[test]
    fn test_nslookup_reverse_lookup_public_ip() {
        // Resolvers return PTR names fully qualified, with the trailing root dot.
        let (mutator, stub) = mutator(
            HashMap::new(),
            StubResolver::default().reverse("8.8.8.8", &["dns.google."]),
        );
        let value = json!("8.8.8.8");
        let result = mutator.apply("destination.ip", &json!({}), &value).unwrap();

        assert_eq!(stub.queries(), vec!["PTR 8.8.8.8"]);
        assert_eq!(
            result,
            json!({
                "_tql_enrichment": {
                    "type": "dns",
                    "domain_field": "destination.domain",
                    "dns_field": "destination.dns",
                    "domain": "dns.google",
                    "dns": {
                        "question": {"name": "8.8.8.8", "type": "PTR"},
                        "answers": ["dns.google"],
                        "response_code": "NOERROR"
                    }
                },
                "_tql_return_value": "dns.google",
                "_tql_preserve_original": true,
                "_tql_original_value": "8.8.8.8"
            })
        );
    }

    #[test]
    fn test_nslookup_forward_lookup() {
        let (mutator, stub) = mutator(
            HashMap::new(),
            StubResolver::default().forward("localhost", &["127.0.0.1", "::1"]),
        );
        let result = mutator
            .apply("hostname", &json!({}), &json!("localhost"))
            .unwrap();

        assert_eq!(stub.queries(), vec!["A localhost"]);
        let enrichment = &result["_tql_enrichment"];
        assert_eq!(enrichment["type"], "dns");
        assert_eq!(enrichment["domain_field"], "domain");
        assert_eq!(enrichment["dns_field"], "dns");
        assert_eq!(
            enrichment["dns"],
            json!({
                "question": {"name": "localhost", "type": "A"},
                "answers": ["127.0.0.1", "::1"],
                "resolved_ip": ["127.0.0.1", "::1"],
                "response_code": "NOERROR"
            })
        );
        // The first answer is what a comparison on the mutated value sees.
        assert_eq!(result["_tql_return_value"], "127.0.0.1");
    }

    #[test]
    fn test_nslookup_forward_lookup_failure_is_nxdomain() {
        let (mutator, stub) = mutator(HashMap::new(), StubResolver::default());
        let value = json!("no-such-host.invalid");
        let result = mutator.apply("hostname", &json!({}), &value).unwrap();

        assert_eq!(stub.queries(), vec!["A no-such-host.invalid"]);
        let enrichment = &result["_tql_enrichment"];
        assert_eq!(enrichment["domain"], JsonValue::Null);
        // No `resolved_ip` key at all when nothing resolved.
        assert_eq!(
            enrichment["dns"],
            json!({
                "question": {"name": "no-such-host.invalid", "type": "A"},
                "answers": [],
                "response_code": "NXDOMAIN"
            })
        );
        assert_eq!(result["_tql_return_value"], value);
    }

    #[test]
    fn test_nslookup_private_ip_no_reverse() {
        // A failed reverse lookup -- the usual outcome for RFC 1918 space.
        let (mutator, stub) = mutator(HashMap::new(), StubResolver::default());
        let value = json!("192.168.1.1");
        let result = mutator.apply("source.ip", &json!({}), &value).unwrap();

        assert_eq!(stub.queries(), vec!["PTR 192.168.1.1"]);
        let enrichment = &result["_tql_enrichment"];
        assert_eq!(enrichment["domain_field"], "source.domain");
        assert_eq!(enrichment["domain"], JsonValue::Null);
        assert_eq!(
            enrichment["dns"],
            json!({
                "question": {"name": "192.168.1.1", "type": "PTR"},
                "answers": [],
                "response_code": "NXDOMAIN"
            })
        );
        // Nothing resolved, so comparisons fall back to the original value.
        assert_eq!(result["_tql_return_value"], value);
        assert_eq!(result["_tql_original_value"], value);
    }

    #[test]
    fn test_nslookup_without_save_enrichment() {
        let mut params = HashMap::new();
        params.insert("save".to_string(), json!(false));
        let (mutator, stub) = mutator(
            params,
            StubResolver::default().reverse("8.8.8.8", &["dns.google."]),
        );

        // 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", &json!({}), &value).unwrap();

        // Should return the original value unchanged
        assert_eq!(
            result, value,
            "With save=false, should return original value"
        );
        // The lookup still happens; only the enrichment is dropped.
        assert_eq!(stub.queries(), vec!["PTR 8.8.8.8"]);
    }

    #[test]
    fn test_nslookup_array_of_ips() {
        // One address resolves and one does not: both outcomes in one array.
        let (mutator, stub) = mutator(
            HashMap::new(),
            StubResolver::default().reverse("8.8.8.8", &["dns.google."]),
        );
        let value = json!(["8.8.8.8", "8.8.4.4"]);
        let result = mutator.apply("destination.ip", &json!({}), &value).unwrap();

        assert_eq!(stub.queries(), vec!["PTR 8.8.8.8", "PTR 8.8.4.4"]);
        let enrichment = &result["_tql_enrichment"];
        assert_eq!(enrichment["domain"], json!(["dns.google"]));
        assert_eq!(
            enrichment["dns"],
            json!([
                {
                    "question": {"name": "8.8.8.8", "type": "PTR"},
                    "answers": ["dns.google"],
                    "response_code": "NOERROR"
                },
                {
                    "question": {"name": "8.8.4.4", "type": "PTR"},
                    "answers": [],
                    "response_code": "NXDOMAIN"
                }
            ])
        );
        // The unresolved element keeps its original value.
        assert_eq!(
            result["_tql_return_value"],
            json!(["dns.google", "8.8.4.4"])
        );
    }

    #[test]
    fn test_nslookup_non_string() {
        let (mutator, stub) = mutator(HashMap::new(), StubResolver::default());
        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)
        );

        // None of those is a name or an address, so nothing is resolved.
        assert!(stub.queries().is_empty());
    }

    #[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"));
    }
}