package webauthn

import (
	"bytes"
	"crypto/rand"
	"errors"
	"io"
	"testing"
)

func testSecret(t *testing.T) []byte {
	t.Helper()
	s := make([]byte, secretLen)
	if _, err := io.ReadFull(rand.Reader, s); err != nil {
		t.Fatalf("rand: %v", err)
	}
	s[0] |= 1 // never all-zero
	return s
}

func testKey(t *testing.T) []byte {
	t.Helper()
	k := make([]byte, keyLen)
	if _, err := io.ReadFull(rand.Reader, k); err != nil {
		t.Fatalf("rand: %v", err)
	}
	return k
}

func TestWrapUnwrapRoundTrip(t *testing.T) {
	key, secret := testKey(t), testSecret(t)

	blob, err := WrapKey(key, secret)
	if err != nil {
		t.Fatalf("WrapKey: %v", err)
	}
	if !bytes.HasPrefix(blob, blobMagicV2) {
		t.Fatalf("blob does not start with the v2 magic: %x", blob[:8])
	}
	// The plaintext key must not be recoverable by reading the file.
	if bytes.Contains(blob, key) {
		t.Fatal("the wrapped blob contains the plaintext key verbatim")
	}

	got, version, err := UnwrapKey(blob, secret)
	if err != nil {
		t.Fatalf("UnwrapKey: %v", err)
	}
	if version != BlobV2 {
		t.Errorf("version = %v, want v2", version)
	}
	if !bytes.Equal(got, key) {
		t.Errorf("unwrapped key differs from the wrapped one")
	}
}

// Fresh salt and nonce per wrap: two blobs of the same key under the same
// secret must not be byte-identical, or the file leaks that nothing changed.
func TestWrapKeyIsNotDeterministic(t *testing.T) {
	key, secret := testKey(t), testSecret(t)
	a, err := WrapKey(key, secret)
	if err != nil {
		t.Fatalf("WrapKey: %v", err)
	}
	b, err := WrapKey(key, secret)
	if err != nil {
		t.Fatalf("WrapKey: %v", err)
	}
	if bytes.Equal(a, b) {
		t.Fatal("two wraps of the same key produced identical blobs")
	}
}

// The failure mode that matters most: a wrong passkey must not unlock.
func TestUnwrapWithWrongSecretFails(t *testing.T) {
	key := testKey(t)
	blob, err := WrapKey(key, testSecret(t))
	if err != nil {
		t.Fatalf("WrapKey: %v", err)
	}
	got, _, err := UnwrapKey(blob, testSecret(t))
	if err == nil {
		t.Fatal("a different secret unwrapped the blob")
	}
	if !errors.Is(err, ErrKeyUnwrap) {
		t.Errorf("err = %v, want ErrKeyUnwrap", err)
	}
	if got != nil {
		t.Error("key material returned alongside an error")
	}
}

// One flipped bit anywhere must fail the GCM tag, including in the salt and
// nonce — those are not authenticated by the tag but they change the
// derivation, so the tag fails anyway.
func TestUnwrapRejectsTamperedBlob(t *testing.T) {
	key, secret := testKey(t), testSecret(t)
	blob, err := WrapKey(key, secret)
	if err != nil {
		t.Fatalf("WrapKey: %v", err)
	}
	for i := range blob {
		bad := bytes.Clone(blob)
		bad[i] ^= 0x01
		if _, _, err := UnwrapKey(bad, secret); err == nil {
			t.Fatalf("byte %d of %d could be flipped and the blob still opened", i, len(blob))
		}
	}
}

func TestUnwrapRejectsTruncatedBlob(t *testing.T) {
	key, secret := testKey(t), testSecret(t)
	blob, err := WrapKey(key, secret)
	if err != nil {
		t.Fatalf("WrapKey: %v", err)
	}
	for _, n := range []int{0, 1, len(blobMagicV2), len(blobMagicV2) + saltLen, len(blob) - 1} {
		if _, _, err := UnwrapKey(blob[:n], secret); err == nil {
			t.Errorf("a %d-byte blob unwrapped", n)
		}
	}
}

