anti_ping/

pinger.rs

//! Main pinger coordinator module
//!
//! This module provides a unified interface for all ping types (ICMP, UDP, TCP).
//! It handles the creation and coordination of the appropriate ping implementation
//! based on the configured mode.

use crate::icmp::IcmpPinger;
use crate::tcp::TcpPinger;
use crate::udp::UdpPinger;
use anti_common::{PingConfig, PingError, PingMode, PingReply, PingResult, PingStatistics};
use std::time::{Duration, Instant};

/// Main pinger that coordinates different ping implementations
pub struct Pinger {
    mode: PingMode,
    config: PingConfig,
    icmp_pinger: Option<IcmpPinger>,
    udp_pinger: Option<UdpPinger>,
    tcp_pinger: Option<TcpPinger>,
}

impl Pinger {
    /// Create a new pinger with the specified configuration and mode
    pub fn new(config: PingConfig, mode: PingMode) -> PingResult<Self> {
        let mut pinger = Self {
            mode,
            config: config.clone(),
            icmp_pinger: None,
            udp_pinger: None,
            tcp_pinger: None,
        };

        // Initialize the appropriate pinger based on mode
        match mode {
            PingMode::Icmp => {
                pinger.icmp_pinger = Some(IcmpPinger::new(config)?);
            }
            PingMode::Udp => {
                pinger.udp_pinger = Some(UdpPinger::new(config)?);
            }
            PingMode::Tcp => {
                pinger.tcp_pinger = Some(TcpPinger::new(config));
            }
        }

        Ok(pinger)
    }

    /// Create a new TCP pinger with a specific port
    pub fn new_tcp_with_port(config: PingConfig, port: u16) -> PingResult<Self> {
        let tcp_pinger = TcpPinger::new_with_port(config.clone(), port);

        Ok(Self {
            mode: PingMode::Tcp,
            config,
            icmp_pinger: None,
            udp_pinger: None,
            tcp_pinger: Some(tcp_pinger),
        })
    }

    /// Create a new UDP pinger with a specific port
    pub fn new_udp_with_port(config: PingConfig, port: u16) -> PingResult<Self> {
        let udp_pinger = UdpPinger::new_with_port(config.clone(), port)?;

        Ok(Self {
            mode: PingMode::Udp,
            config,
            icmp_pinger: None,
            udp_pinger: Some(udp_pinger),
            tcp_pinger: None,
        })
    }

    /// Send a single ping and return the result
    pub fn ping_once(&self) -> PingResult<PingReply> {
        self.ping_sequence(1)
    }

    /// Send a ping with a specific sequence number
    pub fn ping_sequence(&self, sequence: u16) -> PingResult<PingReply> {
        match self.mode {
            PingMode::Icmp => {
                if let Some(ref pinger) = self.icmp_pinger {
                    pinger.ping(sequence)
                } else {
                    Err(PingError::Configuration {
                        message: "ICMP pinger not initialized".to_string(),
                    })
                }
            }
            PingMode::Udp => {
                if let Some(ref pinger) = self.udp_pinger {
                    pinger.ping(sequence)
                } else {
                    Err(PingError::Configuration {
                        message: "UDP pinger not initialized".to_string(),
                    })
                }
            }
            PingMode::Tcp => {
                if let Some(ref pinger) = self.tcp_pinger {
                    pinger.ping(sequence)
                } else {
                    Err(PingError::Configuration {
                        message: "TCP pinger not initialized".to_string(),
                    })
                }
            }
        }
    }

    /// Send multiple pings according to the configuration
    pub fn ping_all(&self) -> PingResult<PingStatistics> {
        let mut stats = PingStatistics::new();
        let mut rtts = Vec::new();
        let mut sequence = 1u16;

        println!("PING {} using {} mode", self.config.target, self.mode);
        println!(
            "Sending {} packets with {}ms interval",
            self.config.count,
            self.config.interval.as_millis()
        );

        for i in 0..self.config.count {
            stats.add_transmitted();

            match self.ping_sequence(sequence) {
                Ok(reply) => {
                    stats.add_reply(&reply);
                    rtts.push(reply.rtt);

                    println!(
                        "Reply from {}: seq={} time={:.2}ms bytes={}",
                        reply.from,
                        reply.sequence,
                        reply.rtt.as_secs_f64() * 1000.0,
                        reply.bytes_received
                    );
                }
                Err(e) => match e {
                    PingError::Timeout { .. } => {
                        println!("Request timeout for seq={}", sequence);
                    }
                    _ => {
                        println!("Error for seq={}: {}", sequence, e);
                    }
                },
            }

            sequence = sequence.wrapping_add(1);

            // Sleep between pings (except for the last one)
            if i < self.config.count - 1 {
                std::thread::sleep(self.config.interval);
            }
        }

        stats.finalize(&rtts);
        Ok(stats)
    }

    /// Send pings with a custom callback for each result
    pub fn ping_with_callback<F>(&self, mut callback: F) -> PingResult<PingStatistics>
    where
        F: FnMut(u16, &PingResult<PingReply>),
    {
        let mut stats = PingStatistics::new();
        let mut rtts = Vec::new();
        let mut sequence = 1u16;

        for _i in 0..self.config.count {
            stats.add_transmitted();

            let start_time = Instant::now();
            let result = self.ping_sequence(sequence);

            // Call the callback with the result
            callback(sequence, &result);

            match result {
                Ok(reply) => {
                    stats.add_reply(&reply);
                    rtts.push(reply.rtt);
                }
                Err(_) => {
                    // Error already handled by callback
                }
            }

            sequence = sequence.wrapping_add(1);

            // Ensure we don't send too quickly
            let elapsed = start_time.elapsed();
            if elapsed < self.config.interval {
                std::thread::sleep(self.config.interval - elapsed);
            }
        }

        stats.finalize(&rtts);
        Ok(stats)
    }

