// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

// The AES encryption profiles of RFC 3962 over the simplified profile
// of RFC 3961: key derivation DK, the CBC-CTS cipher, and the keyed
// checksum, with the key usage constants of RFC 4120 and RFC 4121.

package krb5

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"encoding/binary"
	"errors"
	"hash"
)

// Encryption types of RFC 3962.
const (
	EtypeAES128 = 17
	EtypeAES256 = 18
)

// Key usage numbers of RFC 4120 section 7.5.1 and RFC 4121 section 4.
const (
	UsageTicket        = 2
	UsageInitiatorSign = 7
	UsageAcceptorSign  = 8
	UsageInitiatorSeal = 9
	UsageAcceptorSeal  = 10
	UsageAPReqAuth     = 11
	UsageAPRepAuth     = 12
)

// ErrIntegrity marks a checksum mismatch on decryption.
var ErrIntegrity = errors.New("krb5: integrity check failed")

// deriveKey implements DK of RFC 3961 section 5.1: the constant is
// n-folded to the cipher block size, encrypted with the base key until
// enough bits accumulate, and the first keySize bytes are the key.
func deriveKey(key []byte, constant []byte, keyBytes int) []byte {
	block := 16
	folded := NFold(constant, block)
	var out []byte
	prev := folded
	for len(out) < keyBytes {
		prev = aesECB(key, prev)
		out = append(out, prev...)
	}
	return out[:keyBytes]
}

// DK derives the protocol key for a usage: the constant is the usage
// number as four big endian octets followed by the label octet.
func DK(key []byte, usage uint32, label byte) []byte {
	var constant [5]byte
	binary.BigEndian.PutUint32(constant[:4], usage)
	constant[4] = label
	return deriveKey(key, constant[:], len(key))
}

// deriveKeys answers the three keys of the simplified profile:
// Kc for checksums, Ke for encryption, Ki for the message HMAC.
func deriveKeys(key []byte, usage uint32) (kc, ke, ki []byte) {
	return DK(key, usage, 0x99), DK(key, usage, 0xAA), DK(key, usage, 0x55)
}

// aesECB encrypts a block aligned input in ECB mode, which is how the
// profile uses the cipher for key derivation.
func aesECB(key, block []byte) []byte {
	c, err := aes.NewCipher(key)
	if err != nil {
		panic("krb5: " + err.Error())
	}
	out := make([]byte, len(block))
	for i := 0; i < len(block); i += c.BlockSize() {
		c.Encrypt(out[i:], block[i:])
	}
	return out
}

// ctsEncrypt implements the Kerberos CBC-CTS mode of RFC 3962: plain
// CBC for everything but the tail, with the last two blocks through
// ciphertext stealing. The output carries the input length, except an
// ctsEncrypt implements the Kerberos CBC-CTS mode of RFC 3962, ported
// from the MIT krb5 reference: plain CBC over everything but the last
// two blocks, whose ciphertexts swap places, with the final block
// truncated back to the input length.
func ctsEncrypt(key []byte, plaintext []byte) []byte {
	c, err := aes.NewCipher(key)
	if err != nil {
		panic("krb5: " + err.Error())
	}
	n := len(plaintext)
	if n == 0 {
		panic("krb5: empty plaintext")
	}
	blocks := (n + 15) / 16
	lastLen := n - (blocks-1)*16
	if blocks == 1 {
		block := make([]byte, 16)
		copy(block, plaintext)
		out := make([]byte, 16)
		c.Encrypt(out, block)
		return out
	}
	buf := make([]byte, blocks*16)
	copy(buf, plaintext)
	head := buf[:(blocks-2)*16]
	iv := make([]byte, 16)
	cipher.NewCBCEncrypter(c, iv).CryptBlocks(head, head)
	prev := make([]byte, 16)
	if blocks > 2 {
		copy(prev, head[len(head)-16:])
	}
	p1 := buf[(blocks-2)*16 : (blocks-1)*16]
	p2 := buf[(blocks-1)*16:]
	c1 := make([]byte, 16)
	for i := range c1 {
		c1[i] = p1[i] ^ prev[i]
	}
	c.Encrypt(c1, c1)
	c2 := make([]byte, 16)
	for i := range c2 {
		c2[i] = p2[i] ^ c1[i]
	}
	c.Encrypt(c2, c2)
	out := make([]byte, 0, n)
	out = append(out, head...)
	out = append(out, c2...)
	out = append(out, c1[:lastLen]...)
	return out
}

