package gateway

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/binary"
	"errors"
)

const (
	encryptedRTSPHeaderSize = 24
	encryptedRTSPMaxPayload = 64 << 10
)

var errEncryptedRTSPFrame = errors.New("invalid encrypted RTSP frame")

// encryptedRTSPCodec keeps client and host nonce spaces disjoint. It accepts
// only strictly increasing host sequence numbers, so a replay cannot be fed
// into the RTSP parser after it has already authenticated once.
type encryptedRTSPCodec struct {
	aead           cipher.AEAD
	nextClient     uint32
	lastClient     uint32
	clientReceived bool
	lastHost       uint32
	hostReceived   bool
}

func newEncryptedRTSPCodec(key []byte) (*encryptedRTSPCodec, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return &encryptedRTSPCodec{aead: aead, nextClient: 1}, nil
}

func (codec *encryptedRTSPCodec) SealClient(plaintext []byte) ([]byte, error) {
	if codec == nil || codec.aead == nil || len(plaintext) == 0 || len(plaintext) > encryptedRTSPMaxPayload || codec.nextClient == 0 {
		return nil, errEncryptedRTSPFrame
	}
	sequence := codec.nextClient
	codec.nextClient++
	nonce := encryptedRTSPNonce(sequence, 'C', 'R')
	sealed := codec.aead.Seal(nil, nonce[:], plaintext, nil)
	ciphertext, tag := sealed[:len(plaintext)], sealed[len(plaintext):]
	frame := make([]byte, encryptedRTSPHeaderSize+len(ciphertext))
	binary.BigEndian.PutUint32(frame[:4], uint32(len(ciphertext))|0x80000000)
	binary.BigEndian.PutUint32(frame[4:8], sequence)
	copy(frame[8:24], tag)
	copy(frame[24:], ciphertext)
	return frame, nil
}

func (codec *encryptedRTSPCodec) OpenHost(frame []byte) ([]byte, error) {
	return codec.open(frame, 'H', 'R', &codec.lastHost, &codec.hostReceived)
}

// OpenClient validates the client-originated direction. The native gateway
// only sends this direction, but retaining the inverse lets a bounded fake
// provider verify the negotiated session key and frame shape.
func (codec *encryptedRTSPCodec) OpenClient(frame []byte) ([]byte, error) {
	return codec.open(frame, 'C', 'R', &codec.lastClient, &codec.clientReceived)
}

func (codec *encryptedRTSPCodec) open(frame []byte, origin, protocol byte, last *uint32, received *bool) ([]byte, error) {
	if codec == nil || codec.aead == nil || len(frame) < encryptedRTSPHeaderSize {
		return nil, errEncryptedRTSPFrame
	}
	length := binary.BigEndian.Uint32(frame[:4])
	if length&0x80000000 == 0 || int(length&0x7fffffff) > encryptedRTSPMaxPayload || len(frame) != encryptedRTSPHeaderSize+int(length&0x7fffffff) {
		return nil, errEncryptedRTSPFrame
	}
	sequence := binary.BigEndian.Uint32(frame[4:8])
	if sequence == 0 || (*received && sequence <= *last) {
		return nil, errEncryptedRTSPFrame
	}
	nonce := encryptedRTSPNonce(sequence, origin, protocol)
	sealed := make([]byte, int(length&0x7fffffff)+codec.aead.Overhead())
	copy(sealed, frame[24:])
	copy(sealed[length&0x7fffffff:], frame[8:24])
	plaintext, err := codec.aead.Open(nil, nonce[:], sealed, nil)
	if err != nil {
		return nil, errEncryptedRTSPFrame
	}
	*last, *received = sequence, true
	return plaintext, nil
}

func encryptedRTSPNonce(sequence uint32, origin, protocol byte) [12]byte {
	var nonce [12]byte
	binary.BigEndian.PutUint32(nonce[:4], sequence)
	nonce[10], nonce[11] = origin, protocol
	return nonce
}