drive/backend/handlers/auth.go

package handlers

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"strings"
	"crypto/subtle"
	"sync"

	"git.elijahkuntz.com/Elijah/drive/config"
	"git.elijahkuntz.com/Elijah/drive/database"
	"git.elijahkuntz.com/Elijah/drive/middleware"
	"github.com/gofiber/fiber/v2"
	"github.com/pquerna/otp/totp"
	"golang.org/x/crypto/argon2"
)

type AuthHandler struct {
	DB     *database.DB
	Config *config.Config
	Guard  *middleware.BruteForceGuard
	
	tempTOTP map[string]string
	totpMu   sync.Mutex
}

func NewAuthHandler(db *database.DB, cfg *config.Config, guard *middleware.BruteForceGuard) *AuthHandler {
	return &AuthHandler{
		DB:       db,
		Config:   cfg,
		Guard:    guard,
		tempTOTP: make(map[string]string),
	}
}

type SetupRequest struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

type LoginRequest struct {
	Username string `json:"username"`
	Password string `json:"password"`
	TOTPCode string `json:"totp_code,omitempty"`
}

// Setup handles the first-run account creation wizard.
// POST /api/setup
func (h *AuthHandler) Setup(c *fiber.Ctx) error {
	// Check if setup is already complete
	complete, err := h.DB.IsSetupComplete()
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
	}
	if complete {
		return c.Status(fiber.StatusConflict).JSON(fiber.Map{"error": "setup already completed"})
	}

	var req SetupRequest
	if err := c.BodyParser(&req); err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request body"})
	}

	if len(req.Username) < 3 {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "username must be at least 3 characters"})
	}
	if len(req.Password) < 8 {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "password must be at least 8 characters"})
	}

	hash, err := hashPassword(req.Password)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to hash password"})
	}

	_, err = h.DB.Exec(
		"INSERT INTO users (username, password_hash) VALUES (?, ?)",
		req.Username, hash,
	)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to create user"})
	}

	h.DB.AddAuditLog("setup_complete", fmt.Sprintf("User '%s' created", req.Username), c.IP())

	return c.Status(fiber.StatusCreated).JSON(fiber.Map{
		"message": "setup complete, please log in",
	})
}

// CheckSetup returns whether initial setup has been completed.
// GET /api/setup/status
func (h *AuthHandler) CheckSetup(c *fiber.Ctx) error {
	complete, err := h.DB.IsSetupComplete()
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
	}
	return c.JSON(fiber.Map{"setup_complete": complete})
}

// Login authenticates a user and returns JWT tokens.
// POST /api/auth/login
func (h *AuthHandler) Login(c *fiber.Ctx) error {
	var req LoginRequest
	if err := c.BodyParser(&req); err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request body"})
	}

	var id int
	var passwordHash, totpSecret string
	var totpEnabled bool

	err := h.DB.QueryRow(
		"SELECT id, password_hash, totp_secret, totp_enabled FROM users WHERE username = ?",
		req.Username,
	).Scan(&id, &passwordHash, &totpSecret, &totpEnabled)

	if err != nil {
		h.Guard.RecordFailure(c.IP())
		h.DB.AddAuditLog("login_failed", fmt.Sprintf("Unknown user '%s'", req.Username), c.IP())
		return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid credentials"})
	}

	if !verifyPassword(req.Password, passwordHash) {
		h.Guard.RecordFailure(c.IP())
		h.DB.AddAuditLog("login_failed", fmt.Sprintf("Bad password for '%s'", req.Username), c.IP())
		return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid credentials"})
	}

	// Check TOTP if enabled
	if totpEnabled {
		if req.TOTPCode == "" {
			return c.Status(fiber.StatusPreconditionRequired).JSON(fiber.Map{
				"error":      "2fa_required",
				"message":    "Two-factor authentication code required",
				"totp_required": true,
			})
		}
		if !totp.Validate(req.TOTPCode, totpSecret) {
			h.Guard.RecordFailure(c.IP())
			h.DB.AddAuditLog("login_failed", fmt.Sprintf("Bad TOTP for '%s'", req.Username), c.IP())
			return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid 2FA code"})
		}
	}

	// Generate tokens
	accessToken, err := middleware.GenerateAccessToken(id, req.Username, h.Config.JWTSecret, h.Config.JWTExpirySecs)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate token"})
	}

	refreshToken, err := middleware.GenerateRefreshToken(id, req.Username, h.Config.JWTSecret)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate refresh token"})
	}

	h.Guard.RecordSuccess(c.IP())
	h.DB.AddAuditLog("login_success", fmt.Sprintf("User '%s' logged in", req.Username), c.IP())

	return c.JSON(fiber.Map{
		"access_token":  accessToken,
		"refresh_token": refreshToken,
		"expires_in":    h.Config.JWTExpirySecs,
	})
}

// Refresh issues a new access token from a valid refresh token.
// POST /api/auth/refresh
func (h *AuthHandler) Refresh(c *fiber.Ctx) error {
	// The auth middleware already validated the token
	userID := c.Locals("userID")
	username := c.Locals("username")

	accessToken, err := middleware.GenerateAccessToken(
		int(userID.(float64)),
		username.(string),
		h.Config.JWTSecret,
		h.Config.JWTExpirySecs,
	)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate token"})
	}

	return c.JSON(fiber.Map{
		"access_token": accessToken,
		"expires_in":   h.Config.JWTExpirySecs,
	})
}

// GetDownloadToken issues a short-lived token for file downloads
// GET /api/auth/download-token
func (h *AuthHandler) GetDownloadToken(c *fiber.Ctx) error {
	userID := c.Locals("userID")
	username := c.Locals("username")

	token, err := middleware.GenerateDownloadToken(
		int(userID.(float64)),
		username.(string),
		h.Config.JWTSecret,
	)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate download token"})
	}

	return c.JSON(fiber.Map{
		"token": token,
	})
}

