package security

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"net/http"
	"strings"
)

// CSRF issues and verifies signed double submit tokens. A token is a random nonce plus an HMAC of
// that nonce under a server key, so a value the server did not mint (for example a cookie planted by
// a compromised sibling subdomain) is rejected even though the double submit cookie and header match.
// This composes with the __Host- cookie prefix, SameSite Strict on the session, and the Origin
// allowlist below.
type CSRF struct{ key []byte }

// NewCSRF builds a CSRF signer. A nil or empty key disables verification (dev only); production
// config requires the key.
func NewCSRF(key []byte) *CSRF { return &CSRF{key: key} }

// Issue mints a fresh signed token to set as the __Host-mm_csrf cookie and echo in the header.
func (c *CSRF) Issue() (string, error) {
	nonce := make([]byte, 16)
	if _, err := rand.Read(nonce); err != nil {
		return "", err
	}
	nb := base64.RawURLEncoding.EncodeToString(nonce)
	return nb + "." + c.sign(nb), nil
}

// Valid reports whether token was minted by this signer.
func (c *CSRF) Valid(token string) bool {
	if len(c.key) == 0 {
		return true // verification disabled (dev)
	}
	nb, sig, found := strings.Cut(token, ".")
	if !found || nb == "" {
		return false
	}
	return subtle.ConstantTimeCompare([]byte(sig), []byte(c.sign(nb))) == 1
}

func (c *CSRF) sign(nonceB64 string) string {
	mac := hmac.New(sha256.New, c.key)
	mac.Write([]byte(nonceB64))
	return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}

// Protect is the middleware enforcing CSRF on unsafe verbs: a hard Origin allowlist (reject absent or
// mismatched Origin) plus the signed double submit check (the __Host-mm_csrf cookie and the
// X-CSRF-Token header must match and be a valid signed token). Safe verbs pass through. allowedOrigins
// is matched exactly; an empty list disables the Origin check (dev only).
func (c *CSRF) Protect(cookieName, headerName string, allowedOrigins []string) func(http.Handler) http.Handler {
	allow := make(map[string]bool, len(allowedOrigins))
	for _, o := range allowedOrigins {
		allow[strings.TrimSpace(o)] = true
	}
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if safeMethod(r.Method) {
				next.ServeHTTP(w, r)
				return
			}
			if len(allow) > 0 {
				origin := r.Header.Get("Origin")
				if origin == "" {
					origin = originFromReferer(r.Header.Get("Referer"))
				}
				if !allow[origin] {
					http.Error(w, "forbidden", http.StatusForbidden)
					return
				}
			}
			cookie, err := r.Cookie(cookieName)
			header := r.Header.Get(headerName)
			if err != nil || header == "" ||
				subtle.ConstantTimeCompare([]byte(cookie.Value), []byte(header)) != 1 ||
				!c.Valid(header) {
				http.Error(w, "forbidden", http.StatusForbidden)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

func safeMethod(m string) bool {
	switch m {
	case http.MethodGet, http.MethodHead, http.MethodOptions, http.MethodTrace:
		return true
	}
	return false
}

func originFromReferer(referer string) string {
	if referer == "" {
		return ""
	}
	// Take scheme://host[:port] from a referer URL without importing net/url for this hot path.
	i := strings.Index(referer, "://")
	if i < 0 {
		return ""
	}
	rest := referer[i+3:]
	if j := strings.IndexByte(rest, '/'); j >= 0 {
		rest = rest[:j]
	}
	return referer[:i+3] + rest
}