rusty-ping 0.1.0

use clap::Parser;
use std::net::{IpAddr, ToSocketAddrs};
use tokio::sync::mpsc;
use tokio::sync::mpsc::error::TryRecvError;
use tokio::time::sleep;

mod args;
mod ip;
mod pinger;
mod print;
mod state;
mod stats;
use args::Args;
use ip::parse_ip;
use pinger::{PingResult, spawn_ping};
use print::{print_legend_line, print_ping_line, print_stats, print_status_line};
use state::State;
use stats::compute_stats;

fn parse_target(target_str: &str) -> IpAddr {
    if let Ok(ip) = parse_ip(target_str) {
        ip
    } else {
        // Try DNS resolution
        let dns_str = format!("{}:0", target_str);
        match dns_str.to_socket_addrs() {
            Ok(mut addrs) => {
                if let Some(addr) = addrs.next() {
                    addr.ip()
                } else {
                    eprintln!("Error: no DNS addresses found");
                    std::process::exit(1);
                }
            }
            Err(e) => {
                eprintln!("Error: DNS resolution failed: {}", e);
                std::process::exit(1);
            }
        }
    }
}

fn main_loop(
    args: &Args,
    states: &mut Vec<State>,
    target: &IpAddr,
    tx: &mpsc::Sender<PingResult>,
    rx: &mut mpsc::Receiver<PingResult>,
    packet_id: &mut u32,
) {
    let current_time = State::current_time_ms();
    for state in states.iter_mut() {
        if let State::NotYetReceived { time_sent } = state {
            // TODO: 2000ms Should be a configurable variable
            if current_time - *time_sent > 2000 {
                // 2s = 2000ms
                *state = State::Lost;
            }
        }
    }

    // Send the ping
    states.push(State::NotYetReceived {
        time_sent: current_time,
    });
    spawn_ping(*target, tx.clone(), *packet_id, args.ttl.unwrap_or(64));
    *packet_id += 1;

    // The update state loop
    loop {
        match rx.try_recv() {
            // Update the state of the received packet
            Ok(result) => {
                if args.beep {
                    print!("\x07"); // Rust doesn't support \a ?!
                }
                match result {
                    PingResult::Ping { rtt, id } => {
                        if (id as usize) < states.len() {
                            states[id as usize] = State::Received { rtt };
                        }
                    }
                    PingResult::NoPingErr { id } => {
                        states[id as usize] = State::SendError;
                    }
                }
            }
            Err(TryRecvError::Disconnected) => {
                eprintln!("Critical error, pinger thread killed");
                std::process::exit(2);
            }
            Err(TryRecvError::Empty) => {
                // No more ping to read
                break;
            }
        }
    }

    // Put the cursor back up
    if !std::process::Command::new("tput")
        .arg("rc")
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
    {
        println!("Process failed too");
        std::process::exit(1);
    }

    // Prints lines based on the state of each ping
    print_ping_line(states);

    // Print stats
    if args.stats {
        if let Some(stats) = compute_stats(states) {
            print_stats(&stats);
        }
        // Print stats of the last 30 pings
        let first_index = if (states.len() as i32) - 30 > 0 {
            states.len() - 30
        } else {
            0
        };
        if let Some(stats) = compute_stats(&states[first_index..]) {
            print_stats(&stats);
        }
    }
}

#[tokio::main]
async fn main() {
    let args = Args::parse();

    let target_str: String;
    let target: IpAddr = if let Some(target_arg) = args.target.clone() {
        target_str = target_arg;
        parse_target(&target_str)
    } else if let Some(default_target_str) = args.default_target.clone() {
        target_str = default_target_str;
        parse_target(&target_str)
    } else {
        eprintln!("Error: target address is required");
        std::process::exit(1);
    };

    let (tx, mut rx) = mpsc::channel::<PingResult>(10);

    let mut states: Vec<State> = Vec::new();
    let mut packet_id: u32 = 0;

    if args.legend {
        print_legend_line();
    }

    if args.status {
        print_status_line(&target_str, &target);
    }

    // Save the current cursor position
    if !std::process::Command::new("tput")
        .arg("sc")
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
    {
        println!("Process failed");
        std::process::exit(1);
    }

    let wait_time = args.interval.unwrap_or(1000);
    match args.count {
        Some(c) => {
            for _ in 1..c {
                main_loop(&args, &mut states, &target, &tx, &mut rx, &mut packet_id);
                sleep(tokio::time::Duration::from_millis(wait_time as u64)).await;
            }
        }
        None => loop {
            main_loop(&args, &mut states, &target, &tx, &mut rx, &mut packet_id);
            sleep(tokio::time::Duration::from_millis(wait_time as u64)).await;
        },
    }
}