MiauInv/handlers/account.go

package handlers

import (
	"MiauInv/auth"
	"MiauInv/models"
	"MiauInv/storage"
	utils "MiauInv/util"
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"image/png"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/pquerna/otp/totp"
)

func APIRegister(w http.ResponseWriter, r *http.Request) {
	var user models.User
	if err := json.NewDecoder(r.Body).Decode(&user); err != nil {
		log.Println("POST [api/register] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request body", http.StatusBadRequest)
		return
	}

	if user.Username == "" || user.Password == "" {
		log.Println("POST [api/register] " + r.RemoteAddr + ": Username or Password is empty")
		http.Error(w, "username and password required", http.StatusBadRequest)
		return
	}

	if len(user.Password) > 72 {
		log.Println("POST [api/register] User password too long")
		http.Error(w, "Password exceeds the maximum allowed length of 72 characters", http.StatusUnprocessableEntity)
		return
	}

	hashed, err := auth.HashPassword(user.Password)
	if err != nil {
		log.Println("POST [api/register] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "internal error", http.StatusInternalServerError)
		return
	}
	user.Password = hashed
	user.ID = utils.GenerateUUID()
	user.Role = models.RoleUser

	if err := storage.AddUser(&user); err != nil {
		log.Println("POST [api/register] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "user already exists", http.StatusBadRequest)
		return
	}

	w.WriteHeader(http.StatusCreated)
	log.Println("POST [api/register] " + r.RemoteAddr + ": Successfully created user")
}

func APILogin(w http.ResponseWriter, r *http.Request) {
	var creds struct {
		Username string `json:"username"`
		Password string `json:"password"`
	}
	if err := json.NewDecoder(r.Body).Decode(&creds); err != nil {
		log.Println("POST [api/login] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	user, err := storage.GetUserByUsername(creds.Username)
	if err != nil {
		log.Println("POST [api/login] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid credentials", http.StatusUnauthorized)
		return
	}

	if !auth.CheckPasswordHash(creds.Password, user.Password) {
		log.Println("POST [api/login] " + r.RemoteAddr + ": Invalid credentials")
		http.Error(w, "Invalid credentials", http.StatusUnauthorized)
		return
	}

	secret := []byte(os.Getenv("JWT_SECRET"))
	if len(secret) == 0 {
		log.Println("POST [api/login] " + r.RemoteAddr + ": Server misconfiguration")
		http.Error(w, "Server misconfiguration", http.StatusInternalServerError)
		return
	}

	if user.TwoFactorEnabled {
		twoFactorToken, err := auth.GeneratePurposeJWT(user.ID, "2fa_login", secret, 5*time.Minute)
		if err != nil {
			log.Println("POST [api/login] " + r.RemoteAddr + ": " + err.Error())
			http.Error(w, "Could not generate 2FA challenge", http.StatusInternalServerError)
			return
		}

		writeJSON(w, http.StatusOK, map[string]interface{}{
			"requires_2fa":     true,
			"two_factor_token": twoFactorToken,
		})
		log.Println("POST [api/login] " + r.RemoteAddr + ": Password accepted, 2FA required")
		return
	}

	issueLoginSession(w, r, user)
	log.Println("POST [api/login] " + r.RemoteAddr + ": Successfully logged in")
}

func APILoginTwoFactor(w http.ResponseWriter, r *http.Request) {
	var req struct {
		TwoFactorToken string `json:"two_factor_token"`
		Code           string `json:"code"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		log.Println("POST [api/login/2fa] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	secret := []byte(os.Getenv("JWT_SECRET"))
	claims, err := auth.ValidatePurposeJWT(req.TwoFactorToken, "2fa_login", secret)
	if err != nil {
		log.Println("POST [api/login/2fa] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid or expired 2FA challenge", http.StatusUnauthorized)
		return
	}

	user, err := storage.GetUserById(claims.UserID)
	if err != nil || !user.TwoFactorEnabled || user.TwoFactorSecret == "" {
		log.Println("POST [api/login/2fa] " + r.RemoteAddr + ": 2FA not available for user")
		http.Error(w, "Invalid 2FA state", http.StatusUnauthorized)
		return
	}

	code := strings.TrimSpace(req.Code)
	validTOTP := totp.Validate(code, user.TwoFactorSecret)
	usedRecoveryCode := false

	if !validTOTP {
		recoveryCodeHash := utils.HashToken(normalizeRecoveryCode(code))
		usedRecoveryCode, err = storage.UseUserRecoveryCode(user.ID, recoveryCodeHash)
		if err != nil {
			log.Println("POST [api/login/2fa] " + r.RemoteAddr + ": " + err.Error())
			http.Error(w, "Could not validate recovery code", http.StatusInternalServerError)
			return
		}
	}

	if !validTOTP && !usedRecoveryCode {
		log.Println("POST [api/login/2fa] " + r.RemoteAddr + ": Invalid 2FA or recovery code")
		http.Error(w, "Invalid 2FA or recovery code", http.StatusUnauthorized)
		return
	}

	issueLoginSession(w, r, user)
	if usedRecoveryCode {
		log.Println("POST [api/login/2fa] " + r.RemoteAddr + ": Successfully logged in with recovery code")
		return
	}
	log.Println("POST [api/login/2fa] " + r.RemoteAddr + ": Successfully logged in with 2FA")
}

func TwoFactorSetup(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	user, err := storage.GetUserById(claims.UserID)
	if err != nil {
		log.Println("POST [api/2fa/setup] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "User not found", http.StatusNotFound)
		return
	}

	if user.TwoFactorEnabled {
		http.Error(w, "2FA is already enabled", http.StatusConflict)
		return
	}

	key, err := totp.Generate(totp.GenerateOpts{
		Issuer:      "MiauInv",
		AccountName: user.Username,
		SecretSize:  20,
	})
	if err != nil {
		log.Println("POST [api/2fa/setup] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not generate 2FA secret", http.StatusInternalServerError)
		return
	}

	if err := storage.SetUserTwoFactorSecret(user.ID, key.Secret()); err != nil {
		log.Println("POST [api/2fa/setup] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not save 2FA secret", http.StatusInternalServerError)
		return
	}

	img, err := key.Image(220, 220)
	if err != nil {
		log.Println("POST [api/2fa/setup] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not generate QR code", http.StatusInternalServerError)
		return
	}

	var qr bytes.Buffer
	if err := png.Encode(&qr, img); err != nil {
		log.Println("POST [api/2fa/setup] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not encode QR code", http.StatusInternalServerError)
		return
	}

	writeJSON(w, http.StatusOK, map[string]interface{}{
		"secret":      key.Secret(),
		"otpauth_url": key.URL(),
		"qr_code":     "data:image/png;base64," + base64.StdEncoding.EncodeToString(qr.Bytes()),
	})
	log.Println("POST [api/2fa/setup] " + r.RemoteAddr + ": Created 2FA setup challenge")
}

func TwoFactorEnable(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		Code string `json:"code"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		log.Println("POST [api/2fa/enable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	user, err := storage.GetUserById(claims.UserID)
	if err != nil {
		log.Println("POST [api/2fa/enable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "User not found", http.StatusNotFound)
		return
	}

	if user.TwoFactorSecret == "" {
		http.Error(w, "2FA setup has not been started", http.StatusBadRequest)
		return
	}

	if !totp.Validate(strings.TrimSpace(req.Code), user.TwoFactorSecret) {
		http.Error(w, "Invalid 2FA code", http.StatusUnauthorized)
		return
	}

	recoveryCodes, recoveryCodeHashes, err := generateRecoveryCodes(10)
	if err != nil {
		log.Println("POST [api/2fa/enable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not generate recovery codes", http.StatusInternalServerError)
		return
	}

	if err := storage.EnableUserTwoFactorWithRecoveryCodes(user.ID, recoveryCodeHashes); err != nil {
		log.Println("POST [api/2fa/enable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not enable 2FA", http.StatusInternalServerError)
		return
	}

	writeJSON(w, http.StatusOK, map[string]interface{}{
		"two_factor_enabled": true,
		"recovery_codes":     recoveryCodes,
	})
	log.Println("POST [api/2fa/enable] " + r.RemoteAddr + ": Enabled 2FA and generated recovery codes")
}

func TwoFactorDisable(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		Password string `json:"password"`
		Code     string `json:"code"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		log.Println("POST [api/2fa/disable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	user, err := storage.GetUserById(claims.UserID)
	if err != nil {
		log.Println("POST [api/2fa/disable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "User not found", http.StatusNotFound)
		return
	}

	if !auth.CheckPasswordHash(req.Password, user.Password) {
		http.Error(w, "Invalid password", http.StatusUnauthorized)
		return
	}

	if user.TwoFactorEnabled && !totp.Validate(strings.TrimSpace(req.Code), user.TwoFactorSecret) {
		http.Error(w, "Invalid 2FA code", http.StatusUnauthorized)
		return
	}

	if err := storage.DisableUserTwoFactor(user.ID); err != nil {
		log.Println("POST [api/2fa/disable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not disable 2FA", http.StatusInternalServerError)
		return
	}

	if err := storage.RevokeAllRefreshTokensForUser(user.ID); err != nil {
		log.Println("POST [api/2fa/disable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not revoke sessions", http.StatusInternalServerError)
		return
	}

	clearAuthCookies(w)
	writeJSON(w, http.StatusOK, map[string]interface{}{"two_factor_enabled": false})
	log.Println("POST [api/2fa/disable] " + r.RemoteAddr + ": Disabled 2FA")
}

func TwoFactorRegenerateRecoveryCodes(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		Password string `json:"password"`
		Code     string `json:"code"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		log.Println("POST [api/2fa/recovery-codes/regenerate] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	user, err := storage.GetUserById(claims.UserID)
	if err != nil {
		log.Println("POST [api/2fa/recovery-codes/regenerate] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "User not found", http.StatusNotFound)
		return
	}

	if !user.TwoFactorEnabled || user.TwoFactorSecret == "" {
		http.Error(w, "2FA is not enabled", http.StatusBadRequest)
		return
	}

	if !auth.CheckPasswordHash(req.Password, user.Password) {
		http.Error(w, "Invalid password", http.StatusUnauthorized)
		return
	}

	if !totp.Validate(strings.TrimSpace(req.Code), user.TwoFactorSecret) {
		http.Error(w, "Invalid 2FA code", http.StatusUnauthorized)
		return
	}

	recoveryCodes, recoveryCodeHashes, err := generateRecoveryCodes(10)
	if err != nil {
		log.Println("POST [api/2fa/recovery-codes/regenerate] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not generate recovery codes", http.StatusInternalServerError)
		return
	}

	if err := storage.ReplaceUserRecoveryCodes(user.ID, recoveryCodeHashes); err != nil {
		log.Println("POST [api/2fa/recovery-codes/regenerate] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not save recovery codes", http.StatusInternalServerError)
		return
	}

	writeJSON(w, http.StatusOK, map[string]interface{}{
		"recovery_codes": recoveryCodes,
	})
	log.Println("POST [api/2fa/recovery-codes/regenerate] " + r.RemoteAddr + ": Regenerated recovery codes")
}

func Logout(w http.ResponseWriter, r *http.Request) {
	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	err := storage.RevokeAllRefreshTokensForUser(claims.UserID)
	if err != nil {
		log.Println("GET [api/logout] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}
	clearAuthCookies(w)
	w.WriteHeader(http.StatusNoContent)
}

func TestHandler(w http.ResponseWriter, r *http.Request) {
	claims, _ := utils.IsLoggedIn(w, r)

	w.Header().Set("Content-Type", "application/json")
	err := json.NewEncoder(w).Encode(map[string]interface{}{
		"user_id": claims.UserID,
		"role":    claims.Role,
		"msg":     "Authentication successful",
	})
	if err != nil {
		log.Println("GET [api/ping] " + r.RemoteAddr + ": " + err.Error())
		return
	}
	log.Println("GET [api/ping] " + r.RemoteAddr + ": Successfully tested connection")
}

func RefreshToken(w http.ResponseWriter, r *http.Request) {
	var req struct {
		RefreshToken string `json:"refresh_token"`
	}

	if r.Body != nil {
		_ = json.NewDecoder(r.Body).Decode(&req)
	}
	if req.RefreshToken == "" {
		cookie, err := r.Cookie("refresh_token")
		if err == nil {
			req.RefreshToken = cookie.Value
		}
	}
	if req.RefreshToken == "" {
		log.Println("POST [api/refresh] " + r.RemoteAddr + ": Missing refresh token")
		http.Error(w, "Invalid refresh token", http.StatusUnauthorized)
		return
	}

	hashed := utils.HashToken(req.RefreshToken)

	tokenRow, err := storage.GetRefreshToken(hashed)
	if err != nil || tokenRow.Revoked || tokenRow.ExpiresAt < time.Now().Unix() {
		log.Println("POST [api/refresh] " + r.RemoteAddr + ": Invalid refresh token")
		http.Error(w, "Invalid refresh token", http.StatusUnauthorized)
		return
	}

	if err := storage.RevokeRefreshToken(tokenRow.ID); err != nil {
		log.Println(err)
	}

	user, err := storage.GetUserById(tokenRow.UserID)
	if err != nil {
		log.Println("POST [api/refresh] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}

	issueLoginSession(w, r, user)
	log.Println("POST [api/refresh] " + r.RemoteAddr + ": Successfully refreshed token")
}

func UserInfo(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		log.Println("GET [api/userinfo] " + r.RemoteAddr + ": Method " + r.Method + " not allowed")
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	query := r.URL.Query()
	idParam := query.Get("id")

	if idParam == "" {
		tokenStr := ""
		authHeader := r.Header.Get("Authorization")

		if strings.HasPrefix(authHeader, "Bearer ") {
			tokenStr = strings.TrimPrefix(authHeader, "Bearer ")
		}

		if tokenStr == "" {
			cookie, err := r.Cookie("access_token")
			if err == nil {
				tokenStr = cookie.Value
			}
		}

		if tokenStr == "" {
			http.Error(w, "Missing token", http.StatusUnauthorized)
			return
		}

		claims, err := auth.ValidateJWT(tokenStr, models.JWTSecret)

		if err != nil {
			log.Println("GET [api/userinfo] " + r.RemoteAddr + ": " + err.Error())
			http.Error(w, "Unauthorized: Invalid or expired token", http.StatusUnauthorized)
			return
		}

		idParam = claims.UserID
	}

	user, err := storage.GetUserById(idParam)
	if err != nil {
		log.Println("GET [api/userinfo] " + r.RemoteAddr + ": User " + idParam + " not found")
		http.Error(w, "User not found", http.StatusNotFound)
		return
	}

	recoveryCodesRemaining := 0
	if user.TwoFactorEnabled {
		if count, err := storage.CountUnusedRecoveryCodes(user.ID); err == nil {
			recoveryCodesRemaining = count
		}
	}

	w.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w).Encode(map[string]interface{}{
		"id":                       user.ID,
		"username":                 user.Username,
		"avatar_url":               "",
		"two_factor_enabled":       user.TwoFactorEnabled,
		"recovery_codes_remaining": recoveryCodesRemaining,
	})
	if err != nil {
		log.Println("GET [api/userinfo] " + r.RemoteAddr + ": " + err.Error())
		return
	}
	log.Println("GET [api/userinfo] " + r.RemoteAddr + ": Successfully retrieved user info of " + user.Username + " (" + user.ID + ")")
}

func issueLoginSession(w http.ResponseWriter, r *http.Request, user models.User) {
	secret := []byte(os.Getenv("JWT_SECRET"))
	if len(secret) == 0 {
		log.Println("AUTH " + r.RemoteAddr + ": Server misconfiguration")
		http.Error(w, "Server misconfiguration", http.StatusInternalServerError)
		return
	}

	accessToken, err := auth.GenerateJWT(user.ID, user.Role, secret)
	if err != nil {
		log.Println("AUTH " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not generate token", http.StatusInternalServerError)
		return
	}

	refreshTokenPlain, err := utils.GenerateRefreshToken()
	if err != nil {
		log.Println("AUTH " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "could not generate refresh token", http.StatusInternalServerError)
		return
	}

	refreshExpires := time.Now().Add(7 * 24 * time.Hour).Unix()
	if err := storage.AddRefreshToken(&models.RefreshToken{
		ID:         utils.GenerateUUID(),
		UserID:     user.ID,
		Token:      utils.HashToken(refreshTokenPlain),
		ExpiresAt:  refreshExpires,
		DeviceInfo: r.Header.Get("User-Agent"),
		CreatedAt:  time.Now().Unix(),
		Revoked:    false,
	}); err != nil {
		log.Println("AUTH " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "could not save refresh token", http.StatusInternalServerError)
		return
	}

	setAuthCookies(w, accessToken, refreshTokenPlain)
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"access_token":  accessToken,
		"refresh_token": refreshTokenPlain,
		"user": map[string]interface{}{
			"id":                 user.ID,
			"username":           user.Username,
			"role":               user.Role,
			"two_factor_enabled": user.TwoFactorEnabled,
		},
	})
}

func generateRecoveryCodes(count int) ([]string, []string, error) {
	codes := make([]string, 0, count)
	hashes := make([]string, 0, count)
	seen := make(map[string]struct{}, count)

	for len(codes) < count {
		code, err := generateRecoveryCode()
		if err != nil {
			return nil, nil, err
		}

		normalized := normalizeRecoveryCode(code)
		if _, exists := seen[normalized]; exists {
			continue
		}
		seen[normalized] = struct{}{}

		codes = append(codes, code)
		hashes = append(hashes, utils.HashToken(normalized))
	}

	return codes, hashes, nil
}

func generateRecoveryCode() (string, error) {
	bytes := make([]byte, 10)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}

	raw := hex.EncodeToString(bytes)
	return raw[0:5] + "-" + raw[5:10] + "-" + raw[10:15] + "-" + raw[15:20], nil
}

func normalizeRecoveryCode(code string) string {
	code = strings.TrimSpace(code)
	code = strings.ReplaceAll(code, "-", "")
	code = strings.ReplaceAll(code, " ", "")
	return strings.ToLower(code)
}

func setAuthCookies(w http.ResponseWriter, accessToken, refreshToken string) {
	http.SetCookie(w, &http.Cookie{
		Name:     "access_token",
		Value:    accessToken,
		Path:     "/",
		MaxAge:   15 * 60,
		HttpOnly: true,
		Secure:   true,
		SameSite: http.SameSiteLaxMode,
	})
	http.SetCookie(w, &http.Cookie{
		Name:     "refresh_token",
		Value:    refreshToken,
		Path:     "/",
		MaxAge:   7 * 24 * 60 * 60,
		HttpOnly: true,
		Secure:   true,
		SameSite: http.SameSiteLaxMode,
	})
}

func clearAuthCookies(w http.ResponseWriter) {
	for _, name := range []string{"access_token", "refresh_token"} {
		http.SetCookie(w, &http.Cookie{
			Name:     name,
			Value:    "",
			Path:     "/",
			MaxAge:   -1,
			HttpOnly: true,
			Secure:   true,
			SameSite: http.SameSiteLaxMode,
		})
	}
}

func writeJSON(w http.ResponseWriter, status int, data interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	if err := json.NewEncoder(w).Encode(data); err != nil {
		log.Println("JSON response error: " + err.Error())
	}
}