package leigh.stack.dnsd;

import leigh.mecha.log.MechaLogger;
import leigh.mecha.log.MechaLoggerFactory;
import leigh.mecha.util.VelocityWatch;
import org.xbill.DNS.Message;

import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.UnknownHostException;
import java.util.HashSet;
import java.util.Scanner;

/**
 * This implementation of {@link NameStore} provides an evil name repository. Queries for evil names will be
 * returned, but pointing at 127.0.0.1, resulting in a black hole.
 *
 * @author Alex Leigh
 */
public class EvilNameStore implements NameStore {
    private static final MechaLogger logger = MechaLoggerFactory.getLogger(EvilNameStore.class);
    private final HashSet<String> evil = new HashSet<>();

    public EvilNameStore() throws UnknownHostException, MalformedURLException {
        new Thread(() -> {

            VelocityWatch vw = new VelocityWatch(logger, 1);
            logger.info("Loading EVIL urls.");
            try {
                // We delay here in case we are our own DNS server.
                Thread.sleep(500);
                loadEvilNames(new URL("https://v.firebog.net/hosts/Easyprivacy.txt"), vw);
                loadEvilNames(new URL("https://v.firebog.net/hosts/AdguardDNS.txt"), vw);
                loadEvilNames(new URL("https://v.firebog.net/hosts/Admiral.txt"), vw);
                loadEvilNames(new URL("https://v.firebog.net/hosts/Easylist.txt"), vw);
                loadEvilNames(new URL("https://v.firebog.net/hosts/Prigent-Ads.txt"), vw);
                loadEvilNames(new URL("https://v.firebog.net/hosts/Prigent-Malware.txt"), vw);
                loadEvilNames(new URL("https://v.firebog.net/hosts/Shalla-mal.txt"), vw);
                vw.status();
                logger.info("Loaded evil names. [cnt: {}]", evil.size());
            } catch (MalformedURLException | InterruptedException e) {
                e.printStackTrace();
            }
        }).start();
    }

    public void loadEvilNames(URL url, VelocityWatch vw) {
        String ustr = url.toString();
        try {
            Scanner s = new Scanner(url.openStream());
            while (s.hasNext()) {
                String ln = s.nextLine();
                vw.event(ustr);
                evil.add(ln + ".");
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @Override
    public Message doQuery(Message request) throws IOException {
        String name = request.getQuestion().getName().toString();

        if (evil.contains(name)) {
            logger.info("Evil query blocked. [name: {}]", name);
            // ha.. ha.. ha..
            return DNSUtil.buildEmpty(request);
        }

        return null;
    }
}
