// fleet-enrol-broker — TPM identity enrolment broker.
//
// The ONE unauthenticated, pre-auth bootstrap endpoint of the fleet. An empty-NV
// host generates a key IN its TPM, builds a CSR, POSTs it here, and shows the
// returned fingerprint on its Plymouth splash. An admin cross-checks that
// fingerprint against the pending queue (a separate mTLS web UI) and approves;
// the broker has step-ca sign the CSR and hands the cert back on the host's next
// status poll. The host writes it to TPM NV and authenticates.
//
// See docs/tpm-enrolment-design.md. Two listeners, by design:
//
//	PUBLIC (plain HTTP, behind the serving nginx /enrol/)  -- unauthenticated:
//	    POST /request          submit a CSR -> {enrol_id, fingerprint}
//	    GET  /status/{id}      poll -> pending | approved+cert | denied | expired
//
//	ADMIN (mTLS, management network / loopback only)       -- admin client cert:
//	    GET  /                 the pending-queue web UI
//	    GET  /pending          JSON queue
//	    POST /approve/{id}     OOB-approve -> sign -> store cert
//	    POST /deny/{id}        reject
//
// Stdlib only; signing shells out to `step`. Nothing here is per-host baked state.
package main

import (
	"context"
	_ "embed"
	"encoding/json"
	"log"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

//go:embed wordlist.txt
var wordlistRaw string

type config struct {
	publicAddr string
	adminAddr  string
	ttl        time.Duration
	retention  time.Duration
	maxPending int
	fpWords    int
	ratePerMin int
	trustXFF   bool
	hostBase   string

	signerKind string
	// step signer
	stepBin      string
	caURL        string
	caRoot       string
	provisioner  string
	passwordFile string
	notAfter     string
	// admin mTLS
	adminCert     string
	adminKey      string
	adminClientCA string
	adminInsecure bool
}

func loadConfig() config {
	return config{
		publicAddr:    env("BROKER_PUBLIC_ADDR", "127.0.0.1:9090"),
		adminAddr:     env("BROKER_ADMIN_ADDR", "127.0.0.1:9443"),
		ttl:           envDur("BROKER_PENDING_TTL", 10*time.Minute),
		retention:     envDur("BROKER_RETENTION", time.Hour),
		maxPending:    envInt("BROKER_MAX_PENDING", 256),
		fpWords:       envInt("BROKER_WORDS", 8),
		ratePerMin:    envInt("BROKER_RATE_PER_MIN", 12),
		trustXFF:      envBool("BROKER_TRUST_XFF", true),
		hostBase:      env("HOSTNQN_BASE", "nqn.2026-06.org.example.fleet:host-"),
		signerKind:    env("BROKER_SIGNER", "step"),
		stepBin:       env("STEP_BIN", "step"),
		caURL:         env("STEP_CA_URL", "https://ca.fleet.lab:8443"),
		caRoot:        env("STEP_CA_ROOT", "/pki/root_ca.crt"),
		provisioner:   env("STEP_PROVISIONER", "enrol"),
		passwordFile:  env("STEP_PROVISIONER_PASSWORD_FILE", "/pki/stepca/secrets/password"),
		notAfter:      env("STEP_CERT_NOT_AFTER", "720h"),
		adminCert:     env("BROKER_ADMIN_CERT", ""),
		adminKey:      env("BROKER_ADMIN_KEY", ""),
		adminClientCA: env("BROKER_ADMIN_CLIENT_CA", ""),
		adminInsecure: envBool("BROKER_ADMIN_INSECURE", false),
	}
}

type broker struct {
	cfg    config
	store  *Store
	signer Signer
	rl     *rateLimiter
}

func main() {
	log.SetFlags(log.LstdFlags | log.LUTC)
	cfg := loadConfig()

	words := parseWordlist(wordlistRaw)
	store, err := NewStore(cfg.ttl, cfg.retention, cfg.maxPending, words, cfg.fpWords, cfg.hostBase)
	if err != nil {
		log.Fatalf("store: %v", err)
	}

	var signer Signer
	switch cfg.signerKind {
	case "step":
		signer = &stepSigner{
			bin: cfg.stepBin, caURL: cfg.caURL, root: cfg.caRoot,
			provisioner: cfg.provisioner, passwordFile: cfg.passwordFile, notAfter: cfg.notAfter,
		}
	case "mock":
		signer, err = newMockSigner()
		if err != nil {
			log.Fatalf("mock signer: %v", err)
		}
	default:
		log.Fatalf("unknown BROKER_SIGNER %q (want step|mock)", cfg.signerKind)
	}

	b := &broker{cfg: cfg, store: store, signer: signer, rl: newRateLimiter(cfg.ratePerMin)}

	// background reaper: expire timed-out requests + evict old decided ones.
	go func() {
		t := time.NewTicker(30 * time.Second)
		defer t.Stop()
		for range t.C {
			if n := store.Reap(); n > 0 {
				log.Printf("audit reap expired_pending=%d", n)
			}
			b.rl.sweep()
		}
	}()

	var wg sync.WaitGroup
	wg.Add(2)
	go func() { defer wg.Done(); b.servePublic() }()
	go func() { defer wg.Done(); b.serveAdmin() }()
	log.Printf("enrol-broker up: public=%s admin=%s signer=%s words=%d ttl=%s",
		cfg.publicAddr, cfg.adminAddr, signer.Kind(), len(words), cfg.ttl)
	wg.Wait()
}

// --- public listener (unauthenticated) -------------------------------------

func (b *broker) servePublic() {
	mux := http.NewServeMux()
	mux.HandleFunc("POST /request", b.handleRequest)
	mux.HandleFunc("GET /status/{id}", b.handleStatus)
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})
	srv := &http.Server{
		Addr:              b.cfg.publicAddr,
		Handler:           mux,
		ReadHeaderTimeout: 5 * time.Second,
		ReadTimeout:       10 * time.Second,
		WriteTimeout:      15 * time.Second,
		IdleTimeout:       30 * time.Second,
		MaxHeaderBytes:    8 << 10,
	}
	log.Fatalf("public listener: %v", srv.ListenAndServe())
}

