package main

import (
	"context"
	"errors"
	"net"
	"net/netip"
	"sync"
	"time"

	"codeberg.org/miekg/dns"
	"codeberg.org/miekg/dns/rdata"
	"github.com/jedisct1/dlog"
)

const rfc7050WKN = "ipv4only.arpa."

var (
	rfc7050WKA1 = net.IPv4(192, 0, 0, 170)
	rfc7050WKA2 = net.IPv4(192, 0, 0, 171)
)

type PluginDNS64 struct {
	pref64Mutex    *sync.RWMutex
	pref64         []*net.IPNet
	dns64Resolvers []string
	ipv4Resolver   string
	proxy          *Proxy
}

func (plugin *PluginDNS64) Name() string {
	return "dns64"
}

func (plugin *PluginDNS64) Description() string {
	return "Synthesize DNS64 AAAA responses"
}

func (plugin *PluginDNS64) Init(proxy *Proxy) error {
	if len(proxy.listenAddresses) == 0 {
		return errors.New("At least one listening IP address must be configured for the DNS64 plugin to work")
	}
	plugin.ipv4Resolver = proxy.listenAddresses[0] // query is sent to ourselves
	plugin.pref64Mutex = new(sync.RWMutex)
	plugin.proxy = proxy

	if len(proxy.dns64Prefixes) != 0 {
		plugin.pref64Mutex.Lock()
		defer plugin.pref64Mutex.Unlock()
		for _, prefStr := range proxy.dns64Prefixes {
			_, pref, err := net.ParseCIDR(prefStr)
			if err != nil {
				return err
			}
			if ones, _ := pref.Mask.Size(); ones > 96 {
				return errors.New("DNS64 prefixes must be /96 or shorter")
			}
			dlog.Noticef("Registered DNS64 prefix [%s]", pref.String())
			plugin.pref64 = append(plugin.pref64, pref)
		}
	} else if len(proxy.dns64Resolvers) != 0 {
		plugin.dns64Resolvers = proxy.dns64Resolvers
		if err := plugin.refreshPref64(); err != nil {
			return err
		}
	} else {
		return nil
	}
	dlog.Notice("DNS64 map enabled")

	return nil
}

func (plugin *PluginDNS64) Drop() error {
	return nil
}

func (plugin *PluginDNS64) Reload() error {
	return nil
}

func (plugin *PluginDNS64) Eval(pluginsState *PluginsState, msg *dns.Msg) error {
	if hasAAAAAnswer(msg) {
		return nil
	}

	question := pluginsState.questionMsg.Question[0]
	qtype := dns.RRToType(question)
	if question.Header().Class != dns.ClassINET || qtype != dns.TypeAAAA {
		return nil
	}

	msgA := dns.NewMsg(question.Header().Name, dns.TypeA)
	msgA.ID = pluginsState.questionMsg.ID
	msgA.RecursionDesired = pluginsState.questionMsg.RecursionDesired
	if err := msgA.Pack(); err != nil {
		return err
	}
	msgAPacket := msgA.Data

	if !plugin.proxy.clientsCountInc() {
		return errors.New("Too many concurrent connections to handle DNS64 subqueries")
	}
	respPacket := plugin.proxy.processIncomingQuery(
		"trampoline",
		plugin.proxy.xTransport.mainProto,
		msgAPacket,
		nil,
		nil,
		time.Now(),
		false,
	)
	plugin.proxy.clientsCountDec()

	if len(respPacket) == 0 {
		return errors.New("Empty response from DNS64 trampoline query")
	}

	resp := dns.Msg{Data: respPacket}
	if err := resp.Unpack(); err != nil {
		return err
	}

	if resp.Rcode != dns.RcodeSuccess {
		return nil
	}

	if len(resp.Answer) == 0 {
		return nil
	}

	initialTTL := uint32(600)
	for _, ns := range resp.Ns {
		header := ns.Header()
		if dns.RRToType(ns) == dns.TypeSOA {
			initialTTL = header.TTL
		}
	}

	synth64 := make([]dns.RR, 0)
	for _, answer := range resp.Answer {
		header := answer.Header()
		rrtype := dns.RRToType(answer)
		if rrtype == dns.TypeCNAME {
			synth64 = append(synth64, answer)
		} else if rrtype == dns.TypeA {
			ttl := min(initialTTL, header.TTL)

			ipv4 := answer.(*dns.A).A.Addr.AsSlice()
			if len(ipv4) == 4 {
				plugin.pref64Mutex.RLock()
				for _, prefix := range plugin.pref64 {
					ipv6 := translateToIPv6(net.IP(ipv4), prefix)
					synthAAAA := new(dns.AAAA)
					synthAAAA.Hdr = dns.Header{
						Name:  header.Name,
						Class: header.Class,
						TTL:   ttl,
					}
					synthAAAA.AAAA = rdata.AAAA{Addr: netip.AddrFrom16([16]byte(ipv6.To16()))}
					synth64 = append(synth64, synthAAAA)
				}
				plugin.pref64Mutex.RUnlock()
			}
		}
	}

	msg.Answer = synth64
	msg.AuthenticatedData = false
	msg.UDPSize = uint16(MaxDNSUDPSafePacketSize)

	pluginsState.returnCode = PluginsReturnCodeCloak

	return nil
}

