yps/src/main.rs

mod config;
mod udp;
mod yggdrasil;

use anyhow::Result;
use colored::*;
use config::Config;
use std::{
    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: Vec<String> = Vec::with_capacity(config.index_capacity);
    let mut tcp = if config.tcp {
        Some(Vec::with_capacity(config.index_capacity))
    } else {
        None
    };
    let mut udp = config.udp.as_ref().map(|bind| {
        let server = Udp::init(bind).unwrap();
        let index = Vec::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
    for peer in p.response.unwrap().peers {
        crawl(peer.key, &config, &mut ygg, &mut key, &mut tcp, &mut udp)?;
    }
    Ok(())
}

fn crawl(
    k: String,
    config: &Config,
    ygg: &mut Yggdrasil,
    key: &mut Vec<String>,
    tcp: &mut Option<Vec<SocketAddr>>,
    udp: &mut Option<(Vec<SocketAddr>, Udp)>,
) -> Result<()> {
    if !key.contains(&k) {
        if config.debug {
            println!("get peers for `{k}`...");
        }
        let p = ygg.remote_peers(&k)?;
        if p.status == "success" {
            key.push(k);
            for (host, peers) in p.response.unwrap() {
                for port in &config.port {
                    let address = SocketAddr::new(IpAddr::V6(host), *port);
                    if let Some(index) = tcp {
                        if !index.contains(&address) {
                            let url = format!("tcp://{address}");
                            if TcpStream::connect_timeout(&address, Duration::from_secs(1)).is_ok()
                            {
                                println!("\t{}: {url}", SUCCESS.green())
                            } else if config.debug {
                                println!("\t{}: {url}", FAILURE.red());
                            }
                            index.push(address)
                        }
                    }
                    if let Some((index, server)) = udp {
                        let url = format!("udp://{address}");
                        if !index.contains(&address) {
                            if server.check(address) {
                                println!("\t{}: {url}", SUCCESS.green());
                            } else if config.debug {
                                println!("\t{}: {url}", FAILURE.red());
                            }
                            index.push(address)
                        }
                    }
                }
                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
            );
            key.push(k) // ban
        }
    }
    Ok(())
}

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