diff --git a/src/dnslib.rs b/src/dnslib.rs index b2ea2dc..93345c8 100644 --- a/src/dnslib.rs +++ b/src/dnslib.rs @@ -1,17 +1,15 @@ -use trust_dns_resolver::proto::{rr::RecordType, xfer::DnsRequestOptions}; - use { - crate::structs::DomainData, + crate::structs::{DomainData, LibOptions}, futures::stream::{self, StreamExt}, - std::collections::{HashMap, HashSet}, - trust_dns_resolver::config::ResolverOpts, -}; - -use { - std::net::SocketAddr, + std::{ + collections::{HashMap, HashSet}, + net::SocketAddr, + }, trust_dns_resolver::{ - config::{NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig}, + config::{NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig, ResolverOpts}, + lookup::{Ipv4Lookup, Lookup}, name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime}, + proto::{rr::RecordType, xfer::DnsRequestOptions}, AsyncResolver, TokioAsyncResolver, }, }; @@ -52,63 +50,58 @@ pub fn return_tokio_asyncresolver( .unwrap() } -#[allow(clippy::too_many_arguments)] -pub async fn return_hosts_data( - hosts: HashSet, - resolver: AsyncResolver>, - trustable_resolver: AsyncResolver>, - wildcard_ips: HashSet, - disable_double_check: bool, - mut threads: usize, - show_ip_addresses: bool, - quiet_flag: bool, -) -> HashMap { - if hosts.len() < threads { - threads = hosts.len(); - } +pub async fn return_hosts_data(options: &LibOptions) -> HashMap { + let threads = if options.hosts.len() < options.threads { + options.hosts.len() + } else { + options.threads + }; - stream::iter(hosts) - .map(|host| { - let resolver_fut = resolver.ipv4_lookup(host.trim_end_matches('.').to_owned() + "."); - let trustable_resolver_fut = - trustable_resolver.ipv4_lookup(host.trim_end_matches('.').to_owned() + "."); - let wildcard_ips = wildcard_ips.clone(); + stream::iter(options.hosts.clone().into_iter().map(|host| { + let lookup_host = host.trim_end_matches('.').to_owned() + "."; - let mut domain_data = DomainData::default(); + let resolver_fut = options.resolvers.ipv4_lookup(lookup_host.clone()); + let trustable_resolver_fut = options.trustable_resolver.ipv4_lookup(lookup_host); + let wildcard_ips = options.wildcard_ips.clone(); - async move { - if let Ok(ip) = resolver_fut.await { - if disable_double_check { - domain_data.ipv4_addresses = ip - .into_iter() - .map(|x| x.to_string()) - .collect::>(); - } else if let Ok(ip) = trustable_resolver_fut.await { - domain_data.ipv4_addresses = ip - .into_iter() - .map(|x| x.to_string()) - .collect::>(); - } + let mut domain_data = DomainData::default(); + + let mut ip_lookup = Option::::None; + + async move { + if let Ok(ip) = resolver_fut.await { + if options.disable_double_check { + ip_lookup = Some(ip); + } else if let Ok(ip) = trustable_resolver_fut.await { + ip_lookup = Some(ip); } - domain_data.is_wildcard = domain_data - .ipv4_addresses - .iter() - .all(|ip| wildcard_ips.contains(ip)); - - if !quiet_flag { - if show_ip_addresses && !domain_data.is_wildcard { - println!("{};{:?}", host, domain_data.ipv4_addresses); - } else if !domain_data.is_wildcard { - println!("{}", host) - } - } - - (host, domain_data) } - }) - .buffer_unordered(threads) - .collect::>() - .await + + if let Some(ip_lookup) = ip_lookup { + for ip in ip_lookup.iter() { + domain_data.ipv4_addresses.insert(ip.to_string()); + } + } + + domain_data.is_wildcard = domain_data + .ipv4_addresses + .iter() + .all(|ip| wildcard_ips.contains(ip)); + + if !options.quiet_flag { + if options.show_ip_address && !domain_data.is_wildcard { + println!("{};{:?}", host, domain_data.ipv4_addresses); + } else if !domain_data.is_wildcard { + println!("{}", host) + } + } + + (host, domain_data) + } + })) + .buffer_unordered(threads) + .collect::>() + .await } // Used internally for now @@ -128,41 +121,32 @@ pub async fn return_cname_data( stream::iter(hosts) .map(|host| { - let resolver_fut = resolver.lookup( - host.trim_end_matches('.').to_owned() + ".", - record_type, - request_options, - ); - let trustable_resolver_fut = trustable_resolver.lookup( - host.trim_end_matches('.').to_owned() + ".", - record_type, - request_options, - ); + let lookup_host = host.trim_end_matches('.').to_owned() + "."; + let resolver_fut = resolver.lookup(lookup_host.clone(), record_type, request_options); + let trustable_resolver_fut = + trustable_resolver.lookup(lookup_host, record_type, request_options); let mut domain_data = DomainData::default(); + let mut cname_lookup = Option::::None; + async move { if let Ok(lookup) = resolver_fut.await { if disable_double_check { - domain_data.cname = lookup - .iter() - .filter_map(|rdata| rdata.as_cname()) - .map(|name| { - let name = name.to_string(); - name[..name.len() - 1].to_owned() - }) - .collect(); + cname_lookup = Some(lookup); } else if let Ok(lookup) = trustable_resolver_fut.await { - domain_data.cname = lookup - .iter() - .filter_map(|rdata| rdata.as_cname()) - .map(|name| { - let name = name.to_string(); - name[..name.len() - 1].to_owned() - }) - .collect(); + cname_lookup = Some(lookup); } } + + if let Some(lookup) = cname_lookup { + for record in lookup.iter() { + if let Some(cname) = record.as_cname() { + domain_data.cname = cname.to_string(); + } + } + } + (host, domain_data) } }) diff --git a/src/main.rs b/src/main.rs index c948c0c..e785040 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,6 @@ use { clap::{value_t, App, Arg}, - rusolver::{dnslib, utils}, + rusolver::{dnslib, structs, utils}, std::collections::HashSet, tokio::{ self, @@ -9,7 +9,7 @@ use { trust_dns_resolver::config::{LookupIpStrategy, ResolverOpts}, }; -// Please add support for AAAA, TXT, SRV, NAPTR, PTR, CNAME, DNAME, MX, NS, SOA, LOC, SVCB, HTTPS, SPF, CAA and AVC resource records. +// WIP: add support for AAAA, TXT, SRV, NAPTR, PTR, CNAME, DNAME, MX, NS, SOA, LOC, SVCB, HTTPS, SPF, CAA and AVC resource records. // This could use a new command line option such as -t, e.g. echo www.example.com | rusolver -i -t AAAA. It might also make sense // to change -i/--ip to -d/--data with the text Display the record data. @@ -76,7 +76,7 @@ async fn main() -> Result<(), Box> { .get_matches(); // Assign values or use defaults - let show_ip_adress = matches.is_present("ip"); + let show_ip_address = matches.is_present("ip"); let threads = value_t!(matches.value_of("threads"), usize).unwrap_or_else(|_| 100); let timeout = value_t!(matches.value_of("timeout"), u64).unwrap_or_else(|_| 3); let retries = value_t!(matches.value_of("retries"), usize).unwrap_or_else(|_| 0); @@ -118,7 +118,7 @@ async fn main() -> Result<(), Box> { "84.200.70.40:53", ] .iter() - .map(|x| x.to_string()) + .map(ToString::to_string) .collect(); // Create resolvers @@ -152,17 +152,18 @@ async fn main() -> Result<(), Box> { buffer.lines().map(str::to_owned).collect() }; - dnslib::return_hosts_data( + let options = structs::LibOptions { hosts, resolvers, trustable_resolver, wildcard_ips, disable_double_check, threads, - show_ip_adress, - false, - ) - .await; + show_ip_address, + quiet_flag, + }; + + dnslib::return_hosts_data(&options).await; Ok(()) } diff --git a/src/structs.rs b/src/structs.rs index f281670..f408081 100644 --- a/src/structs.rs +++ b/src/structs.rs @@ -1,4 +1,10 @@ -use std::collections::HashSet; +use { + std::collections::HashSet, + trust_dns_resolver::{ + name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime}, + AsyncResolver, + }, +}; #[derive(Clone, Debug)] pub struct DomainData { @@ -18,3 +24,16 @@ impl Default for DomainData { } } } + +#[derive(Clone, Debug)] +pub struct LibOptions { + pub hosts: HashSet, + pub resolvers: AsyncResolver>, + pub trustable_resolver: + AsyncResolver>, + pub wildcard_ips: HashSet, + pub disable_double_check: bool, + pub threads: usize, + pub show_ip_address: bool, + pub quiet_flag: bool, +}