    /// Get the ping mode
    pub fn mode(&self) -> PingMode {
        self.mode
    }

    /// Get the configuration
    pub fn config(&self) -> &PingConfig {
        &self.config
    }

    /// Check if the pinger is properly initialized
    pub fn is_initialized(&self) -> bool {
        match self.mode {
            PingMode::Icmp => self.icmp_pinger.is_some(),
            PingMode::Udp => self.udp_pinger.is_some(),
            PingMode::Tcp => self.tcp_pinger.is_some(),
        }
    }

    /// Get information about the pinger setup
    pub fn info(&self) -> String {
        let mode_info = match self.mode {
            PingMode::Icmp => {
                if let Some(ref pinger) = self.icmp_pinger {
                    format!(
                        "ICMP (socket type: {})",
                        if pinger.is_raw() { "raw" } else { "dgram" }
                    )
                } else {
                    "ICMP (not initialized)".to_string()
                }
            }
            PingMode::Udp => {
                if let Some(ref pinger) = self.udp_pinger {
                    format!("UDP (base port: {})", pinger.base_port())
                } else {
                    "UDP (not initialized)".to_string()
                }
            }
            PingMode::Tcp => {
                if let Some(ref pinger) = self.tcp_pinger {
                    format!("TCP (port: {})", pinger.target_port())
                } else {
                    "TCP (not initialized)".to_string()
                }
            }
        };

        format!(
            "Pinger: {} -> {} ({})",
            self.config.target,
            mode_info,
            if self.is_initialized() {
                "ready"
            } else {
                "not ready"
            }
        )
    }
}

/// Builder pattern for creating pingers with custom configurations
pub struct PingerBuilder {
    config: PingConfig,
}

impl PingerBuilder {
    /// Create a new pinger builder
    pub fn new() -> Self {
        Self {
            config: PingConfig::default(),
        }
    }

    /// Set the target IP address
    pub fn target(mut self, target: std::net::Ipv4Addr) -> Self {
        self.config.target = target;
        self
    }

    /// Set the number of pings to send
    pub fn count(mut self, count: u16) -> Self {
        self.config.count = count;
        self
    }

    /// Set the timeout for each ping
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.config.timeout = timeout;
        self
    }

    /// Set the interval between pings
    pub fn interval(mut self, interval: Duration) -> Self {
        self.config.interval = interval;
        self
    }

    /// Set the packet size
    pub fn packet_size(mut self, size: usize) -> Self {
        self.config.packet_size = size;
        self
    }

    /// Set a custom identifier
    pub fn identifier(mut self, id: u16) -> Self {
        self.config.identifier = Some(id);
        self
    }

    /// Build an ICMP pinger
    pub fn build_icmp(self) -> PingResult<Pinger> {
        Pinger::new(self.config, PingMode::Icmp)
    }

    /// Build a UDP pinger
    pub fn build_udp(self) -> PingResult<Pinger> {
        Pinger::new(self.config, PingMode::Udp)
    }

    /// Build a UDP pinger with custom port
    pub fn build_udp_with_port(self, port: u16) -> PingResult<Pinger> {
        Pinger::new_udp_with_port(self.config, port)
    }

    /// Build a TCP pinger
    pub fn build_tcp(self) -> PingResult<Pinger> {
        Pinger::new(self.config, PingMode::Tcp)
    }

    /// Build a TCP pinger with custom port
    pub fn build_tcp_with_port(self, port: u16) -> PingResult<Pinger> {
        Pinger::new_tcp_with_port(self.config, port)
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::Ipv4Addr;

    #[test]
    fn test_pinger_builder() {
        let builder = PingerBuilder::new()
            .target(Ipv4Addr::new(8, 8, 8, 8))
            .count(3)
            .timeout(Duration::from_secs(2))
            .interval(Duration::from_millis(500))
            .packet_size(128)
            .identifier(12345);

        assert_eq!(builder.config.target, Ipv4Addr::new(8, 8, 8, 8));
        assert_eq!(builder.config.count, 3);
        assert_eq!(builder.config.timeout, Duration::from_secs(2));
        assert_eq!(builder.config.interval, Duration::from_millis(500));
        assert_eq!(builder.config.packet_size, 128);
        assert_eq!(builder.config.identifier, Some(12345));
    }

    #[test]
    fn test_pinger_info() {
        let config = PingConfig {
            target: Ipv4Addr::new(127, 0, 0, 1),
            ..Default::default()
        };

        // TCP pinger should always initialize successfully
        let tcp_pinger = Pinger::new(config, PingMode::Tcp).unwrap();
        assert!(tcp_pinger.is_initialized());
        assert_eq!(tcp_pinger.mode(), PingMode::Tcp);

        let info = tcp_pinger.info();
        assert!(info.contains("TCP"));
        assert!(info.contains("127.0.0.1"));
    }

    #[test]
    fn test_pinger_config_access() {
        let config = PingConfig {
            target: Ipv4Addr::new(8, 8, 8, 8),
            count: 5,
            timeout: Duration::from_secs(3),
            ..Default::default()
        };

        let pinger = Pinger::new(config, PingMode::Tcp).unwrap();

        assert_eq!(pinger.config().target, Ipv4Addr::new(8, 8, 8, 8));
        assert_eq!(pinger.config().count, 5);
        assert_eq!(pinger.config().timeout, Duration::from_secs(3));
    }
}