// Package simtest validates tpmkey's TPM ceremony against the go-tpm software
// simulator (the MS TPM 2.0 reference). It lives in its own module so the CGO
// simulator never enters the shipped fleet-enrol build; see go.mod.
package simtest

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"strings"
	"testing"

	"fleet-enrol/tpmkey"

	"github.com/google/go-tpm-tools/simulator"
	"github.com/google/go-tpm/tpm2"
	"github.com/google/go-tpm/tpm2/transport"
)

func openSim(t *testing.T) transport.TPM {
	t.Helper()
	sim, err := simulator.Get()
	if err != nil {
		t.Fatalf("simulator: %v", err)
	}
	t.Cleanup(func() { sim.Close() })
	return transport.FromReadWriter(sim)
}

func TestEnsureKeyAndSignCSR(t *testing.T) {
	tpm := openSim(t)
	const h = tpm2.TPMHandle(0x81000100)

	signer, created, err := tpmkey.EnsureKey(tpm, h)
	if err != nil {
		t.Fatalf("EnsureKey: %v", err)
	}
	if !created {
		t.Errorf("first EnsureKey should report created=true")
	}
	if _, ok := signer.Public().(*ecdsa.PublicKey); !ok {
		t.Fatalf("signer public key is not ECDSA")
	}

	// idempotent: second call reuses the persistent key (same pubkey).
	signer2, created2, err := tpmkey.EnsureKey(tpm, h)
	if err != nil || created2 {
		t.Fatalf("second EnsureKey: created=%v err=%v", created2, err)
	}
	if !signer.Public().(*ecdsa.PublicKey).Equal(signer2.Public()) {
		t.Errorf("persistent key changed between calls")
	}

	// the TPM signer must produce a CSR that verifies — this is the exact
	// openssl-tpm2-provider path the Go binary replaces.
	tmpl := &x509.CertificateRequest{Subject: pkix.Name{CommonName: "fleet-host"}}
	der, err := x509.CreateCertificateRequest(rand.Reader, tmpl, signer)
	if err != nil {
		t.Fatalf("CreateCertificateRequest with TPM signer: %v", err)
	}
	csr, err := x509.ParseCertificateRequest(der)
	if err != nil {
		t.Fatalf("parse CSR: %v", err)
	}
	if err := csr.CheckSignature(); err != nil {
		t.Fatalf("CSR signature (TPM-signed) invalid: %v", err)
	}

	// raw signer sanity: sign a digest, verify with the public key.
	dig := sha256.Sum256([]byte("hello"))
	sig, err := signer.Sign(rand.Reader, dig[:], nil)
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	if !ecdsa.VerifyASN1(signer.Public().(*ecdsa.PublicKey), dig[:], sig) {
		t.Fatalf("TPM ECDSA signature did not verify")
	}
}

func TestNVWriteAndPresence(t *testing.T) {
	tpm := openSim(t)
	const nv = tpm2.TPMHandle(0x01800100)

	if tpmkey.NVPresent(tpm, nv) {
		t.Fatalf("NV should be empty on a fresh TPM")
	}
	if _, err := tpmkey.ReadNVCert(tpm, nv); err == nil {
		t.Fatalf("ReadNVCert should error on an undefined index")
	}
	// realistic cert-sized payload that also crosses the 512B read/write chunk.
	want := make([]byte, 700)
	rand.Read(want)
	if err := tpmkey.WriteNVCert(tpm, nv, want); err != nil {
		t.Fatalf("WriteNVCert: %v", err)
	}
	if !tpmkey.NVPresent(tpm, nv) {
		t.Fatalf("NV should be present after write")
	}
	// ReadNVCert must return exactly what was written (across the chunk boundary).
	got, err := tpmkey.ReadNVCert(tpm, nv)
	if err != nil {
		t.Fatalf("ReadNVCert: %v", err)
	}
	if !bytes.Equal(got, want) {
		t.Fatalf("ReadNVCert roundtrip mismatch: got %d bytes, want %d", len(got), len(want))
	}
	// re-write (re-enrolment) must succeed (undefine + redefine).
	if err := tpmkey.WriteNVCert(tpm, nv, want[:400]); err != nil {
		t.Fatalf("WriteNVCert rewrite: %v", err)
	}
	if got, _ := tpmkey.ReadNVCert(tpm, nv); !bytes.Equal(got, want[:400]) {
		t.Fatalf("ReadNVCert after rewrite mismatch: got %d bytes, want 400", len(got))
	}
}

