archivmail/internal/api/totp_handlers.go

package api

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"

	"github.com/archivmail/internal/audit"
	"github.com/archivmail/internal/auth"
)

// ── PROJ-24: TOTP 2FA Handlers ───────────────────────────────────────────

// handleTOTPSetupGet generates a new TOTP secret and QR code for the current user.
// GET /api/auth/totp/setup
func (s *Server) handleTOTPSetupGet(w http.ResponseWriter, r *http.Request) {
	sess := sessionFromCtx(r.Context())
	if sess.UserID == 0 {
		writeError(w, http.StatusUnauthorized, "not authenticated")
		return
	}

	secret, otpauthURL, qrPNG, err := auth.GenerateSecret(sess.Username, "archivmail")
	if err != nil {
		s.logger.Error("totp setup: generate secret failed", "err", err)
		writeError(w, http.StatusInternalServerError, "failed to generate TOTP secret")
		return
	}

	// Encrypt the secret with AES-256-GCM before storing
	encryptedSecret, err := s.authMgr.EncryptAES([]byte(secret))
	if err != nil {
		s.logger.Error("totp setup: encrypt secret failed", "err", err)
		writeError(w, http.StatusInternalServerError, "failed to encrypt TOTP secret")
		return
	}

	// Store encrypted secret in DB (not yet activated)
	if err := s.users.SetTOTPSecret(r.Context(), sess.UserID, encryptedSecret); err != nil {
		s.logger.Error("totp setup: store secret failed", "err", err)
		writeError(w, http.StatusInternalServerError, "failed to store TOTP secret")
		return
	}

	resp := map[string]interface{}{
		"secret":     secret,
		"otpauth_url": otpauthURL,
	}
	if len(qrPNG) > 0 {
		resp["qr_code"] = base64.StdEncoding.EncodeToString(qrPNG)
	}

	writeJSON(w, http.StatusOK, resp)
}