// A v2 blob must not be downgradeable to v1 by stripping its header: the magic
// is GCM additional data, so the tag fails once it is gone.
func TestV2BlobCannotBeStrippedToV1(t *testing.T) {
	key, secret := testKey(t), testSecret(t)
	blob, err := WrapKey(key, secret)
	if err != nil {
		t.Fatalf("WrapKey: %v", err)
	}
	if _, _, err := UnwrapKey(blob[len(blobMagicV2):], secret); err == nil {
		t.Fatal("a header-stripped v2 blob was accepted as v1")
	}
}

// v1 blobs still open, and report themselves as v1 so the daemon can warn.
// wrapV1 reproduces the legacy writer this file no longer has.
func wrapV1(t *testing.T, key, secret []byte) []byte {
	t.Helper()
	salt := make([]byte, saltLen)
	nonce := make([]byte, nonceLen)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		t.Fatalf("rand: %v", err)
	}
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		t.Fatalf("rand: %v", err)
	}
	gcm, err := gcmFor(secret, salt, wrapInfoV1)
	if err != nil {
		t.Fatalf("gcmFor: %v", err)
	}
	out := append([]byte{}, salt...)
	out = append(out, nonce...)
	return append(out, gcm.Seal(nil, nonce, key, nil)...)
}

func TestUnwrapReadsLegacyV1(t *testing.T) {
	key := testKey(t)
	// v1 was keyed on the credential public key: not 32 bytes, and that is
	// deliberately still accepted on the read path.
	pub := make([]byte, 77)
	if _, err := io.ReadFull(rand.Reader, pub); err != nil {
		t.Fatalf("rand: %v", err)
	}
	blob := wrapV1(t, key, pub)

	got, version, err := UnwrapKey(blob, pub)
	if err != nil {
		t.Fatalf("UnwrapKey(v1): %v", err)
	}
	if version != BlobV1 {
		t.Errorf("version = %v, want v1", version)
	}
	if !bytes.Equal(got, key) {
		t.Error("v1 round-trip lost the key")
	}
	if _, _, err := UnwrapKey(blob, pub[:76]); err == nil {
		t.Error("a truncated public key opened the v1 blob")
	}
}

// The structural guard against the bug this replaces: a COSE public key is not
// 32 bytes, so it can never be used to write a new blob.
func TestWrapKeyRefusesNonPRFSecret(t *testing.T) {
	key := testKey(t)
	cases := map[string][]byte{
		"nil":              nil,
		"empty":            {},
		"short":            make([]byte, 16),
		"cose public key":  make([]byte, 77),
		"all-zero 32 byte": make([]byte, 32),
	}
	for name, secret := range cases {
		t.Run(name, func(t *testing.T) {
			if _, err := WrapKey(key, secret); err == nil {
				t.Fatalf("WrapKey accepted a %s secret", name)
			}
		})
	}
}

func TestWrapKeyRefusesWrongKeyLength(t *testing.T) {
	secret := testSecret(t)
	for _, n := range []int{0, 16, 31, 33, 64} {
		if _, err := WrapKey(make([]byte, n), secret); err == nil {
			t.Errorf("WrapKey accepted a %d-byte plaintext key", n)
		}
	}
}

// A v2 blob demands exactly 32 bytes on the read path too, so a caller cannot
// go back to passing a public key.
func TestUnwrapV2RefusesNonPRFSecret(t *testing.T) {
	blob, err := WrapKey(testKey(t), testSecret(t))
	if err != nil {
		t.Fatalf("WrapKey: %v", err)
	}
	if _, _, err := UnwrapKey(blob, make([]byte, 77)); !errors.Is(err, ErrKeyUnwrap) {
		t.Fatalf("err = %v, want ErrKeyUnwrap for a 77-byte secret", err)
	}
	if _, _, err := UnwrapKey(blob, nil); err == nil {
		t.Fatal("an empty secret unwrapped a v2 blob")
	}
}

func TestUnwrapRejectsOversizeBlob(t *testing.T) {
	if _, _, err := UnwrapKey(make([]byte, maxBlobLen+1), testSecret(t)); !errors.Is(err, ErrBlobTooLong) {
		t.Fatalf("err = %v, want ErrBlobTooLong", err)
	}
}