cidr_block 0.1.0

mod extras;
pub use extras::*;
use std::net::{Ipv4Addr, Ipv6Addr};

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("invalid_cidr_block_format")]
    InvalidBlock,
    #[error("invalid_prefix_length")]
    InvalidPrefixLength,
}

/// Calculates the number of IP addresses in an IPv6 CIDR block.
///
/// # Arguments
///
/// * `block` - A string slice that represents the IPv6 CIDR block (e.g., "2001:db8::/32").
///
/// # Returns
///
/// Returns `Ok(u128)` with the number of IPs if the CIDR block is valid.
/// Returns an `Error` if the block format is invalid or if the prefix length is out of range.
///
/// # Errors
///
/// * `Error::InvalidBlock` - If the CIDR block format is invalid.
/// * `Error::InvalidPrefixLength` - If the prefix length is greater than 128 or cannot be parsed.
///
pub fn v6_num_ips(block: &str) -> Result<u128, Error> {
    let parts: Vec<&str> = block.split('/').collect();
    if parts.len() != 2 {
        return Err(Error::InvalidBlock);
    }

    let prefix_length: u32 = parts[1].parse().map_err(|_| Error::InvalidPrefixLength)?;
    if prefix_length > 128 {
        return Err(Error::InvalidPrefixLength);
    }

    Ok(1 << (128 - prefix_length))
}
/// Calculates the number of IP addresses in an IPv4 CIDR block.
///
/// # Arguments
///
/// * `block` - A string slice that represents the IPv4 CIDR block (e.g., "192.168.0.0/24").
///
/// # Returns
///
/// Returns `Ok(u32)` with the number of IPs if the CIDR block is valid.
/// Returns an `Error` if the block format is invalid or if the prefix length is out of range.
///
/// # Errors
///
/// * `Error::InvalidBlock` - If the CIDR block format is invalid.
/// * `Error::InvalidPrefixLength` - If the prefix length is greater than 32 or cannot be parsed.
///
pub fn v4_num_ips(block: &str) -> Result<u32, Error> {
    let parts: Vec<&str> = block.split('/').collect();
    if parts.len() != 2 {
        return Err(Error::InvalidBlock);
    }

    let prefix_length: u32 = parts[1].parse().map_err(|_| Error::InvalidPrefixLength)?;
    if prefix_length > 32 {
        return Err(Error::InvalidPrefixLength);
    }

    Ok(1 << (32 - prefix_length))
}
/// Calculates the next IPv6 address within the same CIDR block, wrapping around if necessary.
///
/// # Arguments
///
/// * `ip` - An `Ipv6Addr` that represents the starting IPv6 address.
/// * `prefix_length` - The prefix length of the CIDR block.
/// * `n` - The number of addresses to skip.
///
/// # Returns
///
/// Returns the next `Ipv6Addr` within the CIDR block.
///
/// This function wraps around to the beginning of the block if the next address exceeds the block size.
pub fn next_ip_v6(ip: Ipv6Addr, prefix_length: u8, n: u128) -> Result<Ipv6Addr, Error> {
    if prefix_length > 128 {
        return Err(Error::InvalidPrefixLength);
    }

    let mask = if prefix_length == 0 {
        u128::MAX
    } else {
        (1 << (128 - prefix_length)) - 1
    };
    let ip: u128 = u128::from(ip);
    let mut host_part = ip & mask;
    let base_part = ip & !mask;

    // Use wrapping_add to increment host_part by n
    host_part = host_part.wrapping_add(n);

    // Ensure host_part is within mask range
    host_part &= mask;

    Ok(Ipv6Addr::from(base_part | host_part))
}

