anti-ping 0.1.2

//! ICMP ping implementation
//!
//! This module provides ICMP echo request/reply functionality for network connectivity testing.
//! It supports both raw sockets (requiring root privileges) and DGRAM sockets where available.

use anti_common::{calculate_checksum, icmp, PingConfig, PingError, PingReply, PingResult};
use bytes::{BufMut, BytesMut};
use socket2::{Domain, Protocol, Socket, Type};
use std::io::Read;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::os::unix::io::{AsRawFd, RawFd};
use std::time::{Duration, Instant};

/// ICMP packet structure for echo requests and replies
#[derive(Debug, Clone)]
pub struct IcmpPacket {
    /// ICMP type (8 for echo request, 0 for echo reply)
    pub icmp_type: u8,
    /// ICMP code (usually 0 for echo)
    pub code: u8,
    /// Internet checksum
    pub checksum: u16,
    /// Identifier to match requests with replies
    pub identifier: u16,
    /// Sequence number for ordering
    pub sequence: u16,
    /// Payload data
    pub data: Vec<u8>,
}

impl IcmpPacket {
    /// Create a new ICMP echo request packet
    pub fn new_echo_request(identifier: u16, sequence: u16, data_size: usize) -> Self {
        let data = vec![0x08; data_size.max(8).min(1024)]; // Reasonable size limits

        Self {
            icmp_type: icmp::ECHO_REQUEST,
            code: 0,
            checksum: 0,
            identifier,
            sequence,
            data,
        }
    }

    /// Create an ICMP packet from raw bytes
    pub fn from_bytes(data: &[u8]) -> PingResult<Self> {
        if data.len() < 8 {
            return Err(PingError::InvalidResponse {
                reason: "ICMP packet too short".to_string(),
            });
        }

        Ok(Self {
            icmp_type: data[0],
            code: data[1],
            checksum: u16::from_be_bytes([data[2], data[3]]),
            identifier: u16::from_be_bytes([data[4], data[5]]),
            sequence: u16::from_be_bytes([data[6], data[7]]),
            data: data[8..].to_vec(),
        })
    }

    /// Convert packet to bytes for transmission
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut buf = BytesMut::new();
        buf.put_u8(self.icmp_type);
        buf.put_u8(self.code);
        buf.put_u16(self.checksum);
        buf.put_u16(self.identifier);
        buf.put_u16(self.sequence);
        buf.extend_from_slice(&self.data);
        buf.to_vec()
    }

    /// Calculate and set the checksum for this packet
    pub fn calculate_checksum(&mut self) {
        self.checksum = 0;
        let bytes = self.to_bytes();
        self.checksum = calculate_checksum(&bytes);
    }

    /// Check if this is an echo reply packet
    pub fn is_echo_reply(&self) -> bool {
        self.icmp_type == icmp::ECHO_REPLY
    }

    /// Check if this packet matches the given identifier and sequence
    pub fn matches(&self, identifier: u16, sequence: u16) -> bool {
        self.identifier == identifier && self.sequence == sequence
    }
}

/// ICMP socket wrapper that handles both raw and DGRAM sockets
pub struct IcmpSocket {
    socket: Socket,
    is_raw: bool,
}

