// fleet-enrol — inline empty-NV first-boot TPM enrolment client.
//
// Invoked from the dracut hook (95dot1x-macsec/tpm-enrol.sh) ONLY when the
// secure build boots a host whose TPM NV 0x01800100 is empty. It:
//  1. ensures an ECDSA signing key exists at a non-reserved owner handle
//     (default 0x81000100 — deliberately NOT the EK slot 0x81010001),
//  2. builds a CSR with it (key never leaves the TPM),
//  3. submits CSR + device context to the enrolment broker, shows the returned
//     fingerprint on the Plymouth splash, and polls until an admin approves,
//  4. writes the issued cert (DER) to NV 0x01800100.
//
// The hook then reboots into the normal secure path. Everything host-specific
// (key, cert, hostnqn) is created here at runtime — nothing is baked (invariant #1).
//
// Exit codes (default/enrol mode): 0 = enrolled (cert written to NV; caller
// should reboot); 10 = a VALID fleet cert is already present (normal boot —
// caller proceeds); 1 = failure (caller leaves NV empty -> boot fails -> emergency).
//
// Exit codes (-check-only mode, no network/no key-gen): 0 = a valid fleet cert is
// present (caller: normal boot); 20 = enrolment needed (NV empty/garbage/foreign/
// expired); 1 = error (e.g. TPM open failed). The dracut hook uses this as the
// single content-aware gate instead of a blind "is the NV index defined" check.
package main

