meowping 2.0.18

A flexible ping utility Tool written in Rust, that is focused on being size efficient and fast.
use super::watch::perform_subnet_watch;
use super::{Ipv4Subnet, Ipv6Subnet};
use crate::colors::Colorize;
use crate::icmp::ping_host_once;
use crate::output::{color_time, micros_to_ms, print_with_prefix};
use crate::tcp::tcp_connect_once;
use crate::udp::{ProbeOutcome, udp_probe_once};
use std::collections::HashSet;
use std::fmt::Write as _;
use std::net::{IpAddr, SocketAddr};
use std::thread;
use std::time::Duration;

#[derive(Clone, Copy)]
pub enum ProbeKind {
    Tcp {
        port: u16,
        timeout_ms: u64,
    },
    Udp {
        port: u16,
        timeout: Duration,
    },
    Icmp {
        timeout: Duration,
        ttl: u8,
        ident: u16,
        payload: [u8; 24],
    },
}

impl ProbeKind {
    pub const fn tcp(port: u16, timeout_ms: u64) -> Self {
        Self::Tcp { port, timeout_ms }
    }

    pub const fn udp(port: u16, timeout_ms: u64) -> Self {
        Self::Udp {
            port,
            timeout: Duration::from_millis(timeout_ms),
        }
    }

    pub const fn icmp(timeout_ms: u64, ttl: u8, ident: u16, payload: [u8; 24]) -> Self {
        Self::Icmp {
            timeout: Duration::from_millis(timeout_ms),
            ttl,
            ident,
            payload,
        }
    }

    pub(super) const fn header_protocol(&self) -> &'static str {
        match self {
            Self::Tcp { .. } => "TCP",
            Self::Udp { .. } => "UDP",
            Self::Icmp { .. } => "ICMP",
        }
    }

    pub(super) const fn header_port(&self) -> Option<u16> {
        match self {
            Self::Tcp { port, .. } | Self::Udp { port, .. } => Some(*port),
            Self::Icmp { .. } => None,
        }
    }
}

#[derive(Clone, Copy)]
enum ScanVerdict {
    Open { rtt_us: u128 },
    Down,
    UdpClosed,
    UdpNoResponse,
}

#[derive(Clone, Copy)]
pub(super) struct ScanResult {
    pub(super) host: IpAddr,
    verdict: ScanVerdict,
}

impl ScanResult {
    pub(super) const fn latency_us(&self) -> Option<u128> {
        match self.verdict {
            ScanVerdict::Open { rtt_us } => Some(rtt_us),
            _ => None,
        }
    }

    pub(super) const fn is_responsive(&self) -> bool {
        matches!(self.verdict, ScanVerdict::Open { .. })
    }

    pub(super) fn format(&self, minimal: bool) -> String {
        match self.verdict {
            ScanVerdict::Open { rtt_us } => {
                let colored_ip = self.host.to_string().green();
                if minimal {
                    colored_ip
                } else {
                    format!("{} {}", colored_ip, color_time(micros_to_ms(rtt_us)))
                }
            }
            ScanVerdict::Down => self.host.to_string().red(),
            ScanVerdict::UdpClosed => format!("{} closed", self.host.to_string().red()),
            ScanVerdict::UdpNoResponse => {
                format!("{} open|filtered", self.host.to_string().orange())
            }
        }
    }
}

const MAX_CONCURRENT_PROBES: usize = 512;

pub(super) fn probe_chunk(hosts: &[IpAddr], kind: &ProbeKind, seq: &mut u16) -> Vec<ScanResult> {
    let mut results = Vec::with_capacity(hosts.len());
    for batch in hosts.chunks(MAX_CONCURRENT_PROBES) {
        results.extend(probe_batch(batch, kind, seq));
    }
    results
}

fn probe_batch(hosts: &[IpAddr], kind: &ProbeKind, seq: &mut u16) -> Vec<ScanResult> {
    let mut handles = Vec::with_capacity(hosts.len());
    for &host in hosts {
        let verdict = match *kind {
            ProbeKind::Tcp { port, timeout_ms } => thread::spawn(move || {
                tcp_connect_once(host, port, timeout_ms).map_or(ScanVerdict::Down, |latency| {
                    ScanVerdict::Open {
                        rtt_us: latency.as_micros(),
                    }
                })
            }),
            ProbeKind::Udp { port, timeout } => {
                let payload = crate::udp::probe_payload(port);
                thread::spawn(move || {
                    let addr = SocketAddr::new(host, port);
                    match udp_probe_once(addr, &payload, timeout) {
                        ProbeOutcome::Open { rtt, .. } => ScanVerdict::Open {
                            rtt_us: rtt.as_micros(),
                        },
                        ProbeOutcome::Closed => ScanVerdict::UdpClosed,
                        ProbeOutcome::NoResponse => ScanVerdict::UdpNoResponse,
                    }
                })
            }
            ProbeKind::Icmp {
                timeout,
                ttl,
                ident,
                payload,
            } => {
                let current_seq = *seq;
                *seq = seq.wrapping_add(1);
                thread::spawn(move || {
                    match ping_host_once(host, current_seq, timeout, ttl, ident, &payload) {
                        Ok((_bytes, rtt)) => ScanVerdict::Open {
                            rtt_us: rtt.as_micros(),
                        },
                        Err(_) => ScanVerdict::Down,
                    }
                })
            }
        };
        handles.push((host, verdict));
    }

    handles
        .into_iter()
        .map(|(host, handle)| {
            let fallback = match kind {
                ProbeKind::Udp { .. } => ScanVerdict::UdpNoResponse,
                _ => ScanVerdict::Down,
            };
            ScanResult {
                host,
                verdict: handle.join().unwrap_or(fallback),
            }
        })
        .collect()
}