impl IcmpSocket {
    /// Create a new ICMP socket
    ///
    /// This attempts to create a DGRAM socket first (non-root), then falls back to
    /// a raw socket (requires root privileges) if needed.
    pub fn new() -> PingResult<Self> {
        // Try DGRAM first (non-root on some systems like macOS)
        match Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::ICMPV4)) {
            Ok(socket) => {
                socket
                    .set_nonblocking(false)
                    .map_err(|e| PingError::SocketCreation(e.to_string()))?;
                socket
                    .set_read_timeout(Some(Duration::from_secs(5)))
                    .map_err(|e| PingError::SocketCreation(e.to_string()))?;
                socket.set_broadcast(true).ok(); // Enable broadcast for DGRAM sockets

                Ok(Self {
                    socket,
                    is_raw: false,
                })
            }
            Err(_) => {
                // Fall back to raw socket (requires root)
                let socket = Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::ICMPV4))
                    .map_err(|e| {
                        if e.kind() == std::io::ErrorKind::PermissionDenied {
                            PingError::PermissionDenied {
                                context: "ICMP ping requires root privileges. Try running with sudo or use UDP/TCP ping instead.".to_string(),
                            }
                        } else {
                            PingError::SocketCreation(e.to_string())
                        }
                    })?;

                socket
                    .set_nonblocking(false)
                    .map_err(|e| PingError::SocketCreation(e.to_string()))?;
                socket
                    .set_read_timeout(Some(Duration::from_secs(5)))
                    .map_err(|e| PingError::SocketCreation(e.to_string()))?;

                Ok(Self {
                    socket,
                    is_raw: true,
                })
            }
        }
    }

    /// Connect the socket to a target address (only for raw sockets)
    pub fn connect(&self, target: Ipv4Addr) -> PingResult<()> {
        if self.is_raw {
            let addr = SocketAddr::new(IpAddr::V4(target), 0);
            self.socket
                .connect(&addr.into())
                .map_err(|e| PingError::SocketCreation(e.to_string()))
        } else {
            // DGRAM sockets don't need to be connected for ICMP
            Ok(())
        }
    }

    /// Send an ICMP packet
    pub fn send(&self, packet: &IcmpPacket, target: Option<Ipv4Addr>) -> PingResult<usize> {
        let mut packet = packet.clone();
        packet.calculate_checksum();
        let bytes = packet.to_bytes();

        let result = if self.is_raw {
            self.socket.send(&bytes)
        } else {
            // For DGRAM sockets, always use send_to with target address
            let target_addr = target.unwrap_or(Ipv4Addr::new(127, 0, 0, 1));
            let addr = SocketAddr::new(IpAddr::V4(target_addr), 0);
            self.socket.send_to(&bytes, &addr.into())
        };

        result.map_err(|e| PingError::SocketCreation(e.to_string()))
    }

    /// Receive an ICMP packet
    pub fn recv(&self, timeout: Duration) -> PingResult<(IcmpPacket, Ipv4Addr, Option<u8>)> {
        self.socket
            .set_read_timeout(Some(timeout))
            .map_err(|e| PingError::SocketCreation(e.to_string()))?;

        let mut buf = [0u8; 1024];
        let start = Instant::now();

        loop {
            if start.elapsed() >= timeout {
                return Err(PingError::Timeout { duration: timeout });
            }

            let size = if self.is_raw {
                match (&self.socket).read(&mut buf) {
                    Ok(n) => n,
                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                        std::thread::sleep(Duration::from_millis(1));
                        continue;
                    }
                    Err(e) => return Err(PingError::SocketCreation(e.to_string())),
                }
            } else {
                // For DGRAM sockets, use recv_from with proper buffer handling
                let mut uninit_buffer = [std::mem::MaybeUninit::<u8>::uninit(); 1024];
                match self.socket.recv_from(&mut uninit_buffer) {
                    Ok((n, _from_addr)) => {
                        // Copy from MaybeUninit to regular buffer
                        for i in 0..n {
                            buf[i] = unsafe { uninit_buffer[i].assume_init() };
                        }
                        n
                    }
                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                        std::thread::sleep(Duration::from_millis(1));
                        continue;
                    }
                    Err(e) => return Err(PingError::SocketCreation(e.to_string())),
                }
            };

            // On macOS, both RAW and DGRAM sockets include IP headers
            let icmp_data = if size > 20 {
                // Check if this looks like an IP packet (version 4)
                if (buf[0] >> 4) == 4 {
                    let ip_header_len = ((buf[0] & 0x0F) * 4) as usize;
                    if size > ip_header_len {
                        &buf[ip_header_len..size]
                    } else {
                        continue;
                    }
                } else if size >= 8 {
                    // Direct ICMP data
                    &buf[..size]
                } else {
                    continue;
                }
            } else if size >= 8 {
                // Direct ICMP data for smaller packets
                &buf[..size]
            } else {
                continue;
            };

            match IcmpPacket::from_bytes(icmp_data) {
                Ok(packet) => {
                    // Extract source IP from IP header if available
                    let (source_ip, ttl) = if size >= 20 && (buf[0] >> 4) == 4 {
                        (
                            Ipv4Addr::new(buf[12], buf[13], buf[14], buf[15]),
                            Some(buf[8]),
                        )
                    } else {
                        (Ipv4Addr::new(0, 0, 0, 0), None)
                    };
                    return Ok((packet, source_ip, ttl));
                }
                Err(_) => continue, // Invalid ICMP packet, keep trying
            }
        }
    }

    /// Check if this socket is using raw mode
    pub fn is_raw(&self) -> bool {
        self.is_raw
    }
}

