package dnstt

import (
	"encoding/binary"
	"strings"
)

// ProcessIPv4Packet records DNSTT DNS traffic from a raw IPv4 packet.
func ProcessIPv4Packet(data []byte, port int, collector *Collector) {
	if len(data) < 20 || data[0]>>4 != 4 {
		return
	}

	ihl := int(data[0]&0x0f) * 4
	if ihl < 20 || len(data) < ihl {
		return
	}
	if data[9] != 17 {
		return
	}

	udpData := data[ihl:]
	if len(udpData) < 8 {
		return
	}

	srcPort := binary.BigEndian.Uint16(udpData[0:2])
	dstPort := binary.BigEndian.Uint16(udpData[2:4])
	udpLen := int(binary.BigEndian.Uint16(udpData[4:6]))
	dnsPayload := udpData[8:]
	if len(dnsPayload) < 12 {
		return
	}

	if int(dstPort) == port {
		domain, clientID := ExtractQuery(dnsPayload, collector.Domains())
		if domain != "" {
			collector.RecordQuery(domain, clientID, udpLen)
		}
		return
	}

	if int(srcPort) == port {
		domain := extractQueryDomain(dnsPayload)
		if domain != "" {
			collector.RecordResponse(domain, udpLen)
		}
	}
}

func extractQueryDomain(dns []byte) string {
	if len(dns) < 12 {
		return ""
	}
	labels := parseDNSLabels(dns[12:])
	if len(labels) == 0 {
		return ""
	}
	return normalizeDomain(strings.Join(labels, "."))
}