pub(super) fn ensure_rounds(rounds: usize) -> usize {
    rounds.max(1)
}

pub(super) fn format_scan_header(
    verb: &str,
    protocol: &str,
    subnet_notation: &str,
    host_count: u128,
    port: Option<u16>,
    suffix: Option<&str>,
) -> String {
    let mut message = format!(
        "{verb} {} ({} hosts) via {}",
        subnet_notation.bright_blue(),
        host_count,
        protocol
    );

    if let Some(p) = port {
        write!(&mut message, " port {p}").expect("writing to String should not fail");
    }

    if let Some(s) = suffix {
        write!(&mut message, " {s}").expect("writing to String should not fail");
    }

    message
}

pub(super) fn print_host_summary(total: usize, responsive: usize, minimal: bool) {
    let summary = format!(
        "Hosts responsive: {}/{}",
        responsive.to_string().green(),
        total
    );
    print_with_prefix(minimal, &summary);
}

pub(super) fn print_responsive_minimal(responsive_hosts: &HashSet<IpAddr>, minimal: bool) {
    let mut responsive_list: Vec<IpAddr> = responsive_hosts.iter().copied().collect();
    responsive_list.sort_by_key(|ip| match ip {
        IpAddr::V4(v4) => (0, u128::from(u32::from(*v4))),
        IpAddr::V6(v6) => (1, u128::from(*v6)),
    });
    if responsive_list.is_empty() {
        return;
    }
    let entries = responsive_list
        .iter()
        .map(|ip| ip.to_string().green())
        .collect::<Vec<_>>()
        .join(", ");
    let message = format!("[{entries}]");
    print_with_prefix(minimal, &message);
}

pub(super) struct ScanConfig {
    pub(super) notation: String,
    pub(super) host_count: u128,
    too_large: bool,
    pub(super) kind: ProbeKind,
    pub(super) rounds: Option<usize>,
    pub(super) minimal: bool,
}

fn print_empty_or_too_large(cfg: &ScanConfig, hosts: &[IpAddr]) -> bool {
    if cfg.host_count == 0 {
        let message = format!("{} has no usable host addresses", cfg.notation.yellow());
        print_with_prefix(cfg.minimal, &message);
        return true;
    }
    if cfg.too_large {
        let message = format!(
            "{} has too many addresses to scan (max /112 supported)",
            cfg.notation.yellow()
        );
        print_with_prefix(cfg.minimal, &message);
        return true;
    }
    if hosts.is_empty() {
        let message = format!("{} has no usable host addresses", cfg.notation.yellow());
        print_with_prefix(cfg.minimal, &message);
        return true;
    }
    false
}

fn scan_hosts(hosts: &[IpAddr], cfg: &ScanConfig) {
    if print_empty_or_too_large(cfg, hosts) {
        return;
    }

    perform_subnet_watch(hosts, cfg);
}

#[derive(Clone, Copy)]
pub enum Subnet {
    V4(Ipv4Subnet),
    V6(Ipv6Subnet),
}

pub fn perform_subnet_scan(subnet: Subnet, kind: ProbeKind, rounds: Option<usize>, minimal: bool) {
    let (hosts, notation, host_count, too_large) = match subnet {
        Subnet::V4(s) => (
            s.iter_hosts().map(IpAddr::V4).collect::<Vec<_>>(),
            s.notation(),
            s.host_count(),
            false,
        ),
        Subnet::V6(s) => {
            let host_count = s.host_count();
            (
                s.iter_hosts().map(IpAddr::V6).collect::<Vec<_>>(),
                s.notation(),
                host_count,
                host_count == u128::MAX,
            )
        }
    };

    scan_hosts(
        &hosts,
        &ScanConfig {
            notation,
            host_count,
            too_large,
            kind,
            rounds,
            minimal,
        },
    );
}