yps/src/main.rs

mod config;
mod udp;
mod yggdrasil;

use anyhow::Result;
use colored::*;
use config::Config;
use std::{
    collections::{HashMap, HashSet},
    net::{IpAddr, SocketAddr, TcpStream},
    time::Duration,
};
use udp::Udp;
use yggdrasil::Yggdrasil;

fn main() -> Result<()> {
    use clap::Parser;
    let config = Config::parse();
    if config.port.is_empty() {
        panic!("at least one port must be specified for scan!")
    }
    if !config.tcp && config.udp.is_none() {
        panic!("at least one TCP or UDP protocol is required for scan!")
    }
    let mut ygg = Yggdrasil::init(&config.socket)?;
    let mut key: HashSet<String> = HashSet::with_capacity(config.index_capacity);
    let mut tcp = if config.tcp {
        Some(HashMap::with_capacity(config.index_capacity))
    } else {
        None
    };
    let mut udp = config.udp.as_ref().map(|bind| {
        let server = Udp::init(bind).unwrap();
        let index = HashMap::with_capacity(config.index_capacity);
        (index, server)
    });
    println!("crawler started...");
    // get initial peers to crawl
    let p = ygg.peers()?;
    if p.status != "success" {
        todo!()
    }
    // start crawler
    match p.response {
        Some(response) => {
            for peer in response.peers {
                crawl(peer.key, &config, &mut ygg, &mut key, &mut tcp, &mut udp)?;
            }
        }
        None => println!("node has no peers to connect."),
    }
    println!("operation completed!");
    if !key.is_empty() {
        println!("\tkeys crawled: {}", key.len());
    }
    if let Some(tcp) = tcp {
        println!("\tTCP: {}", tcp.len());
        println!("\t\tfound: {}", tcp.values().filter(|v| **v).count());
        println!("\t\ttotal: {}", tcp.len());
    }
    if let Some((udp, _)) = udp {
        println!("\tUDP: {}", udp.len());
        println!("\t\tfound: {}", udp.values().filter(|v| **v).count());
        println!("\t\ttotal: {}", udp.len());
    }
    Ok(())
}

fn crawl(
    k: String,
    config: &Config,
    ygg: &mut Yggdrasil,
    key: &mut HashSet<String>,
    tcp: &mut Option<HashMap<SocketAddr, bool>>,
    udp: &mut Option<(HashMap<SocketAddr, bool>, Udp)>,
) -> Result<()> {
    if key.contains(&k) {
        return Ok(());
    }
    if config.debug {
        println!("get peers for `{k}`...")
    }
    let p = ygg.remote_peers(&k)?;
    if p.status == "success" {
        assert!(key.insert(k));
        if let Some(response) = p.response {
            for (host, peers) in response {
                for port in &config.port {
                    let address = SocketAddr::new(IpAddr::V6(host), *port);
                    if let Some(index) = tcp
                        && !index.contains_key(&address)
                    {
                        let url = format!("tcp://{address}");
                        let result =
                            TcpStream::connect_timeout(&address, Duration::from_secs(1)).is_ok();
                        if result {
                            println!("\t{}: {url}", SUCCESS.green())
                        } else if config.debug {
                            println!("\t{}: {url}", FAILURE.red())
                        }
                        assert!(index.insert(address, result).is_none())
                    }
                    if let Some((index, server)) = udp
                        && !index.contains_key(&address)
                    {
                        let url = format!("udp://{address}");
                        let result = server.check(address);
                        if result {
                            println!("\t{}: {url}", SUCCESS.green())
                        } else if config.debug {
                            println!("\t{}: {url}", FAILURE.red())
                        }
                        assert!(index.insert(address, result).is_none())
                    }
                }
                for k in peers.keys {
                    if let Some(l) = config.latency {
                        std::thread::sleep(Duration::from_secs(l))
                    }
                    crawl(k, config, ygg, key, tcp, udp)?;
                }
            }
        }
    } else if config.debug {
        println!(
            "\t{}: peer `{k}` return status `{}`, skip.",
            WARNING.yellow(),
            p.status
        );
        assert!(key.insert(k)) // ban
    }
    Ok(())
}

const FAILURE: &str = "FAILURE";
const SUCCESS: &str = "SUCCESS";
const WARNING: &str = "WARNING";