harmony/opnsense-config/src/modules/dnsmasq.rs

// dnsmasq.rs
use crate::modules::dhcp::DhcpError;
use log::{debug, info, warn};
use opnsense_config_xml::dnsmasq::{DhcpRange, DnsMasq, DnsmasqHost}; // Assuming DhcpRange is defined in opnsense_config_xml::dnsmasq
use opnsense_config_xml::{MaybeString, StaticMap};
use std::collections::HashSet;
use std::net::Ipv4Addr;
use std::sync::Arc;
use uuid::Uuid;

use opnsense_config_xml::OPNsense;

use crate::config::OPNsenseShell;
use crate::Error;

pub struct DhcpConfigDnsMasq<'a> {
    opnsense: &'a mut OPNsense,
    opnsense_shell: Arc<dyn OPNsenseShell>,
}

const DNS_MASQ_PXE_CONFIG_FILE: &str = "/usr/local/etc/dnsmasq.conf.d/pxe.conf";

impl<'a> DhcpConfigDnsMasq<'a> {
    pub fn new(opnsense: &'a mut OPNsense, opnsense_shell: Arc<dyn OPNsenseShell>) -> Self {
        Self {
            opnsense,
            opnsense_shell,
        }
    }

    /// Removes a MAC address from a static mapping.
    /// If the mapping has no other MAC addresses associated with it, the entire host entry is removed.
    pub fn remove_static_mapping(&mut self, mac_to_remove: &str) {
        let dnsmasq = self.get_dnsmasq();

        // Update hwaddr fields for hosts that contain the MAC, removing it from the comma-separated list.
        for host in dnsmasq.hosts.iter_mut() {
            let mac = host.hwaddr.content_string();
            let original_macs: Vec<&str> = mac.split(',').collect();
            if original_macs
                .iter()
                .any(|m| m.eq_ignore_ascii_case(mac_to_remove))
            {
                let updated_macs: Vec<&str> = original_macs
                    .into_iter()
                    .filter(|m| !m.eq_ignore_ascii_case(mac_to_remove))
                    .collect();
                host.hwaddr = updated_macs.join(",").into();
            }
        }

        // Remove any host entries that no longer have any MAC addresses.
        dnsmasq
            .hosts
            .retain(|host_entry| !host_entry.hwaddr.content_string().is_empty());
    }

    /// Retrieves a mutable reference to the DnsMasq configuration.
    /// This is located in the <dnsmasq> section of the OPNsense config.
    fn get_dnsmasq(&mut self) -> &mut DnsMasq {
        self.opnsense
            .dnsmasq
            .as_mut()
            .expect("Dnsmasq config must be initialized")
    }

