mcias/internal/ui/handlers_totp.go

package ui

import (
	"encoding/base32"
	"encoding/base64"
	"fmt"
	"net/http"

	qrcode "github.com/skip2/go-qrcode"

	"git.wntrmute.dev/kyle/mcias/internal/audit"
	"git.wntrmute.dev/kyle/mcias/internal/auth"
	"git.wntrmute.dev/kyle/mcias/internal/crypto"
	"git.wntrmute.dev/kyle/mcias/internal/model"
)

// handleTOTPEnrollStart processes the password re-auth step and generates
// the TOTP secret + QR code for the user to scan.
//
// Security (SEC-01): the current password is required to prevent a stolen
// session from enrolling attacker-controlled TOTP. Lockout is checked and
// failures are recorded to prevent brute-force use as a password oracle.
func (u *UIServer) handleTOTPEnrollStart(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, maxFormBytes)
	if err := r.ParseForm(); err != nil {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "invalid form submission"})
		return
	}

	claims := claimsFromContext(r.Context())
	if claims == nil {
		u.renderError(w, r, http.StatusUnauthorized, "unauthorized")
		return
	}

	acct, err := u.db.GetAccountByUUID(claims.Subject)
	if err != nil {
		u.renderError(w, r, http.StatusUnauthorized, "account not found")
		return
	}

	// Already enrolled — show enabled status.
	if acct.TOTPRequired {
		u.renderTOTPSection(w, r, ProfileData{TOTPEnabled: true})
		return
	}

	password := r.FormValue("password")
	if password == "" {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "password is required"})
		return
	}

	// Security: check lockout before verifying password.
	locked, lockErr := u.db.IsLockedOut(acct.ID)
	if lockErr != nil {
		u.logger.Error("lockout check (UI TOTP enroll)", "error", lockErr)
	}
	if locked {
		u.writeAudit(r, model.EventTOTPEnrolled, &acct.ID, &acct.ID, `{"result":"locked"}`)
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "account temporarily locked, please try again later"})
		return
	}

	// Security: verify current password with constant-time Argon2id path.
	ok, verifyErr := auth.VerifyPassword(password, acct.PasswordHash)
	if verifyErr != nil || !ok {
		_ = u.db.RecordLoginFailure(acct.ID)
		u.writeAudit(r, model.EventTOTPEnrolled, &acct.ID, &acct.ID, `{"result":"wrong_password"}`)
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "password is incorrect"})
		return
	}

	// Generate TOTP secret.
	rawSecret, b32Secret, err := auth.GenerateTOTPSecret()
	if err != nil {
		u.logger.Error("generate TOTP secret", "error", err)
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}

	// Encrypt and store as pending (totp_required stays 0 until confirmed).
	masterKey, err := u.vault.MasterKey()
	if err != nil {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}
	secretEnc, secretNonce, err := crypto.SealAESGCM(masterKey, rawSecret)
	if err != nil {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}

	// Security: use StorePendingTOTP (not SetTOTP) so that totp_required
	// remains 0 until the user proves possession via ConfirmTOTP.
	if err := u.db.StorePendingTOTP(acct.ID, secretEnc, secretNonce); err != nil {
		u.logger.Error("store pending TOTP", "error", err)
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}

	otpURI := fmt.Sprintf("otpauth://totp/MCIAS:%s?secret=%s&issuer=MCIAS", acct.Username, b32Secret)

	// Generate QR code PNG.
	png, err := qrcode.Encode(otpURI, qrcode.Medium, 200)
	if err != nil {
		u.logger.Error("generate QR code", "error", err)
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}
	qrDataURI := "data:image/png;base64," + base64.StdEncoding.EncodeToString(png)

	// Issue enrollment nonce for the confirm step.
	nonce, err := u.issueTOTPEnrollNonce(acct.ID)
	if err != nil {
		u.logger.Error("issue TOTP enroll nonce", "error", err)
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}

	csrfToken, _ := u.setCSRFCookies(w)
	u.render(w, "totp_enroll_qr", ProfileData{
		PageData:        PageData{CSRFToken: csrfToken},
		TOTPSecret:      b32Secret,
		TOTPQR:          qrDataURI,
		TOTPEnrollNonce: nonce,
	})
}

