tellaro-query-language 2.0.0

//! GeoIP lookup mutators for TQL.
//!
//! Provides IP address geolocation using MaxMind and DB-IP databases.
//!
//! Supports environment variables for database configuration:
//! - TQL_GEOIP_DB_PATH: Full combined database path
//! - TQL_GEOIP_DB_CITY_PATH: City database path
//! - TQL_GEOIP_DB_COUNTRY_PATH: Country database path
//! - TQL_GEOIP_DB_ASN_PATH: ASN database path
//! - TQL_GEOIP_MMDB_PATH: Base directory for auto-detection

use super::{Mutator, MutatorParams};
use crate::error::{Result, TqlError};
use log::warn;
use maxminddb::{geoip2, Reader};
use memmap2::Mmap;
use serde_json::{json, Value as JsonValue};
use std::env;
use std::fs::File;
use std::net::IpAddr;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

/// Supported GeoIP database types
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatabaseType {
    MaxMind,
    DbIp,
}

/// Mutator that performs GeoIP lookups to enrich IP addresses with location data
pub struct GeoIPMutator {
    _params: MutatorParams,
    reader: Option<Arc<Reader<Mmap>>>,
    db_type: DatabaseType,
}

/// Expand tilde (~) in file paths to user home directory
fn expand_tilde(path: &str) -> PathBuf {
    if let Some(stripped) = path.strip_prefix("~/") {
        if let Ok(home) = env::var("HOME") {
            return Path::new(&home).join(stripped);
        }
    }
    PathBuf::from(path)
}

impl GeoIPMutator {
    /// Create a new GeoIP mutator
    ///
    /// Supports parameters:
    /// - "db_path": Path to the GeoIP database file
    /// - "database_path": Alternative parameter name (legacy)
    /// - "database_type": "maxmind" or "dbip" (default: "maxmind")
    ///
    /// Environment variables (checked in order):
    /// 1. TQL_GEOIP_DB_PATH - Full database path
    /// 2. TQL_GEOIP_DB_CITY_PATH - City database path
    /// 3. TQL_GEOIP_MMDB_PATH - Base directory for database files
    ///
    /// All paths support tilde (~) expansion for home directory
    pub fn new(params: MutatorParams) -> Self {
        // Priority: parameter > environment variable
        let db_path = params
            .get("db_path")
            .or_else(|| params.get("database_path"))
            .and_then(|v| v.as_str())
            .map(expand_tilde)
            .or_else(|| {
                // Check environment variables in priority order
                env::var("TQL_GEOIP_DB_PATH")
                    .ok()
                    .or_else(|| env::var("TQL_GEOIP_DB_CITY_PATH").ok())
                    .or_else(|| {
                        // Check base path and look for common filenames
                        env::var("TQL_GEOIP_MMDB_PATH").ok().and_then(|base| {
                            let base_path = expand_tilde(&base);
                            // Try common database filenames
                            let candidates = vec![
                                "GeoLite2-City.mmdb",
                                "GeoIP2-City.mmdb",
                                "dbip-city-lite.mmdb",
                                "dbip-full.mmdb",
                            ];
                            for filename in candidates {
                                let path = base_path.join(filename);
                                if path.exists() {
                                    return Some(path.to_string_lossy().to_string());
                                }
                            }
                            None
                        })
                    })
                    .map(|s| expand_tilde(&s))
            });

        let db_type_str = params
            .get("database_type")
            .and_then(|v| v.as_str())
            .unwrap_or("maxmind");

        let db_type = match db_type_str.to_lowercase().as_str() {
            "dbip" | "db-ip" | "db_ip" => DatabaseType::DbIp,
            _ => DatabaseType::MaxMind,
        };

        let reader = db_path.and_then(|path| {
            // Use memory-mapped I/O for efficient access to large database files
            File::open(&path)
                .ok()
                .and_then(|file| {
                    // SAFETY: Mmap::map requires that the underlying file is not modified or
                    // truncated while the mapping is alive, as that would cause undefined
                    // behavior (SIGBUS / access violation). The calling application must ensure
                    // that the .mmdb file is not replaced, truncated, or written to while this
                    // GeoIPMutator instance (and its Arc<Reader<Mmap>>) is in use. In practice,
                    // GeoIP database updates should create a new GeoIPMutator after the file is
                    // fully written, rather than modifying the file in-place.
                    // nosemgrep: unsafe-usage
                    unsafe { Mmap::map(&file).ok() }
                })
                .and_then(|mmap| Reader::from_source(mmap).ok())
                .map(Arc::new)
        });

        Self {
            _params: params,
            reader,
            db_type,
        }
    }