    /// Adds or updates a static DHCP mapping.
    ///
    /// This function implements specific logic to handle existing entries:
    /// - If no host exists for the given IP or hostname, a new entry is created.
    /// - If exactly one host exists for the IP and/or hostname, the new MAC is appended to it.
    /// - It will panic if the IP and hostname exist but point to two different host entries,
    ///   as this represents an unresolvable conflict.
    /// - It will also panic if multiple entries are found for the IP or hostname, indicating an
    ///   ambiguous state.
    pub fn add_static_mapping(
        &mut self,
        mac: &str,
        ipaddr: Ipv4Addr,
        hostname: &str,
    ) -> Result<(), DhcpError> {
        if !Self::is_valid_mac(mac) {
            return Err(DhcpError::InvalidMacAddress(mac.to_string()));
        }

        let ip_str = ipaddr.to_string();
        let hosts = &mut self.get_dnsmasq().hosts;

        let ip_indices: Vec<usize> = hosts
            .iter()
            .enumerate()
            .filter(|(_, h)| h.ip.content_string() == ip_str)
            .map(|(i, _)| i)
            .collect();

        let hostname_indices: Vec<usize> = hosts
            .iter()
            .enumerate()
            .filter(|(_, h)| h.host == hostname)
            .map(|(i, _)| i)
            .collect();

        let ip_set: HashSet<usize> = ip_indices.iter().cloned().collect();
        let hostname_set: HashSet<usize> = hostname_indices.iter().cloned().collect();

        if !ip_indices.is_empty()
            && !hostname_indices.is_empty()
            && ip_set.intersection(&hostname_set).count() == 0
        {
            panic!(
                "Configuration conflict: IP {} and hostname '{}' exist, but in different static host entries.",
                ipaddr, hostname
            );
        }

        let mut all_indices: Vec<&usize> = ip_set.union(&hostname_set).collect();
        all_indices.sort();

        match all_indices.len() {
            0 => {
                info!(
                    "Creating new static host for {} ({}) with MAC {}",
                    hostname, ipaddr, mac
                );
                let new_host = DnsmasqHost {
                    uuid: Uuid::new_v4().to_string(),
                    host: hostname.to_string(),
                    ip: ip_str.into(),
                    hwaddr: mac.to_string().into(),
                    local: MaybeString::from("1"),
                    ignore: Some(0),
                    ..Default::default()
                };
                hosts.push(new_host);
            }
            1 => {
                let host_index = *all_indices[0];
                let host_to_modify = &mut hosts[host_index];
                let host_to_modify_ip = host_to_modify.ip.content_string();
                if host_to_modify_ip != ip_str {
                    warn!(
                        "Hostname '{}' already exists with a different IP ({}). Appending MAC {}.",
                        hostname, host_to_modify_ip, mac
                    );
                } else if host_to_modify.host != hostname {
                    warn!(
                        "IP {} already exists with a different hostname ('{}'). Appending MAC {}.",
                        ipaddr, host_to_modify.host, mac
                    );
                }

                if !host_to_modify
                    .hwaddr
                    .content_string()
                    .split(',')
                    .any(|m| m.eq_ignore_ascii_case(mac))
                {
                    info!(
                        "Appending MAC {} to existing static host for {} ({})",
                        mac, host_to_modify.host, host_to_modify_ip
                    );
                    let mut updated_macs = host_to_modify.hwaddr.content_string().to_string();
                    updated_macs.push(',');
                    updated_macs.push_str(mac);
                    host_to_modify.hwaddr = updated_macs.into();
                } else {
                    info!(
                        "MAC {} already present in static host entry for {} ({}). No changes made.",
                        mac, host_to_modify.host, host_to_modify_ip
                    );
                }
            }
            _ => {
                panic!(
                    "Configuration conflict: Found multiple host entries matching IP {} and/or hostname '{}'. Cannot resolve automatically.",
                    ipaddr, hostname
                );
            }
        }

        Ok(())
    }

    /// Helper function to validate a MAC address format.
    fn is_valid_mac(mac: &str) -> bool {
        let parts: Vec<&str> = mac.split(':').collect();
        if parts.len() != 6 {
            return false;
        }
        parts
            .iter()
            .all(|part| part.len() <= 2 && part.chars().all(|c| c.is_ascii_hexdigit()))
    }

    /// Retrieves the list of current static mappings by shelling out to `configctl`.
    /// This provides the real-time state from the running system.
    pub async fn get_static_mappings(&self) -> Result<Vec<StaticMap>, Error> {
        // Note: This command is for the 'dhcpd' service. If dnsmasq uses a different command
        // or key, this will need to be adjusted.
        let list_static_output = self
            .opnsense_shell
            .exec("configctl dhcpd list static")
            .await?;

        let value: serde_json::Value = serde_json::from_str(&list_static_output)
            .unwrap_or_else(|_| panic!("Got invalid json from configctl {list_static_output}"));

        // The JSON output key might be 'dhcpd' even when dnsmasq is the backend.
        let static_maps = value["dhcpd"]
            .as_array()
            .ok_or(Error::Command(format!(
                "Invalid DHCP data from configctl command, got {list_static_output}"
            )))?
            .iter()
            .map(|entry| StaticMap {
                mac: entry["mac"].as_str().unwrap_or_default().to_string(),
                ipaddr: entry["ipaddr"].as_str().unwrap_or_default().to_string(),
                hostname: entry["hostname"].as_str().unwrap_or_default().to_string(),
                descr: entry["descr"].as_str().map(MaybeString::from),
                ..Default::default()
            })
            .collect();

        Ok(static_maps)
    }