// Regression for the bench failure "key at handle is not ECC ... selector value
// was 1": an older provisioning path left an RSA SIGNING key at the enrol handle,
// and EnsureKey fatally failed trying to read it as ECC. It must now evict the
// stale key and create a fresh ECC one.
func TestEnsureKeyEvictsStaleNonECC(t *testing.T) {
	tpm := openSim(t)
	const h = tpm2.TPMHandle(0x81000100)

	// Pre-populate the handle with an RSA SIGNING key (mimics tpm2-tools/openssl
	// default) — SignEncrypt, NOT restricted/decrypt, so it is safe to evict.
	rsaTmpl := tpm2.TPMTPublic{
		Type:    tpm2.TPMAlgRSA,
		NameAlg: tpm2.TPMAlgSHA256,
		ObjectAttributes: tpm2.TPMAObject{
			FixedTPM: true, FixedParent: true, SensitiveDataOrigin: true,
			UserWithAuth: true, SignEncrypt: true, NoDA: true,
		},
		Parameters: tpm2.NewTPMUPublicParms(tpm2.TPMAlgRSA, &tpm2.TPMSRSAParms{
			Symmetric: tpm2.TPMTSymDefObject{Algorithm: tpm2.TPMAlgNull},
			Scheme: tpm2.TPMTRSAScheme{
				Scheme: tpm2.TPMAlgRSASSA,
				Details: tpm2.NewTPMUAsymScheme(tpm2.TPMAlgRSASSA, &tpm2.TPMSSigSchemeRSASSA{
					HashAlg: tpm2.TPMAlgSHA256,
				}),
			},
			KeyBits: 2048,
		}),
	}
	cp, err := tpm2.CreatePrimary{
		PrimaryHandle: tpm2.TPMRHOwner, InPublic: tpm2.New2B(rsaTmpl),
	}.Execute(tpm)
	if err != nil {
		t.Fatalf("create RSA primary: %v", err)
	}
	if _, err := (tpm2.EvictControl{
		Auth:             tpm2.TPMRHOwner,
		ObjectHandle:     &tpm2.NamedHandle{Handle: cp.ObjectHandle, Name: cp.Name},
		PersistentHandle: h,
	}).Execute(tpm); err != nil {
		t.Fatalf("persist stale RSA key: %v", err)
	}
	_, _ = (tpm2.FlushContext{FlushHandle: cp.ObjectHandle}).Execute(tpm)

	// EnsureKey must self-heal: evict the RSA key, create + persist a fresh ECC key.
	signer, created, err := tpmkey.EnsureKey(tpm, h)
	if err != nil {
		t.Fatalf("EnsureKey over stale RSA key: %v", err)
	}
	if !created {
		t.Errorf("EnsureKey should have recreated the key (created=true)")
	}
	if _, ok := signer.Public().(*ecdsa.PublicKey); !ok {
		t.Fatalf("recreated signer is not ECDSA")
	}
	// and it must be usable for signing (the whole point of the recovery).
	dig := sha256.Sum256([]byte("post-heal"))
	sig, err := signer.Sign(rand.Reader, dig[:], nil)
	if err != nil {
		t.Fatalf("sign after heal: %v", err)
	}
	if !ecdsa.VerifyASN1(signer.Public().(*ecdsa.PublicKey), dig[:], sig) {
		t.Fatalf("post-heal ECDSA signature did not verify")
	}
}

