From 00579b38b21c9e9eed158f7d32a08a46548d2e1a Mon Sep 17 00:00:00 2001 From: Uncle Stretch Date: Sat, 28 Dec 2024 19:23:42 +0300 Subject: [PATCH] pretty print and ability to save file added Signed-off-by: Uncle Stretch --- Cargo.lock | 10 ++++++++++ Cargo.toml | 1 + src/locator.rs | 45 +++++++++++++++++++++++++++++++++++---------- src/main.rs | 4 ++++ src/walker.rs | 43 ++++++++++++++++--------------------------- 5 files changed, 66 insertions(+), 37 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 03c2114..135fa31 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -991,6 +991,7 @@ dependencies = [ "maxminddb", "parity-scale-codec", "pin-project", + "pretty-table", "primitive-types", "thiserror 2.0.9", "tokio", @@ -2432,6 +2433,15 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "pretty-table" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ebabdeeb9dde45278a6fd4caf2ec417c8c8d81b75d4450dc8d8305c25a4ac1d" +dependencies = [ + "thiserror 1.0.69", +] + [[package]] name = "primitive-types" version = "0.13.1" diff --git a/Cargo.toml b/Cargo.toml index 833c755..d1c4d63 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ ip_network = "0.4.1" libp2p = { version = "0.52.0", features = ["dns", "identify", "kad", "macros", "mdns", "noise", "ping", "tcp", "tokio", "yamux", "websocket", "request-response"] } maxminddb = "0.24.0" pin-project = "1.1.7" +pretty-table = "0.1.3" primitive-types = { version = "0.13.1", default-features = false, features = ["codec", "scale-info", "serde"] } thiserror = "2.0.9" tokio = { version = "1.42.0", features = ["macros", "time", "rt-multi-thread"] } diff --git a/src/locator.rs b/src/locator.rs index 6b72aa2..adda12d 100644 --- a/src/locator.rs +++ b/src/locator.rs @@ -1,16 +1,18 @@ use std::{collections::HashMap, net::IpAddr}; +use pretty_table::prelude::*; +use ip_network::IpNetwork; use trust_dns_resolver::{ config::{ResolverConfig, ResolverOpts}, TokioAsyncResolver, }; use maxminddb::{geoip2::City, Reader as GeoIpReader}; use libp2p::{ - identify::Info, multiaddr::Protocol, - PeerId + identify::Info, multiaddr::Protocol, Multiaddr, PeerId }; pub struct Locator<'a> { + maybe_filename: &'a Option, peers: &'a HashMap<&'a PeerId, &'a Info>, db: maxminddb::Reader<&'static [u8]>, } @@ -18,24 +20,39 @@ pub struct Locator<'a> { impl<'a> Locator<'a> { const CITY_DATA: &'static [u8] = include_bytes!("../artifacts/GeoLite2-City.mmdb"); - pub fn new(peers: &'a HashMap<&'a PeerId, &'a Info>) -> Self { + pub fn new( + peers: &'a HashMap<&'a PeerId, &'a Info>, + maybe_filename: &'a Option, + ) -> Self { Self { peers, + maybe_filename, db: GeoIpReader::from_source(Self::CITY_DATA) .expect("City data is always valid; qed"), } } + fn is_public_adress(&self, addr: &Multiaddr) -> bool { + let ip = match addr.iter().next() { + Some(Protocol::Ip4(ip)) => IpNetwork::from(ip), + Some(Protocol::Ip6(ip)) => IpNetwork::from(ip), + Some(Protocol::Dns(_)) | Some(Protocol::Dns4(_)) | Some(Protocol::Dns6(_)) => return true, + _ => return false, + }; + ip.is_global() + } + pub async fn locate_peers(&self) { - let mut geo_peers: HashMap = HashMap::new(); + let mut geo_peers = Vec::new(); let resolver = TokioAsyncResolver::tokio( ResolverConfig::default(), ResolverOpts::default(), ); - for (peer, info) in self.peers { + for (_, info) in self.peers { for addr in &info.listen_addrs { + if !self.is_public_adress(addr) { continue; } let located = match addr.iter().next() { Some(Protocol::Ip4(ip)) => self.locate(IpAddr::V4(ip)), Some(Protocol::Ip6(ip)) => self.locate(IpAddr::V6(ip)), @@ -50,14 +67,22 @@ impl<'a> Locator<'a> { _ => continue, }; - let Some(located) = located else { continue; }; - geo_peers.insert(**peer, located); + geo_peers.push(vec![ + addr.to_string(), + located.unwrap_or_default(), + ]); } } - for (peer, location) in geo_peers { - println!("\tPeer={peer} location={location}"); - } + match &self.maybe_filename { + Some(filename) => { + if let Err(err) = write_table_to_file(filename.clone(), geo_peers) { + tracing::error!("Could not write to file {} because of: {err}", &filename); + } + println!("Results saved to {}", &filename); + }, + None => print_table!(geo_peers), + }; } fn locate(&self, ip: IpAddr) -> Option { diff --git a/src/main.rs b/src/main.rs index 5ec0691..7ab3531 100644 --- a/src/main.rs +++ b/src/main.rs @@ -29,6 +29,9 @@ struct Opts { value_parser = parse_duration, )] timeout: std::time::Duration, + + #[clap(short, long, default_value = None, value_parser)] + filename: Option, } fn parse_duration(arg: &str) -> Result { @@ -48,6 +51,7 @@ async fn main() -> Result<(), Box> { .with_genesis(args.genesis) .with_bootnodes(args.bootnodes) .with_timeout(args.timeout) + .with_filename(args.filename) .build()? .walk_around() .await diff --git a/src/walker.rs b/src/walker.rs index 32b22b9..25ec63c 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -6,14 +6,12 @@ use std::{ use futures::StreamExt; use codec::Decode; -use ip_network::IpNetwork; use primitive_types::H256; use libp2p::{ identify::Info, identity, kad::{ Event as KademliaEvent, GetClosestPeersError, GetClosestPeersOk, QueryId, QueryResult, }, - multiaddr::Protocol, swarm::{Config as SwarmConfig, SwarmEvent}, Multiaddr, PeerId, Swarm }; @@ -37,6 +35,7 @@ use super::p2p::{ pub struct Walker { genesis: String, timeout: Duration, + maybe_filename: Option, swarm: Swarm, queries: HashSet, discovered_with_address: HashMap>, @@ -159,22 +158,12 @@ impl Walker { } } - fn is_public_adress(&self, addr: &Multiaddr) -> bool { - let ip = match addr.iter().next() { - Some(Protocol::Ip4(ip)) => IpNetwork::from(ip), - Some(Protocol::Ip6(ip)) => IpNetwork::from(ip), - Some(Protocol::Dns(_)) | Some(Protocol::Dns4(_)) | Some(Protocol::Dns6(_)) => return true, - _ => return false, - }; - ip.is_global() - } - pub async fn walk_around(&mut self) -> Result<(), Box> { let _ = tokio::time::timeout(self.timeout, self.inner_walk()).await; - println!("Number of dialed peers: {}", self.dialed_peers.len()); - println!("Total peers discovered: {}", self.discovered_with_address.len()); - println!("Peers with local iden.: {}", self.peer_details.len()); + println!("Number of dialed peers : {}", self.dialed_peers.len()); + println!("Total peers discovered : {}", self.discovered_with_address.len()); + println!("Peers with local identity : {}", self.peer_details.len()); let infos: HashMap<_, _> = self .peer_details @@ -185,19 +174,10 @@ impl Walker { .any(|stream_proto| stream_proto.as_ref().contains(&self.genesis)) }) .collect(); - println!("Support correct genesis {}: {}", &self.genesis, infos.len()); - - let peers_with_public: HashMap<_, _> = infos - .iter() - .filter(|(_, info)| info.listen_addrs - .iter() - .any(|a| self.is_public_adress(&a))) - .collect(); - println!("Peers with public address: {}", peers_with_public.len()); - println!("Peers with private address: {}", infos.len() - peers_with_public.len()); + println!("Support correct genesis : {}", infos.len()); println!("Peers with associated role: {}", self.peer_role.len()); - Locator::new(&infos).locate_peers().await; + Locator::new(&infos, &self.maybe_filename).locate_peers().await; Ok(()) } @@ -207,6 +187,7 @@ pub struct WalkerBuilder { genesis: String, bootnodes: Vec, timeout: Duration, + maybe_filename: Option, } impl Default for WalkerBuilder { @@ -215,6 +196,7 @@ impl Default for WalkerBuilder { genesis: Default::default(), bootnodes: Default::default(), timeout: Default::default(), + maybe_filename: None, } } } @@ -255,7 +237,8 @@ impl WalkerBuilder { let discovery = DiscoveryBuilder::new() .record_ttl(Some(Duration::from_secs(0))) .provider_ttl(Some(Duration::from_secs(0))) - .query_timeout(Duration::from_secs(5 * 60)) + //.query_timeout(Duration::from_secs(5 * 60)) + .query_timeout(self.timeout) .build(local_peer_id, &self.genesis); let peer_info = PeerBehaviour::new(local_key.public()); @@ -296,10 +279,16 @@ impl WalkerBuilder { self } + pub fn with_filename(mut self, maybe_filename: Option) -> Self { + self.maybe_filename = maybe_filename; + self + } + pub fn build(&self) -> Result> { let swarm = self.create_swarm()?; Ok(Walker { swarm, + maybe_filename: self.maybe_filename.clone(), genesis: self.genesis.clone(), timeout: self.timeout, queries: HashSet::with_capacity(1024),