impl AsRawFd for IcmpSocket {
    fn as_raw_fd(&self) -> RawFd {
        self.socket.as_raw_fd()
    }
}

/// ICMP-specific pinger implementation
pub struct IcmpPinger {
    socket: IcmpSocket,
    config: PingConfig,
    identifier: u16,
}

impl IcmpPinger {
    /// Create a new ICMP pinger
    pub fn new(config: PingConfig) -> PingResult<Self> {
        let socket = IcmpSocket::new()?;
        socket.connect(config.target)?;

        let identifier = config.identifier.unwrap_or_else(|| rand::random::<u16>());

        Ok(Self {
            socket,
            config,
            identifier,
        })
    }

    /// Send a single ICMP ping and wait for reply
    pub fn ping(&self, sequence: u16) -> PingResult<PingReply> {
        let packet = IcmpPacket::new_echo_request(
            self.identifier,
            sequence,
            self.config.packet_size.saturating_sub(8), // Account for ICMP header
        );

        let start = Instant::now();

        // Send the packet
        self.socket.send(&packet, Some(self.config.target))?;

        // Wait for reply
        loop {
            let elapsed = start.elapsed();
            if elapsed >= self.config.timeout {
                return Err(PingError::Timeout {
                    duration: self.config.timeout,
                });
            }

            let remaining = self.config.timeout - elapsed;
            match self.socket.recv(remaining) {
                Ok((reply_packet, source, ttl)) => {
                    if reply_packet.is_echo_reply()
                        && reply_packet.matches(self.identifier, sequence)
                    {
                        let rtt = start.elapsed();
                        return Ok(PingReply {
                            sequence,
                            rtt,
                            bytes_received: reply_packet.to_bytes().len(),
                            from: if source.is_unspecified() {
                                self.config.target
                            } else {
                                source
                            },
                            ttl,
                        });
                    }
                    // Wrong packet, continue waiting
                }
                Err(PingError::Timeout { .. }) => {
                    return Err(PingError::Timeout {
                        duration: self.config.timeout,
                    });
                }
                Err(e) => return Err(e),
            }
        }
    }

    /// Check if the underlying socket is using raw mode
    pub fn is_raw(&self) -> bool {
        self.socket.is_raw()
    }
}

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

    #[test]
    fn test_icmp_packet_creation() {
        let packet = IcmpPacket::new_echo_request(12345, 1, 56);
        assert_eq!(packet.icmp_type, icmp::ECHO_REQUEST);
        assert_eq!(packet.code, 0);
        assert_eq!(packet.identifier, 12345);
        assert_eq!(packet.sequence, 1);
        assert_eq!(packet.data.len(), 56);
    }

    #[test]
    fn test_icmp_packet_serialization() {
        let mut packet = IcmpPacket::new_echo_request(12345, 1, 8);
        packet.calculate_checksum();

        let bytes = packet.to_bytes();
        assert!(bytes.len() >= 16); // 8 header + 8 data

        let parsed = IcmpPacket::from_bytes(&bytes).unwrap();
        assert_eq!(parsed.icmp_type, packet.icmp_type);
        assert_eq!(parsed.identifier, packet.identifier);
        assert_eq!(parsed.sequence, packet.sequence);
    }

    #[test]
    fn test_packet_matching() {
        let packet = IcmpPacket {
            icmp_type: icmp::ECHO_REPLY,
            code: 0,
            checksum: 0,
            identifier: 12345,
            sequence: 42,
            data: vec![],
        };

        assert!(packet.is_echo_reply());
        assert!(packet.matches(12345, 42));
        assert!(!packet.matches(12345, 41));
        assert!(!packet.matches(12344, 42));
    }

    #[test]
    fn test_checksum_calculation() {
        let mut packet = IcmpPacket::new_echo_request(1, 1, 8);
        packet.calculate_checksum();
        assert_ne!(packet.checksum, 0);
    }
}