package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/miekg/dns"
	"gopkg.in/yaml.v3"
)

type zoneView struct {
	rr  rrMap
	acl []string
}

type zoneMap map[string][]zoneView

type rrMap map[uint16]map[string][]dns.RR

type config struct {
	Zones   []configZone  `yaml:"zones"`
	ACL     []configACL   `yaml:"acl"`
	Forward configForward `yaml:"forward"`
	Address string        `yaml:"address"`
}

type configForward struct {
	ACL    []string `yaml:"acl"`
	Server string   `yaml:"server"`
}

type configACL struct {
	Name string `yaml:"name"`
	CIDR string `yaml:"cidr"`
}

type configZone struct {
	Zone string   `yaml:"zone"`
	File string   `yaml:"file"`
	ACL  []string `yaml:"acl"`
}

var anyRecordTypes = []uint16{
	dns.TypeSOA,
	dns.TypeA,
	dns.TypeAAAA,
	dns.TypeNS,
	dns.TypeCNAME,
	dns.TypeMX,
	dns.TypeTXT,
	dns.TypeSRV,
	dns.TypeCAA,
}

func loadConfig(configPath string) (*config, error) {
	file, err := ioutil.ReadFile(configPath)
	if err != nil {
		return nil, err
	}

	var loadedConfig config

	err = yaml.Unmarshal(file, &loadedConfig)
	if err != nil {
		return nil, err
	}

	return &loadedConfig, nil
}

func loadZones(configZones []configZone) (zoneMap, error) {
	zones := make(zoneMap)
	for _, z := range configZones {
		rrs, err := loadZonefile(z.File, z.Zone)
		if err != nil {
			return nil, err
		}
		if zones[z.Zone] == nil {
			zones[z.Zone] = make([]zoneView, 0)
		}
		zones[z.Zone] = append(zones[z.Zone], zoneView{
			rr:  createRRMap(rrs),
			acl: z.ACL,
		})

		log.Printf("Loaded zone %s\n", z.Zone)
	}

	return zones, nil
}

func createRRMap(rrs []dns.RR) rrMap {
	rrMap := make(rrMap)
	for _, rr := range rrs {
		if rrMap[rr.Header().Rrtype] == nil {
			rrMap[rr.Header().Rrtype] = make(map[string][]dns.RR)
		}

		if rrMap[rr.Header().Rrtype][rr.Header().Name] == nil {
			rrMap[rr.Header().Rrtype][rr.Header().Name] = make([]dns.RR, 0)
		}
		rrMap[rr.Header().Rrtype][rr.Header().Name] = append(rrMap[rr.Header().Rrtype][rr.Header().Name], rr)
	}

	return rrMap
}

func loadZonefile(filepath, origin string) ([]dns.RR, error) {
	file, err := os.Open(filepath)

	if err != nil {
		return nil, err
	}

	parser := dns.NewZoneParser(file, origin, filepath)

	var rrs = make([]dns.RR, 0)

	for rr, ok := parser.Next(); ok; rr, ok = parser.Next() {
		rrs = append(rrs, rr)
	}

	if err := parser.Err(); err != nil {
		log.Println(err)
	}

	return rrs, nil
}

func createACLList(config []configACL) (map[string]*net.IPNet, error) {
	acls := make(map[string]*net.IPNet)

	for _, aclRule := range config {
		_, mask, err := net.ParseCIDR(aclRule.CIDR)

		if err != nil {
			return nil, err
		}

		acls[aclRule.Name] = mask
	}

	return acls, nil
}

func createServer(zones zoneMap, config config, aclList map[string]*net.IPNet) *dns.ServeMux {
	srv := dns.NewServeMux()
	c := new(dns.Client)

	for zoneName, zones := range zones {
		srv.HandleFunc(zoneName, func(w dns.ResponseWriter, r *dns.Msg) {

			// Parse IP
			remoteIP, _, err := net.SplitHostPort(w.RemoteAddr().String())
			ip := net.ParseIP(remoteIP)
			if err != nil && ip != nil {
				log.Printf("Faild to parse remote IP WTF? :%s\n", err.Error())
				return
			}

			// find out what view to handle the request
			zoneIndex := -1

			for i, zone := range zones {
				if (len(zone.acl) == 0 && zoneIndex == -1) || checkACL(zone.acl, aclList, ip) {
					zoneIndex = i
				}
			}

			if zoneIndex == -1 {
				rcodeRequest(w, r, dns.RcodeRefused)
				return
			}

			handleRequest(w, r, zones[zoneIndex])
		})
	}

	// Handle any other request
	srv.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
		remoteIP, _, err := net.SplitHostPort(w.RemoteAddr().String())
		ip := net.ParseIP(remoteIP)

		if err != nil && ip != nil {
			log.Printf("Faild to parse remote IP WTF? :%s\n", err.Error())
			return
		}

		// Check ACL rules
		if !checkACL(config.Forward.ACL, aclList, ip) {
			rcodeRequest(w, r, dns.RcodeRefused)
			return
		}

		// Forward request
		in, _, err := c.Exchange(r, config.Forward.Server)

		if err != nil {
			rcodeRequest(w, r, dns.RcodeServerFailure)
			return
		}

		w.WriteMsg(in)
	})

	return srv
}

