harmony/harmony/src/domain/topology/network.rs

use std::{net::Ipv4Addr, str::FromStr, sync::Arc};

use async_trait::async_trait;
use harmony_types::net::{IpAddress, MacAddress};
use serde::Serialize;

use crate::executors::ExecutorError;

use super::{LogicalHost, k8s::K8sClient};

#[derive(Debug)]
pub struct DHCPStaticEntry {
    pub name: String,
    pub mac: Vec<MacAddress>,
    pub ip: Ipv4Addr,
}

impl std::fmt::Display for DHCPStaticEntry {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mac = self
            .mac
            .iter()
            .map(|m| m.to_string())
            .collect::<Vec<String>>()
            .join(",");
        f.write_fmt(format_args!(
            "DHCPStaticEntry : name {}, mac {}, ip {}",
            self.name, mac, self.ip
        ))
    }
}

pub trait Firewall: Send + Sync {
    fn add_rule(&mut self, rule: FirewallRule) -> Result<(), ExecutorError>;
    fn remove_rule(&mut self, rule_id: &str) -> Result<(), ExecutorError>;
    fn list_rules(&self) -> Vec<FirewallRule>;
    fn get_ip(&self) -> IpAddress;
    fn get_host(&self) -> LogicalHost;
}

impl std::fmt::Debug for dyn Firewall {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("Firewall {}", self.get_ip()))
    }
}

pub struct NetworkDomain {
    pub name: String,
}

#[async_trait]
pub trait K8sclient: Send + Sync {
    async fn k8s_client(&self) -> Result<Arc<K8sClient>, String>;
}

pub struct PxeOptions {
    pub ipxe_filename: String,
    pub bios_filename: String,
    pub efi_filename: String,
    pub tftp_ip: Option<IpAddress>,
}

#[async_trait]
pub trait DhcpServer: Send + Sync + std::fmt::Debug {
    async fn add_static_mapping(&self, entry: &DHCPStaticEntry) -> Result<(), ExecutorError>;
    async fn remove_static_mapping(&self, mac: &MacAddress) -> Result<(), ExecutorError>;
    async fn list_static_mappings(&self) -> Vec<(MacAddress, IpAddress)>;
    async fn set_pxe_options(&self, pxe_options: PxeOptions) -> Result<(), ExecutorError>;
    async fn set_dhcp_range(&self, start: &IpAddress, end: &IpAddress)
    -> Result<(), ExecutorError>;
    fn get_ip(&self) -> IpAddress;
    fn get_host(&self) -> LogicalHost;
    async fn commit_config(&self) -> Result<(), ExecutorError>;
}

#[async_trait]
pub trait DnsServer: Send + Sync {
    async fn register_dhcp_leases(&self, register: bool) -> Result<(), ExecutorError>;
    async fn register_hosts(&self, hosts: Vec<DnsRecord>) -> Result<(), ExecutorError>;
    fn remove_record(&self, name: &str, record_type: DnsRecordType) -> Result<(), ExecutorError>;
    async fn list_records(&self) -> Vec<DnsRecord>;
    fn get_ip(&self) -> IpAddress;
    fn get_host(&self) -> LogicalHost;
    async fn commit_config(&self) -> Result<(), ExecutorError>;
    async fn ensure_hosts_registered(&self, hosts: Vec<DnsRecord>) -> Result<(), ExecutorError> {
        let current_hosts = self.list_records().await;
        let mut hosts_to_register = vec![];

        for host in hosts {
            if !current_hosts.iter().any(|h| h == &host) {
                hosts_to_register.push(host);
            }
        }

        if !hosts_to_register.is_empty() {
            self.register_hosts(hosts_to_register).await?;
        }

        Ok(())
    }
}

impl std::fmt::Debug for dyn DnsServer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("DnsServer {}", self.get_ip()))
    }
}
#[derive(Clone, Debug)]
pub struct FirewallRule {
    pub id: String,
    pub source: IpAddress,
    pub destination: IpAddress,
    pub port: u16,
    pub protocol: Protocol,
    pub action: Action,
}

#[derive(Clone, Debug)]
pub enum Protocol {
    TCP,
    UDP,
    ICMP,
}