    pub async fn set_dhcp_range(&mut self, start: &str, end: &str) -> Result<(), DhcpError> {
        let dnsmasq = self.get_dnsmasq();
        let ranges = &mut dnsmasq.dhcp_ranges;

        // Assuming DnsMasq has dhcp_ranges: Vec<DhcpRange>
        // Find existing range for "lan" interface
        if let Some(range) = ranges
            .iter_mut()
            .find(|r| r.interface == Some("lan".to_string()))
        {
            // Update existing range
            range.start_addr = Some(start.to_string());
            range.end_addr = Some(end.to_string());
        } else {
            // Create new range
            let new_range = DhcpRange {
                uuid: Some(Uuid::new_v4().to_string()),
                interface: Some("lan".to_string()),
                start_addr: Some(start.to_string()),
                end_addr: Some(end.to_string()),
                domain_type: Some("range".to_string()),
                nosync: Some(0),
                ..Default::default()
            };
            ranges.push(new_range);
        }

        Ok(())
    }

    pub async fn set_pxe_options(
        &self,
        tftp_ip: Option<String>,
        bios_filename: String,
        efi_filename: String,
        ipxe_filename: String,
    ) -> Result<(), DhcpError> {
        // OPNsense does not support negative tags via its API for dnsmasq, and the required
        // logic is complex. Therefore, we write a configuration file directly to the
        // dnsmasq.conf.d directory to achieve the desired PXE boot behavior.
        let tftp_str = tftp_ip.map_or(String::new(), |i| format!(",{i},{i}"));

        let config = format!(
            "
# Add tag ipxe to dhcp requests with user class (77) = iPXE
dhcp-match=set:ipxe,77,iPXE
# Add tag bios to dhcp requests with arch (93) = 0
dhcp-match=set:bios,93,0
# Add tag efi to dhcp requests with arch (93) = 7
dhcp-match=set:efi,93,7

# Provide ipxe efi file to uefi but NOT ipxe clients
dhcp-boot=tag:efi,tag:!ipxe,{efi_filename}{tftp_str}

# Provide ipxe boot script to ipxe clients
dhcp-boot=tag:ipxe,{ipxe_filename}{tftp_str}

# Provide undionly to legacy bios clients
dhcp-boot=tag:bios,{bios_filename}{tftp_str}
"
        );
        info!("Writing configuration file to {DNS_MASQ_PXE_CONFIG_FILE}");
        debug!("Content:\n{config}");
        self.opnsense_shell
            .write_content_to_file(&config, DNS_MASQ_PXE_CONFIG_FILE)
            .await
            .map_err(|e| {
                DhcpError::Configuration(format!(
                    "Could not configure pxe for dhcp because of : {e}"
                ))
            })?;

        info!("Restarting dnsmasq to apply changes");
        self.opnsense_shell
            .exec("configctl dnsmasq restart")
            .await
            .map_err(|e| DhcpError::Configuration(format!("Restarting dnsmasq failed : {e}")))?;
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use crate::config::DummyOPNSenseShell;

    use super::*;
    use opnsense_config_xml::OPNsense;
    use std::net::Ipv4Addr;
    use std::sync::Arc;

    /// Helper function to create a DnsmasqHost with minimal boilerplate.
    fn create_host(uuid: &str, host: &str, ip: &str, hwaddr: &str) -> DnsmasqHost {
        DnsmasqHost {
            uuid: uuid.to_string(),
            host: host.to_string(),
            ip: ip.into(),
            hwaddr: hwaddr.into(),
            local: MaybeString::from("1"),
            ignore: Some(0),
            ..Default::default()
        }
    }

    /// Helper to set up the test environment with an initial OPNsense configuration.
    fn setup_test_env(initial_hosts: Vec<DnsmasqHost>) -> DhcpConfigDnsMasq<'static> {
        let opnsense_config = Box::leak(Box::new(OPNsense {
            dnsmasq: Some(DnsMasq {
                hosts: initial_hosts,
                ..Default::default()
            }),
            ..Default::default()
        }));

        DhcpConfigDnsMasq::new(opnsense_config, Arc::new(DummyOPNSenseShell {}))
    }

    #[test]
    fn test_add_first_static_mapping() {
        let mut dhcp_config = setup_test_env(vec![]);
        let ip = Ipv4Addr::new(192, 168, 1, 10);
        let mac = "00:11:22:33:44:55";
        let hostname = "new-host";

        dhcp_config.add_static_mapping(mac, ip, hostname).unwrap();

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert_eq!(hosts.len(), 1);
        let host = &hosts[0];
        assert_eq!(host.host, hostname);
        assert_eq!(host.ip, ip.to_string().into());
        assert_eq!(host.hwaddr.content_string(), mac);
        assert!(Uuid::parse_str(&host.uuid).is_ok());
    }