// handleTOTPConfirm validates the TOTP code and activates enrollment.
//
// Security (CRIT-01): the counter is recorded to prevent replay of the same
// code within its validity window.
func (u *UIServer) handleTOTPConfirm(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, maxFormBytes)
	if err := r.ParseForm(); err != nil {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "invalid form submission"})
		return
	}

	claims := claimsFromContext(r.Context())
	if claims == nil {
		u.renderError(w, r, http.StatusUnauthorized, "unauthorized")
		return
	}

	nonce := r.FormValue("totp_enroll_nonce")
	totpCode := r.FormValue("totp_code")

	// Security: consume the nonce (single-use); reject if unknown or expired.
	accountID, ok := u.consumeTOTPEnrollNonce(nonce)
	if !ok {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "session expired, please start enrollment again"})
		return
	}

	acct, err := u.db.GetAccountByID(accountID)
	if err != nil {
		u.logger.Error("get account for TOTP confirm", "error", err, "account_id", accountID)
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}

	// Security: verify nonce accountID matches session claims.
	if acct.UUID != claims.Subject {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "session mismatch"})
		return
	}

	if acct.TOTPSecretEnc == nil {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "enrollment not started"})
		return
	}

	// Decrypt and validate TOTP code.
	masterKey, err := u.vault.MasterKey()
	if err != nil {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}
	secret, err := crypto.OpenAESGCM(masterKey, acct.TOTPSecretNonce, acct.TOTPSecretEnc)
	if err != nil {
		u.logger.Error("decrypt TOTP secret for confirm", "error", err, "account_id", acct.ID)
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}

	valid, totpCounter, err := auth.ValidateTOTP(secret, totpCode)
	if err != nil || !valid {
		// Re-issue a fresh nonce so the user can retry without restarting.
		u.reissueTOTPEnrollQR(w, r, acct, secret, "invalid TOTP code")
		return
	}

	// Security (CRIT-01): reject replay of a code already used.
	if err := u.db.CheckAndUpdateTOTPCounter(acct.ID, totpCounter); err != nil {
		u.reissueTOTPEnrollQR(w, r, acct, secret, "invalid TOTP code")
		return
	}

	// Activate TOTP (sets totp_required=1).
	if err := u.db.SetTOTP(acct.ID, acct.TOTPSecretEnc, acct.TOTPSecretNonce); err != nil {
		u.logger.Error("set TOTP", "error", err, "account_id", acct.ID)
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}

	u.writeAudit(r, model.EventTOTPEnrolled, &acct.ID, nil, "")

	u.renderTOTPSection(w, r, ProfileData{
		TOTPEnabled: true,
		TOTPSuccess: "Two-factor authentication enabled successfully.",
	})
}

// reissueTOTPEnrollQR re-renders the QR code page with a fresh nonce after
// a failed code confirmation, so the user can retry without restarting.
func (u *UIServer) reissueTOTPEnrollQR(w http.ResponseWriter, r *http.Request, acct *model.Account, secret []byte, errMsg string) {
	b32Secret := base32.StdEncoding.EncodeToString(secret)
	otpURI := fmt.Sprintf("otpauth://totp/MCIAS:%s?secret=%s&issuer=MCIAS", acct.Username, b32Secret)

	png, err := qrcode.Encode(otpURI, qrcode.Medium, 200)
	if err != nil {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}
	qrDataURI := "data:image/png;base64," + base64.StdEncoding.EncodeToString(png)

	newNonce, nonceErr := u.issueTOTPEnrollNonce(acct.ID)
	if nonceErr != nil {
		u.renderTOTPSection(w, r, ProfileData{TOTPError: "internal error"})
		return
	}

	csrfToken, _ := u.setCSRFCookies(w)
	u.render(w, "totp_enroll_qr", ProfileData{
		PageData:        PageData{CSRFToken: csrfToken},
		TOTPSecret:      b32Secret,
		TOTPQR:          qrDataURI,
		TOTPEnrollNonce: newNonce,
		TOTPError:       errMsg,
	})
}

// handleAdminTOTPRemove removes TOTP from an account (admin only).
func (u *UIServer) handleAdminTOTPRemove(w http.ResponseWriter, r *http.Request) {
	accountUUID := r.PathValue("id")
	if accountUUID == "" {
		u.renderError(w, r, http.StatusBadRequest, "missing account ID")
		return
	}

	acct, err := u.db.GetAccountByUUID(accountUUID)
	if err != nil {
		u.renderError(w, r, http.StatusNotFound, "account not found")
		return
	}

	if err := u.db.ClearTOTP(acct.ID); err != nil {
		u.logger.Error("clear TOTP (admin)", "error", err, "account_id", acct.ID)
		u.renderError(w, r, http.StatusInternalServerError, "internal error")
		return
	}

	claims := claimsFromContext(r.Context())
	var actorID *int64
	if claims != nil {
		if actor, err := u.db.GetAccountByUUID(claims.Subject); err == nil {
			actorID = &actor.ID
		}
	}
	u.writeAudit(r, model.EventTOTPRemoved, actorID, &acct.ID,
		audit.JSON("admin", "true"))

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	_, _ = fmt.Fprint(w, `Disabled <span class="text-muted text-small">(removed)</span>`)
}

// renderTOTPSection is a helper to render the totp_section fragment with
// common page data fields populated.
func (u *UIServer) renderTOTPSection(w http.ResponseWriter, r *http.Request, data ProfileData) {
	csrfToken, _ := u.setCSRFCookies(w)
	data.CSRFToken = csrfToken
	data.ActorName = u.actorName(r)
	data.IsAdmin = isAdmin(r)
	u.render(w, "totp_section", data)
}