// Copyright (c) 2005 Brian Wellington (bwelling@xbill.org)

package org.xbill.DNS.spi;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.StringTokenizer;

import org.xbill.DNS.AAAARecord;
import org.xbill.DNS.ARecord;
import org.xbill.DNS.ExtendedResolver;
import org.xbill.DNS.Lookup;
import org.xbill.DNS.Name;
import org.xbill.DNS.PTRRecord;
import org.xbill.DNS.Record;
import org.xbill.DNS.Resolver;
import org.xbill.DNS.ReverseMap;
import org.xbill.DNS.TextParseException;
import org.xbill.DNS.Type;

/**
 * This class implements a Name Service Provider, which Java can use 
 * (starting with version 1.4), to perform DNS resolutions instead of using 
 * the standard calls. 
 * <p>
 * This Name Service Provider uses dnsjava.
 * <p>
 * To use this provider, you must set the following system property:
 * <b>sun.net.spi.nameservice.provider.1=dns,dnsjava</b>
 *
 * @author Brian Wellington
 * @author Paul Cowan (pwc21@yahoo.com)
 */

public class DNSJavaNameService implements InvocationHandler {

private static final String nsProperty = "sun.net.spi.nameservice.nameservers";
private static final String domainProperty = "sun.net.spi.nameservice.domain";
private static final String v6Property = "java.net.preferIPv6Addresses";

private boolean preferV6 = false;

/**
 * Creates a DNSJavaNameService instance.
 * <p>
 * Uses the
 * <b>sun.net.spi.nameservice.nameservers</b>,
 * <b>sun.net.spi.nameservice.domain</b>, and
 * <b>java.net.preferIPv6Addresses</b> properties for configuration.
 */
protected
DNSJavaNameService() {
	String nameServers = System.getProperty(nsProperty);
	String domain = System.getProperty(domainProperty);
	String v6 = System.getProperty(v6Property);

	if (nameServers != null) {
		StringTokenizer st = new StringTokenizer(nameServers, ",");
		String [] servers = new String[st.countTokens()];
		int n = 0;
		while (st.hasMoreTokens())
			servers[n++] = st.nextToken();
		try {
			Resolver res = new ExtendedResolver(servers);
			Lookup.setDefaultResolver(res);
		}
		catch (UnknownHostException e) {
			System.err.println("DNSJavaNameService: invalid " +
					   nsProperty);
		}
	}

	if (domain != null) {
		try {
			Lookup.setDefaultSearchPath(new String[] {domain});
		}
		catch (TextParseException e) {
			System.err.println("DNSJavaNameService: invalid " +
					   domainProperty);
		}
	}

	if (v6 != null && v6.equalsIgnoreCase("true"))
		preferV6 = true;
}


public Object
invoke(Object proxy, Method method, Object[] args) throws Throwable {
	try {
		if (method.getName().equals("getHostByAddr")) {
			return this.getHostByAddr((byte[]) args[0]);
		} else if (method.getName().equals("lookupAllHostAddr")) {
			InetAddress[] addresses;
			addresses = this.lookupAllHostAddr((String) args[0]);
			Class returnType = method.getReturnType();
			if (returnType.equals(InetAddress[].class)) {
				// method for Java >= 1.6
				return addresses;
			} else if (returnType.equals(byte[][].class)) {
				// method for Java <= 1.5
				int naddrs = addresses.length;
				byte [][] byteAddresses = new byte[naddrs][];
				byte [] addr;
				for (int i = 0; i < naddrs; i++) {
					addr = addresses[i].getAddress();
					byteAddresses[i] = addr;
				}
				return byteAddresses;
			}
		}		
	} catch (Throwable e) {
		System.err.println("DNSJavaNameService: Unexpected error.");
		e.printStackTrace();
		throw e;
	}
	throw new IllegalArgumentException(
					"Unknown function name or arguments.");
}

/**
 * Performs a forward DNS lookup for the host name.
 * @param host The host name to resolve.
 * @return All the ip addresses found for the host name.
 */
public InetAddress []
lookupAllHostAddr(String host) throws UnknownHostException {
	Name name = null;

	try {
		name = new Name(host);
	}
	catch (TextParseException e) {
		throw new UnknownHostException(host);
	}

	Record [] records = null;
	if (preferV6)
		records = new Lookup(name, Type.AAAA).run();
	if (records == null)
		records = new Lookup(name, Type.A).run();
	if (records == null && !preferV6)
		records = new Lookup(name, Type.AAAA).run();
	if (records == null)
		throw new UnknownHostException(host);

	InetAddress[] array = new InetAddress[records.length];
	for (int i = 0; i < records.length; i++) {
		Record record = records[i];
		if (records[i] instanceof ARecord) {
			ARecord a = (ARecord) records[i];
			array[i] = a.getAddress();
		} else {
			AAAARecord aaaa = (AAAARecord) records[i];
			array[i] = aaaa.getAddress();
		}
	}
	return array;
}

/**
 * Performs a reverse DNS lookup.
 * @param addr The ip address to lookup.
 * @return The host name found for the ip address.
 */
public String
getHostByAddr(byte [] addr) throws UnknownHostException {
	Name name = ReverseMap.fromAddress(InetAddress.getByAddress(addr));
	Record [] records = new Lookup(name, Type.PTR).run();
	if (records == null)
		throw new UnknownHostException();
	return ((PTRRecord) records[0]).getTarget().toString();
}
}