type enrolRequest struct {
	CSR     string        `json:"csr"`
	Context DeviceContext `json:"context"`
}

type enrolRequestResp struct {
	EnrolID     string `json:"enrol_id"`
	Fingerprint string `json:"fingerprint"`
	State       string `json:"state"`
	PollAfter   int    `json:"poll_after_seconds"`
	ExpiresIn   int    `json:"expires_in_seconds"`
}

type statusResp struct {
	State      string `json:"state"`
	Cert       string `json:"cert,omitempty"`
	Hostnqn    string `json:"hostnqn,omitempty"`
	CertSerial string `json:"cert_serial,omitempty"`
}

func (b *broker) handleRequest(w http.ResponseWriter, r *http.Request) {
	ip := b.clientIP(r)
	if !b.rl.allow(ip) {
		writeJSONErr(w, http.StatusTooManyRequests, "rate limit exceeded")
		return
	}
	r.Body = http.MaxBytesReader(w, r.Body, 16<<10) // a CSR is ~1KB
	var req enrolRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeJSONErr(w, http.StatusBadRequest, "invalid JSON body")
		return
	}
	csr, err := parseCSR(req.CSR)
	if err != nil {
		writeJSONErr(w, http.StatusBadRequest, "invalid CSR: "+err.Error())
		return
	}
	e, err := b.store.Add(csr.RawSubjectPublicKeyInfo, csr.Subject.CommonName, req.CSR, req.Context)
	if err != nil {
		if err == errTooManyPending {
			writeJSONErr(w, http.StatusServiceUnavailable, "broker busy: too many pending enrolments")
			return
		}
		writeJSONErr(w, http.StatusInternalServerError, "internal error")
		return
	}
	log.Printf("audit request id=%s cn=%q pubkey=%s mac=%q fp=%q from=%s",
		e.ID, e.CN, short(e.PubKeyHash), e.Context.MAC, e.Fingerprint, ip)
	writeJSON(w, http.StatusAccepted, enrolRequestResp{
		EnrolID:     e.ID,
		Fingerprint: e.Fingerprint,
		State:       string(e.State),
		PollAfter:   5,
		ExpiresIn:   int(time.Until(e.ExpiresAt).Seconds()),
	})
}

func (b *broker) handleStatus(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	e := b.store.Get(id)
	if e == nil {
		writeJSONErr(w, http.StatusNotFound, "unknown enrol-id")
		return
	}
	resp := statusResp{State: string(e.State)}
	if e.State == StateApproved {
		resp.Cert = e.certPEM
		resp.Hostnqn = e.Hostnqn
		resp.CertSerial = e.CertSerial
	}
	writeJSON(w, http.StatusOK, resp)
}

func (b *broker) clientIP(r *http.Request) string {
	if b.cfg.trustXFF {
		if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
			return strings.TrimSpace(strings.Split(xff, ",")[0])
		}
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// --- rate limiter (per-IP fixed window) ------------------------------------

type rateLimiter struct {
	mu      sync.Mutex
	perMin  int
	windows map[string]*window
}

type window struct {
	start time.Time
	count int
}

func newRateLimiter(perMin int) *rateLimiter {
	return &rateLimiter{perMin: perMin, windows: make(map[string]*window)}
}

func (rl *rateLimiter) allow(ip string) bool {
	if rl.perMin <= 0 {
		return true
	}
	rl.mu.Lock()
	defer rl.mu.Unlock()
	now := time.Now()
	w := rl.windows[ip]
	if w == nil || now.Sub(w.start) >= time.Minute {
		rl.windows[ip] = &window{start: now, count: 1}
		return true
	}
	if w.count >= rl.perMin {
		return false
	}
	w.count++
	return true
}

func (rl *rateLimiter) sweep() {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	for ip, w := range rl.windows {
		if time.Since(w.start) >= time.Minute {
			delete(rl.windows, ip)
		}
	}
}

// --- small helpers ---------------------------------------------------------

func parseWordlist(s string) []string {
	var out []string
	for _, line := range strings.Split(s, "\n") {
		if w := strings.TrimSpace(line); w != "" {
			out = append(out, w)
		}
	}
	return out
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(v)
}

func writeJSONErr(w http.ResponseWriter, code int, msg string) {
	writeJSON(w, code, map[string]string{"error": msg})
}

func short(s string) string {
	if len(s) > 16 {
		return s[:16]
	}
	return s
}

func env(k, def string) string {
	if v := os.Getenv(k); v != "" {
		return v
	}
	return def
}

func envInt(k string, def int) int {
	if v := os.Getenv(k); v != "" {
		if n, err := strconv.Atoi(v); err == nil {
			return n
		}
	}
	return def
}

func envBool(k string, def bool) bool {
	if v := os.Getenv(k); v != "" {
		b, err := strconv.ParseBool(v)
		if err == nil {
			return b
		}
	}
	return def
}

func envDur(k string, def time.Duration) time.Duration {
	if v := os.Getenv(k); v != "" {
		if d, err := time.ParseDuration(v); err == nil {
			return d
		}
	}
	return def
}

// ctxWithTimeout is used by admin approve to bound the step-ca call.
func ctxWithTimeout(d time.Duration) (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), d)
}