/// Calculates the next IPv4 address within the same CIDR block, wrapping around if necessary.
///
/// # Arguments
///
/// * `ip` - An `Ipv4Addr` that represents the starting IPv4 address.
/// * `prefix_length` - The prefix length of the CIDR block.
/// * `n` - The number of addresses to skip.
///
/// # Returns
///
/// Returns the next `Ipv4Addr` within the CIDR block.
///
/// This function wraps around to the beginning of the block if the next address exceeds the block size.
pub fn next_ip_v4(ip: Ipv4Addr, prefix_length: u8, n: u32) -> Result<Ipv4Addr, Error> {
    if prefix_length > 32 {
        return Err(Error::InvalidPrefixLength);
    }

    let mask = (1 << (32 - prefix_length)) - 1;
    let ip: u32 = u32::from(ip);
    let mut host_part = ip & mask;
    let base_part = ip & !mask;

    // Use wrapping_add to increment host_part by n
    host_part = host_part.wrapping_add(n);

    // Ensure host_part doesn't exceed mask
    if host_part > mask {
        host_part &= mask; // Reset to within the mask range
    }

    Ok(Ipv4Addr::from(base_part | host_part))
}
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_v4_num_ips() {
        assert_eq!(v4_num_ips("192.168.0.0/24").unwrap(), 256);
        assert_eq!(v4_num_ips("10.0.0.0/8").unwrap(), 16_777_216);
        assert!(v4_num_ips("192.168.0.0/33").is_err());
        assert!(v4_num_ips("invalid").is_err());
    }

    #[test]
    fn test_v6_num_ips() {
        assert_eq!(
            v6_num_ips("2001:db8::/32").unwrap(),
            79_228_162_514_264_337_593_543_950_336
        );
        assert_eq!(v6_num_ips("::/128").unwrap(), 1);
        assert!(v6_num_ips("2001:db8::/129").is_err());
        assert!(v6_num_ips("invalid").is_err());
    }
    #[test]
    fn test_next_ip_v6() {
        let ip = "2001:db8::1".parse::<Ipv6Addr>().unwrap();
        let next_ip = next_ip_v6(ip, 32, 10).unwrap();
        assert_eq!(next_ip, "2001:db8::b".parse::<Ipv6Addr>().unwrap());

        // Test wrapping behavior - when we reach the end of the CIDR block
        let ip = "2001:db8::ffff:ffff".parse::<Ipv6Addr>().unwrap();
        let next_ip = next_ip_v6(ip, 32, 2).unwrap();
        assert_eq!(next_ip, "2001:db8::1:0:1".parse::<Ipv6Addr>().unwrap());

        // Test overflow with large n value
        let ip = "2001:db8::ffff:ffff".parse::<Ipv6Addr>().unwrap();
        let next_ip = next_ip_v6(ip, 32, 100).unwrap();
        assert_eq!(next_ip, "2001:db8::1:0:63".parse::<Ipv6Addr>().unwrap());

        // Test with different prefix length (64)
        let ip = "2001:db8::1".parse::<Ipv6Addr>().unwrap();
        let next_ip = next_ip_v6(ip, 64, 10).unwrap();
        assert_eq!(next_ip, "2001:db8::b".parse::<Ipv6Addr>().unwrap());
    }

    #[test]
    fn test_next_ip_v4() {
        // 192.168.1.0/24
        let ip = "192.168.1.1".parse::<Ipv4Addr>().unwrap();
        let next_ip = next_ip_v4(ip, 24, 10).unwrap();
        assert_eq!(next_ip, "192.168.1.11".parse::<Ipv4Addr>().unwrap());
        // Test last IP to check if it wraps around
        let last_ip = "192.168.1.255".parse::<Ipv4Addr>().unwrap();
        let next_ip = next_ip_v4(last_ip, 24, 2).unwrap();
        assert_eq!(next_ip, "192.168.1.1".parse::<Ipv4Addr>().unwrap());
        // Add an overflowing test case
        let last_ip = "192.168.1.255".parse::<Ipv4Addr>().unwrap();
        let next_ip = next_ip_v4(last_ip, 24, 100).unwrap();
        assert_eq!(next_ip, "192.168.1.99".parse::<Ipv4Addr>().unwrap());
    }
}