// Package simtest validates the agent's reused-TPM mTLS identity 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-agent build; see go.mod.
//
// The crux it proves: crypto/tls can drive tpmkey's go-tpm crypto.Signer to sign
// the TLS CertificateVerify, so a host holding ONLY a TPM key (no exportable
// private key) authenticates to a server doing RequireAndVerifyClientCert against
// the fleet CA — the exact posture of the fleet-manager agent plane (agent.go).
package simtest

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"math/big"
	"testing"
	"time"

	"fleet-agent/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)
}

// testCA is a throwaway fleet CA for the handshake.
type testCA struct {
	cert *x509.Certificate
	key  *ecdsa.PrivateKey
	pool *x509.CertPool
}

func newTestCA(t *testing.T) *testCA {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	tmpl := &x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{CommonName: "test fleet CA"},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().Add(time.Hour),
		IsCA:                  true,
		KeyUsage:              x509.KeyUsageCertSign,
		BasicConstraintsValid: true,
	}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, key.Public(), key)
	if err != nil {
		t.Fatal(err)
	}
	cert, _ := x509.ParseCertificate(der)
	pool := x509.NewCertPool()
	pool.AddCert(cert)
	return &testCA{cert: cert, key: key, pool: pool}
}

// issue signs a leaf cert for pub with the CA (as the enrol-broker/step-ca does).
func (ca *testCA) issue(t *testing.T, tmpl *x509.Certificate, pub crypto.PublicKey) []byte {
	t.Helper()
	der, err := x509.CreateCertificate(rand.Reader, tmpl, ca.cert, pub, ca.key)
	if err != nil {
		t.Fatalf("CA issue: %v", err)
	}
	return der
}

func TestTPMClientCertHandshake(t *testing.T) {
	tpm := openSim(t)
	const (
		keyHandle = tpm2.TPMHandle(0x81000100)
		nvIndex   = tpm2.TPMHandle(0x01800100)
	)

	// Enrolment side: create the persistent signing key (once).
	if _, _, err := tpmkey.EnsureKey(tpm, keyHandle); err != nil {
		t.Fatalf("EnsureKey: %v", err)
	}
	// Agent side: load it read-only (the function the agent actually uses).
	signer, err := tpmkey.LoadSigner(tpm, keyHandle)
	if err != nil {
		t.Fatalf("LoadSigner: %v", err)
	}

	ca := newTestCA(t)

	// The host's identity cert (CA-signed over the TPM public key). Its SERIAL is
	// the value the manager derives as the inventory key.
	wantSerial := big.NewInt(0xC0FFEE)
	leafDER := ca.issue(t, &x509.Certificate{
		SerialNumber: wantSerial,
		Subject:      pkix.Name{CommonName: "fleet-host"},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, signer.Public())

	// Exercise the NV round-trip the agent's identity path uses.
	if err := tpmkey.WriteNVCert(tpm, nvIndex, leafDER); err != nil {
		t.Fatalf("WriteNVCert: %v", err)
	}
	got, err := tpmkey.ReadNVCert(tpm, nvIndex)
	if err != nil {
		t.Fatalf("ReadNVCert: %v", err)
	}
	if string(got) != string(leafDER) {
		t.Fatalf("NV cert round-trip mismatch")
	}
	leaf, _ := x509.ParseCertificate(leafDER)

	// TPM-backed client cert (private key = the TPM signer; never exported).
	clientCert := tls.Certificate{Certificate: [][]byte{leafDER}, PrivateKey: signer, Leaf: leaf}

	// Server standing in for the manager agent plane: verify the client to the CA.
	serverKey, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	serverDER := ca.issue(t, &x509.Certificate{
		SerialNumber: big.NewInt(2),
		Subject:      pkix.Name{CommonName: "fleet-manager"},
		DNSNames:     []string{"fleet-manager"},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}, serverKey.Public())
	serverCfg := &tls.Config{
		Certificates: []tls.Certificate{{Certificate: [][]byte{serverDER}, PrivateKey: serverKey}},
		ClientAuth:   tls.RequireAndVerifyClientCert,
		ClientCAs:    ca.pool,
		MinVersion:   tls.VersionTLS12,
	}
	clientCfg := &tls.Config{
		Certificates: []tls.Certificate{clientCert},
		RootCAs:      ca.pool,
		ServerName:   "fleet-manager",
		MinVersion:   tls.VersionTLS12,
	}

	ln, err := tls.Listen("tcp", "127.0.0.1:0", serverCfg)
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	serialCh := make(chan *big.Int, 1)
	errCh := make(chan error, 1)
	go func() {
		conn, err := ln.Accept()
		if err != nil {
			errCh <- err
			return
		}
		defer conn.Close()
		tc := conn.(*tls.Conn)
		tc.SetDeadline(time.Now().Add(5 * time.Second))
		if err := tc.Handshake(); err != nil {
			errCh <- err
			return
		}
		st := tc.ConnectionState()
		if len(st.PeerCertificates) == 0 {
			errCh <- err
			return
		}
		serialCh <- st.PeerCertificates[0].SerialNumber
	}()

	conn, err := tls.Dial("tcp", ln.Addr().String(), clientCfg)
	if err != nil {
		t.Fatalf("client handshake with TPM cert failed: %v", err)
	}
	defer conn.Close()

	select {
	case err := <-errCh:
		t.Fatalf("server rejected TPM client cert: %v", err)
	case serial := <-serialCh:
		if serial.Cmp(wantSerial) != 0 {
			t.Fatalf("server saw serial %x, want %x", serial, wantSerial)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("handshake timed out")
	}
}

// TestLoadSignerRejectsEmptyHandle proves the agent fails clearly on a host that
// isn't enrolled (no key at the handle) rather than creating one.
func TestLoadSignerRejectsEmptyHandle(t *testing.T) {
	tpm := openSim(t)
	if _, err := tpmkey.LoadSigner(tpm, tpm2.TPMHandle(0x81000100)); err == nil {
		t.Fatal("LoadSigner on an empty handle should error (agent must not create a key)")
	}
}