package protocol

import (
	"log"

	"golang.org/x/crypto/cryptobyte"
)

func init() {
	registerSTUN()
}

func registerSTUN() {
	Register(Both, "????\x21\x12\xa4\x42", detectSTUN)
}

func detectSTUN(dir Direction, data []byte, srcPort, dstPort int) (proto *Protocol, confidence float64) {
	// All STUN messages comprise a 20-byte header followed by zero or more attributes.
	var (
		stream        = cryptobyte.String(data)
		messageType   uint16
		messageLength uint16
		magicCookie   uint32
		transactionID []byte
	)
	if !stream.ReadUint16(&messageType) ||
		!stream.ReadUint16(&messageLength) ||
		!stream.ReadUint32(&magicCookie) ||
		!stream.ReadBytes(&transactionID, 12) {
		return
	}

	if len(stream) != int(messageLength) {
		confidence -= .2
	}

	// The most significant 2 bits of every STUN message MUST be zeroes.
	if messageType&0b11000000 != 0 {
		log.Printf("type %#08b", messageType)
		return
	}

	// The Magic Cookie field MUST contain the fixed value 0x2112A442 in network byte order.
	if magicCookie != 0x2112A442 {
		log.Printf("cookie %#04x", magicCookie)
		return
	}
	confidence += .5

	// Decode the message class & type.
	class := (messageType & 0b00000100000000) >> 7
	class |= (messageType & 0b00000000010000) >> 4
	method := (messageType & 0b11111000000000) >> 2
	method |= (messageType & 0b00000011100000) >> 1
	method |= (messageType & 0b00000000001111) >> 0

	// Validate most common methods.
	// TODO(maze): check method arguments to improve accuracy
	switch class {
	case stunRequest: // request
		switch method {
		case stunBinding: // binding
			confidence += .2
		}

	case stunSuccesResponse: // success response
		switch method {
		case stunBinding: // binding
			confidence += .2
		}

	case stunErrorResponse:
		switch method {
		case stunBinding: // binding error
			confidence += .2
		}
	}

	if stunPort[srcPort] || stunPort[dstPort] {
		confidence += .2
	}

	return &Protocol{
		Type: TypeSTUN,
	}, confidence
}

const (
	stunRequest        = 0b00
	stunSuccesResponse = 0b10
	stunErrorResponse  = 0b11
	stunBinding        = 0x0001
)

// IANA has updated the reference from RFC 5389 to RFC 8489 for the
// following ports in the "Service Name and Transport Protocol Port
// Number Registry".
var stunPort = map[int]bool{
	3478: true,
	5349: true,
}