// Command fleet-ssh-hostcert obtains a fleet-CA-signed SSH *host* certificate for
// this host's per-boot SSH host key and installs it, so administrators trust the
// host via a single `@cert-authority` line instead of trust-on-first-use on a key
// that changes every boot (the ephemeral RAM overlay regenerates host keys each
// boot — see scripts/extract-artifacts.sh). Closes F-02 / POA&M P-17.
//
// Flow (Design A, broker-mediated — docs/ssh-host-certificate-design.md):
//  1. load the host's mTLS identity from the TPM (persistent key 0x81000100 +
//     leaf cert in NV 0x01800100) — the same identity used for 802.1X, the NVMe
//     hostnqn and agent mTLS. The private key never leaves the TPM.
//  2. read the freshly generated public host key (/etc/ssh/ssh_host_ecdsa_key.pub).
//  3. POST it to the enrol-broker's host plane over mTLS. The broker verifies the
//     client cert to the fleet CA (proving an enrolled host) and has step-ca sign
//     an SSH host certificate — the broker, not the host, holds the CA provisioner
//     credential, keeping step-ca behind the single front door.
//  4. write the returned cert + a HostCertificate sshd drop-in, then reload sshd.
//
// FAIL-OPEN: any failure (no network, broker/CA down, host not enrolled) exits
// non-zero WITHOUT touching sshd — the host keeps serving its bare per-boot key
// (today's behaviour). It never installs a missing/broken HostCertificate that
// would stop sshd from starting. A systemd timer retries. Fleet-generic: no
// per-host data in the image (invariant #1).
package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"

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

func main() {
	log.SetFlags(0)
	log.SetPrefix("fleet-ssh-hostcert: ")

	var (
		tpmPath    = flag.String("tpm", envOr("SSH_HOSTCERT_TPM", "/dev/tpmrm0"), "TPM resource-manager device")
		handleStr  = flag.String("handle", envOr("SSH_HOSTCERT_HANDLE", "0x81000100"), "persistent handle of the TPM identity key")
		nvindexStr = flag.String("nvindex", envOr("SSH_HOSTCERT_NVINDEX", "0x01800100"), "NV index of the identity cert (DER)")
		brokerURL  = flag.String("broker-url", envOr("SSH_HOSTCERT_BROKER_URL", "https://boot.fleet.lab:9444/ssh/hostcert"), "enrol-broker host-plane SSH host-cert endpoint (mTLS)")
		caFile     = flag.String("cacert", envOr("SSH_HOSTCERT_CACERT", "/etc/dot1x/radius-ca.pem"), "fleet CA bundle (verifies the broker's server cert)")
		hostKey    = flag.String("hostkey", envOr("SSH_HOSTCERT_HOSTKEY", "/etc/ssh/ssh_host_ecdsa_key.pub"), "SSH host public key to certify")
		outCert    = flag.String("out", envOr("SSH_HOSTCERT_OUT", "/etc/ssh/ssh_host_ecdsa_key-cert.pub"), "path to write the host certificate")
		dropIn     = flag.String("dropin", envOr("SSH_HOSTCERT_DROPIN", "/etc/ssh/sshd_config.d/60-hostcert.conf"), "sshd drop-in that references the host certificate")
		principals = flag.String("principals", envOr("SSH_HOSTCERT_PRINCIPALS", ""), "comma-separated cert principals (default: hostname)")
		reload     = flag.Bool("reload", envOr("SSH_HOSTCERT_RELOAD", "true") == "true", "reload sshd after installing the certificate")
	)
	flag.Parse()

	handle, err := parseHandle(*handleStr)
	if err != nil {
		log.Fatalf("bad -handle: %v", err)
	}
	nvindex, err := parseHandle(*nvindexStr)
	if err != nil {
		log.Fatalf("bad -nvindex: %v", err)
	}

	// --- mTLS identity from the TPM ------------------------------------------
	id, err := loadTPMIdentity(*tpmPath, handle, nvindex)
	if err != nil {
		log.Fatalf("TPM identity: %v", err)
	}
	defer id.Close()

	// --- request inputs -------------------------------------------------------
	pubLine, err := readHostPubLine(*hostKey)
	if err != nil {
		log.Fatalf("read host key %s: %v", *hostKey, err)
	}
	princ := splitPrincipals(*principals)
	if len(princ) == 0 {
		princ = defaultPrincipals()
	}
	keyID := ""
	if id.cert.Leaf != nil {
		keyID = id.cert.Leaf.Subject.CommonName
	}
	if keyID == "" {
		keyID = princ[0]
	}

	// --- ask the broker to sign (over mTLS) ----------------------------------
	client, err := brokerClient(id, *brokerURL, *caFile)
	if err != nil {
		log.Fatalf("broker client: %v", err)
	}
	certLine, err := requestHostCert(client, *brokerURL, pubLine, princ, keyID)
	if err != nil {
		log.Fatalf("request SSH host cert: %v", err)
	}

	// --- install (only now do we touch sshd's config surface) -----------------
	if err := writeFileAtomic(*outCert, []byte(strings.TrimRight(certLine, "\n")+"\n"), 0o644); err != nil {
		log.Fatalf("write cert %s: %v", *outCert, err)
	}
	dropContent := fmt.Sprintf("# Managed by fleet-ssh-hostcert. Host cert signed by the fleet SSH CA.\nHostCertificate %s\n", *outCert)
	if err := writeFileAtomic(*dropIn, []byte(dropContent), 0o644); err != nil {
		log.Fatalf("write drop-in %s: %v", *dropIn, err)
	}
	log.Printf("installed SSH host certificate for principals %v (keyID %q)", princ, keyID)

	if *reload {
		if out, err := exec.Command("systemctl", "reload", "sshd.service").CombinedOutput(); err != nil {
			// Non-fatal: the cert is on disk and takes effect on the next sshd (re)start.
			log.Printf("warning: sshd reload failed (%v): %s", err, strings.TrimSpace(string(out)))
		}
	}
}