    fn lookup_ip(&self, ip_str: &str) -> Result<JsonValue> {
        // Parse IP address
        let ip = IpAddr::from_str(ip_str).map_err(|e| {
            TqlError::MutatorError(format!("Invalid IP address '{}': {}", ip_str, e))
        })?;

        // Check if database is available
        let reader = self
            .reader
            .as_ref()
            .ok_or_else(|| TqlError::MutatorError("GeoIP database not configured".to_string()))?;

        // Perform lookup
        match self.db_type {
            DatabaseType::MaxMind => self.lookup_maxmind(reader, ip),
            DatabaseType::DbIp => self.lookup_dbip(reader, ip),
        }
    }

    fn lookup_maxmind(&self, reader: &Reader<Mmap>, ip: IpAddr) -> Result<JsonValue> {
        let city: geoip2::City = reader
            .lookup(ip)
            .map_err(|e| TqlError::MutatorError(format!("GeoIP lookup error: {}", e)))?
            .decode()
            .map_err(|e| TqlError::MutatorError(format!("GeoIP decode error: {}", e)))?
            .ok_or_else(|| {
                TqlError::MutatorError(format!("IP address {} not found in GeoIP database", ip))
            })?;

        // Build ECS-compliant structure with separate geo and as objects
        let mut geo = json!({});
        let as_obj = json!({});

        // Geographic data (ECS format)
        // In maxminddb 0.27+, struct fields are direct values with named language fields
        let country = &city.country;
        if let Some(iso_code) = country.iso_code {
            geo["country_iso_code"] = json!(iso_code);
        }
        if let Some(name) = country.names.english {
            geo["country_name"] = json!(name);
        }

        let city_data = &city.city;
        if let Some(name) = city_data.names.english {
            geo["city_name"] = json!(name);
        }

        let location = &city.location;
        let mut location_obj = json!({});
        if let Some(lat) = location.latitude {
            location_obj["lat"] = json!(lat);
        }
        if let Some(lon) = location.longitude {
            location_obj["lon"] = json!(lon);
        }
        if !location_obj.is_null() && location_obj.as_object().is_some_and(|o| !o.is_empty()) {
            geo["location"] = location_obj;
        }
        if let Some(tz) = location.time_zone {
            geo["timezone"] = json!(tz);
        }

        let postal = &city.postal;
        if let Some(code) = postal.code {
            geo["postal_code"] = json!(code);
        }

        let subdivisions = &city.subdivisions;
        if let Some(subdivision) = subdivisions.first() {
            if let Some(iso_code) = subdivision.iso_code {
                geo["region_iso_code"] = json!(iso_code);
            }
            if let Some(name) = subdivision.names.english {
                geo["region_name"] = json!(name);
            }
        }

        let continent = &city.continent;
        if let Some(code) = continent.code {
            geo["continent_code"] = json!(code);
        }
        if let Some(name) = continent.names.english {
            geo["continent_name"] = json!(name);
        }

        // Add database type marker
        geo["mmdb_type"] = json!(match self.db_type {
            DatabaseType::MaxMind => "maxmind",
            DatabaseType::DbIp => "dbip_pro",
        });

        // Try to get ASN data (requires separate ASN database or combined database)
        // For now, we'll check if the reader has ASN data
        // Note: MaxMind City databases don't include ASN by default
        // This would need a separate ASN database lookup

        // Return ECS-style structure
        let mut result = json!({
            "geo": geo
        });

        // Only add 'as' object if we have ASN data
        if !as_obj.is_null() && as_obj.as_object().is_some_and(|o| !o.is_empty()) {
            result["as"] = as_obj;
        }

        Ok(result)
    }

