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

package krb5

import (
	"bytes"
	"encoding/hex"
	"testing"
)

func unhex(t *testing.T, s string) []byte {
	t.Helper()
	b, err := hex.DecodeString(s)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

// The n-fold test vectors of RFC 3961 appendix A.
func TestNFoldVectors(t *testing.T) {
	cases := []struct {
		in   string
		bits int
		want string
	}{
		{"303132333435", 64, "be072631276b1955"},
		{"70617373776f7264", 56, "78a07b6caf85fa"},
		{"526f75676820436f6e73656e7375732c20616e642052756e" +
			"6e696e6720436f6465", 64, "bb6ed30870b7f0e0"},
		{"70617373776f7264", 168, "59e4a8ca7c0385c3c37b3f6d2000247cb6e6bd5b3e"},
		{"4d41535341434856534554545320494e5354495456544520" +
			"4f4620544543484e4f4c4f4759", 192,
			"db3b0d8f0b061e603282b308a50841229ad798fab9540c1b"},
		{"51", 168, "518a54a215a8452a518a54a215a8452a518a54a215"},
		{"6261", 168, "fb25d531ae8974499f52fd92ea9857c4ba24cf297e"},
	}
	for _, c := range cases {
		got := NFold(unhex(t, c.in), c.bits/8)
		if !bytes.Equal(got, unhex(t, c.want)) {
			t.Errorf("nfold %d bits of %s: % x, want %s", c.bits, c.in, got, c.want)
		}
	}
}

// The key derivation test values of the MIT krb5 reference suite: the
// AES-128 key with the checksum and encryption constants of usage two.
func TestDeriveVector(t *testing.T) {
	key := unhex(t, "42263c6e89f4fc28b8df68ee09799f15")
	kc := DK(key, 2, 0x99)
	if !bytes.Equal(kc, unhex(t, "34280a382bc92769b2da2f9ef066854b")) {
		t.Fatalf("Kc % x", kc)
	}
}

// The PBKDF2 string-to-key vectors of RFC 3962 appendix B, and the
// resulting protocol keys.
func TestStringToKeyVectors(t *testing.T) {
	salt := []byte("ATHENA.MIT.EDUraeburn")
	k128 := StringToKey(EtypeAES128, []byte("password"), salt, 1, 16)
	if !bytes.Equal(k128, unhex(t, "42263c6e89f4fc28b8df68ee09799f15")) {
		t.Fatalf("aes128 key % x", k128)
	}
	k256 := StringToKey(EtypeAES256, []byte("password"), salt, 1, 32)
	if !bytes.Equal(k256, unhex(t, "fe697b52bc0d3ce14432ba036a92e65bbb52280990a2fa27883998d72af30161")) {
		t.Fatalf("aes256 key % x", k256)
	}
	k2 := StringToKey(EtypeAES256, []byte("password"), salt, 2, 32)
	if !bytes.Equal(k2, unhex(t, "a2e16d16b36069c135d5e9d2e25f896102685618b95914b467c67622225824ff")) {
		t.Fatalf("aes256 two rounds % x", k2)
	}
}

// The checksum test value of the MIT krb5 reference suite: HMAC-SHA1-96
// under the derived checksum key of usage three.
func TestChecksumVector(t *testing.T) {
	key := unhex(t, "9062430c8cda3388922e6d6a509f5b7a")
	sum, err := Checksum(EtypeAES128, key, 3, []byte("eight nine ten eleven twelve thirteen"))
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(sum, unhex(t, "01a4b088d45628f6946614e3")) {
		t.Fatalf("checksum % x", sum)
	}
}

// The profile round trips at both key sizes and over lengths that walk
// the CTS edge cases.
func TestEncryptRoundTrip(t *testing.T) {
	key := unhex(t, "fe697b52bc0d3ce14432ba036a92e65bbb52280990a2fa27883998d72af30161")
	for _, size := range []int{0, 1, 15, 16, 17, 31, 32, 33, 100, 1000} {
		plain := make([]byte, size)
		for i := range plain {
			plain[i] = byte(i)
		}
		ct, err := Encrypt(EtypeAES256, key, UsageInitiatorSeal, plain)
		if err != nil {
			t.Fatalf("size %d: %v", size, err)
		}
		got, err := Decrypt(EtypeAES256, key, UsageInitiatorSeal, ct)
		if err != nil {
			t.Fatalf("size %d: %v", size, err)
		}
		if !bytes.Equal(got, plain) {
			t.Fatalf("size %d: round trip differs", size)
		}
		// One flipped byte must break the integrity check.
		ct[len(ct)/2] ^= 0xff
		if _, err := Decrypt(EtypeAES256, key, UsageInitiatorSeal, ct); err == nil {
			t.Fatalf("size %d: tampering passed", size)
		}
	}
}