arpa_enum.go

package main

import (
	"bufio"
	"database/sql"
	"fmt"
	"iter"
	"net/netip"
	"os"
	"regexp"
	"slices"
	"strings"
	"sync"

	"github.com/miekg/dns"
	"github.com/monoidic/cidranger"
)

// database ID for a partial arpa PTR name in the and an array of its children
type arpaResults struct {
	rootID  int64
	results []arpaResult
}

// arpa PTR name, whether it gave a NXDOMAIN response, and its nameservers
type arpaResult struct {
	name     string
	NSes     []dns.NS
	nxdomain bool
}

// simple cidranger.RangerEntry implementation, with a boolean marking whether it matches
type rangerEntry struct {
	subnet netip.Prefix
	match  bool
}

// implement cidranger.RangerEntry interface
func (entry rangerEntry) Network() netip.Prefix {
	return entry.subnet
}

// numbers of dots in full-length PTR names
const (
	ARPAV4DOTS = 6
	ARPAV6DOTS = 34
)

var ptrV6Pattern = regexp.MustCompile("....")

// return writerF function for performing arpa queries and writing the results,
// with a function for generating new subnames given as an argument
func arpaWriter(translator func(iter.Seq[fieldData]) iter.Seq[arpaResults]) writerF[fieldData] {
	return func(db *sql.DB, seq iter.Seq[fieldData]) {
		tablesFields := map[string]string{
			"name":     "name",
			"rr_type":  "name",
			"rr_name":  "name",
			"rr_value": "value",
		}
		namesStmts := map[string]string{
			"delete_root": "DELETE FROM unwalked_root WHERE id=?",
			"add_root":    "INSERT INTO unwalked_root (name, ent) VALUES (?, ?)",
			"zone2rr":     "INSERT OR IGNORE INTO zone2rr (zone_id, rr_type_id, rr_name_id, rr_value_id) VALUES (?, ?, ?, ?)",
		}

		results := translator(seq)
		results = filterTranslates(results)

		netWriter(db, results, tablesFields, namesStmts, arpaWorker, arpaWrite)
	}
}

// read network file generated by github.com/monoidic/rir and insert matching and non-matching subnets
func getCCRanger() cidranger.Ranger {
	fd := check1(os.Open(networksFile))
	scanner := bufio.NewScanner(fd)
	ranger := cidranger.NewPCTrieRanger()
	for scanner.Scan() {
		text := scanner.Text()
		split := strings.SplitN(text, "\t", 2)
		if len(split) != 2 {
			panic(fmt.Sprintf("invalid line: %q", text))
		}
		subnet := check1(netip.ParsePrefix(split[1]))
		check(ranger.Insert(rangerEntry{subnet: subnet, match: strings.ToUpper(split[0]) == netCC}))
	}

	check(fd.Close())
	return ranger
}

// filter arpa names so that only names which have a matching subnet as its child or its most precise covering subnet
func filterTranslates(seq iter.Seq[arpaResults]) iter.Seq[arpaResults] {
	if netCC == "" {
		return seq
	}
	return func(yield func(arpaResults) bool) {
		ranger := getCCRanger()

		for results := range seq {
			var newList []arpaResult
			for _, result := range results.results {
				net, err := partialPtrToNet(result.name)
				if err != nil {
					continue
				}

				var match bool

				if coveredNets := check1(ranger.CoveredNetworks(net)); anyRangerMatch(coveredNets) {
					// fmt.Printf("%s covered in %s\n", net, netCC)
					match = true
				} else if coveringNets := check1(ranger.CoveringNetworks(net)); len(coveringNets) > 0 && coveringNets[len(coveringNets)-1].(rangerEntry).match {
					// fmt.Printf("%s covering in %s\n", net, netCC)
					match = true
				}

				if match {
					newList = append(newList, result)
					//} else {
					//	fmt.Printf("%s NOT contained in %s\n", net, netCC)
				}
			}

			if len(newList) < len(results.results) {
				results.results = newList
			}
			if !yield(results) {
				return
			}
		}
	}
}

// takes a list of ranger entries and returns true if any of them match
func anyRangerMatch(in []cidranger.RangerEntry) bool {
	for _, v := range in {
		if v.(rangerEntry).match {
			return true
		}
	}
	return false
}

// generate addresses for IPv4 PTR records
func arpaV4Translate(seq iter.Seq[fieldData]) iter.Seq[arpaResults] {
	return func(yield func(arpaResults) bool) {
		for fd := range seq {
			res := arpaResults{rootID: fd.id, results: make([]arpaResult, 256)}
			for i := range 256 {
				res.results[i] = arpaResult{name: fmt.Sprintf("%d.%s", i, fd.name)}
			}
			if !yield(res) {
				return
			}
		}
	}
}

// generate addresses for IPv6 PTR records
func arpaV6Translate(seq iter.Seq[fieldData]) iter.Seq[arpaResults] {
	return func(yield func(arpaResults) bool) {
		for fd := range seq {
			res := arpaResults{rootID: fd.id, results: make([]arpaResult, 16)}
			for i, c := range "0123456789abcdef" {
				res.results[i] = arpaResult{name: fmt.Sprintf("%c.%s", c, fd.name)}
			}
			if !yield(res) {
				return
			}
		}
	}
}