    fn lookup_dbip(&self, reader: &Reader<Mmap>, ip: IpAddr) -> Result<JsonValue> {
        // DB-IP uses MaxMind format, so we can reuse the same lookup
        // The difference is mainly in the database file format and fields available
        self.lookup_maxmind(reader, ip)
    }
}

/// Global flag to ensure the "no database configured" warning fires only once,
/// preventing log flooding in high-throughput pipelines.
static WARNED_NO_DB: AtomicBool = AtomicBool::new(false);

impl Mutator for GeoIPMutator {
    fn apply(
        &self,
        _field_name: &str,
        _record: &JsonValue,
        value: &JsonValue,
    ) -> Result<JsonValue> {
        // Warn once when no database is configured. This distinguishes a
        // misconfiguration (no DB path / file not found) from the expected case where
        // a valid IP simply isn't present in the database.
        if self.reader.is_none() {
            if !WARNED_NO_DB.swap(true, Ordering::Relaxed) {
                warn!(
                    "GeoIP database not configured — skipping lookup. \
                     Set TQL_GEOIP_DB_PATH or pass db_path parameter."
                );
            }
            return Ok(value.clone());
        }

        match value {
            JsonValue::String(s) => {
                // Try to lookup the IP — errors here are expected (e.g., IP not in DB)
                match self.lookup_ip(s) {
                    Ok(geo_data) => Ok(geo_data),
                    Err(_) => {
                        // IP not found in database; return original value silently
                        Ok(value.clone())
                    }
                }
            }
            JsonValue::Array(arr) => {
                let results: Vec<JsonValue> = arr
                    .iter()
                    .map(|item| {
                        if let JsonValue::String(s) = item {
                            self.lookup_ip(s).unwrap_or_else(|_| item.clone())
                        } else {
                            item.clone()
                        }
                    })
                    .collect();
                Ok(JsonValue::Array(results))
            }
            _ => Ok(value.clone()),
        }
    }

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

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

    #[test]
    fn test_geoip_no_database() {
        // Test without database configured
        let mutator = GeoIPMutator::new(HashMap::new());
        let record = json!({});

        let value = json!("8.8.8.8");
        let result = mutator.apply("field", &record, &value).unwrap();

        // Should return original value when lookup fails
        assert_eq!(result, json!("8.8.8.8"));
    }

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

        // Test invalid IP
        let value = json!("not-an-ip");
        let result = mutator.apply("field", &record, &value).unwrap();

        // Should return original value on invalid IP
        assert_eq!(result, json!("not-an-ip"));
    }

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

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

        // Should return original array when no database
        assert_eq!(result, json!(["8.8.8.8", "1.1.1.1"]));
    }

    #[test]
    fn test_geoip_non_string() {
        let mutator = GeoIPMutator::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)
        );
    }

    #[test]
    fn test_database_type_parsing() {
        let mut params = HashMap::new();
        params.insert("database_type".to_string(), json!("maxmind"));
        let mutator = GeoIPMutator::new(params);
        assert_eq!(mutator.db_type, DatabaseType::MaxMind);

        let mut params = HashMap::new();
        params.insert("database_type".to_string(), json!("dbip"));
        let mutator = GeoIPMutator::new(params);
        assert_eq!(mutator.db_type, DatabaseType::DbIp);

        let mut params = HashMap::new();
        params.insert("database_type".to_string(), json!("db-ip"));
        let mutator = GeoIPMutator::new(params);
        assert_eq!(mutator.db_type, DatabaseType::DbIp);

        // Default should be MaxMind
        let mutator = GeoIPMutator::new(HashMap::new());
        assert_eq!(mutator.db_type, DatabaseType::MaxMind);
    }

    // Note: Integration tests with actual GeoIP databases would go in tests/integration_test.rs
    // Those tests would require downloading test databases from MaxMind or DB-IP
}