183 lines
5.2 KiB
Rust
183 lines
5.2 KiB
Rust
/*
|
|
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)
|
|
}
|
|
}
|
|
|
|
// _ => (),
|
|
}
|
|
}
|