package webui

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"net/http"
)

// csrfSessionCookieNames are tried in order to find whatever session-identifying
// cookie the current request carries — the full admin session, the full mailbox
// session, or either's pending-MFA cookie (covers the MFA step of login, which
// happens before the full session exists but after a pending cookie is set). The
// bare initial login POST, before any of these cookies exist yet, has no source
// token and so is intentionally not CSRF-checked — standard practice, since there's
// no session yet for a forged request to act against.
var csrfSessionCookieNames = []string{
	sessionCookieName, mailboxSessionCookieName, pendingMFACookieName, mailboxPendingMFACookieName,
}

func csrfSourceToken(r *http.Request) (string, bool) {
	for _, name := range csrfSessionCookieNames {
		if c, err := r.Cookie(name); err == nil && c.Value != "" {
			return c.Value, true
		}
	}
	return "", false
}

// csrfTokenFor derives this request's expected CSRF token: an HMAC over whatever
// session-identifying cookie is present, keyed by the app secret. Deterministic and
// unstored — recomputed fresh on both render (render.go injects it into every page)
// and validation (CSRFProtect below), so there's no server-side token table to
// manage or expire.
func (a *App) csrfTokenFor(r *http.Request) string {
	token, ok := csrfSourceToken(r)
	if !ok {
		return ""
	}
	mac := hmac.New(sha256.New, a.appSecret)
	mac.Write([]byte(token))
	return hex.EncodeToString(mac.Sum(nil))
}

// csrfProtectedMethods are the only ones CSRFProtect checks — GET/HEAD/OPTIONS never
// mutate state in this app (see the M2 fix that removed the one exception that used
// to exist) so they're exempt, matching standard CSRF-defense scope.
var csrfProtectedMethods = map[string]bool{http.MethodPost: true, http.MethodPut: true, http.MethodPatch: true, http.MethodDelete: true}

// CSRFProtect rejects state-changing requests whose csrf_token doesn't match what
// csrfTokenFor computes for the request's own session cookie. Every authenticated
// HTML form gets the token auto-injected as a hidden field, and every same-origin
// fetch() call gets it auto-attached as an X-CSRF-Token header — both via the shared
// csrf_script.html partial parsed into every page (see render.go/loadTemplates) —
// so no individual handler or template needed to change for this to apply
// uniformly across the whole app, admin and webmail alike.
func (a *App) CSRFProtect(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !csrfProtectedMethods[r.Method] {
			next.ServeHTTP(w, r)
			return
		}
		expected := a.csrfTokenFor(r)
		if expected == "" {
			// No session cookie at all — an unauthenticated route (e.g. the login POST
			// itself). Nothing to protect yet.
			next.ServeHTTP(w, r)
			return
		}
		got := r.Header.Get("X-CSRF-Token")
		if got == "" {
			// Only fall back to parsing the body if the header wasn't already present —
			// keeps the common fetch()-with-header path from ever triggering an implicit
			// multipart parse here (the handler still parses it normally afterward).
			got = r.FormValue("csrf_token")
		}
		if got == "" || !hmac.Equal([]byte(got), []byte(expected)) {
			http.Error(w, "Forbidden: missing or invalid CSRF token", http.StatusForbidden)
			return
		}
		next.ServeHTTP(w, r)
	})
}