DNS

package org.kuali.common.devops.logic;

import static com.google.common.collect.Iterables.filter;
import static com.google.common.collect.Lists.newArrayList;
import static com.google.common.collect.Maps.newHashMap;
import static com.google.common.collect.Maps.newTreeMap;
import static com.google.common.collect.Sets.newHashSet;
import static java.lang.String.format;
import static org.apache.commons.io.FileUtils.readFileToString;
import static org.apache.commons.io.FileUtils.write;
import static org.apache.commons.lang3.StringUtils.removeEnd;
import static org.kuali.common.devops.archive.sweep.Functions.hostnameToKey;
import static org.kuali.common.devops.logic.Auth.getDNSMECredentials;
import static org.kuali.common.dns.dnsme.DNSME.DNSME_REST_API_URL_PRODUCTION;
import static org.kuali.common.dns.model.DnsRecordType.CNAME;
import static org.kuali.common.util.FormatUtils.getCount;
import static org.kuali.common.util.FormatUtils.getTime;
import static org.kuali.common.util.base.Exceptions.illegalState;
import static org.kuali.common.util.log.Loggers.newLogger;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;

import org.kuali.common.core.json.api.JsonService;
import org.kuali.common.core.json.jackson.JacksonJsonService;
import org.kuali.common.dns.api.DnsService;
import org.kuali.common.dns.dnsme.DNSMadeEasyDnsService;
import org.kuali.common.dns.dnsme.model.DNSMadeEasyCredentials;
import org.kuali.common.dns.dnsme.model.DNSMadeEasyServiceContext;
import org.kuali.common.dns.model.SimpleDnsRecord;
import org.kuali.common.util.file.CanonicalFile;
import org.kuali.common.util.property.ImmutableProperties;
import org.slf4j.Logger;

import com.google.common.base.Function;
import com.google.common.base.Predicate;
import com.google.common.base.Stopwatch;
import com.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;
import com.google.common.collect.HashMultiset;
import com.google.common.collect.ImmutableBiMap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Multiset;
import com.google.common.collect.Ordering;
import com.google.common.collect.Sets;

public class DNS {

	private static final Logger logger = newLogger();
	private static final String DOMAIN = "kuali.org";
	private static final File CACHE = new CanonicalFile("./target/cache/dns/records.json");

	/**
	 * Return a mapping of aliases to canonical name records. There can be multiple aliases pointing to the same canonical name record.
	 *
	 * <pre>
	 * alias                  canonical
	 * env1.rice.kuali.org -> ec2-174-129-109-246.compute-1.amazonaws.com
	 * </pre>
	 */
	public static Map<String, String> getAliasMap(boolean refresh) {
		List<SimpleDnsRecord> records = getDnsRecords(refresh);
		info("records -> %s", getCount(records));
		return asCNAMEAliasFQDNMap(records);
	}

	private static List<SimpleDnsRecord> getDnsRecords(boolean refresh) {
		if (refresh || !CACHE.exists()) {
			info("domain  -> %s", DOMAIN);
			List<SimpleDnsRecord> records = queryProvider();
			store(records);
			return records;
		} else {
			return load();
		}
	}

	/**
	 * Return a mapping of canonical name records to aliases. The only canonical name records present in this map are those that contain one (and only one) alias.
	 *
	 * This allows you to deterministically find the correct alias if you have the canonical name record.
	 *
	 * <pre>
	 * canonical                                      alias
	 * ec2-174-129-109-246.compute-1.amazonaws.com -> env1.rice.kuali.org
	 * </pre>
	 */
	public static BiMap<String, String> getCanonicalMap(boolean refresh) {
		Map<String, String> map = newHashMap(getAliasMap(refresh));
		removeAllKeysWithDuplicateValues(map);
		BiMap<String, String> aliases = HashBiMap.create(map);
		BiMap<String, String> canonical = aliases.inverse();
		info("using   -> %s uniquely aliased CNAME records", getCount(canonical));
		return ImmutableBiMap.copyOf(canonical);
	}

