Add --reverse, --file and --interval to util/resolve

crates/resolver/src/lookup.rs

@@ -429,6 +429,12 @@ macro_rules! lookup_type {
             }
         }
 
+        impl From<$l> for Lookup {
+            fn from(revlookup: $l) -> Self {
+                revlookup.0
+            }
+        }
+
         /// An iterator over the Lookup type
         pub struct $i<'i>(LookupIter<'i>);
 

util/Cargo.toml

@@ -85,7 +85,6 @@ trust-dns-client = { version = "0.22.0", path = "../crates/client" }
 trust-dns-proto = { version = "0.22.0", path = "../crates/proto" }
 trust-dns-recursor = { version = "0.22.0", path = "../crates/recursor" }
 trust-dns-resolver = { version = "0.22.0", path = "../crates/resolver" }
-tokio = { version = "1.0", features = ["rt-multi-thread", "macros"] }
+tokio = { version = "1.0", features = ["rt-multi-thread", "macros", "time"] }
 webpki = { version = "0.22.0", optional = true }
 webpki-roots = { version = "0.22.1", optional = true }
-

util/src/resolve.rs

@@ -20,14 +20,24 @@
     unreachable_pub
 )]
 
+use std::fs::File;
+use std::io::{BufRead, BufReader};
 use std::net::{IpAddr, SocketAddr};
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::time::Duration;
 
-use clap::Parser;
+use clap::{ArgGroup, Parser};
 use console::style;
+use tokio::task::JoinSet;
 
+use tokio::time::MissedTickBehavior;
+use trust_dns_client::rr::Record;
 use trust_dns_resolver::config::{
     NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig, ResolverOpts,
 };
+use trust_dns_resolver::error::{ResolveError, ResolveErrorKind};
+use trust_dns_resolver::lookup::Lookup;
 use trust_dns_resolver::proto::rr::RecordType;
 use trust_dns_resolver::TokioAsyncResolver;
 
@@ -40,19 +50,36 @@ use trust_dns_resolver::TokioAsyncResolver;
 /// used with the `--system` FLAG. Other nameservers, as many as desired, can
 /// be configured directly with the `--nameserver` OPTION.
 #[derive(Debug, Parser)]
-#[clap(name = "resolve")]
+#[clap(name = "resolve",
+    group(ArgGroup::new("qtype").args(&["happy", "reverse", "ty"])),
+    group(ArgGroup::new("input").required(true).args(&["domainname", "inputfile"]))
+)]
 struct Opts {
     /// Name to attempt to resolve, if followed by a '.' then it's a fully-qualified-domain-name.
-    domainname: String,
+    domainname: Option<String>,
+
+    /// File containing one domainname to resolve per line
+    #[clap(
+        short = 'f',
+        long = "file",
+        value_parser,
+        value_name = "FILE",
+        conflicts_with("domainname")
+    )]
+    inputfile: Option<PathBuf>,
 
     /// Type of query to issue, e.g. A, AAAA, NS, etc.
     #[clap(short = 't', long = "type", default_value = "A")]
     ty: RecordType,
 
     /// Happy eye balls lookup, ipv4 and ipv6
-    #[clap(short = 'e', long = "happy", conflicts_with("ty"))]
+    #[clap(short = 'e', long = "happy", conflicts_with_all(&["reverse", "ty"]))]
     happy: bool,
 
+    /// Reverse DNS lookup
+    #[clap(short = 'r', long = "reverse", conflicts_with_all(&["happy", "ty"]))]
+    reverse: bool,
+
     /// Use system configuration, e.g. /etc/resolv.conf, instead of defaults
     #[clap(short = 's', long = "system")]
     system: bool,
