<commit_before>package main

import (
	"fmt"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
)

type dnsServerHandler struct {
	ds *dns.Server
}

// TODO: Move in-memory mapping toa real datastore
var records = map[string]string{
	"foo.com.": "192.168.0.1",
}

func parseQuery(m *dns.Msg) {
	rl := NewLogger(log.New().Writer())

	for _, q := range m.Question {
		switch q.Qtype {
		case dns.TypeA:
			rl.Info("Query for: ", q.Name)
			ip := records[q.Name]
			if ip != "" {
				rr, err := dns.NewRR(fmt.Sprintf("%s A %s", q.Name, ip))
				if err == nil {
					m.Answer = append(m.Answer, rr)
				}
			} else {
				rl.Error("Record: ", q.Name, " not found")
			}
		default:
			rl.Info("Query type: ", q.Qtype, " is currently not supported")
		}
	}
}

func serve(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)
	// TODO: Enable over the wire compression
	m.Compress = false

	switch r.Opcode {
	case dns.OpcodeQuery:
		parseQuery(m)
	}
	w.WriteMsg(m)
}

func newDNSServerHandler(addr string, port int, conntype string) *dns.Server {
	return &dns.Server{
		Addr: fmt.Sprintf("%s:%d", addr, port),
		Net:  conntype,
	}
}

func newDefaultDNSServerHandler() *dnsServerHandler {
	return &dnsServerHandler{
		ds: newDNSServerHandler(DNS_SERVER_ADDR, DNS_SERVER_PORT, "udp"),
	}
}

func (h *dnsServerHandler) RouteDNS() error {
	rl := NewLogger(log.New().Writer())

	// TODO: Add more domains to the serving list
	dns.HandleFunc("com.", serve)
	err := h.ds.ListenAndServe()
	if err != nil {
		rl.Error("Failed to start DNS server", err)
		return err
	}
	return nil
}
<commit_msg>dns: Support DNS with NS1<commit_after>package main

import (
	"fmt"

	"github.com/miekg/dns"
	"math/rand"
	log "github.com/sirupsen/logrus"
	ns1 "gopkg.in/ns1/ns1-go.v2/rest"
	"net/http"
	"time"
	"os"
	"strings"
)

var ns1Client *ns1.Client

type dnsServerHandler struct {
	ds *dns.Server
}

func getRandomAnswer(max int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(max - 0) + 0
}

func parseQuery(m *dns.Msg) {
	rl := NewLogger(log.New().Writer())

	for _, q := range m.Question {
		switch q.Qtype {
		case dns.TypeA:
			zone := q.Name[strings.IndexAny(q.Name, ".")+1:len(q.Name)-1]
			domain := q.Name[:strings.LastIndex(q.Name, ".")]

			rl.Info("Query for > ", "domain: ", domain, " zone: ", zone, " type: ", q.Qtype)
			record, _, err := ns1Client.Records.Get(zone, domain, "A")
			if err != nil {
				rl.Error("Record: ", q.Name, err)
			} else {
				rl.Info(record.Answers)
				var answers []string
				for _, answer := range record.Answers {
					answers = append(answers, answer.String())
				}
				// TODO: Implement RR answer retrieval
				rr, err := dns.NewRR(fmt.Sprintf("%s A %s", q.Name, answers[getRandomAnswer(len(answers))]))
				if err == nil {
					m.Answer = append(m.Answer, rr)
				} else {
					rl.Error(err)
				}
			}

		// TODO: Add IPv6 support
		default:
			rl.Info("Query type: ", q.Qtype, " is currently not supported")
		}
	}
}

func serve(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)
	// TODO: Enable over the wire compression
	m.Compress = false

	switch r.Opcode {
	case dns.OpcodeQuery:
		parseQuery(m)
	}
	w.WriteMsg(m)
}

func newDNSServerHandler(addr string, port int, conntype string) *dns.Server {
	return &dns.Server{
		Addr: fmt.Sprintf("%s:%d", addr, port),
		Net:  conntype,
	}
}

func newDefaultDNSServerHandler() *dnsServerHandler {
	return &dnsServerHandler{
		ds: newDNSServerHandler(DNS_SERVER_ADDR, DNS_SERVER_PORT, "udp"),
	}
}
func init() {
	k := os.Getenv("NS1_APIKEY")
	if k == "" {
		fmt.Println("NS1_APIKEY environment variable is not set, stopping DNS server")
	}
	httpClient := &http.Client{Timeout: time.Second * 10}
	ns1Client = ns1.NewClient(httpClient, ns1.SetAPIKey(k))
}

func (h *dnsServerHandler) RouteDNS() error {
	rl := NewLogger(log.New().Writer())

	dns.HandleFunc(".", serve)
	err := h.ds.ListenAndServe()
	if err != nil {
		rl.Error("Failed to start DNS server", err)
		return err
	}
	return nil
}

{'fraction_non_alphanumeric': 0.16255297930690601, 'fraction_numerical': 0.006980802792321117, 'mean_word_length': 3.507865168539326, 'pattern_counts': {'":': 1, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6994102', 'n_tokens_mistral': 1728, 'n_tokens_neox': 1572, 'n_words': 358}