// Enable2FA generates a TOTP secret and returns the provisioning URI.
// POST /api/auth/2fa/enable
func (h *AuthHandler) Enable2FA(c *fiber.Ctx) error {
	username := c.Locals("username").(string)

	key, err := totp.Generate(totp.GenerateOpts{
		Issuer:      "Drive",
		AccountName: username,
	})
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate TOTP secret"})
	}

	// Store the secret temporarily in memory
	h.totpMu.Lock()
	h.tempTOTP[username] = key.Secret()
	h.totpMu.Unlock()

	return c.JSON(fiber.Map{
		"secret": key.Secret(),
		"url":    key.URL(),
	})
}

// Confirm2FA verifies a TOTP code and enables 2FA.
// POST /api/auth/2fa/confirm
func (h *AuthHandler) Confirm2FA(c *fiber.Ctx) error {
	username := c.Locals("username").(string)

	var body struct {
		Code string `json:"code"`
	}
	if err := c.BodyParser(&body); err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request"})
	}

	h.totpMu.Lock()
	secret, exists := h.tempTOTP[username]
	h.totpMu.Unlock()

	if !exists {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "no pending 2FA setup found"})
	}

	if !totp.Validate(body.Code, secret) {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid TOTP code"})
	}

	_, err := h.DB.Exec("UPDATE users SET totp_enabled = 1, totp_secret = ? WHERE username = ?", secret, username)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to enable 2FA"})
	}

	h.totpMu.Lock()
	delete(h.tempTOTP, username)
	h.totpMu.Unlock()

	h.DB.AddAuditLog("2fa_enabled", fmt.Sprintf("User '%s' enabled 2FA", username), c.IP())

	return c.JSON(fiber.Map{"message": "2FA enabled successfully"})
}

// Disable2FA turns off two-factor authentication.
// POST /api/auth/2fa/disable
func (h *AuthHandler) Disable2FA(c *fiber.Ctx) error {
	username := c.Locals("username").(string)

	var body struct {
		Code string `json:"code"`
	}
	if err := c.BodyParser(&body); err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request"})
	}

	var secret string
	err := h.DB.QueryRow("SELECT totp_secret FROM users WHERE username = ?", username).Scan(&secret)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
	}

	if !totp.Validate(body.Code, secret) {
		h.Guard.RecordFailure(c.IP())
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid TOTP code"})
	}

	_, err = h.DB.Exec("UPDATE users SET totp_enabled = 0, totp_secret = '' WHERE username = ?", username)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to disable 2FA"})
	}

	h.DB.AddAuditLog("2fa_disabled", fmt.Sprintf("User '%s' disabled 2FA", username), c.IP())

	return c.JSON(fiber.Map{"message": "2FA disabled successfully"})
}

// ChangePassword allows an authenticated user to change their password.
// PUT /api/auth/password
func (h *AuthHandler) ChangePassword(c *fiber.Ctx) error {
	username := c.Locals("username").(string)

	var req struct {
		OldPassword string `json:"old_password"`
		NewPassword string `json:"new_password"`
	}
	if err := c.BodyParser(&req); err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request body"})
	}

	if len(req.NewPassword) < 8 {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "new password must be at least 8 characters"})
	}

	var currentHash string
	if err := h.DB.QueryRow("SELECT password_hash FROM users WHERE username = ?", username).Scan(&currentHash); err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
	}

	if !verifyPassword(req.OldPassword, currentHash) {
		h.Guard.RecordFailure(c.IP())
		return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "incorrect current password"})
	}

	newHash, err := hashPassword(req.NewPassword)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to hash new password"})
	}

	if _, err := h.DB.Exec("UPDATE users SET password_hash = ? WHERE username = ?", newHash, username); err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to update password"})
	}

	h.DB.AddAuditLog("password_changed", fmt.Sprintf("User '%s' changed their password", username), c.IP())
	return c.JSON(fiber.Map{"message": "password changed successfully"})
}

// --- Password Hashing with Argon2 ---

func hashPassword(password string) (string, error) {
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return "", err
	}

	hash := argon2.IDKey([]byte(password), salt, 2, 64*1024, 4, 32)
	// Output PHC format
	encoded := fmt.Sprintf("$argon2id$v=19$m=65536,t=2,p=4$%s$%s",
		base64.RawStdEncoding.EncodeToString(salt),
		base64.RawStdEncoding.EncodeToString(hash),
	)
	return encoded, nil
}

func verifyPassword(password, encoded string) bool {
	if strings.HasPrefix(encoded, "$argon2id$") {
		// PHC Format
		parts := strings.Split(encoded, "$")
		if len(parts) != 6 {
			return false
		}
		salt, err := base64.RawStdEncoding.DecodeString(parts[4])
		if err != nil { return false }
		expectedHash, err := base64.RawStdEncoding.DecodeString(parts[5])
		if err != nil { return false }

		hash := argon2.IDKey([]byte(password), salt, 2, 64*1024, 4, 32)
		return subtle.ConstantTimeCompare(hash, expectedHash) == 1
	}

	// Legacy format: base64(salt)$base64(hash)
	parts := strings.SplitN(encoded, "$", 2)
	if len(parts) != 2 {
		return false
	}

	salt, err := base64.RawStdEncoding.DecodeString(parts[0])
	if err != nil {
		return false
	}

	expectedHash, err := base64.RawStdEncoding.DecodeString(parts[1])
	if err != nil {
		return false
	}

	hash := argon2.IDKey([]byte(password), salt, 2, 64*1024, 4, 32)
	return subtle.ConstantTimeCompare(hash, expectedHash) == 1
}