	protected static <T> void removeAllKeysWithDuplicateValues(Map<?, T> map) {
		int oldSize = map.size();
		Set<T> duplicates = getDuplicateValues(map);
		Set<?> keys = newHashSet(map.keySet());
		for (Object key : keys) {
			T value = map.get(key);
			if (duplicates.contains(value)) {
				debug("remove  -> %s", key);
				map.remove(key);
			}
		}
		int newSize = map.size();
		int removed = oldSize - newSize;
		info("removed -> %s CNAME records with duplicate aliases", getCount(removed));
	}

	protected static <T> Set<T> getDuplicateValues(Map<?, T> map) {
		Multiset<T> multi = HashMultiset.create();
		multi.addAll(map.values());
		Set<T> duplicates = Sets.newHashSet();
		Set<T> elements = multi.elementSet();
		for (T element : elements) {
			if (multi.count(element) > 1) {
				duplicates.add(element);
			}
		}
		return duplicates;
	}

	/**
	 * Grab all the DNS records known to this provider
	 */
	protected static List<SimpleDnsRecord> queryProvider() {
		DNSMadeEasyCredentials credentials = getDNSMECredentials();
		String url = DNSME_REST_API_URL_PRODUCTION;
		String domain = DOMAIN;
		DNSMadeEasyServiceContext context = new DNSMadeEasyServiceContext(credentials, url, domain);
		DnsService dns = new DNSMadeEasyDnsService(context);
		return dns.getRecords();
	}

	/**
	 * Filter the list down to just CNAME records.<br>
	 * Convert the DNS record name into a valid FQDN by appending the domain name.<br>
	 * Convert the DNS record value into a valid FQDN by remove the trailing dot (if there is one)<br>
	 */
	protected static Map<String, String> asCNAMEAliasFQDNMap(List<SimpleDnsRecord> records) {
		List<SimpleDnsRecord> filtered = newArrayList(filter(records, CNAMEPredicate.INSTANCE));
		Map<String, String> map = newTreeMap();
		for (SimpleDnsRecord record : filtered) {
			String alias = record.getName() + "." + DOMAIN;
			String cname = removeEnd(record.getValue(), ".");
			map.put(alias, cname);
		}
		int removed = records.size() - filtered.size();
		info("removed -> %s records that were not CNAME's", getCount(removed));
		return ImmutableMap.copyOf(map);
	}

	protected static List<SimpleDnsRecord> load() {
		JsonService json = new JacksonJsonService();
		try {
			info("load    -> %s", CACHE);
			String content = readFileToString(CACHE);
			return ImmutableList.copyOf(json.readString(content, SimpleDnsRecord[].class));
		} catch (IOException e) {
			throw illegalState(e);
		}
	}

	protected static void store(List<SimpleDnsRecord> records) {
		JsonService json = new JacksonJsonService();
		Function<SimpleDnsRecord, String> sorter = ReverseHostnameFunction.INSTANCE;
		List<SimpleDnsRecord> sorted = Ordering.natural().onResultOf(sorter).sortedCopy(records);
		String data = json.writeString(sorted);
		try {
			info("create  -> %s", CACHE);
			write(CACHE, data);
		} catch (IOException e) {
			throw illegalState(e);
		}
	}

	protected static Properties convert(Map<String, String> map) {
		Properties props = new Properties();
		props.putAll(map);
		return ImmutableProperties.copyOf(props);
	}

	protected static Map<String, String> convert(Properties props) {
		Map<String, String> map = newTreeMap();
		for (String key : props.stringPropertyNames()) {
			map.put(key, props.getProperty(key));
		}
		return ImmutableMap.copyOf(map);
	}

	private enum CNAMEPredicate implements Predicate<SimpleDnsRecord> {
		INSTANCE;

		@Override
		public boolean apply(SimpleDnsRecord record) {
			return record.getType().equals(CNAME);
		}
	}

	private enum ReverseHostnameFunction implements Function<SimpleDnsRecord, String> {
		INSTANCE;

		@Override
		public String apply(SimpleDnsRecord record) {
			return hostnameToKey().apply(record.getName());
		}
	}

	protected static void elapsed(Stopwatch sw) {
		info("elapsed -> %s", getTime(sw));
	}

	protected static void debug(String msg, Object... args) {
		logger.debug((args == null || args.length == 0) ? msg : format(msg, args));
	}

	protected static void info(String msg, Object... args) {
		logger.info((args == null || args.length == 0) ? msg : format(msg, args));
	}

}