@@ -113,6 +140,106 @@ struct Opts {
     /// Enable error logging
     #[clap(long)]
     error: bool,
+
+    /// Set the time interval between requests (in seconds, useful with --file)
+    #[clap(long, default_value = "1.0")]
+    interval: f32,
+}
+
+fn print_record(r: &Record) {
+    print!(
+        "\t{name} {ttl} {class} {ty}",
+        name = style(r.name()).blue(),
+        ttl = style(r.ttl()).blue(),
+        class = style(r.dns_class()).blue(),
+        ty = style(r.record_type()).blue(),
+    );
+    if let Some(rdata) = r.data() {
+        println!(" {rdata}", rdata = rdata);
+    } else {
+        println!("NULL")
+    }
+}
+
+fn print_ok(lookup: Lookup) {
+    println!(
+        "{} for query {}",
+        style("Success").green(),
+        style(lookup.query()).blue()
+    );
+
+    for r in lookup.record_iter() {
+        print_record(r);
+    }
+}
+
+fn print_error(error: ResolveError) {
+    match error.kind() {
+        ResolveErrorKind::NoRecordsFound { query, soa, .. } => {
+            println!(
+                "{} for query {}",
+                style("NoRecordsFound").red(),
+                style(query).blue()
+            );
+            if let Some(r) = soa {
+                print_record(r);
+            }
+        }
+        &_ => {
+            println!("{:?}", error);
+        }
+    }
+}
+
+fn print_result(result: Result<Lookup, ResolveError>) {
+    match result {
+        Ok(lookup) => print_ok(lookup),
+        Err(re) => print_error(re),
+    }
+}
+
+fn log_query(name: &str, ty: RecordType, name_servers: &str, opts: &Opts) {
+    if opts.happy {
+        println!(
+            "Querying for {name} {ty} from {ns}",
+            name = style(name).yellow(),
+            ty = style("A+AAAA").yellow(),
+            ns = style(name_servers).blue()
+        );
+    } else if opts.reverse {
+        println!(
+            "Querying {reverse} for {name} from {ns}",
+            reverse = style("reverse").yellow(),
+            name = style(name).yellow(),
+            ns = style(name_servers).blue()
+        );
+    } else {
+        println!(
+            "Querying for {name} {ty} from {ns}",
+            name = style(name).yellow(),
+            ty = style(ty).yellow(),
+            ns = style(name_servers).blue()
+        );
+    }
+}
+
+async fn execute_query(
+    resolver: Arc<TokioAsyncResolver>,
+    name: String,
+    happy: bool,
+    reverse: bool,
+    ty: RecordType,
+) -> Result<Lookup, ResolveError> {
+    if happy {
+        Ok(resolver.lookup_ip(name.to_string()).await?.into())
+    } else if reverse {
+        let v4addr = name
+            .parse::<IpAddr>()
+            .unwrap_or_else(|_| panic!("Could not parse {} into an IP address", name));
+        Ok(resolver.reverse_lookup(v4addr).await?.into())
+    } else {
+        Ok(resolver.lookup(name.to_string(), ty).await?)
+    }
 }
 
 /// Run the resolve program
@@ -200,68 +327,56 @@ pub async fn main() -> Result<(), Box<dyn std::error::Error>> {
         config.add_name_server(ns.clone());
     }
 
-    let name_servers = config.name_servers().iter().map(|n| format!("{}", n)).fold(
-        String::new(),
-        |mut names, n| {
-            if !names.is_empty() {
-                names.push_str(", ")
-            }
-
-            names.push_str(&n);
-            names
-        },
-    );
-
-    // query parameters
-    let name = &opts.domainname;
-    let ty = opts.ty;
+    let name_servers = config
+        .name_servers()
+        .iter()
+        .map(|ns| format!("{}", ns))
+        .collect::<Vec<String>>()
+        .join(", ");
 
     // configure the resolver options
     let mut options = sys_options.unwrap_or_default();
     if opts.happy {
         options.ip_strategy = trust_dns_resolver::config::LookupIpStrategy::Ipv4AndIpv6;
     }
 
-    let resolver = TokioAsyncResolver::tokio(config, options)?;
-
-    // execute query
-    println!(
-        "Querying for {name} {ty} from {ns}",
-        name = style(name).yellow(),
-        ty = style(ty).yellow(),
-        ns = style(name_servers).blue()
-    );
-
-    let lookup = if opts.happy {
-        let lookup = resolver.lookup_ip(name.to_string()).await?;
-
-        lookup.into()
+    let resolver_arc = Arc::new(TokioAsyncResolver::tokio(config, options)?);
+
+    if let Some(domainname) = &opts.domainname {
+        log_query(domainname, opts.ty, &name_servers, &opts);
+        let lookup = execute_query(
+            resolver_arc,
+            domainname.to_owned(),
+            opts.happy,
+            opts.reverse,
+            opts.ty,
+        )
+        .await;
+        print_result(lookup);
     } else {
-        resolver.lookup(name.to_string(), ty).await?
-    };
-
-    // report response, TODO: better display of errors
-    println!(
-        "{} for query {}",
-        style("Success").green(),
-        style(lookup.query()).blue()
-    );
-
-    for r in lookup.record_iter() {
-        print!(
-            "\t{name} {ttl} {class} {ty}",
-            name = style(r.name()).blue(),
-            ttl = style(r.ttl()).blue(),
-            class = style(r.dns_class()).blue(),
-            ty = style(r.record_type()).blue(),
-        );
-
-        if let Some(rdata) = r.data() {
-            println!(" {rdata}", rdata = rdata);
-        } else {
-            println!("NULL")
+        let duration = Duration::from_secs_f32(opts.interval);
+        let fd = File::open(&opts.inputfile.as_ref().unwrap())?;
+        let reader = BufReader::new(fd);
+        let mut taskset = JoinSet::new();
+        let mut timer = tokio::time::interval(duration);
+        timer.set_missed_tick_behavior(MissedTickBehavior::Burst);
+        for name in reader.lines().filter_map(|line| line.ok()) {
+            let (happy, reverse, ty) = (opts.happy, opts.reverse, opts.ty);
+            log_query(&name, ty, &name_servers, &opts);
+            let resolver = resolver_arc.clone();
+            taskset.spawn(async move { execute_query(resolver, name, happy, reverse, ty).await });
+            loop {
+                tokio::select! {
+                    _ = timer.tick() => break,
+                    lookup_opt = taskset.join_next() => match lookup_opt {
+                        Some(lookup_rr) => {
+                            print_result(lookup_rr?);
+                        },
+                        None => { timer.tick().await; break; }
+                    }
+                };
+            }
         }
     }
-
     Ok(())
 }