wg/internal/auth/auth.go

package auth

import (
	"context"
	"encoding/gob"
	"errors"
	"fmt"
	"net/http"
	"time"

	"github.com/evilfox/wg-panel/internal/config"
	"github.com/evilfox/wg-panel/internal/models"
	"github.com/gorilla/sessions"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"golang.org/x/crypto/bcrypt"
)

func init() {
	gob.Register(int64(0))
}

const SessionName = "wg_session"

type Store struct {
	Sessions *sessions.CookieStore
	DB       *pgxpool.Pool
	Cfg      config.Config
}

func NewStore(cfg config.Config, db *pgxpool.Pool) *Store {
	store := sessions.NewCookieStore([]byte(cfg.SessionSecret))
	store.Options = &sessions.Options{
		Path:     "/",
		HttpOnly: true,
		Secure:   cfg.SecureCookies,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   0,
	}
	return &Store{Sessions: store, DB: db, Cfg: cfg}
}

func (s *Store) EnsureAdmin(ctx context.Context) error {
	if s.Cfg.AdminPassword == "" {
		return errors.New("ADMIN_PASSWORD is required in environment")
	}
	hash, err := bcrypt.GenerateFromPassword([]byte(s.Cfg.AdminPassword), bcrypt.DefaultCost)
	if err != nil {
		return err
	}

	var id int64
	err = s.DB.QueryRow(ctx, `SELECT id FROM users WHERE role = 'admin' ORDER BY id LIMIT 1`).Scan(&id)
	if errors.Is(err, pgx.ErrNoRows) {
		_, err = s.DB.Exec(ctx, `
			INSERT INTO users (username, email, password_hash, role, is_approved)
			VALUES ($1, $2, $3, 'admin', TRUE)
			ON CONFLICT (username) DO UPDATE
			SET email = EXCLUDED.email,
			    password_hash = EXCLUDED.password_hash,
			    role = 'admin',
			    is_approved = TRUE,
			    updated_at = NOW()`,
			s.Cfg.AdminUsername, s.Cfg.AdminEmail, string(hash),
		)
		return err
	}
	if err != nil {
		return err
	}

	_, err = s.DB.Exec(ctx, `
		UPDATE users
		SET username = $1,
		    email = $2,
		    password_hash = $3,
		    role = 'admin',
		    is_approved = TRUE,
		    updated_at = NOW()
		WHERE id = $4`,
		s.Cfg.AdminUsername, s.Cfg.AdminEmail, string(hash), id,
	)
	return err
}

func (s *Store) Authenticate(ctx context.Context, login, password string) (*models.User, error) {
	var u models.User
	err := s.DB.QueryRow(ctx, `
		SELECT id, username, email, password_hash, role, is_approved, created_at
		FROM users
		WHERE username = $1 OR email = $1
		LIMIT 1`, login,
	).Scan(&u.ID, &u.Username, &u.Email, &u.PasswordHash, &u.Role, &u.IsApproved, &u.CreatedAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, errors.New("invalid credentials")
	}
	if err != nil {
		return nil, err
	}
	if !u.IsApproved {
		return nil, errors.New("account pending approval")
	}
	if err := bcrypt.CompareHashAndPassword([]byte(u.PasswordHash), []byte(password)); err != nil {
		return nil, errors.New("invalid credentials")
	}
	return &u, nil
}

func (s *Store) Login(w http.ResponseWriter, r *http.Request, u *models.User, remember bool) error {
	sess, _ := s.Sessions.Get(r, SessionName)
	sess.Values["user_id"] = u.ID
	sess.Values["username"] = u.Username
	sess.Values["role"] = u.Role
	sess.Values["remember"] = remember
	if remember {
		sess.Options.MaxAge = int(s.Cfg.RememberDuration().Seconds())
	} else {
		sess.Options.MaxAge = 0
	}
	return sess.Save(r, w)
}

func (s *Store) Logout(w http.ResponseWriter, r *http.Request) error {
	sess, _ := s.Sessions.Get(r, SessionName)
	sess.Options.MaxAge = -1
	sess.Values = map[interface{}]interface{}{}
	return sess.Save(r, w)
}

func (s *Store) CurrentUser(r *http.Request) *models.SessionUser {
	sess, err := s.Sessions.Get(r, SessionName)
	if err != nil {
		return nil
	}
	id, ok := sess.Values["user_id"].(int64)
	if !ok {
		// gorilla may decode numbers as int
		switch v := sess.Values["user_id"].(type) {
		case int:
			id = int64(v)
			ok = true
		case float64:
			id = int64(v)
			ok = true
		}
	}
	if !ok || id == 0 {
		return nil
	}
	username, _ := sess.Values["username"].(string)
	role, _ := sess.Values["role"].(string)
	return &models.SessionUser{ID: id, Username: username, Role: role}
}

func (s *Store) RequireAdmin(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		u := s.CurrentUser(r)
		if u == nil {
			http.Redirect(w, r, "/login", http.StatusSeeOther)
			return
		}
		if u.Role != "admin" {
			http.Error(w, "Доступ запрещён", http.StatusForbidden)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func HashPassword(password string) (string, error) {
	b, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return "", err
	}
	return string(b), nil
}

func FormatDuration(d time.Duration) string {
	return fmt.Sprintf("%d", int(d.Seconds()))
}