feat(cli): add custom dns resolver option

Cargo.toml

@@ -34,7 +34,7 @@ serde = "1.0.124"
 serde_derive = "1.0.116"
 cidr-utils = "0.5.1"
 itertools = "0.9.0"
-trust-dns-resolver = { version = "0.23.2", features = ["dns-over-rustls"] }
+hickory-resolver = { version = "0.24.0", features = ["dns-over-rustls"] }
 anyhow = "1.0.40"
 subprocess = "0.2.6"
 text_placeholder = { version = "0.5", features = ["struct_context"] }

src/address.rs

@@ -1,15 +1,16 @@
 //! Provides functions to parse input IP addresses, CIDRs or files.
-use std::fs::File;
+use std::fs::{self, File};
 use std::io::{prelude::*, BufReader};
-use std::net::{IpAddr, ToSocketAddrs};
+use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
 use std::path::Path;
+use std::str::FromStr;
 
 use cidr_utils::cidr::IpCidr;
-use log::debug;
-use trust_dns_resolver::{
-    config::{ResolverConfig, ResolverOpts},
+use hickory_resolver::{
+    config::{NameServerConfig, Protocol, ResolverConfig, ResolverOpts},
     Resolver,
 };
+use log::debug;
 
 use crate::input::Opts;
 use crate::warning;
@@ -29,8 +30,7 @@ use crate::warning;
 pub fn parse_addresses(input: &Opts) -> Vec<IpAddr> {
     let mut ips: Vec<IpAddr> = Vec::new();
     let mut unresolved_addresses: Vec<&str> = Vec::new();
-    let backup_resolver =
-        Resolver::new(ResolverConfig::cloudflare_tls(), ResolverOpts::default()).unwrap();
+    let backup_resolver = get_resolver(&input.resolver);
 
     for address in &input.addresses {
         let parsed_ips = parse_address(address, &backup_resolver);
@@ -72,11 +72,14 @@ pub fn parse_addresses(input: &Opts) -> Vec<IpAddr> {
 /// Given a string, parse it as a host, IP address, or CIDR.
 ///
 /// This allows us to pass files as hosts or cidr or IPs easily
-/// Call this every time you have a possible IP_or_host
+/// Call this every time you have a possible IP-or-host.
+///
+/// If the address is a domain, we can self-resolve the domain locally
+/// or resolve it by dns resolver list.
 ///
 /// ```rust
 /// # use rustscan::address::parse_address;
-/// # use trust_dns_resolver::Resolver;
+/// # use hickory_resolver::Resolver;
 /// let ips = parse_address("127.0.0.1", &Resolver::default().unwrap());
 /// ```
 pub fn parse_address(address: &str, resolver: &Resolver) -> Vec<IpAddr> {
@@ -94,7 +97,7 @@ pub fn parse_address(address: &str, resolver: &Resolver) -> Vec<IpAddr> {
 
 /// Uses DNS to get the IPS associated with host
 fn resolve_ips_from_host(source: &str, backup_resolver: &Resolver) -> Vec<IpAddr> {
-    let mut ips: Vec<std::net::IpAddr> = Vec::new();
+    let mut ips: Vec<IpAddr> = Vec::new();
 
     if let Ok(addrs) = source.to_socket_addrs() {
         for ip in addrs {
@@ -107,16 +110,64 @@ fn resolve_ips_from_host(source: &str, backup_resolver: &Resolver) -> Vec<IpAddr
     ips
 }
 
+/// Derive a DNS resolver.
+///
+/// 1. if the `resolver` parameter has been set:
+///     1. assume the parameter is a path and attempt to read IPs.
+///     2. parse the input as a comma-separated list of IPs.
+/// 2. if `resolver` is not set:
+///    1. attempt to derive a resolver from the system config. (e.g.
+///       `/etc/resolv.conf` on *nix).
+///    2. finally, build a CloudFlare-based resolver (default
+///       behaviour).
+fn get_resolver(resolver: &Option<String>) -> Resolver {
+    match resolver {
+        Some(r) => {
+            let mut config = ResolverConfig::new();
+            let resolver_ips = match read_resolver_from_file(r) {
+                Ok(ips) => ips,
+                Err(_) => r
+                    .split(',')
+                    .filter_map(|r| IpAddr::from_str(r).ok())
+                    .collect::<Vec<_>>(),
+            };
+            for ip in resolver_ips {
+                config.add_name_server(NameServerConfig::new(
+                    SocketAddr::new(ip, 53),
+                    Protocol::Udp,
+                ));
+            }
+            Resolver::new(config, ResolverOpts::default()).unwrap()
+        }
+        None => match Resolver::from_system_conf() {
+            Ok(resolver) => resolver,
+            Err(_) => {
+                Resolver::new(ResolverConfig::cloudflare_tls(), ResolverOpts::default()).unwrap()
+            }
+        },
+    }
+}
+
+/// Parses and input file of IPs for use in DNS resolution.
+fn read_resolver_from_file(path: &str) -> Result<Vec<IpAddr>, std::io::Error> {
+    let ips = fs::read_to_string(path)?
+        .lines()
+        .filter_map(|line| IpAddr::from_str(line.trim()).ok())
+        .collect();
+
+    Ok(ips)
+}
+
 #[cfg(not(tarpaulin_include))]
 /// Parses an input file of IPs and uses those
 fn read_ips_from_file(
     ips: &std::path::Path,
     backup_resolver: &Resolver,
-) -> Result<Vec<std::net::IpAddr>, std::io::Error> {
+) -> Result<Vec<IpAddr>, std::io::Error> {
     let file = File::open(ips)?;
     let reader = BufReader::new(file);
 
-    let mut ips: Vec<std::net::IpAddr> = Vec::new();
+    let mut ips: Vec<IpAddr> = Vec::new();
 
     for address_line in reader.lines() {
         if let Ok(address) = address_line {
@@ -131,7 +182,7 @@ fn read_ips_from_file(
 
 #[cfg(test)]
 mod tests {
-    use super::{parse_addresses, Opts};
+    use super::{get_resolver, parse_addresses, Opts};
     use std::net::Ipv4Addr;
 
     #[test]
@@ -204,4 +255,27 @@ mod tests {
         let ips = parse_addresses(&opts);
         assert_eq!(ips.len(), 0);
     }
+
+    #[test]
+    fn resolver_default_cloudflare() {
+        let opts = Opts::default();
+
+        let resolver = get_resolver(&opts.resolver);
+        let lookup = resolver.lookup_ip("www.example.com.").unwrap();
+
+        assert!(opts.resolver.is_none());
+        assert!(lookup.iter().next().is_some());
+    }
+
+    #[test]
+    fn resolver_args_google_dns() {
+        let mut opts = Opts::default();
+        // https://developers.google.com/speed/public-dns
+        opts.resolver = Some("8.8.8.8,8.8.4.4".to_owned());
+
+        let resolver = get_resolver(&opts.resolver);
+        let lookup = resolver.lookup_ip("www.example.com.").unwrap();
+
+        assert!(lookup.iter().next().is_some());
+    }
 }

src/input.rs

@@ -100,6 +100,10 @@ pub struct Opts {
     #[structopt(long)]
     pub accessible: bool,
 
+    /// A comma-delimited list or file of DNS resolvers.
+    #[structopt(long)]
+    pub resolver: Option<String>,
+
     /// The batch size for port scanning, it increases or slows the speed of
     /// scanning. Depends on the open file limit of your OS.  If you do 65535
     /// it will do every port at the same time. Although, your OS may not
@@ -208,7 +212,7 @@ impl Opts {
             self.ports = Some(ports);
         }
 
-        merge_optional!(range, ulimit, exclude_ports);
+        merge_optional!(range, resolver, ulimit, exclude_ports);
     }
 }
 
@@ -225,6 +229,7 @@ impl Default for Opts {
             ulimit: None,
             command: vec![],
             accessible: false,
+            resolver: None,
             scan_order: ScanOrder::Serial,
             no_config: true,
             top: false,
@@ -250,6 +255,7 @@ pub struct Config {
     timeout: Option<u32>,
     tries: Option<u8>,
     ulimit: Option<u64>,
+    resolver: Option<String>,
     scan_order: Option<ScanOrder>,
     command: Option<Vec<String>>,
     scripts: Option<ScriptsRequired>,
@@ -317,6 +323,7 @@ mod tests {
                 ulimit: None,
                 command: Some(vec!["-A".to_owned()]),
                 accessible: Some(true),
+                resolver: None,
                 scan_order: Some(ScanOrder::Random),
                 scripts: None,
                 exclude_ports: None,
@@ -364,10 +371,12 @@ mod tests {
             end: 1_000,
         });
         config.ulimit = Some(1_000);
+        config.resolver = Some("1.1.1.1".to_owned());
 
         opts.merge_optional(&config);
 
         assert_eq!(opts.range, config.range);
         assert_eq!(opts.ulimit, config.ulimit);
+        assert_eq!(opts.resolver, config.resolver);
     }
 }