package protocol

import (
	"log"

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

func init() {
	Register(Both, "\x10?*MQTT", detectMQTT)
}

func detectMQTT(dir Direction, data []byte, srcPort, dstPort int) (proto *Protocol, confidence float64) {
	stream := cryptobyte.String(data)

	var packetType byte
	if !stream.ReadUint8(&packetType) || packetType != 0x10 {
		// Not an MQTT CONNECT control packet.
		return
	}

	// We read the value but only check in the end if we have bytes remaining.
	var remainingLength uint32
	if !readMQTTVariableLengthUint32(&stream, &remainingLength) {
		return
	}

	var protocolName cryptobyte.String
	if !stream.ReadUint16LengthPrefixed(&protocolName) || string(protocolName) != "MQTT" {
		return
	}

	// We are reasonabily sure this is MQTT now.
	proto = &Protocol{
		Name: ProtocolMQTT,
	}
	confidence = 0.5

	var protocolVersion uint8
	if !stream.ReadUint8(&protocolVersion) {
		return
	}
	switch protocolVersion {
	case 4:
		proto.Version = Version{Major: 3, Minor: 1, Patch: 1}
		confidence += .2
	case 5:
		proto.Version = Version{Major: 5, Minor: 0, Patch: -1}
		confidence += .2
	}

	var connectFlags uint8
	if !stream.ReadUint8(&connectFlags) { // read connect flags
		return
	}
	confidence += .05
	var (
		hasUsername = (connectFlags&0b10000000)>>7 == 1
		hasPassword = (connectFlags&0b01000000)>>6 == 1
		willFlag    = (connectFlags&0b00000100)>>2 == 1
		reserved    = (connectFlags&0b00000001)>>0 == 1
		discard     cryptobyte.String
	)
	if reserved {
		// Reserved bit is supposed to be not set.
		confidence -= .2
	}

	var keepAlive uint16
	if !stream.ReadUint16(&keepAlive) {
		return
	}

	if !stream.ReadUint16LengthPrefixed(&discard) { // read client ID
		return
	}
	confidence += .05
	discard = discard[:0]

	if willFlag {
		if !stream.ReadUint16LengthPrefixed(&discard) { // read will topic
			log.Println("will topic fail")
			return
		}
		discard = discard[:0]
		if !stream.ReadUint16LengthPrefixed(&discard) { // read will message
			log.Println("will message fail")
			return
		}
		discard = discard[:0]
		confidence += .05
	} else {
		confidence += .05
	}

	if hasUsername {
		if !stream.ReadUint16LengthPrefixed(&discard) { // read username
			log.Println("user fail")
			return
		}
		discard = discard[:0]
		confidence += .05
	} else {
		confidence += .05
	}

	if hasPassword {
		if !stream.ReadUint16LengthPrefixed(&discard) { // read password
			log.Println("pass fail")
			return
		}
		confidence += .05
	} else {
		confidence += .05
	}

	if !stream.Empty() {
		confidence -= .2
	} else {
		confidence += .04
	}

	return
}

func readMQTTVariableLengthUint32(stream *cryptobyte.String, value *uint32) bool {
	var (
		multiplier uint32 = 1
		read       int
	)
	for {
		var encodedByte byte
		if !stream.ReadUint8(&encodedByte) {
			return false
		}
		read++

		*value += uint32(encodedByte&0x7F) * multiplier
		if encodedByte&0x80 == 0 {
			// Last byte
			break
		}

		if read >= 4 {
			return false
		}
		multiplier *= 128
	}
	return true
}