// ctsDecrypt reverses ctsEncrypt.
func ctsDecrypt(key []byte, ciphertext []byte) []byte {
	c, err := aes.NewCipher(key)
	if err != nil {
		panic("krb5: " + err.Error())
	}
	n := len(ciphertext)
	blocks := (n + 15) / 16
	lastLen := n - (blocks-1)*16
	if blocks == 1 {
		block := make([]byte, 16)
		copy(block, ciphertext)
		out := make([]byte, 16)
		c.Decrypt(out, block)
		return out
	}
	buf := make([]byte, blocks*16)
	copy(buf, ciphertext)
	head := buf[:(blocks-2)*16]
	prev := make([]byte, 16)
	if blocks > 2 {
		copy(prev, head[len(head)-16:])
	}
	if len(head) > 0 {
		cipher.NewCBCDecrypter(c, make([]byte, 16)).CryptBlocks(head, head)
	}
	blockN2 := append([]byte{}, buf[(blocks-2)*16:(blocks-1)*16]...)
	blockN1 := append([]byte{}, buf[(blocks-1)*16:]...)
	blockN1 = append(blockN1, make([]byte, 16-len(blockN1))...)
	// The final plaintext block decrypts with the truncated ciphertext
	// block as the CBC vector.
	p2 := make([]byte, 16)
	c.Decrypt(p2, blockN2)
	for i := range p2 {
		p2[i] ^= blockN1[i]
	}
	// The stolen tail of the decrypted final block restores the truncated
	// second-to-last ciphertext block.
	for i := lastLen; i < 16; i++ {
		blockN1[i] = p2[i]
	}
	p1 := make([]byte, 16)
	c.Decrypt(p1, blockN1)
	for i := range p1 {
		p1[i] ^= prev[i]
	}
	out := make([]byte, 0, n)
	out = append(out, head...)
	out = append(out, p1...)
	out = append(out, p2[:lastLen]...)
	return out
}

// Encrypt applies the RFC 3962 simplified profile over AES: a random
// confounder, the plaintext, and a truncated HMAC under Ki, the whole
// confounder and plaintext encrypted under Ke by CBC-CTS.
func Encrypt(etype uint32, key []byte, usage uint32, plaintext []byte) ([]byte, error) {
	_, ke, ki := deriveKeys(key, usage)
	conf := make([]byte, 16)
	if _, err := rand.Read(conf); err != nil {
		return nil, err
	}
	data := append(append([]byte{}, conf...), plaintext...)
	body := ctsEncrypt(ke, data)
	mac := hmac.New(sha1.New, ki)
	mac.Write(data)
	return append(body, mac.Sum(nil)[:12]...), nil
}

// Decrypt verifies and removes the protection Encrypt applied.
func Decrypt(etype uint32, key []byte, usage uint32, ciphertext []byte) ([]byte, error) {
	_, ke, ki := deriveKeys(key, usage)
	if len(ciphertext) < 16+12 {
		return nil, ErrIntegrity
	}
	body := ciphertext[:len(ciphertext)-12]
	mac := ciphertext[len(ciphertext)-12:]
	data := ctsDecrypt(ke, body)
	expect := hmac.New(sha1.New, ki)
	expect.Write(data)
	if !hmac.Equal(expect.Sum(nil)[:12], mac) {
		return nil, ErrIntegrity
	}
	return data[16:], nil
}

// Checksum answers the keyed checksum of RFC 3961 section 5.4: the
// leading twelve bytes of HMAC-SHA1 under Kc.
func Checksum(etype uint32, key []byte, usage uint32, message []byte) ([]byte, error) {
	kc, _, _ := deriveKeys(key, usage)
	mac := hmac.New(sha1.New, kc)
	mac.Write(message)
	return mac.Sum(nil)[:12], nil
}

// StringToKey implements the PBKDF2+DK function of RFC 3962 section 4:
// PBKDF2 over HMAC-SHA1 with the given iteration count, then one
// application of the key derivation with the well known constant
// "kerberos" of the AES profile.
func StringToKey(etype uint32, password, salt []byte, iterations uint32, keyBytes int) []byte {
	_ = etype
	seed := pbkdf2(sha1.New, password, salt, int(iterations), keyBytes)
	pepper := []byte("kerberos")
	return deriveKey(seed, pepper, keyBytes)
}

// pbkdf2 is the PBKDF2 of RFC 2898 over the given hash.
func pbkdf2(h func() hash.Hash, password, salt []byte, rounds, length int) []byte {
	out := make([]byte, 0, length)
	var block [4]byte
	for blockIndex := 1; len(out) < length; blockIndex++ {
		block[0] = byte(blockIndex >> 24)
		block[1] = byte(blockIndex >> 16)
		block[2] = byte(blockIndex >> 8)
		block[3] = byte(blockIndex)
		mac := hmac.New(sha1.New, password)
		mac.Write(salt)
		mac.Write(block[:])
		u := mac.Sum(nil)
		t := make([]byte, len(u))
		copy(t, u)
		for i := 1; i < rounds; i++ {
			mac.Reset()
			mac.Write(u)
			u = mac.Sum(nil)
			for j := range t {
				t[j] ^= u[j]
			}
		}
		out = append(out, t...)
	}
	return out[:length]
}