    #[test]
    fn test_add_mac_to_existing_host_by_ip_and_hostname() {
        let initial_host = create_host(
            "uuid-1",
            "existing-host",
            "192.168.1.20",
            "AA:BB:CC:DD:EE:FF",
        );
        let mut dhcp_config = setup_test_env(vec![initial_host]);
        let ip = Ipv4Addr::new(192, 168, 1, 20);
        let new_mac = "00:11:22:33:44:55";
        let hostname = "existing-host";

        dhcp_config
            .add_static_mapping(new_mac, ip, hostname)
            .unwrap();

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert_eq!(hosts.len(), 1);
        let host = &hosts[0];
        assert_eq!(
            host.hwaddr.content_string(),
            "AA:BB:CC:DD:EE:FF,00:11:22:33:44:55"
        );
    }

    #[test]
    fn test_add_mac_to_existing_host_by_ip_only() {
        let initial_host = create_host(
            "uuid-1",
            "existing-host",
            "192.168.1.20",
            "AA:BB:CC:DD:EE:FF",
        );
        let mut dhcp_config = setup_test_env(vec![initial_host]);
        let ip = Ipv4Addr::new(192, 168, 1, 20);
        let new_mac = "00:11:22:33:44:55";

        // Using a different hostname should still find the host by IP and log a warning.
        dhcp_config
            .add_static_mapping(new_mac, ip, "different-host-name")
            .unwrap();

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert_eq!(hosts.len(), 1);
        let host = &hosts[0];
        assert_eq!(
            host.hwaddr.content_string(),
            "AA:BB:CC:DD:EE:FF,00:11:22:33:44:55"
        );
        assert_eq!(host.host, "existing-host"); // Original hostname should be preserved.
    }

    #[test]
    fn test_add_mac_to_existing_host_by_hostname_only() {
        let initial_host = create_host(
            "uuid-1",
            "existing-host",
            "192.168.1.20",
            "AA:BB:CC:DD:EE:FF",
        );
        let mut dhcp_config = setup_test_env(vec![initial_host]);
        let new_mac = "00:11:22:33:44:55";
        let hostname = "existing-host";

        // Using a different IP should still find the host by hostname and log a warning.
        dhcp_config
            .add_static_mapping(new_mac, Ipv4Addr::new(192, 168, 1, 99), hostname)
            .unwrap();

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert_eq!(hosts.len(), 1);
        let host = &hosts[0];
        assert_eq!(
            host.hwaddr.content_string(),
            "AA:BB:CC:DD:EE:FF,00:11:22:33:44:55"
        );
        assert_eq!(host.ip.content_string(), "192.168.1.20"); // Original IP should be preserved.
    }

    #[test]
    fn test_add_duplicate_mac_to_host() {
        let initial_mac = "AA:BB:CC:DD:EE:FF";
        let initial_host = create_host("uuid-1", "host-1", "192.168.1.20", initial_mac);
        let mut dhcp_config = setup_test_env(vec![initial_host]);

        dhcp_config
            .add_static_mapping(initial_mac, Ipv4Addr::new(192, 168, 1, 20), "host-1")
            .unwrap();

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert_eq!(hosts.len(), 1);
        assert_eq!(hosts[0].hwaddr.content_string(), initial_mac); // No change, no duplication.
    }

    #[test]
    fn test_add_invalid_mac_address() {
        let mut dhcp_config = setup_test_env(vec![]);
        let result =
            dhcp_config.add_static_mapping("invalid-mac", Ipv4Addr::new(10, 0, 0, 1), "host");
        assert!(matches!(result, Err(DhcpError::InvalidMacAddress(_))));
    }

