package leigh.stack.dnsd;

import leigh.mecha.log.MechaLogger;
import leigh.mecha.log.MechaLoggerFactory;
import org.xbill.DNS.Message;
import org.xbill.DNS.SimpleResolver;
import org.xbill.DNS.Type;

import java.io.IOException;
import java.net.*;
import java.util.LinkedList;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * Simple implementation of a DNS server. This implementation attempts to lookup queries in a {@link NameStore},
 * and, failing that, forwards the query to <code>8.8.8.8</code>.
 *
 * @author Alex Leigh
 */
public class DNSService implements Runnable {
    private static final MechaLogger logger = MechaLoggerFactory.getLogger(DNSService.class);
    private static final int UDP_SIZE = 512;
    private final int port;
    private final SimpleResolver resolver;
    private final LinkedList<NameStore> stores = new LinkedList<>();
    private final ExecutorService executorService = Executors.newFixedThreadPool(128);

    DNSService(int port) throws UnknownHostException, MalformedURLException {
        this.port = port;
        resolver = new SimpleResolver("8.8.8.8");

        stores.add(new LeighCoZone());
        stores.add(new EvilNameStore());
    }

    public void run() {
        try {
            DatagramSocket socket = new DatagramSocket(port);
            //noinspection InfiniteLoopStatement
            while (true) {
                byte[] in = new byte[UDP_SIZE];
                DatagramPacket indp = new DatagramPacket(in, UDP_SIZE);
                socket.receive(indp);

                executorService.submit(() -> {
                    try {
                        query(socket, in, indp.getAddress(), indp.getPort());
                    } catch (Exception e) {
                        //   e.printStackTrace();
                    }
                });
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private void query(DatagramSocket socket, byte[] in, InetAddress dstAddr, int dstPort) throws IOException {
        Message request = new Message(in);
        Message reply = null;

        for (NameStore store : stores) {
            reply = store.doQuery(request);
            if (reply != null) break;
        }

        logger.info("query [name: {}] [type: {}] [relay: {}]", request.getQuestion().getName().toString(),
                Type.string(request.getQuestion().getType()), reply == null);

        if (reply == null)
            reply = resolver.send(request);

        if (reply != null) {
            byte[] resp = reply.toWire();
            DatagramPacket outdp = new DatagramPacket(resp, resp.length,
                    dstAddr, dstPort);
            socket.send(outdp);
        }
    }
}