diff --git a/Cargo.toml b/Cargo.toml index e9b6669..fec06cd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "yps" -version = "0.1.1" +version = "0.1.2" edition = "2024" license = "MIT" readme = "README.md" @@ -12,5 +12,6 @@ repository = "https://github.com/YGGverse/yps" [dependencies] anyhow = "1.0" clap = { version = "4.5", features = ["derive"] } +colored = "3.0" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" diff --git a/src/main.rs b/src/main.rs index ff1afb6..643b777 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ mod udp; mod yggdrasil; use anyhow::Result; +use colored::*; use config::Config; use std::{ net::{IpAddr, SocketAddr, TcpStream}, @@ -20,82 +21,92 @@ fn main() -> Result<()> { if !config.tcp && config.udp.is_none() { panic!("at least one TCP or UDP protocol is required for scan!") } - let latency = config.latency.map(Duration::from_secs); // parse once let mut ygg = Yggdrasil::init(&config.socket)?; - let mut tcp: Option> = if config.tcp { + let mut key: Vec = 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 = Vec::with_capacity(config.index_capacity); + let index = Vec::with_capacity(config.index_capacity); (index, server) }); + println!("crawler started..."); // get initial peers to crawl - if config.debug { - println!("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, latency, &config, &mut ygg, &mut tcp, &mut udp)?; + crawl(peer.key, &config, &mut ygg, &mut key, &mut tcp, &mut udp)?; } Ok(()) } fn crawl( - key: String, - latency: Option, + k: String, config: &Config, ygg: &mut Yggdrasil, + key: &mut Vec, tcp: &mut Option>, udp: &mut Option<(Vec, Udp)>, ) -> Result<()> { - if config.debug { - println!("get peers for `{}`...", &key); - } - let p = ygg.remote_peers(key)?; - if p.status != "success" { - todo!() - } - 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 config.debug { - println!("try `{url}`..."); - } - if TcpStream::connect_timeout(&address, Duration::from_secs(1)).is_ok() { - println!("{url}"); - } - index.push(address) - } - } - if let Some((index, server)) = udp { - let url = format!("udp://{address}"); - if !index.contains(&address) { - if config.debug { - println!("try `{url}`..."); - } - if server.check(address) { - println!("{url}"); - } - index.push(address) - } - } + if !key.contains(&k) { + if config.debug { + println!("get peers for `{k}`..."); } - for k in peers.keys { - if let Some(l) = latency { - std::thread::sleep(l); + 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(i) = tcp { + if !i.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()); + } + i.push(address) + } + } + if let Some((i, server)) = udp { + let url = format!("udp://{address}"); + if !i.contains(&address) { + if server.check(address) { + println!("\t{}: {url}", SUCCESS.green()); + } else if config.debug { + println!("\t{}: {url}", FAILURE.red()); + } + i.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)?; + } } - crawl(k, latency, config, ygg, 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"; diff --git a/src/yggdrasil.rs b/src/yggdrasil.rs index 89bab8b..ed2a694 100644 --- a/src/yggdrasil.rs +++ b/src/yggdrasil.rs @@ -25,7 +25,7 @@ impl Yggdrasil { pub fn remote_peers( &mut self, - public_key: String, + public_key: &str, ) -> Result>> { let r: Response> = serde_json::from_slice(&self.request( format!("{{\"keepalive\":true,\"request\":\"debug_remotegetpeers\",\"arguments\":{{\"key\":\"{public_key}\"}}}}")