package ueps

import (
	"bufio"
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"io"
)

// packet, err := ReadAndVerify(bufio.NewReader(conn), secret)
// if err == errMissingHMAC { /* no HMAC tag found in frame */ }
var errMissingHMAC = sentinelError("UEPS packet missing HMAC signature")

// packet, err := ReadAndVerify(bufio.NewReader(conn), wrongSecret)
// if err == errIntegrityViolation { header.ThreatScore += 100; /* reject packet */ }
var errIntegrityViolation = sentinelError("integrity violation: HMAC mismatch")

// dispatch(packet.Header.IntentID, packet.Header.ThreatScore, packet.Payload)
type ParsedPacket struct {
	Header  UEPSHeader // packet.Header.Version == 0x09; packet.Header.IntentID == 0x01; packet.Header.ThreatScore == 0
	Payload []byte     // packet.Payload == []byte("hello world")
}

// packet, err := ueps.ReadAndVerify(bufio.NewReader(conn), []byte("my-shared-secret"))
// if err == errMissingHMAC { return } // unauthenticated: no HMAC tag in stream
// if err == errIntegrityViolation { return } // tampered: HMAC mismatch; reject and raise threat score
// if err == nil { dispatch(packet.Header.IntentID, packet.Header.ThreatScore, packet.Payload) }
func ReadAndVerify(reader *bufio.Reader, sharedSecret []byte) (*ParsedPacket, error) {
	var (
		hmacInputBuffer bytes.Buffer
		packetHeader    UEPSHeader
		hmacSignature   []byte
		payload         []byte
	)

	for {
		tagType, err := reader.ReadByte()
		if err != nil {
			return nil, err
		}

		if tagType == TagPayload {
			payload, err = io.ReadAll(reader)
			if err != nil {
				return nil, err
			}
			break
		}

		tagLength, err := reader.ReadByte()
		if err != nil {
			return nil, err
		}

		tagValue := make([]byte, int(tagLength))
		if _, err := io.ReadFull(reader, tagValue); err != nil {
			return nil, err
		}

		switch tagType {
		case TagVersion:
			packetHeader.Version = tagValue[0]
			hmacInputBuffer.WriteByte(tagType)
			hmacInputBuffer.WriteByte(tagLength)
			hmacInputBuffer.Write(tagValue)
		case TagCurrentLayer:
			packetHeader.CurrentLayer = tagValue[0]
			hmacInputBuffer.WriteByte(tagType)
			hmacInputBuffer.WriteByte(tagLength)
			hmacInputBuffer.Write(tagValue)
		case TagTargetLayer:
			packetHeader.TargetLayer = tagValue[0]
			hmacInputBuffer.WriteByte(tagType)
			hmacInputBuffer.WriteByte(tagLength)
			hmacInputBuffer.Write(tagValue)
		case TagIntent:
			packetHeader.IntentID = tagValue[0]
			hmacInputBuffer.WriteByte(tagType)
			hmacInputBuffer.WriteByte(tagLength)
			hmacInputBuffer.Write(tagValue)
		case TagThreatScore:
			packetHeader.ThreatScore = binary.BigEndian.Uint16(tagValue)
			hmacInputBuffer.WriteByte(tagType)
			hmacInputBuffer.WriteByte(tagLength)
			hmacInputBuffer.Write(tagValue)
		case TagHMAC:
			hmacSignature = tagValue
		default:
			hmacInputBuffer.WriteByte(tagType)
			hmacInputBuffer.WriteByte(tagLength)
			hmacInputBuffer.Write(tagValue)
		}
	}

	if len(hmacSignature) == 0 {
		return nil, errMissingHMAC
	}

	messageAuthCode := hmac.New(sha256.New, sharedSecret)
	messageAuthCode.Write(hmacInputBuffer.Bytes())
	messageAuthCode.Write(payload)
	expectedMessageAuthCode := messageAuthCode.Sum(nil)

	if !hmac.Equal(hmacSignature, expectedMessageAuthCode) {
		return nil, errIntegrityViolation
	}

	return &ParsedPacket{
		Header:  packetHeader,
		Payload: payload,
	}, nil
}