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};
pub struct Pinger {
mode: PingMode,
config: PingConfig,
icmp_pinger: Option<IcmpPinger>,
udp_pinger: Option<UdpPinger>,
tcp_pinger: Option<TcpPinger>,
}
impl Pinger {
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,
};
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)
}
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),
})
}
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,
})
}
pub fn ping_once(&self) -> PingResult<PingReply> {
self.ping_sequence(1)
}
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(),
})
}
}
}
}
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);
if i < self.config.count - 1 {
std::thread::sleep(self.config.interval);
}
}
stats.finalize(&rtts);
Ok(stats)
}
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);
callback(sequence, &result);
match result {
Ok(reply) => {
stats.add_reply(&reply);
rtts.push(reply.rtt);
}
Err(_) => {
}
}
sequence = sequence.wrapping_add(1);
let elapsed = start_time.elapsed();
if elapsed < self.config.interval {
std::thread::sleep(self.config.interval - elapsed);
}
}
stats.finalize(&rtts);
Ok(stats)
}
pub fn mode(&self) -> PingMode {
self.mode
}
pub fn config(&self) -> &PingConfig {
&self.config
}
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(),
}
}
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"
}
)
}
}
pub struct PingerBuilder {
config: PingConfig,
}
impl PingerBuilder {
pub fn new() -> Self {
Self {
config: PingConfig::default(),
}
}
pub fn target(mut self, target: std::net::Ipv4Addr) -> Self {
self.config.target = target;
self
}
pub fn count(mut self, count: u16) -> Self {
self.config.count = count;
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.config.timeout = timeout;
self
}
pub fn interval(mut self, interval: Duration) -> Self {
self.config.interval = interval;
self
}
pub fn packet_size(mut self, size: usize) -> Self {
self.config.packet_size = size;
self
}
pub fn identifier(mut self, id: u16) -> Self {
self.config.identifier = Some(id);
self
}
pub fn build_icmp(self) -> PingResult<Pinger> {
Pinger::new(self.config, PingMode::Icmp)
}
pub fn build_udp(self) -> PingResult<Pinger> {
Pinger::new(self.config, PingMode::Udp)
}
pub fn build_udp_with_port(self, port: u16) -> PingResult<Pinger> {
Pinger::new_udp_with_port(self.config, port)
}
pub fn build_tcp(self) -> PingResult<Pinger> {
Pinger::new(self.config, PingMode::Tcp)
}
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()
};
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));
}
}