#[derive(Clone, Debug)]
pub enum Action {
    Allow,
    Deny,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub enum DnsRecordType {
    A,
    AAAA,
    CNAME,
    MX,
    TXT,
}

impl std::fmt::Display for DnsRecordType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DnsRecordType::A => write!(f, "A"),
            DnsRecordType::AAAA => write!(f, "AAAA"),
            DnsRecordType::CNAME => write!(f, "CNAME"),
            DnsRecordType::MX => write!(f, "MX"),
            DnsRecordType::TXT => write!(f, "TXT"),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct DnsRecord {
    pub host: String,
    pub domain: String,
    pub record_type: DnsRecordType,
    pub value: IpAddress,
}

impl FromStr for DnsRecordType {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "A" => Ok(DnsRecordType::A),
            "AAAA" => Ok(DnsRecordType::AAAA),
            "CNAME" => Ok(DnsRecordType::CNAME),
            "MX" => Ok(DnsRecordType::MX),
            "TXT" => Ok(DnsRecordType::TXT),
            _ => Err(format!("Unknown DNSRecordType {s}")),
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use tokio::sync::RwLock;

    use super::*;

    #[tokio::test]
    async fn test_ensure_hosts_registered_no_new_hosts() {
        let server = DummyDnsServer::default();
        let existing_host = DnsRecord {
            host: "existing".to_string(),
            domain: "example.com".to_string(),
            record_type: DnsRecordType::A,
            value: IpAddress::V4(Ipv4Addr::new(192, 168, 1, 2)),
        };

        server
            .register_hosts(vec![existing_host.clone()])
            .await
            .unwrap();

        let new_hosts = vec![
            existing_host, // already exists
        ];

        server.ensure_hosts_registered(new_hosts).await.unwrap();
        assert_eq!(server.list_records().await.len(), 1);
    }

    #[tokio::test]
    async fn test_ensure_hosts_registered_with_new_hosts() {
        let server = DummyDnsServer::default();

        let existing_host = DnsRecord {
            host: "existing".to_string(),
            domain: "example.com".to_string(),
            record_type: DnsRecordType::A,
            value: IpAddress::V4(Ipv4Addr::new(192, 168, 1, 2)),
        };

        server
            .register_hosts(vec![existing_host.clone()])
            .await
            .unwrap();

        let new_hosts = vec![
            existing_host.clone(), // already exists
            DnsRecord {
                host: "new".to_string(),
                domain: "example.com".to_string(),
                record_type: DnsRecordType::A,
                value: IpAddress::V4(Ipv4Addr::new(192, 168, 1, 3)),
            },
        ];

        server.ensure_hosts_registered(new_hosts).await.unwrap();
        assert_eq!(server.list_records().await.len(), 2);
    }

    #[tokio::test]
    async fn test_ensure_hosts_registered_no_hosts() {
        let server = DummyDnsServer::default();

        let new_hosts = vec![];

        server.ensure_hosts_registered(new_hosts).await.unwrap();
        assert_eq!(server.list_records().await.len(), 0);
    }

    #[tokio::test]
    async fn test_ensure_existing_host_kept_no_new_host() {
        let server = DummyDnsServer::default();

        let new_hosts = vec![];

        server.ensure_hosts_registered(new_hosts).await.unwrap();
        assert_eq!(server.list_records().await.len(), 0);
    }

    #[async_trait::async_trait]
    impl DnsServer for DummyDnsServer {
        async fn register_dhcp_leases(&self, _register: bool) -> Result<(), ExecutorError> {
            Ok(())
        }

        async fn register_hosts(&self, hosts: Vec<DnsRecord>) -> Result<(), ExecutorError> {
            self.hosts.write().await.extend(hosts);
            Ok(())
        }

        fn remove_record(
            &self,
            _name: &str,
            _record_type: DnsRecordType,
        ) -> Result<(), ExecutorError> {
            Ok(())
        }

        async fn list_records(&self) -> Vec<DnsRecord> {
            self.hosts.read().await.clone()
        }

        fn get_ip(&self) -> IpAddress {
            IpAddress::V4(Ipv4Addr::new(192, 168, 0, 1))
        }

        fn get_host(&self) -> LogicalHost {
            LogicalHost {
                ip: self.get_ip(),
                name: "dummy-host".to_string(),
            }
        }

        async fn commit_config(&self) -> Result<(), ExecutorError> {
            Ok(())
        }
    }

    struct DummyDnsServer {
        hosts: Arc<RwLock<Vec<DnsRecord>>>,
    }

    impl Default for DummyDnsServer {
        fn default() -> Self {
            DummyDnsServer {
                hosts: Arc::new(RwLock::new(vec![])),
            }
        }
    }
}