diff --git a/Cargo.lock b/Cargo.lock index 135fa31..de817d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -539,6 +539,27 @@ dependencies = [ "typenum", ] +[[package]] +name = "csv" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acdc4883a9c96732e4733212c01447ebd805833b7275a73ca3ee080fd77afdaf" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "csv-core" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efa2b3d7902f4b634a20cae3c9c4e6209dc4779feb6863329607560143efa70" +dependencies = [ + "memchr", +] + [[package]] name = "ctr" version = "0.9.2" @@ -983,6 +1004,7 @@ dependencies = [ "asynchronous-codec 0.7.0", "bytes", "clap", + "csv", "either", "futures", "hex", @@ -2833,6 +2855,12 @@ dependencies = [ "static_assertions", ] +[[package]] +name = "ryu" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" + [[package]] name = "scale-info" version = "2.11.6" diff --git a/Cargo.toml b/Cargo.toml index d1c4d63..41ebdf7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,7 @@ asynchronous-codec = "0.7.0" bytes = "1.9.0" clap = { version = "4.5.23", features = ["derive"] } codec = { version = "3.6.12", package = "parity-scale-codec", features = ["derive"] } +csv = "1.3.1" either = "1.13.0" futures = "0.3.31" hex = "0.4.3" diff --git a/src/locator.rs b/src/locator.rs index adda12d..159e29a 100644 --- a/src/locator.rs +++ b/src/locator.rs @@ -1,4 +1,9 @@ -use std::{collections::HashMap, net::IpAddr}; +use std::{ + error::Error, + fs::File, + collections::{HashMap, HashSet}, + net::IpAddr, +}; use pretty_table::prelude::*; use ip_network::IpNetwork; @@ -43,7 +48,8 @@ impl<'a> Locator<'a> { } pub async fn locate_peers(&self) { - let mut geo_peers = Vec::new(); + let mut geo_peers: Vec> = Vec::new(); + let used_addrs: HashSet = HashSet::new(); let resolver = TokioAsyncResolver::tokio( ResolverConfig::default(), @@ -67,17 +73,25 @@ impl<'a> Locator<'a> { _ => continue, }; - geo_peers.push(vec![ - addr.to_string(), - located.unwrap_or_default(), - ]); + if let None = used_addrs.get(addr) { + geo_peers.push(vec![ + addr.to_string(), + located.unwrap_or_default(), + ]); + } } } 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); + if filename.ends_with(".csv") { + if let Err(err) = self.write_to_csv(&filename, &geo_peers) { + tracing::error!("Could not write to file {} because of: {err}", &filename); + } + } else { + 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); }, @@ -85,6 +99,22 @@ impl<'a> Locator<'a> { }; } + fn write_to_csv( + &self, + filename: &str, + data: &Vec>, + ) -> Result<(), Box> { + let file = File::create(filename)?; + let mut writer = csv::Writer::from_writer(file); + + for record in data { + writer.write_record(record)?; + } + + writer.flush()?; + Ok(()) + } + fn locate(&self, ip: IpAddr) -> Option { let City { city, .. } = self.db.lookup(ip).ok()?;