func hasAAAAAnswer(msg *dns.Msg) bool {
	for _, answer := range msg.Answer {
		if dns.RRToType(answer) == dns.TypeAAAA {
			return true
		}
	}
	return false
}

func translateToIPv6(ipv4 net.IP, prefix *net.IPNet) net.IP {
	ipv6 := make(net.IP, net.IPv6len)
	copy(ipv6, prefix.IP)
	n, _ := prefix.Mask.Size()
	ipShift := n / 8
	for i := range net.IPv4len {
		if ipShift+i == 8 {
			ipShift++
		}
		ipv6[ipShift+i] = ipv4[i]
	}
	return ipv6
}

func (plugin *PluginDNS64) fetchPref64(resolver string) error {
	msg := dns.NewMsg(rfc7050WKN, dns.TypeAAAA)

	client := new(dns.Client)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	resp, _, err := client.Exchange(ctx, msg, "udp", resolver)
	if err != nil {
		return err
	}

	if resp == nil || resp.Rcode != dns.RcodeSuccess {
		return errors.New("Unable to fetch Pref64")
	}

	uniqPrefixes := make(map[string]struct{})
	prefixes := make([]*net.IPNet, 0)
	for _, answer := range resp.Answer {
		if dns.RRToType(answer) == dns.TypeAAAA {
			ipv6 := answer.(*dns.AAAA).AAAA.Addr.AsSlice()
			if ipv6 != nil && len(ipv6) == net.IPv6len {
				prefEnd := 0

				if wka := net.IPv4(ipv6[12], ipv6[13], ipv6[14], ipv6[15]); wka.Equal(rfc7050WKA1) ||
					wka.Equal(rfc7050WKA2) { // 96
					prefEnd = 12
				} else if wka := net.IPv4(ipv6[9], ipv6[10], ipv6[11], ipv6[12]); wka.Equal(rfc7050WKA1) || wka.Equal(rfc7050WKA2) { // 64
					prefEnd = 8
				} else if wka := net.IPv4(ipv6[7], ipv6[9], ipv6[10], ipv6[11]); wka.Equal(rfc7050WKA1) || wka.Equal(rfc7050WKA2) { // 56
					prefEnd = 7
				} else if wka := net.IPv4(ipv6[6], ipv6[7], ipv6[9], ipv6[10]); wka.Equal(rfc7050WKA1) || wka.Equal(rfc7050WKA2) { // 48
					prefEnd = 6
				} else if wka := net.IPv4(ipv6[5], ipv6[6], ipv6[7], ipv6[9]); wka.Equal(rfc7050WKA1) || wka.Equal(rfc7050WKA2) { // 40
					prefEnd = 5
				} else if wka := net.IPv4(ipv6[4], ipv6[5], ipv6[6], ipv6[7]); wka.Equal(rfc7050WKA1) || wka.Equal(rfc7050WKA2) { // 32
					prefEnd = 4
				}

				if prefEnd > 0 {
					prefix := new(net.IPNet)
					prefix.IP = append(ipv6[:prefEnd], net.IPv6zero[prefEnd:]...)
					prefix.Mask = net.CIDRMask(prefEnd*8, 128)
					if _, ok := uniqPrefixes[prefix.String()]; !ok {
						prefixes = append(prefixes, prefix)
						uniqPrefixes[prefix.String()] = struct{}{}
						dlog.Infof("Registered DNS64 prefix [%s]", prefix.String())
					}
				}
			}
		}
	}

	if len(prefixes) == 0 {
		return errors.New("Empty Pref64 list")
	}

	plugin.pref64Mutex.Lock()
	defer plugin.pref64Mutex.Unlock()
	plugin.pref64 = prefixes
	return nil
}

func (plugin *PluginDNS64) refreshPref64() error {
	for _, resolver := range plugin.dns64Resolvers {
		if err := plugin.fetchPref64(resolver); err == nil {
			break
		}
	}

	plugin.pref64Mutex.RLock()
	defer plugin.pref64Mutex.RUnlock()
	if len(plugin.pref64) == 0 {
		return errors.New("Empty Pref64 list")
	}

	return nil
}
