// based on https://github.com/aler9/landiscover/blob/main/layer-nbns.go (MIT License)

package layer_nbns

import (
	"encoding/binary"
	"fmt"
	"strings"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
)

type layerNbns struct {
	layers.BaseLayer
	TransactionID   uint16
	IsResponse      bool
	Opcode          uint8
	Truncated       bool
	Recursion       bool
	Broadcast       bool
	Questions       []nbnsQuestion
	Answers         []nbnsAnswer
	AuthorityCount  uint16
	AdditionalCount uint16
}

func NewLayerNbns() *layerNbns {
	return &layerNbns{}
}

type nbnsQuestion struct {
	Query string
	Type  uint16
	Class uint16
}

type nbnsAnswer struct {
	Query string
	Type  uint16
	Class uint16
	TTL   uint32
	Names []nbnsAnswerName
}

type nbnsAnswerName struct {
	Name  string
	Type  uint8
	Flags uint16
}

func (l *layerNbns) DecodeFromBytes(data []byte) error {
	l.BaseLayer = layers.BaseLayer{Contents: data}

	if len(data) < 12 {
		return fmt.Errorf("invalid packet")
	}

	l.TransactionID = binary.BigEndian.Uint16(data[0:2])
	l.IsResponse = (data[3] >> 7) == 0x01
	l.Opcode = (data[3] >> 3) & 0x0F
	l.Truncated = (data[3] >> 1) == 0x01
	l.Recursion = data[3] == 0x01
	l.Broadcast = (data[4] >> 4) == 0x01
	questionCount := binary.BigEndian.Uint16(data[4:6])
	answerCount := binary.BigEndian.Uint16(data[6:8])
	l.AuthorityCount = binary.BigEndian.Uint16(data[8:10])
	l.AdditionalCount = binary.BigEndian.Uint16(data[10:12])
	pos := 12

	if questionCount > 0 {
		return fmt.Errorf("is question, unsupported")
	}

	l.Questions = nil
	for i := uint16(0); i < questionCount; i++ {
		q := nbnsQuestion{}

		var read int
		q.Query, read = dnsQueryDecode(data, pos)
		if read <= 0 {
			return fmt.Errorf("invalid string")
		}
		pos += read

		q.Type = binary.BigEndian.Uint16(data[pos+0 : pos+2])
		q.Class = binary.BigEndian.Uint16(data[pos+2 : pos+4])
		pos += 4
		l.Questions = append(l.Questions, q)
	}

	l.Answers = nil
	for i := uint16(0); i < answerCount; i++ {
		a := nbnsAnswer{}

		var read int
		a.Query, read = dnsQueryDecode(data, pos)
		if read <= 0 {
			return fmt.Errorf("invalid string")
		}
		pos += read

		a.Type = binary.BigEndian.Uint16(data[pos+0 : pos+2])
		a.Class = binary.BigEndian.Uint16(data[pos+2 : pos+4])
		a.TTL = binary.BigEndian.Uint32(data[pos+4 : pos+8])
		dataLen := binary.BigEndian.Uint16(data[pos+8 : pos+10])
		pos += 10

		if a.Type == 0x21 { // NB_STAT
			pos2 := pos
			nameCount := data[pos2]
			pos2++

			for j := uint8(0); j < nameCount; j++ {
				a.Names = append(a.Names, nbnsAnswerName{
					Name:  strings.TrimSuffix(string(data[pos2:pos2+15]), " "),
					Type:  data[pos2+15],
					Flags: binary.BigEndian.Uint16(data[pos2+16 : pos2+18]),
				})
				pos2 += 18
			}
		}

		pos += int(dataLen)
		l.Answers = append(l.Answers, a)
	}

	// TODO: Decode Authorities and Additional Information

	return nil
}

func (l *layerNbns) SerializeTo(b gopacket.SerializeBuffer) error {
	data, err := b.AppendBytes(12)
	if err != nil {
		panic(err)
	}

	binary.BigEndian.PutUint16(data[0:2], l.TransactionID)
	if l.IsResponse {
		data[3] |= 0x01 << 7
	}
	data[3] |= l.Opcode << 3
	if l.Truncated {
		data[3] |= 0x01 << 1
	}
	if l.Recursion {
		data[3] |= 0x01
	}
	if l.Broadcast {
		data[4] |= 0x01 << 4
	}
	binary.BigEndian.PutUint16(data[4:6], uint16(len(l.Questions)))
	binary.BigEndian.PutUint16(data[6:8], uint16(len(l.Answers)))
	binary.BigEndian.PutUint16(data[8:10], l.AuthorityCount)
	binary.BigEndian.PutUint16(data[10:12], l.AdditionalCount)

	for _, q := range l.Questions {
		enc := dnsQueryEncode(q.Query)

		data, err := b.AppendBytes(len(enc) + 4)
		if err != nil {
			panic(err)
		}

		copy(data[:len(enc)], enc)
		data = data[len(enc):]
		binary.BigEndian.PutUint16(data[0:2], q.Type)
		binary.BigEndian.PutUint16(data[2:4], q.Class)
	}

	// TODO: encode answers, authorities and additional informations

	return nil
}

// Source: https://github.com/aler9/landiscover/blob/main/utils.go
// <size>part<size>part
func dnsQueryDecode(data []byte, start int) (string, int) {
	var read []byte
	toread := uint8(0)
	pos := start

	for ; true; pos++ {
		if pos >= len(data) { // decoding terminated before null character
			return "", -1
		}
		if data[pos] == 0x00 {
			if toread > 0 { // decoding terminated before part parsing
				return "", -1
			}
			break // query correctly decoded
		}

		if toread == 0 { // we need a size or pointer
			if len(read) > 0 { // add separator
				read = append(read, '.')
			}

			if (data[pos] & 0xC0) == 0xC0 { // pointer
				ptr := int(binary.BigEndian.Uint16(data[pos:pos+2]) & 0x3FFF)
				pos++ // skip next byte

				substr, subread := dnsQueryDecode(data, ptr)
				if subread <= 0 {
					return "", -1
				}
				read = append(read, []byte(substr)...)
				break // query correctly decoded
			} else { // size
				toread = data[pos]
			}
		} else { // byte inside part
			read = append(read, data[pos])
			toread--
		}
	}
	return string(read), pos + 1 - start
}

func dnsQueryEncode(in string) []byte {
	tmp := strings.Split(in, ".")

	l := 0
	for _, part := range tmp {
		bpart := []byte(part)
		l++
		l += len(bpart)
	}
	l++

	ret := make([]byte, l)
	i := 0

	for _, part := range tmp {
		bpart := []byte(part)
		ret[i] = uint8(len(bpart))
		i++
		copy(ret[i:], bpart)
		i += len(bpart)
	}

	ret[i] = uint8(0)

	return ret
}