package updater

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"strings"
	"testing"
)

// generateTestKey produces an ECDSA P-256 keypair in PEM form,
// matching the shape `cosign generate-key-pair` emits for the public
// half. The private half stays in-test for signing.
func generateTestKey(t *testing.T) (privKey *ecdsa.PrivateKey, pubPEM string) {
	t.Helper()
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}
	der, err := x509.MarshalPKIXPublicKey(&priv.PublicKey)
	if err != nil {
		t.Fatalf("marshal public key: %v", err)
	}
	pubPEM = string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: der}))
	return priv, pubPEM
}

// signBlob mimics `cosign sign-blob`'s output: base64-encoded ASN.1-DER
// ECDSA signature over SHA256(body).
func signBlob(t *testing.T, priv *ecdsa.PrivateKey, body []byte) string {
	t.Helper()
	digest := sha256.Sum256(body)
	sig, err := ecdsa.SignASN1(rand.Reader, priv, digest[:])
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	return base64.StdEncoding.EncodeToString(sig)
}

func TestVerifyBlobSignaturePlaceholderRefuses(t *testing.T) {
	// A build that hasn't replaced the placeholder key must refuse
	// every verify call with ErrSignatureRequired so an un-rotated
	// build can't silently accept anything. Swap the embedded key
	// out for the placeholder shape and assert that.
	prev := BangerReleasePublicKey
	BangerReleasePublicKey = `-----BEGIN PUBLIC KEY-----
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEPLACEHOLDER0000000000000000000
000000000000000000000000000000000000000000000000000000000000PLACE
-----END PUBLIC KEY-----`
	defer func() { BangerReleasePublicKey = prev }()
	err := VerifyBlobSignature([]byte("body"), []byte("sig"))
	if !errors.Is(err, ErrSignatureRequired) {
		t.Fatalf("err = %v, want ErrSignatureRequired", err)
	}
}

func TestVerifyBlobSignatureHappyPath(t *testing.T) {
	priv, pubPEM := generateTestKey(t)
	prev := BangerReleasePublicKey
	BangerReleasePublicKey = pubPEM
	defer func() { BangerReleasePublicKey = prev }()

	body := []byte("SHA256SUMS body bytes")
	sig := signBlob(t, priv, body)
	if err := VerifyBlobSignature(body, []byte(sig)); err != nil {
		t.Fatalf("VerifyBlobSignature: %v", err)
	}
}

func TestVerifyBlobSignatureRejectsTamperedBody(t *testing.T) {
	priv, pubPEM := generateTestKey(t)
	prev := BangerReleasePublicKey
	BangerReleasePublicKey = pubPEM
	defer func() { BangerReleasePublicKey = prev }()

	body := []byte("original body")
	sig := signBlob(t, priv, body)
	tampered := []byte("tampered body")
	err := VerifyBlobSignature(tampered, []byte(sig))
	if err == nil || !strings.Contains(err.Error(), "does not verify") {
		t.Fatalf("err = %v, want signature-mismatch", err)
	}
}

func TestVerifyBlobSignatureRejectsWrongKey(t *testing.T) {
	// Sign with one key, verify with a different one.
	signingPriv, _ := generateTestKey(t)
	_, otherPubPEM := generateTestKey(t)
	prev := BangerReleasePublicKey
	BangerReleasePublicKey = otherPubPEM
	defer func() { BangerReleasePublicKey = prev }()

	body := []byte("body")
	sig := signBlob(t, signingPriv, body)
	err := VerifyBlobSignature(body, []byte(sig))
	if err == nil || !strings.Contains(err.Error(), "does not verify") {
		t.Fatalf("err = %v, want wrong-key rejection", err)
	}
}

func TestVerifyBlobSignatureRejectsMalformed(t *testing.T) {
	_, pubPEM := generateTestKey(t)
	prev := BangerReleasePublicKey
	BangerReleasePublicKey = pubPEM
	defer func() { BangerReleasePublicKey = prev }()
	for _, tc := range []struct {
		name string
		sig  string
	}{
		{name: "not_base64", sig: "!!!not_b64!!!"},
		{name: "empty", sig: ""},
		{name: "garbage_bytes", sig: base64.StdEncoding.EncodeToString([]byte{0x01, 0x02, 0x03})},
	} {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			err := VerifyBlobSignature([]byte("body"), []byte(tc.sig))
			if err == nil {
				t.Fatalf("expected error for %s; got success", tc.name)
			}
		})
	}
}