func listenAndServer(server *dns.ServeMux, address string) {
	go func() {
		if err := dns.ListenAndServe(address, "udp", server); err != nil {
			log.Fatalf("Failed to set udp listener %s\n", err.Error())
		}
	}()

	go func() {
		if err := dns.ListenAndServe(address, "tcp", server); err != nil {
			log.Fatalf("Failed to set tcp listener %s\n", err.Error())
		}
	}()
}

func checkACL(alcRules []string, aclList map[string]*net.IPNet, ip net.IP) bool {
	if len(alcRules) != 0 {
		passed := false
		for _, rule := range alcRules {

			if aclList[rule].Contains(ip) {
				passed = true
			}
		}
		return passed
	}
	return true
}

func rcodeRequest(w dns.ResponseWriter, r *dns.Msg, rcode int) {
	m := new(dns.Msg)
	m.SetReply(r)
	m.SetRcode(r, rcode)
	w.WriteMsg(m)
}

func handleRequest(w dns.ResponseWriter, r *dns.Msg, zone zoneView) {
	m := new(dns.Msg)
	m.SetReply(r)
	m.Authoritative = true

	// maybe only support one question per query like most servers do it ???
	for _, q := range r.Question {
		rrs := zone.rr[q.Qtype]

		// Handle ANY
		if q.Qtype == dns.TypeANY {
			for _, rrType := range anyRecordTypes {
				m.Answer = append(m.Answer, zone.rr[rrType][q.Name]...)
			}
		} else {
			// Handle any other type
			m.Answer = append(m.Answer, rrs[q.Name]...)

			// Check for wildcard
			if len(m.Answer) == 0 {
				parts := dns.SplitDomainName(q.Name)[1:]
				searchDomain := "*." + dns.Fqdn(strings.Join(parts, "."))
				foundDomain := rrs[searchDomain]
				for _, rr := range foundDomain {
					newRR := rr
					newRR.Header().Name = q.Name
					m.Answer = append(m.Answer, newRR)
				}
			}
		}

		// Handle extras
		switch q.Qtype {
		// Dont handle extra stuff when answering ANY request
		// case dns.TypeANY:
		// 	fallthrough
		case dns.TypeMX:
			// Resolve MX domains
			for _, mxRR := range m.Answer {
				if t, ok := mxRR.(*dns.MX); ok {
					m.Extra = append(m.Extra, zone.rr[dns.TypeA][t.Mx]...)
					m.Extra = append(m.Extra, zone.rr[dns.TypeAAAA][t.Mx]...)
				}
			}
		case dns.TypeA, dns.TypeAAAA:
			if len(m.Answer) == 0 {
				// no A or AAAA found. Look for CNAME
				m.Answer = append(m.Answer, zone.rr[dns.TypeCNAME][q.Name]...)
				if len(m.Answer) != 0 {
					// Resolve CNAME
					for _, nameRR := range m.Answer {
						if t, ok := nameRR.(*dns.CNAME); ok {
							m.Answer = append(m.Answer, zone.rr[q.Qtype][t.Target]...)
						}
					}
				}
			}
		case dns.TypeNS:
			// Resove NS records
			for _, nsRR := range m.Answer {
				if t, ok := nsRR.(*dns.NS); ok {
					m.Extra = append(m.Extra, zone.rr[dns.TypeA][t.Ns]...)
					m.Extra = append(m.Extra, zone.rr[dns.TypeAAAA][t.Ns]...)
				}
			}
		}
	}

	w.WriteMsg(m)
}

func main() {

	configPath := flag.String("c", "/etc/cool-dns/config.yaml", "path to the config file")
	flag.Parse()

	config, err := loadConfig(*configPath)
	if err != nil {
		log.Fatalf("Failed to load config: %s\n", err.Error())
	}

	zones, err := loadZones(config.Zones)
	if err != nil {
		log.Fatalf("Failed to load zones: %s\n", err.Error())
	}

	aclList, err := createACLList(config.ACL)
	if err != nil {
		log.Fatalf("Failed to parse ACL rules: %s\n", err.Error())
	}

	server := createServer(zones, *config, aclList)

	listenAndServer(server, config.Address)

	log.Printf("Start listening on udp %s and tcp %s\n", config.Address, config.Address)

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	s := <-sig
	log.Printf("Signal (%v) received, stopping\n", s)
	os.Exit(0)
}