rust-lookup/src/main.rs

#[macro_use]
extern crate clap;
extern crate resolv;

use resolv::{
	Resolver,
	Class,
	Section,
	error::{
		Error,
		ResolutionError,
	},
	record::{
		A,
		AAAA,
	},
};

use std::process::exit;
use std::ffi::OsStr;
use std::os::unix::ffi::OsStrExt;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::fmt;

/// Workaround https://github.com/mikedilger/resolv-rs/issues/1
///
/// When it gets fixed, this will simply call the inner `.next()` a few
/// times too often, which should hurt much.
struct FixedIterator<'a, T: resolv::record::RecordData>(resolv::RecordItems<'a, T>, usize);

impl<'a, T: resolv::record::RecordData> Iterator for FixedIterator<'a, T> {
	type Item = resolv::Record<T>;
	
	fn next(&mut self) -> Option<Self::Item> {
		while self.1 > 0 {
			self.1 -= 1;
			if let Some(r) = self.0.next() {
				return Some(r);
			}
		}
		return None;
	}
}

impl<'a, T: resolv::record::RecordData> FixedIterator<'a, T> {
	pub fn answers(response: &'a mut resolv::Response) -> Self {
		let count = response.get_section_count(Section::Answer);
		FixedIterator(response.answers::<T>(), count)
	}
}

pub trait AddressRecord: resolv::record::RecordData {
	type Address: Into<IpAddr> + fmt::Display + fmt::Debug;

	fn address(&self) -> Self::Address;
}

impl AddressRecord for A {
	type Address = Ipv4Addr;

	fn address(&self) -> Ipv4Addr {
		self.address
	}
}

impl AddressRecord for AAAA {
	type Address = Ipv6Addr;

	fn address(&self) -> Ipv6Addr {
		self.address
	}
}

fn run<R: AddressRecord>(resolver: &mut resolv::Resolver, name: &OsStr) {
	match resolver.query(name.as_bytes(), Class::IN, R::get_record_type()) {
		Ok(mut response) => {
			for answer in FixedIterator::<R>::answers(&mut response) {
				println!("{}", answer.data.address());
			}
		},
		Err(Error::Resolver(ResolutionError::HostNotFound)) => {
			eprintln!("Host not found: {:?}", name);
			exit(0);
		},
		Err(Error::Resolver(ResolutionError::NoData)) => {
			// empty answer is perfectly fine
		},
		Err(e) => {
			eprintln!("Failed looking up {:?}: {}", name, e);
			exit(1);
		},
	}
}

fn main() {
	let app = clap_app!(("lookup") =>
		(version: crate_version!())
		(author: crate_authors!("\n"))
		(about: crate_description!())
		(@arg IPv4: short("4") conflicts_with("IPv6") "Query only IPv4 records (A)")
		(@arg IPv6: short("6") "Query only IPv6 records (AAAA)")
		(@arg NAME: +required "Name to lookup")
	);
	let matches = app.get_matches();

	let ipv4_only = matches.is_present("IPv4");
	let ipv6_only = matches.is_present("IPv6");
	let name = matches.value_of_os("NAME").unwrap();

	let mut resolver = Resolver::new().unwrap_or_else(|| {
		eprintln!("Couldn't initialize resolver");
		exit(1);
	});

	if !ipv6_only {
		run::<A>(&mut resolver, name);
	}

	if !ipv4_only {
		run::<AAAA>(&mut resolver, name);
	}
}