package main

import (
	"log"
	"strings"

	"github.com/miekg/dns"
)

const upstream = "1.1.1.1:53"

var blocklist = []string{
	"x.com.",
	"twitter.com.",
	"example.com.",
}

func isBlocked(domain string) bool {
	domain = strings.ToLower(domain)
	for _, blocked := range blocklist {
		// Check exact match or subdomain
		if domain == blocked || strings.HasSuffix(domain, "."+blocked) {
			return true
		}
	}
	return false
}

func handleDNS(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)

	c := new(dns.Client)
	for _, q := range r.Question {
		log.Printf("question: %s", q.Name)
		if isBlocked(q.Name) {
			res, err := dns.NewRR(q.Name + " 60 IN A 0.0.0.0")
			if err != nil {
				log.Fatal(err)
			}
			m.Answer = append(m.Answer, res)
		} else {
			req := new(dns.Msg)
			req.Question = append(req.Question, q)
			res, _, err := c.Exchange(req, upstream)
			if err != nil {
				log.Fatal(err)
			}
			m.Answer = append(m.Answer, res.Answer...)

		}
	}
	w.WriteMsg(m)
}

func main() {
	dns.HandleFunc(".", handleDNS)
	server := &dns.Server{Addr: ":53", Net: "udp"}
	log.Println("Starting DNS Server on :53")
	log.Fatal(server.ListenAndServe())
}