// Package tpmkey — TPM 2.0 operations via go-tpm (pure Go, no CGO, no tpm2-tools
// / no openssl tpm2 provider). Provides: ensure a persistent ECDSA signing key
// at a non-reserved owner handle (default 0x81000100 — NOT the EK slot
// 0x81010001, see EnsureKey), a crypto.Signer backed by it (for CSR signing),
// and NV cert read-presence + write at 0x01800100. Talks to any transport.TPM,
// so the same
// code runs against /dev/tpmrm0 (production) and the software simulator (tests).
//
// Exported so the ceremony can be tested against the go-tpm simulator from a
// SEPARATE module (simtest/), keeping this module's vendored deps lean (go-tpm
// only) — the simulator (CGO + the MS TPM reference C) never enters the image build.
//
// NOTE: this is a COPY of files/enrol/tpmkey (via files/agent/tpmkey) — kept
// byte-identical apart from the added read-only LoadSigner — so every reused-
// identity client (agent, ssh-hostcert) has the exact same, bench-verified TPM
// signing path without importing across the load-bearing fleet-enrol module.
// ssh-hostcert only READS the key (LoadSigner) + NV cert (ReadNVCert); it never
// creates/evicts. A future refactor could promote tpmkey to a shared module.
package tpmkey

import (
	"crypto"
	"crypto/ecdsa"
	"encoding/asn1"
	"fmt"
	"io"
	"math/big"

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

// eccSigningTemplate: a NON-restricted ECC P-256 ECDSA-SHA256 signing key.
// Non-restricted so it can sign an externally-computed digest (the CSR TBS) with
// a NULL validation ticket; Symmetric/KDF are explicitly NULL (sign-only key).
func eccSigningTemplate() tpm2.TPMTPublic {
	return tpm2.TPMTPublic{
		Type:    tpm2.TPMAlgECC,
		NameAlg: tpm2.TPMAlgSHA256,
		ObjectAttributes: tpm2.TPMAObject{
			FixedTPM:            true,
			FixedParent:         true,
			SensitiveDataOrigin: true,
			UserWithAuth:        true,
			SignEncrypt:         true,
			NoDA:                true,
		},
		Parameters: tpm2.NewTPMUPublicParms(tpm2.TPMAlgECC, &tpm2.TPMSECCParms{
			Symmetric: tpm2.TPMTSymDefObject{Algorithm: tpm2.TPMAlgNull},
			Scheme: tpm2.TPMTECCScheme{
				Scheme: tpm2.TPMAlgECDSA,
				Details: tpm2.NewTPMUAsymScheme(tpm2.TPMAlgECDSA, &tpm2.TPMSSigSchemeECDSA{
					HashAlg: tpm2.TPMAlgSHA256,
				}),
			},
			CurveID: tpm2.TPMECCNistP256,
			KDF:     tpm2.TPMTKDFScheme{Scheme: tpm2.TPMAlgNull},
		}),
	}
}

// Signer implements crypto.Signer using a persistent TPM key.
type Signer struct {
	tpm    transport.TPM
	handle tpm2.TPMHandle
	name   tpm2.TPM2BName
	pub    *ecdsa.PublicKey
}

func (s *Signer) Public() crypto.PublicKey { return s.pub }

func (s *Signer) Sign(_ io.Reader, digest []byte, _ crypto.SignerOpts) ([]byte, error) {
	rsp, err := tpm2.Sign{
		KeyHandle: tpm2.NamedHandle{Handle: s.handle, Name: s.name},
		Digest:    tpm2.TPM2BDigest{Buffer: digest},
		InScheme: tpm2.TPMTSigScheme{
			Scheme: tpm2.TPMAlgECDSA,
			Details: tpm2.NewTPMUSigScheme(tpm2.TPMAlgECDSA, &tpm2.TPMSSchemeHash{
				HashAlg: tpm2.TPMAlgSHA256,
			}),
		},
		Validation: tpm2.TPMTTKHashCheck{Tag: tpm2.TPMSTHashCheck},
	}.Execute(s.tpm)
	if err != nil {
		return nil, fmt.Errorf("tpm sign: %w", err)
	}
	ecc, err := rsp.Signature.Signature.ECDSA()
	if err != nil {
		return nil, fmt.Errorf("tpm sign: not ECDSA: %w", err)
	}
	r := new(big.Int).SetBytes(ecc.SignatureR.Buffer)
	sInt := new(big.Int).SetBytes(ecc.SignatureS.Buffer)
	// x509 expects an ASN.1 DER ECDSA-Sig-Value{ r, s }.
	return asn1.Marshal(struct{ R, S *big.Int }{r, sInt})
}

func signerFromPub(tpm transport.TPM, handle tpm2.TPMHandle, name tpm2.TPM2BName, pub2b tpm2.TPM2BPublic) (*Signer, error) {
	pub, err := pub2b.Contents()
	if err != nil {
		return nil, err
	}
	det, err := pub.Parameters.ECCDetail()
	if err != nil {
		return nil, fmt.Errorf("key at handle is not ECC: %w", err)
	}
	uniq, err := pub.Unique.ECC()
	if err != nil {
		return nil, err
	}
	pk, err := tpm2.ECDSAPub(det, uniq)
	if err != nil {
		return nil, err
	}
	return &Signer{tpm: tpm, handle: handle, name: name, pub: pk}, nil
}

// isProtectedObject reports whether a persistent object is a restricted
// decryption key — i.e. an Endorsement Key or Storage Root Key. We must NEVER
// evict one of those (it would destroy the TPM's endorsement/storage identity,
// including the EK the future EK/AK attestation path depends on). Fail-safe: if
// the public area can't be parsed, treat it as protected.
func isProtectedObject(pub2b tpm2.TPM2BPublic) bool {
	pub, err := pub2b.Contents()
	if err != nil {
		return true
	}
	a := pub.ObjectAttributes
	return a.Restricted && a.Decrypt
}

// ensureKey returns a signer for the persistent key at handle, creating +
// persisting a fresh ECDSA key there if none exists. Returns created=true when
// a new key was generated. handle MUST be a non-reserved owner handle (the fleet
// default is 0x81000100) — NOT the EK (0x81010001/2) or SRK (0x81000001/2) slots.
//
// If the handle already holds an object that is NOT our ECC signing key (e.g. an
// RSA signing key left by an older tpm2-tools/openssl provisioning path — go-tpm
// reports "key at handle is not ECC"), it is treated as STALE and evicted, then a
// fresh ECC key is created. Safe on the enrol path: this only runs on an empty-NV
// host, so no issued cert references the old key. Only a valid ECC key is reused.
// A restricted-decrypt object (EK/SRK) is NEVER evicted — it errors instead, which
// only happens if handle was misconfigured to a reserved slot.
func EnsureKey(tpm transport.TPM, handle tpm2.TPMHandle) (signer *Signer, created bool, err error) {
	if rp, e := (tpm2.ReadPublic{ObjectHandle: handle}).Execute(tpm); e == nil {
		// Guard FIRST: never touch a restricted-decrypt object (EK/SRK). An ECC
		// EK/SRK would otherwise parse as a usable ECC key below and be "reused".
		if isProtectedObject(rp.OutPublic) {
			return nil, false, fmt.Errorf("refusing to evict a restricted-decrypt object (EK/SRK?) at handle %#x — configure a non-reserved owner handle (e.g. 0x81000100)", uint32(handle))
		}
		if s, e2 := signerFromPub(tpm, handle, rp.Name, rp.OutPublic); e2 == nil {
			return s, false, nil // our ECC signing key is already there — reuse it
		}
		// Stale/wrong-type signing key occupying the handle — evict so we recreate.
		if _, e2 := (tpm2.EvictControl{
			Auth:             tpm2.TPMRHOwner,
			ObjectHandle:     &tpm2.NamedHandle{Handle: handle, Name: rp.Name},
			PersistentHandle: handle,
		}).Execute(tpm); e2 != nil {
			return nil, false, fmt.Errorf("evict stale (non-ECC) key at handle: %w", e2)
		}
	}

	cp, err := tpm2.CreatePrimary{
		PrimaryHandle: tpm2.TPMRHOwner,
		InPublic:      tpm2.New2B(eccSigningTemplate()),
	}.Execute(tpm)
	if err != nil {
		return nil, false, fmt.Errorf("createprimary: %w", err)
	}
	// Persist at the well-known handle, then flush the transient object.
	_, err = tpm2.EvictControl{
		Auth:             tpm2.TPMRHOwner,
		ObjectHandle:     &tpm2.NamedHandle{Handle: cp.ObjectHandle, Name: cp.Name},
		PersistentHandle: handle,
	}.Execute(tpm)
	_, _ = (tpm2.FlushContext{FlushHandle: cp.ObjectHandle}).Execute(tpm)
	if err != nil {
		return nil, false, fmt.Errorf("evictcontrol: %w", err)
	}
	rp, err := (tpm2.ReadPublic{ObjectHandle: handle}).Execute(tpm)
	if err != nil {
		return nil, false, fmt.Errorf("readpublic after persist: %w", err)
	}
	s, err := signerFromPub(tpm, handle, rp.Name, rp.OutPublic)
	return s, true, err
}

// LoadSigner returns a crypto.Signer for the persistent key ALREADY at handle,
// WITHOUT creating one. This is the reused-identity path (normal boot / the on-host
// agent): the fleet signing key was created at enrolment, so here we only sign the
// TLS handshake with it — the private key never leaves the TPM. Errors (never
// creates/evicts) if the handle is empty, holds a non-ECC object, or holds a
// restricted-decrypt EK/SRK. Contrast EnsureKey, which is the enrol-time create path.
func LoadSigner(tpm transport.TPM, handle tpm2.TPMHandle) (*Signer, error) {
	rp, err := (tpm2.ReadPublic{ObjectHandle: handle}).Execute(tpm)
	if err != nil {
		return nil, fmt.Errorf("no persistent key at handle %#x: %w", uint32(handle), err)
	}
	if isProtectedObject(rp.OutPublic) {
		return nil, fmt.Errorf("object at handle %#x is a restricted-decrypt key (EK/SRK?), not the fleet signing key", uint32(handle))
	}
	return signerFromPub(tpm, handle, rp.Name, rp.OutPublic)
}

// nvPresent reports whether a defined NV index exists. NOTE: "defined" does NOT
// mean "holds our fleet cert" — leftover/foreign NV (e.g. an earlier
// tpm-provision run, or another OS) also reports present. Prefer ReadNVCert +
// content validation for the enrol decision; this stays for cheap existence checks.
func NVPresent(tpm transport.TPM, index tpm2.TPMHandle) bool {
	_, err := (tpm2.NVReadPublic{NVIndex: index}).Execute(tpm)
	return err == nil
}

// ReadNVCert reads the full contents of NV index (chunked under the NV buffer
// limit). Returns an error if the index is undefined or unreadable — callers
// treat that as "no cert" (enrol). The bytes are the caller's to parse/validate;
// this does not assume they are a certificate.
func ReadNVCert(tpm transport.TPM, index tpm2.TPMHandle) ([]byte, error) {
	rp, err := (tpm2.NVReadPublic{NVIndex: index}).Execute(tpm)
	if err != nil {
		return nil, fmt.Errorf("nv readpublic: %w", err)
	}
	pub, err := rp.NVPublic.Contents()
	if err != nil {
		return nil, err
	}
	size := int(pub.DataSize)
	out := make([]byte, 0, size)
	const chunk = 512 // conservatively under TPM_MAX_NV_BUFFER_SIZE
	for off := 0; off < size; off += chunk {
		end := min(off+chunk, size)
		r, err := (tpm2.NVRead{
			AuthHandle: tpm2.AuthHandle{Handle: index, Name: rp.NVName, Auth: tpm2.PasswordAuth(nil)},
			NVIndex:    tpm2.NamedHandle{Handle: index, Name: rp.NVName},
			Size:       uint16(end - off),
			Offset:     uint16(off),
		}).Execute(tpm)
		if err != nil {
			return nil, fmt.Errorf("nv read @%d: %w", off, err)
		}
		out = append(out, r.Data.Buffer...)
	}
	return out, nil
}

// nvWriteCert (re)defines the NV index sized to der and writes it under owner
// auth. Chunks the write to stay within the TPM's NV buffer limit.
func WriteNVCert(tpm transport.TPM, index tpm2.TPMHandle, der []byte) error {
	// Undefine a stale index first (re-enrolment / resized cert).
	if rp, err := (tpm2.NVReadPublic{NVIndex: index}).Execute(tpm); err == nil {
		// A PLATFORM-created index at our handle is OEM firmware NV (bench-seen on
		// a Dell 820: 0x01800001 was platformcreate|policydelete|writelocked). We
		// cannot undefine it with owner auth, and TPM2_Clear/BIOS won't remove it —
		// so fail LOUDLY with guidance rather than a confusing "nv undefine" error.
		// The remedy is a different, free owner index via rd.dot1x.nvindex.
		if pub, e := rp.NVPublic.Contents(); e == nil && pub.Attributes.PlatformCreate {
			return fmt.Errorf("NV index %#x is a platform-created (firmware-owned) index — cannot be written by enrolment; pick a free owner index via rd.dot1x.nvindex / -nvindex", uint32(index))
		}
		if _, err := (tpm2.NVUndefineSpace{
			AuthHandle: tpm2.TPMRHOwner,
			NVIndex:    tpm2.NamedHandle{Handle: index, Name: rp.NVName},
		}).Execute(tpm); err != nil {
			return fmt.Errorf("nv undefine: %w", err)
		}
	}

	def := tpm2.NVDefineSpace{
		AuthHandle: tpm2.TPMRHOwner,
		PublicInfo: tpm2.New2B(tpm2.TPMSNVPublic{
			NVIndex: index,
			NameAlg: tpm2.TPMAlgSHA256,
			Attributes: tpm2.TPMANV{
				OwnerWrite: true,
				OwnerRead:  true,
				AuthWrite:  true,
				AuthRead:   true,
				NT:         tpm2.TPMNTOrdinary,
				NoDA:       true,
			},
			DataSize: uint16(len(der)),
		}),
	}
	if _, err := def.Execute(tpm); err != nil {
		return fmt.Errorf("nv define: %w", err)
	}
	rp, err := (tpm2.NVReadPublic{NVIndex: index}).Execute(tpm)
	if err != nil {
		return fmt.Errorf("nv readpublic: %w", err)
	}

	const chunk = 512 // conservatively under TPM_MAX_NV_BUFFER_SIZE
	for off := 0; off < len(der); off += chunk {
		end := min(off+chunk, len(der))
		if _, err := (tpm2.NVWrite{
			AuthHandle: tpm2.AuthHandle{Handle: index, Name: rp.NVName, Auth: tpm2.PasswordAuth(nil)},
			NVIndex:    tpm2.NamedHandle{Handle: index, Name: rp.NVName},
			Data:       tpm2.TPM2BMaxNVBuffer{Buffer: der[off:end]},
			Offset:     uint16(off),
		}).Execute(tpm); err != nil {
			return fmt.Errorf("nv write @%d: %w", off, err)
		}
	}
	return nil
}