import (
	"context"
	"crypto/x509"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"fleet-enrol/tpmkey"

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

const (
	exitAlreadyEnrolled = 10 // enrol mode: a valid fleet cert is already present
	exitNeedEnrol       = 20 // -check-only: no valid fleet cert → enrolment needed
)

func main() {
	var (
		url       = flag.String("url", "http://boot.fleet.lab/enrol", "enrolment broker base URL")
		tpmPath   = flag.String("tpm", "/dev/tpmrm0", "TPM device path")
		nvHex     = flag.String("nvindex", "0x01800100", "TPM NV index for the client cert")
		handleHex = flag.String("handle", "0x81000100", "TPM persistent handle for the signing key (non-reserved owner handle; NOT the EK slot 0x81010001)")
		cn        = flag.String("cn", "fleet-host", "CSR CommonName (generic; identity is the cert serial)")
		iface     = flag.String("iface", "", "network interface for the device-context MAC")
		caFile    = flag.String("cacert", "/etc/dot1x/radius-ca.pem", "fleet CA bundle: verifies an https broker URL AND the NV cert's chain")
		waitFor   = flag.Duration("timeout", 30*time.Minute, "max wait for admin approval")
		pollEvery = flag.Duration("poll", 5*time.Second, "status poll interval")
		checkOnly = flag.Bool("check-only", false, "classify the NV cert and exit (no network, no key-gen): 0=valid fleet cert, 20=enrol needed, 1=error")
	)
	flag.Parse()
	log.SetFlags(0)
	log.SetPrefix("fleet-enrol: ")

	nvIndex, err := parseHandle(*nvHex)
	if err != nil {
		log.Fatalf("bad -nvindex %q: %v", *nvHex, err)
	}
	keyHandle, err := parseHandle(*handleHex)
	if err != nil {
		log.Fatalf("bad -handle %q: %v", *handleHex, err)
	}

	// /dev/tpmrm0 (the kernel TPM resource manager) is a simple command/response
	// char device: write one command, read one complete response. Wrapping the
	// *os.File with transport.FromReadWriter (as the simulator tests do) avoids
	// the linuxtpm transport and its golang.org/x/sys dependency, keeping the
	// vendored, reproducible initramfs binary tiny (go-tpm only).
	f, err := os.OpenFile(*tpmPath, os.O_RDWR, 0)
	if err != nil {
		log.Fatalf("open TPM %s: %v", *tpmPath, err)
	}
	defer f.Close()
	tpm := transport.FromReadWriter(f)

	// Content-aware enrolment decision: read the NV cert, parse it, and (when the
	// fleet CA is available) verify it chains — so leftover/foreign/garbage NV does
	// NOT masquerade as "enrolled" (a bare "index defined" check let an old
	// tpm-provision cert or random bytes silently suppress enrolment).
	enrolled, detail := nvEnrolled(tpm, nvIndex, *caFile)

	if *checkOnly {
		if enrolled {
			log.Printf("check-only: %s — normal boot", detail)
			os.Exit(0)
		}
		log.Printf("check-only: enrolment needed — %s", detail)
		os.Exit(exitNeedEnrol)
	}

	// Full enrol mode: a valid fleet cert means this is a normal boot.
	if enrolled {
		log.Printf("%s — normal boot", detail)
		os.Exit(exitAlreadyEnrolled)
	}

	ply("enrol-active")
	plyMsg("TPM enrolment: contacting fleet…")

	signer, created, err := tpmkey.EnsureKey(tpm, keyHandle)
	if err != nil {
		plyMsg("TPM enrolment FAILED (key)")
		log.Fatalf("ensure key at %s: %v", *handleHex, err)
	}
	if created {
		log.Printf("created persistent signing key at %s", *handleHex)
	} else {
		log.Printf("reusing persistent signing key at %s", *handleHex)
	}

	csrPEM, err := buildCSR(signer, *cn)
	if err != nil {
		plyMsg("TPM enrolment FAILED (CSR)")
		log.Fatalf("build CSR: %v", err)
	}

	client, err := newBrokerClient(*url, *caFile)
	if err != nil {
		log.Fatalf("broker client: %v", err)
	}
	dctx := gatherContext(*iface)

	ctx, cancel := context.WithTimeout(context.Background(), *waitFor)
	defer cancel()
	log.Printf("submitting CSR to %s (waiting up to %s for approval)", *url, *waitFor)
	st, err := enrol(ctx, client, csrPEM, dctx, *pollEvery, func(fp string) {
		ply("enrol-wait")
		plyMsg("APPROVE THIS DEVICE — code: " + fp)
		log.Printf("fingerprint: %s", fp)
	})
	if err != nil {
		plyMsg("TPM enrolment FAILED — reboot to retry")
		log.Fatalf("enrolment: %v", err)
	}

	der, err := leafDER(st.Cert)
	if err != nil {
		plyMsg("TPM enrolment FAILED (bad cert)")
		log.Fatalf("issued cert: %v", err)
	}
	if err := tpmkey.WriteNVCert(tpm, nvIndex, der); err != nil {
		plyMsg("TPM enrolment FAILED (NV write)")
		log.Fatalf("write cert to NV %s: %v (TPM owner auth?)", *nvHex, err)
	}

	ply("enrol-done")
	plyMsg("Enrolment complete — rebooting")
	log.Printf("ENROLLED: hostnqn %s (serial %s); cert written to NV %s",
		st.Hostnqn, st.CertSerial, *nvHex)
	// The caller (dracut hook) reboots into the secure path.
	os.Exit(0)
}

func parseHandle(s string) (tpm2.TPMHandle, error) {
	v, err := strconv.ParseUint(strings.TrimPrefix(strings.ToLower(s), "0x"), 16, 32)
	if err != nil {
		return 0, err
	}
	return tpm2.TPMHandle(v), nil
}

func newBrokerClient(url, caFile string) (*brokerClient, error) {
	c := &brokerClient{baseURL: strings.TrimRight(url, "/"), http: httpClient(nil)}
	if strings.HasPrefix(url, "https:") {
		pem, err := os.ReadFile(caFile)
		if err != nil {
			return nil, fmt.Errorf("read CA %s: %w", caFile, err)
		}
		pool := x509.NewCertPool()
		if !pool.AppendCertsFromPEM(pem) {
			return nil, fmt.Errorf("no certs in %s", caFile)
		}
		c.http = httpClient(pool)
	}
	return c, nil
}

// gatherContext collects advisory device context for the admin to recognise the
// device. Best-effort — every field is optional.
func gatherContext(iface string) deviceContext {
	d := deviceContext{
		DMIUUID:   strings.TrimSpace(readFile("/sys/class/dmi/id/product_uuid")),
		DMISerial: strings.TrimSpace(readFile("/sys/class/dmi/id/product_serial")),
	}
	if iface != "" {
		d.MAC = strings.TrimSpace(readFile("/sys/class/net/" + iface + "/address"))
	}
	return d
}

func readFile(p string) string {
	b, err := os.ReadFile(p)
	if err != nil {
		return ""
	}
	return string(b)
}

// Plymouth best-effort splash updates (states only; the fingerprint is public).
func ply(status string)  { _ = exec.Command("plymouth", "update", "--status="+status).Run() }
func plyMsg(text string) { _ = exec.Command("plymouth", "message", "--text="+text).Run() }