package main

import (
	"fmt"
	"log"
	"net/http"
	"net/http/httptest"
	"strings"

	goidentity "github.com/jcmturner/goidentity/v6"
	"github.com/jcmturner/gokrb5/v8/keytab"
	"github.com/jcmturner/gokrb5/v8/service"
	"github.com/jcmturner/gokrb5/v8/spnego"
)

// KerberosAuth validates Kerberos SPNEGO tokens using a service keytab.
//
// The client must obtain a service ticket for the SPN (default: HTTP/<fqdn>@REALM)
// and present it via the standard "Authorization: Negotiate <base64>" header.
// Any GSSAPI-aware client works: curl --negotiate, kinit + curl, python-requests-gssapi, etc.
type KerberosAuth struct {
	kt  *keytab.Keytab
	spn string // e.g. "HTTP/k3s.example.com" — empty = auto-select from ticket SName
}

// NewKerberosAuth loads the keytab from keytabPath.
// spn is written into the service settings so gokrb5 validates the correct principal;
// pass an empty string to accept any principal present in the keytab.
func NewKerberosAuth(keytabPath, spn string) (*KerberosAuth, error) {
	kt, err := keytab.Load(keytabPath)
	if err != nil {
		return nil, fmt.Errorf("loading keytab %s: %w", keytabPath, err)
	}
	return &KerberosAuth{kt: kt, spn: spn}, nil
}

// Authenticate validates the SPNEGO token from the request's Authorization header.
// On success it returns the authenticated username (principal primary, without @REALM).
// On failure it sets a WWW-Authenticate: Negotiate header on w (for mutual-auth
// continuation tokens) and returns an error; the caller is responsible for the 401.
func (k *KerberosAuth) Authenticate(w http.ResponseWriter, r *http.Request) (string, error) {
	if !strings.HasPrefix(r.Header.Get("Authorization"), "Negotiate ") {
		return "", fmt.Errorf("no Negotiate token in request")
	}

	// We use gokrb5's SPNEGO middleware but call it synchronously via a
	// closure — ServeHTTP is not goroutine-concurrent so the captured
	// variables are safe to read after the call returns.
	var (
		authed   bool
		username string
	)
	inner := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		id := goidentity.FromHTTPRequestContext(r)
		if id != nil && id.Authenticated() {
			authed = true
			username = id.UserName()
		}
	})

	var opts []func(*service.Settings)
	if k.spn != "" {
		opts = append(opts, service.KeytabPrincipal(k.spn))
	}
	opts = append(opts, service.DecodePAC(false))
	opts = append(opts, service.Logger(log.Default()))

	// Route through a recorder so we can capture the WWW-Authenticate
	// continuation token (mutual authentication) without writing it yet.
	rec := httptest.NewRecorder()
	spnego.SPNEGOKRB5Authenticate(inner, k.kt, opts...).ServeHTTP(rec, r)

	if !authed {
		// Forward any Negotiate continuation token from the middleware.
		if v := rec.Header().Get("WWW-Authenticate"); v != "" {
			w.Header().Set("WWW-Authenticate", v)
		}
		return "", fmt.Errorf("kerberos authentication failed")
	}
	return username, nil
}