resdns/src/main.rs

/*
Copyright 2023 0xf8.dev@proton.me

This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
*/

use argparse::{ArgumentParser, List, Store};
use resolve::{record, DnsConfig, DnsResolver};
use std::time::Duration;

#[derive(Debug, Clone)]
struct Record {
    pub host: String,
    pub addr: String,
}

#[derive(Debug)]
enum OutputParser {
    Plain,
    JSON { pretty: bool },
    JSONPretty { pretty: bool },
    CSV,
}

struct Resolver {
    resolver: DnsResolver,
}

impl Resolver {
    pub fn init() -> Resolver {
        let config = DnsConfig {
            timeout: Duration::from_secs(3),
            attempts: 2,
            use_inet6: false,
            ..DnsConfig::load_default().expect("Couldn't get default DnsConfig")
        };

        Resolver {
            resolver: DnsResolver::new(config).expect("Couldn't create DNS resolver"),
        }
    }

    pub fn resolve(&self, host: &str) -> Vec<Record> {
        let rec4 = self
            .resolver
            .resolve_record::<record::A>(host)
            .unwrap_or(vec![]);
        let rec6 = self
            .resolver
            .resolve_record::<record::AAAA>(host)
            .unwrap_or(vec![]);

        let mut recs: Vec<Record> = vec![];
        for r in rec4 {
            recs.insert(
                0,
                Record {
                    host: host.to_owned(),
                    addr: r.address.to_string(),
                },
            )
        }
        for r in rec6 {
            recs.insert(
                0,
                Record {
                    host: host.to_owned(),
                    addr: r.address.to_string(),
                },
            )
        }

        recs.sort_unstable_by(|a, b| a.addr.len().cmp(&b.addr.len()));

        recs
    }
}

fn main() {
    let mut output_flag = String::from("Default");
    let mut hosts: Vec<String> = vec![];

    // Argument parser
    {
        let mut args = ArgumentParser::new();
        args.set_description("Resolve Host(s) to IP Address(es).");

        // -o, --output, --parser
        args.refer(&mut output_flag).add_option(
            &["-o", "--output", "--parser"],
            Store,
            "Select how the response is outputted: [Default, Plain, JSON, JSON Pretty, CSV]",
        );

        // Positional arg: hosts
        args.refer(&mut hosts)
            .add_argument("hosts", List, "Hosts to resolve")
            .required();

        args.parse_args_or_exit();
    }

    // OutputParser
    let output = match output_flag.to_lowercase().as_str() {
        "default" | "plain" | "" => OutputParser::Plain,
        "json" => OutputParser::JSON { pretty: false },
        "json-pretty" | "json_pretty" | "json pretty" | "jsonpretty" => {
            OutputParser::JSONPretty { pretty: true }
        }
        "csv" => OutputParser::CSV,
        _ => {
            eprintln!("[FATAL] Invalid output parser: \"{output_flag}\"");
            std::process::exit(1)
        }
    };

    // No hosts
    if hosts.is_empty() {
        eprintln!("[Fatal] No hosts given");
        std::process::exit(1)
    }

    // Resolver
    let resolver: Resolver = Resolver::init();

    let mut record_pile: Vec<Record> = vec![];
    for h in hosts {
        record_pile.append(&mut resolver.resolve(h.as_str()));
    }

    match output {
        // Plain / Default output
        OutputParser::Plain => {
            for r in record_pile {
                println!("{}: {}", r.host, r.addr)
            }
        }

        // JSON / JSON Pretty output
        OutputParser::JSON { pretty } | OutputParser::JSONPretty { pretty } => {
            let mut rj = json::object! { records: {} };

            for r in record_pile {
                if rj["records"][&r.host].is_null() {
                    rj["records"]
                        .insert(
                            &r.host,
                            json::parse(&format!("[\"{}\"]", r.addr))
                                .expect("Unable to parse JSON"),
                        )
                        .expect("Unable to insert record into JSON");
                } else {
                    rj["records"][&r.host]
                        .push(r.addr)
                        .expect("Unable to insert record into JSON");
                }
            }

            println!(
                "{}",
                match pretty {
                    true => rj.pretty(2),
                    false => rj.dump(),
                }
            );
        }

        // CSV output
        OutputParser::CSV => {
            println!("\"Host\", \"Address\"");
            for r in record_pile {
                println!("\"{}\", \"{}\"", r.host, r.addr)
            }
        }

//        _ => (),
    }
}