edu4rdshl/rusolver
1
0
Fork 0
mirror of https://github.com/edu4rdshl/rusolver.git synced 2026-07-17 23:24:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use LibOptions for library options.

Browse source 
Signed-off-by: Eduard Tolosa <edu4rdshl@protonmail.com>
This commit is contained in:
Eduard Tolosa 2022-04-29 12:08:48 -05:00
parent 7755533e1e
commit cf75cafee7
No known key found for this signature in database
GPG key ID: 3A574A4009F553E5
3 changed files with 102 additions and 98 deletions
Download patch file Download diff file Expand all files Collapse all files

160
src/dnslib.rs
View file

@ -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<String>,
resolver: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
trustable_resolver: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
wildcard_ips: HashSet<String>,
disable_double_check: bool,
mut threads: usize,
show_ip_addresses: bool,
quiet_flag: bool,
) -> HashMap<String, DomainData> {
if hosts.len() < threads {
threads = hosts.len();
}
pub async fn return_hosts_data(options: &LibOptions) -> HashMap<String, DomainData> {
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::<HashSet<String>>();
} else if let Ok(ip) = trustable_resolver_fut.await {
domain_data.ipv4_addresses = ip
.into_iter()
.map(|x| x.to_string())
.collect::<HashSet<String>>();
}
let mut domain_data = DomainData::default();
let mut ip_lookup = Option::<Ipv4Lookup>::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::<HashMap<String, DomainData>>()
.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::<HashMap<String, DomainData>>()
.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::<Lookup>::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)
}
})

19
src/main.rs
View file

@ -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<dyn std::error::Error>> {
.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<dyn std::error::Error>> {
"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<dyn std::error::Error>> {
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(())
}

21
src/structs.rs
View file

@ -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<String>,
pub resolvers: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
pub trustable_resolver:
AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
pub wildcard_ips: HashSet<String>,
pub disable_double_check: bool,
pub threads: usize,
pub show_ip_address: bool,
pub quiet_flag: bool,
}