// handleTOTPSetupPost confirms TOTP setup by verifying a code, then activates TOTP.
// POST /api/auth/totp/setup  { "code": "123456" }
func (s *Server) handleTOTPSetupPost(w http.ResponseWriter, r *http.Request) {
	sess := sessionFromCtx(r.Context())
	if sess.UserID == 0 {
		writeError(w, http.StatusUnauthorized, "not authenticated")
		return
	}

	var req struct {
		Code string `json:"code"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.Code == "" {
		writeError(w, http.StatusBadRequest, "missing or invalid code")
		return
	}

	// Load encrypted secret from DB
	encSecret, _, err := s.users.GetTOTPSecret(r.Context(), sess.UserID)
	if err != nil || len(encSecret) == 0 {
		writeError(w, http.StatusBadRequest, "no TOTP secret found, run setup first")
		return
	}

	// Decrypt
	plainSecret, err := s.authMgr.DecryptAESForHandler(encSecret)
	if err != nil {
		s.logger.Error("totp setup confirm: decrypt failed", "err", err)
		writeError(w, http.StatusInternalServerError, "failed to decrypt TOTP secret")
		return
	}

	// Validate code
	if !auth.ValidateTOTP(string(plainSecret), req.Code) {
		writeError(w, http.StatusBadRequest, "invalid TOTP code")
		return
	}

	// Activate TOTP
	if err := s.users.EnableTOTP(r.Context(), sess.UserID); err != nil {
		s.logger.Error("totp setup confirm: enable failed", "err", err)
		writeError(w, http.StatusInternalServerError, "failed to enable TOTP")
		return
	}

	s.audlog.Log(audit.Entry{
		EventType: audit.EventUserMgmt,
		Username:  sess.Username,
		IPAddress: remoteIP(r),
		Detail:    "totp_enabled",
		Success:   true,
	})

	writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
}

// handleTOTPDisable deactivates TOTP for the current user (requires a valid code).
// DELETE /api/auth/totp  { "code": "123456" }
func (s *Server) handleTOTPDisable(w http.ResponseWriter, r *http.Request) {
	sess := sessionFromCtx(r.Context())
	if sess.UserID == 0 {
		writeError(w, http.StatusUnauthorized, "not authenticated")
		return
	}

	var req struct {
		Code string `json:"code"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.Code == "" {
		writeError(w, http.StatusBadRequest, "missing or invalid code")
		return
	}

	// Load and decrypt secret
	encSecret, enabled, err := s.users.GetTOTPSecret(r.Context(), sess.UserID)
	if err != nil || !enabled || len(encSecret) == 0 {
		writeError(w, http.StatusBadRequest, "TOTP is not enabled")
		return
	}

	plainSecret, err := s.authMgr.DecryptAESForHandler(encSecret)
	if err != nil {
		s.logger.Error("totp disable: decrypt failed", "err", err)
		writeError(w, http.StatusInternalServerError, "failed to decrypt TOTP secret")
		return
	}

	if !auth.ValidateTOTP(string(plainSecret), req.Code) {
		writeError(w, http.StatusBadRequest, "invalid TOTP code")
		return
	}

	if err := s.users.DisableTOTP(r.Context(), sess.UserID); err != nil {
		s.logger.Error("totp disable: failed", "err", err)
		writeError(w, http.StatusInternalServerError, "failed to disable TOTP")
		return
	}

	s.audlog.Log(audit.Entry{
		EventType: audit.EventUserMgmt,
		Username:  sess.Username,
		IPAddress: remoteIP(r),
		Detail:    "totp_disabled",
		Success:   true,
	})

	writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
}

// handleTOTPLogin completes a TOTP-pending login by validating the code.
// POST /api/auth/totp  { "session_token": "...", "code": "123456" }
func (s *Server) handleTOTPLogin(w http.ResponseWriter, r *http.Request) {
	var req struct {
		SessionToken string `json:"session_token"`
		Code         string `json:"code"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.SessionToken == "" || req.Code == "" {
		writeError(w, http.StatusBadRequest, "missing session_token or code")
		return
	}

	token, user, err := s.authMgr.ValidateTOTPLogin(req.SessionToken, req.Code)
	if err != nil {
		s.audlog.Log(audit.Entry{
			EventType: audit.EventLogin,
			IPAddress: remoteIP(r),
			Success:   false,
			Detail:    "totp_login_failed: " + err.Error(),
		})
		writeError(w, http.StatusUnauthorized, "invalid TOTP code or expired session")
		return
	}

	_ = s.users.UpdateLastLogin(user.ID)

	s.audlog.Log(audit.Entry{
		EventType: audit.EventLogin,
		Username:  user.Username,
		IPAddress: remoteIP(r),
		Success:   true,
		Detail:    "totp_login_completed",
	})

	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookieName,
		Value:    token,
		Path:     "/",
		MaxAge:   8 * 3600,
		HttpOnly: true,
		SameSite: http.SameSiteStrictMode,
	})

	writeJSON(w, http.StatusOK, map[string]interface{}{
		"user": map[string]interface{}{
			"id":       user.ID,
			"username": user.Username,
			"email":    user.Email,
			"role":     user.Role,
		},
	})
}

// handleTOTPReset allows an admin to reset TOTP for a user.
// POST /api/admin/users/{id}/totp/reset
func (s *Server) handleTOTPReset(w http.ResponseWriter, r *http.Request) {
	id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid user id")
		return
	}

	sess := sessionFromCtx(r.Context())

	// Fetch target user
	target, err := s.users.GetByID(id)
	if err != nil {
		writeError(w, http.StatusNotFound, "user not found")
		return
	}

	// SEC-02: Tenant isolation — domain_admin can only reset TOTP for users in their own tenant.
	if sess.TenantID != nil {
		if target.TenantID == nil || *target.TenantID != *sess.TenantID {
			writeError(w, http.StatusForbidden, "access denied")
			return
		}
	}

	// Reset TOTP
	if err := s.users.ResetTOTP(r.Context(), id, sess.Username); err != nil {
		s.logger.Error("totp reset: failed", "err", err, "target_user", id, "admin", sess.Username)
		writeError(w, http.StatusInternalServerError, "failed to reset TOTP")
		return
	}

	s.audlog.Log(audit.Entry{
		EventType: audit.EventUserMgmt,
		Username:  sess.Username,
		IPAddress: remoteIP(r),
		Detail:    fmt.Sprintf("totp_reset_by_admin: TOTP reset by %s for user %s (id=%d)", sess.Username, target.Username, id),
		Success:   true,
	})

	writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
}