type hostCertRequest struct {
	PublicKey  string   `json:"publicKey"`
	Principals []string `json:"principals"`
	KeyID      string   `json:"keyID"`
}

type hostCertResponse struct {
	Cert  string `json:"cert"`
	Error string `json:"error,omitempty"`
}

// requestHostCert POSTs the host public key to the broker's mTLS host plane and
// returns the signed SSH host-certificate line (`ecdsa-sha2-…-cert-v01@… AAAA…`).
func requestHostCert(client *http.Client, url, pubLine string, principals []string, keyID string) (string, error) {
	body, err := json.Marshal(hostCertRequest{PublicKey: pubLine, Principals: principals, KeyID: keyID})
	if err != nil {
		return "", err
	}
	resp, err := client.Post(url, "application/json", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode/100 != 2 {
		return "", fmt.Errorf("broker %s: %s", resp.Status, strings.TrimSpace(string(raw)))
	}
	var out hostCertResponse
	if err := json.Unmarshal(raw, &out); err != nil || out.Cert == "" {
		return "", fmt.Errorf("unexpected broker response: %s", strings.TrimSpace(string(raw)))
	}
	return out.Cert, nil
}

// readHostPubLine returns the trimmed single-line OpenSSH public key (type +
// base64 + optional comment) that the broker signs.
func readHostPubLine(path string) (string, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	line := strings.TrimSpace(string(b))
	if len(strings.Fields(line)) < 2 {
		return "", fmt.Errorf("malformed public key (want at least 2 fields)")
	}
	return line, nil
}

func splitPrincipals(s string) []string {
	var out []string
	for p := range strings.SplitSeq(s, ",") {
		if p = strings.TrimSpace(p); p != "" {
			out = append(out, p)
		}
	}
	return out
}

func defaultPrincipals() []string {
	h, err := os.Hostname()
	if err != nil || h == "" {
		return []string{"localhost"}
	}
	return []string{h}
}

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

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 writeFileAtomic(path string, data []byte, mode os.FileMode) error {
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, data, mode); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}