// worker performing arpa queries
func arpaWorker(dataChan <-chan retryWrap[arpaResults, arpaResults], refeedChan chan<- retryWrap[arpaResults, arpaResults], outChan chan<- arpaResults, wg, retryWg *sync.WaitGroup) {
	msg := dns.Msg{
		MsgHdr: dns.MsgHdr{
			Opcode:           dns.OpcodeQuery,
			RecursionDesired: true,
			Rcode:            dns.RcodeSuccess,
		},
		Question: []dns.Question{{
			Qclass: dns.ClassINET,
			Qtype:  dns.TypeNS,
		}},
	}
	msgSetSize(&msg)

	resolverWorker(dataChan, refeedChan, outChan, msg, arpaResolve, wg, retryWg)
}

// resolve .arpa PTR records
func arpaResolve(connCache connCache, msg dns.Msg, search *retryWrap[arpaResults, arpaResults]) (results arpaResults, err error) {
	if search.stage == 0 {
		search.tmp.results = make([]arpaResult, len(search.val.results))
	}

	for i := search.stage; i < len(search.val.results); i++ {
		search.stage = i
		result := search.val.results[i]
		msg.Question[0].Name = result.name
		var response *dns.Msg

		response, err = plainResolveRandom(msg, connCache)
		if err != nil {
			if search.retriesLeft == 0 {
				// skip to next
				search.retriesLeft = RETRIES
				err = nil
			} else {
				return
			}
		}

		if response != nil {
			result.nxdomain = response.Rcode != dns.RcodeSuccess
			fmt.Printf("rcode for %s: %s\n", result.name, dns.RcodeToString[response.Rcode])
			for _, rr := range response.Answer {
				switch rrT := rr.(type) {
				case *dns.NS:
					normalizeRR(rrT)
					result.NSes = append(result.NSes, *rrT)
				}
			}
		} else {
			result.nxdomain = true
		}

		search.tmp.results[i] = result
	}
	results = search.tmp
	results.rootID = search.val.rootID
	return
}

// write recursed results
func arpaWrite(tableMap TableMap, stmtMap StmtMap, datum arpaResults) {
	for _, result := range datum.results {
		if !result.nxdomain {
			fmt.Printf("new root %s\n", result.name)
			emptyNonTerminal := len(result.NSes) == 0
			stmtMap.exec("add_root", result.name, emptyNonTerminal)

			if !emptyNonTerminal {
				rrTypeID := tableMap.get("rr_type", "NS")
				for _, ns := range result.NSes {
					rrNameID := tableMap.get("rr_name", ns.Hdr.Name)
					zoneID := tableMap.get("name", ns.Hdr.Name)
					rrValueID := tableMap.get("rr_value", ns.String())

					stmtMap.exec("zone2rr", zoneID, rrTypeID, rrNameID, rrValueID)
				}
			}
		}
	}
	stmtMap.exec("delete_root", datum.rootID)
}

// convert parital PTR name (v4 or v6), e.g 3.2.1.in-addr.arpa, to the netip.Prefix for the corresponding subnet, (1.2.3.0/24 for the given example)
func partialPtrToNet(s string) (netip.Prefix, error) {
	count := strings.Count(s, ".")
	var limit, bits, sectionBits int
	if strings.HasSuffix(s, ".ip6.arpa.") {
		limit = ARPAV6DOTS
		bits = 128
		sectionBits = 4
	} else if strings.HasSuffix(s, ".in-addr.arpa.") {
		limit = ARPAV4DOTS
		bits = 32
		sectionBits = 8
	} else {
		return netip.Prefix{}, Error{s: fmt.Sprintf("invalid addr: %q", s)}
	}

	if count > limit {
		return netip.Prefix{}, Error{s: fmt.Sprintf("invalid addr: %q", s)}
	} else if count == limit {
		return netip.Prefix{}, Error{s: "not partial ptr"}
	}

	s = strings.Repeat("0.", limit-count) + s
	ip, err := ptrToIP(s)
	if err != nil {
		return netip.Prefix{}, err
	}
	subnet := bits - (limit-count)*sectionBits
	return netip.PrefixFrom(ip, subnet), nil
}

// convert full-length PTR name (v4 or v6), e.g 4.3.2.1.in-addr.arpa, to the netip.Addr for the corresponding IP address (1.2.3.4 for the given example)
func ptrToIP(s string) (netip.Addr, error) {
	if strings.HasSuffix(s, ".ip6.arpa.") {
		s = s[:len(s)-len(".ip6.arpa.")]
		s = strings.ReplaceAll(s, ".", "")
		s = reverseASCII(s)
		s = strings.Join(ptrV6Pattern.FindAllString(s, -1), ":")
	} else if strings.HasSuffix(s, ".in-addr.arpa.") {
		s = s[:len(s)-len(".in-addr.arpa.")]
		l := strings.Split(s, ".")
		slices.Reverse(l)
		s = strings.Join(l, ".")
	} else {
		return netip.Addr{}, Error{s: fmt.Sprintf("invalid addr: %q", s)}
	}
	return netip.ParseAddr(s)
}

// initialize recursive enumeration with root
func setupArpa(db *sql.DB, root string) {
	tx := check1(db.Begin())
	check1(tx.Exec(`DELETE FROM unwalked_root`))
	check1(tx.Exec(fmt.Sprintf(`INSERT INTO unwalked_root (name) VALUES ('%s')`, root)))
	check(tx.Commit())
}

func recurseArpaV4(db *sql.DB) {
	setupArpa(db, "in-addr.arpa.")
	readerWriterRecurse("recursing through in-addr.arpa.", db, getUnqueriedArpaRoots, arpaWriter(arpaV4Translate))
}

func recurseArpaV6(db *sql.DB) {
	setupArpa(db, "ip6.arpa.")
	readerWriterRecurse("recursing through ip6.arpa.", db, getUnqueriedArpaRoots, arpaWriter(arpaV6Translate))
}