// Safety guard: if the handle ever holds a restricted-decrypt object (an EK/SRK —
// e.g. handle misconfigured to the EK slot 0x81010001), EnsureKey must REFUSE to
// evict it and error, never clobber the TPM's endorsement/storage identity.
func TestEnsureKeyRefusesToEvictProtectedKey(t *testing.T) {
	tpm := openSim(t)
	const h = tpm2.TPMHandle(0x81000100)

	// A restricted decryption primary — the shape of an EK/SRK (storage key).
	skTmpl := tpm2.TPMTPublic{
		Type:    tpm2.TPMAlgECC,
		NameAlg: tpm2.TPMAlgSHA256,
		ObjectAttributes: tpm2.TPMAObject{
			FixedTPM: true, FixedParent: true, SensitiveDataOrigin: true,
			UserWithAuth: true, Restricted: true, Decrypt: true, NoDA: true,
		},
		Parameters: tpm2.NewTPMUPublicParms(tpm2.TPMAlgECC, &tpm2.TPMSECCParms{
			Symmetric: tpm2.TPMTSymDefObject{
				Algorithm: tpm2.TPMAlgAES,
				KeyBits:   tpm2.NewTPMUSymKeyBits(tpm2.TPMAlgAES, tpm2.TPMKeyBits(128)),
				Mode:      tpm2.NewTPMUSymMode(tpm2.TPMAlgAES, tpm2.TPMAlgCFB),
			},
			Scheme:  tpm2.TPMTECCScheme{Scheme: tpm2.TPMAlgNull},
			CurveID: tpm2.TPMECCNistP256,
			KDF:     tpm2.TPMTKDFScheme{Scheme: tpm2.TPMAlgNull},
		}),
	}
	cp, err := tpm2.CreatePrimary{
		PrimaryHandle: tpm2.TPMRHOwner, InPublic: tpm2.New2B(skTmpl),
	}.Execute(tpm)
	if err != nil {
		t.Fatalf("create storage-like primary: %v", err)
	}
	if _, err := (tpm2.EvictControl{
		Auth:             tpm2.TPMRHOwner,
		ObjectHandle:     &tpm2.NamedHandle{Handle: cp.ObjectHandle, Name: cp.Name},
		PersistentHandle: h,
	}).Execute(tpm); err != nil {
		t.Fatalf("persist storage-like key: %v", err)
	}
	_, _ = (tpm2.FlushContext{FlushHandle: cp.ObjectHandle}).Execute(tpm)

	// EnsureKey must refuse (error), and must NOT have evicted the object.
	if _, _, err := tpmkey.EnsureKey(tpm, h); err == nil {
		t.Fatalf("EnsureKey should refuse to evict a restricted-decrypt (EK/SRK) object")
	}
	if _, err := (tpm2.ReadPublic{ObjectHandle: h}).Execute(tpm); err != nil {
		t.Fatalf("protected object was evicted (should have been left intact): %v", err)
	}
}

// Regression for the Dell 820 finding: our NV index collided with an OEM
// PLATFORM-created index (platformcreate|policydelete) that no owner clear can
// remove. WriteNVCert must REFUSE such an index with a clear error (the remedy
// is a different rd.dot1x.nvindex), not a confusing undefine failure.
func TestWriteNVCertRefusesPlatformIndex(t *testing.T) {
	tpm := openSim(t)
	const nv = tpm2.TPMHandle(0x01800100)

	def := tpm2.NVDefineSpace{
		AuthHandle: tpm2.TPMRHPlatform,
		PublicInfo: tpm2.New2B(tpm2.TPMSNVPublic{
			NVIndex: nv,
			NameAlg: tpm2.TPMAlgSHA256,
			Attributes: tpm2.TPMANV{
				PlatformCreate: true, AuthRead: true, AuthWrite: true,
				NoDA: true, NT: tpm2.TPMNTOrdinary,
			},
			DataSize: 8,
		}),
	}
	if _, err := def.Execute(tpm); err != nil {
		t.Fatalf("define platform NV: %v", err)
	}
	err := tpmkey.WriteNVCert(tpm, nv, make([]byte, 64))
	if err == nil || !strings.Contains(err.Error(), "platform-created") {
		t.Fatalf("WriteNVCert should refuse a platform-created index; got: %v", err)
	}
}