    #[test]
    #[should_panic(
        expected = "Configuration conflict: IP 192.168.1.10 and hostname 'host-b' exist, but in different static host entries."
    )]
    fn test_panic_on_conflicting_ip_and_hostname() {
        let host_a = create_host("uuid-a", "host-a", "192.168.1.10", "AA:AA:AA:AA:AA:AA");
        let host_b = create_host("uuid-b", "host-b", "192.168.1.20", "BB:BB:BB:BB:BB:BB");
        let mut dhcp_config = setup_test_env(vec![host_a, host_b]);

        // This IP belongs to host-a, but the hostname belongs to host-b.
        dhcp_config
            .add_static_mapping(
                "CC:CC:CC:CC:CC:CC",
                Ipv4Addr::new(192, 168, 1, 10),
                "host-b",
            )
            .unwrap();
    }

    #[test]
    #[should_panic(
        expected = "Configuration conflict: Found multiple host entries matching IP 192.168.1.30 and/or hostname 'new-host'."
    )]
    fn test_panic_on_multiple_ip_matches() {
        let host_a = create_host("uuid-a", "host-a", "192.168.1.30", "AA:AA:AA:AA:AA:AA");
        let host_b = create_host("uuid-b", "host-b", "192.168.1.30", "BB:BB:BB:BB:BB:BB");
        let mut dhcp_config = setup_test_env(vec![host_a, host_b]);

        // This IP is ambiguous.
        dhcp_config
            .add_static_mapping(
                "CC:CC:CC:CC:CC:CC",
                Ipv4Addr::new(192, 168, 1, 30),
                "new-host",
            )
            .unwrap();
    }

    #[test]
    fn test_remove_mac_from_multi_mac_host() {
        let host = create_host("uuid-1", "host-1", "192.168.1.50", "mac-1,mac-2,mac-3");
        let mut dhcp_config = setup_test_env(vec![host]);

        dhcp_config.remove_static_mapping("mac-2");

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert_eq!(hosts.len(), 1);
        assert_eq!(hosts[0].hwaddr.content_string(), "mac-1,mac-3");
    }

    #[test]
    fn test_remove_last_mac_from_host() {
        let host = create_host("uuid-1", "host-1", "192.168.1.50", "mac-1");
        let mut dhcp_config = setup_test_env(vec![host]);

        dhcp_config.remove_static_mapping("mac-1");

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert!(hosts.is_empty());
    }

    #[test]
    fn test_remove_non_existent_mac() {
        let host = create_host("uuid-1", "host-1", "192.168.1.50", "mac-1,mac-2");
        let mut dhcp_config = setup_test_env(vec![host.clone()]);

        dhcp_config.remove_static_mapping("mac-nonexistent");

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert_eq!(hosts.len(), 1);
        assert_eq!(hosts[0], host); // The host should be unchanged.
    }

    #[test]
    fn test_remove_mac_case_insensitively() {
        let host = create_host("uuid-1", "host-1", "192.168.1.50", "AA:BB:CC:DD:EE:FF");
        let mut dhcp_config = setup_test_env(vec![host]);

        dhcp_config.remove_static_mapping("aa:bb:cc:dd:ee:ff");

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert!(hosts.is_empty());
    }

    #[test]
    fn test_remove_mac_from_correct_host_only() {
        let host1 = create_host(
            "uuid-1",
            "host-1",
            "192.168.1.50",
            "AA:AA:AA:AA:AA:AA,BB:BB:BB:BB:BB:BB",
        );
        let host2 = create_host(
            "uuid-2",
            "host-2",
            "192.168.1.51",
            "CC:CC:CC:CC:CC:CC,DD:DD:DD:DD:DD:DD",
        );
        let mut dhcp_config = setup_test_env(vec![host1.clone(), host2.clone()]);

        dhcp_config.remove_static_mapping("AA:AA:AA:AA:AA:AA");

        let hosts = &dhcp_config.opnsense.dnsmasq.as_ref().unwrap().hosts;
        assert_eq!(hosts.len(), 2);
        let updated_host1 = hosts.iter().find(|h| h.uuid == "uuid-1").unwrap();
        let unchanged_host2 = hosts.iter().find(|h| h.uuid == "uuid-2").unwrap();

        assert_eq!(updated_host1.hwaddr.content_string(), "BB:BB:BB:BB:BB:BB");
        assert_eq!(
            unchanged_host2.hwaddr.content_string(),
            "CC:CC:CC:CC:CC:CC,DD:DD:DD:DD:DD:DD"
        );
    }
}