am-downloader/auth.go

package main

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/crypto/bcrypt"
)

const sessionCookie = "am_session"

var loginLimiter = struct {
	sync.Mutex
	attempts map[string][]time.Time
}{attempts: map[string][]time.Time{}}

func loginAllowed(key string) bool {
	now := time.Now()
	loginLimiter.Lock()
	defer loginLimiter.Unlock()
	items := loginLimiter.attempts[key]
	cutoff := now.Add(-15 * time.Minute)
	kept := items[:0]
	for _, t := range items {
		if t.After(cutoff) {
			kept = append(kept, t)
		}
	}
	if len(kept) >= 10 {
		return false
	}
	loginLimiter.attempts[key] = append(kept, now)
	return true
}

func hashToken(token string) string {
	sum := sha256.Sum256([]byte(token))
	return hex.EncodeToString(sum[:])
}

func randomToken(n int) (string, error) {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}

func hashPassword(password string) (string, error) {
	b, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	return string(b), err
}

func checkPassword(hash, password string) error {
	return bcrypt.CompareHashAndPassword([]byte(hash), []byte(password))
}

func (a *App) currentUser(c *gin.Context) *User {
	value, ok := c.Get("user")
	if !ok {
		return nil
	}
	u, _ := value.(*User)
	return u
}

func (a *App) requireAuth(c *gin.Context) {
	token, err := c.Cookie(sessionCookie)
	if err != nil || token == "" {
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "authentication required"})
		return
	}
	var session Session
	if a.db.Where("token_hash = ? AND expires_at > ?", hashToken(token), time.Now()).First(&session).Error != nil {
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid or expired session"})
		return
	}
	var user User
	if a.db.First(&user, session.UserID).Error != nil || user.Status != UserActive {
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "user unavailable"})
		return
	}
	c.Set("user", &user)
	c.Next()
}

func (a *App) requireAdmin(c *gin.Context) {
	u := a.currentUser(c)
	if u == nil || u.Role != RoleAdmin {
		c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "admin access required"})
		return
	}
	c.Next()
}

func (a *App) requirePermission(permission string) gin.HandlerFunc {
	return func(c *gin.Context) {
		u := a.currentUser(c)
		if u == nil {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "authentication required"})
			return
		}
		if u.Role == RoleAdmin || a.hasPermission(u.ID, permission) {
			c.Next()
			return
		}
		c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "permission denied"})
	}
}

func (a *App) hasPermission(userID uint, permission string) bool {
	var p UserPermission
	return a.db.Where("user_id = ? AND permission = ? AND enabled = ?", userID, permission, true).First(&p).Error == nil
}

func (a *App) login(c *gin.Context) {
	var in struct {
		Username string `json:"username"`
		Password string `json:"password"`
	}
	if c.BindJSON(&in) != nil || strings.TrimSpace(in.Username) == "" || in.Password == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "username and password required"})
		return
	}
	key := c.ClientIP() + ":" + strings.ToLower(strings.TrimSpace(in.Username))
	if !loginAllowed(key) {
		c.JSON(http.StatusTooManyRequests, gin.H{"error": "too many login attempts; try again later"})
		return
	}
	var user User
	if a.db.Where("username = ?", strings.TrimSpace(in.Username)).First(&user).Error != nil || user.Status != UserActive || checkPassword(user.PasswordHash, in.Password) != nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid username or password"})
		return
	}
	token, err := randomToken(32)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "session generation failed"})
		return
	}
	session := Session{UserID: user.ID, TokenHash: hashToken(token), ExpiresAt: time.Now().Add(7 * 24 * time.Hour)}
	if err := a.db.Create(&session).Error; err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "session creation failed"})
		return
	}
	http.SetCookie(c.Writer, &http.Cookie{Name: sessionCookie, Value: token, Path: "/", HttpOnly: true, SameSite: http.SameSiteLaxMode, Secure: c.Request.TLS != nil, Expires: session.ExpiresAt})
	user.PasswordHash = ""
	c.JSON(http.StatusOK, gin.H{"user": user})
}

func (a *App) logout(c *gin.Context) {
	if token, err := c.Cookie(sessionCookie); err == nil && token != "" {
		a.db.Where("token_hash = ?", hashToken(token)).Delete(&Session{})
	}
	http.SetCookie(c.Writer, &http.Cookie{Name: sessionCookie, Value: "", Path: "/", MaxAge: -1, HttpOnly: true, SameSite: http.SameSiteLaxMode})
	c.Status(http.StatusNoContent)
}

func (a *App) me(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{"user": a.currentUser(c)})
}

func (a *App) bootstrapAdmin() {
	var count int64
	a.db.Model(&User{}).Count(&count)
	if count != 0 {
		return
	}
	username := getenv("ADMIN_USERNAME", "admin")
	password := getenv("ADMIN_PASSWORD", "")
	if password == "" {
		generated, err := randomToken(12)
		if err != nil {
			panic(err)
		}
		password = generated
		fmt.Printf("AM Downloader initial admin password: %s\n", password)
	}
	hash, err := hashPassword(password)
	if err != nil {
		panic(err)
	}
	u := User{Username: username, PasswordHash: hash, Role: RoleAdmin, Status: UserActive}
	if err := a.db.Create(&u).Error; err != nil {
		panic(err)
	}
	if err := a.db.Model(&Task{}).Where("user_id = 0").Update("user_id", u.ID).Error; err != nil {
		panic(err)
	}
	if err := a.db.Create(&CreditAccount{UserID: u.ID}).Error; err != nil {
		panic(err)
	}
}