rmm2/backend/db/users.go

package db

import (
	"database/sql"

	"github.com/cynfo/rmm-backend/models"
	"golang.org/x/crypto/bcrypt"
)

func (d *Database) CreateUser(username, password, displayName, role string) (*models.User, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return nil, err
	}

	var u models.User
	err = d.db.QueryRow(
		"INSERT INTO users (username, password_hash, display_name, role) VALUES ($1, $2, $3, $4) RETURNING id, username, display_name, role, created_at",
		username, string(hash), displayName, role,
	).Scan(&u.ID, &u.Username, &u.DisplayName, &u.Role, &u.CreatedAt)
	if err != nil {
		return nil, err
	}
	return &u, nil
}

func (d *Database) GetUserByUsername(username string) (*models.User, error) {
	var u models.User
	var lastLogin sql.NullTime
	var totpSecret sql.NullString
	err := d.db.QueryRow(
		"SELECT id, username, password_hash, display_name, role, totp_enabled, totp_secret, created_at, last_login FROM users WHERE username = $1",
		username,
	).Scan(&u.ID, &u.Username, &u.PasswordHash, &u.DisplayName, &u.Role, &u.TOTPEnabled, &totpSecret, &u.CreatedAt, &lastLogin)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	if lastLogin.Valid {
		u.LastLogin = &lastLogin.Time
	}
	if totpSecret.Valid {
		u.TOTPSecret = totpSecret.String
	}
	return &u, nil
}

func (d *Database) GetUserByID(id int) (*models.User, error) {
	var u models.User
	var lastLogin sql.NullTime
	var totpSecret sql.NullString
	err := d.db.QueryRow(
		"SELECT id, username, password_hash, display_name, role, totp_enabled, totp_secret, created_at, last_login FROM users WHERE id = $1",
		id,
	).Scan(&u.ID, &u.Username, &u.PasswordHash, &u.DisplayName, &u.Role, &u.TOTPEnabled, &totpSecret, &u.CreatedAt, &lastLogin)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	if lastLogin.Valid {
		u.LastLogin = &lastLogin.Time
	}
	if totpSecret.Valid {
		u.TOTPSecret = totpSecret.String
	}
	return &u, nil
}

func (d *Database) SetTOTPSecret(userID int, secret string) error {
	_, err := d.db.Exec("UPDATE users SET totp_secret = $1, totp_enabled = false WHERE id = $2", secret, userID)
	return err
}

func (d *Database) EnableTOTP(userID int) error {
	_, err := d.db.Exec("UPDATE users SET totp_enabled = true WHERE id = $1", userID)
	return err
}

func (d *Database) DisableTOTP(userID int) error {
	_, err := d.db.Exec("UPDATE users SET totp_enabled = false, totp_secret = NULL WHERE id = $1", userID)
	return err
}

func (d *Database) GetUsers() ([]models.User, error) {
	rows, err := d.db.Query("SELECT id, username, display_name, role, created_at, last_login FROM users ORDER BY username")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var users []models.User
	for rows.Next() {
		var u models.User
		var lastLogin sql.NullTime
		if err := rows.Scan(&u.ID, &u.Username, &u.DisplayName, &u.Role, &u.CreatedAt, &lastLogin); err != nil {
			return nil, err
		}
		if lastLogin.Valid {
			u.LastLogin = &lastLogin.Time
		}
		users = append(users, u)
	}
	if users == nil {
		users = []models.User{}
	}
	return users, rows.Err()
}

func (d *Database) UpdateUserLastLogin(id int) error {
	_, err := d.db.Exec("UPDATE users SET last_login = NOW() WHERE id = $1", id)
	return err
}

func (d *Database) ChangePassword(id int, newPassword string) error {
	hash, err := bcrypt.GenerateFromPassword([]byte(newPassword), bcrypt.DefaultCost)
	if err != nil {
		return err
	}
	_, err = d.db.Exec("UPDATE users SET password_hash = $1 WHERE id = $2", string(hash), id)
	return err
}

func (d *Database) DeleteUser(id int) (bool, error) {
	res, err := d.db.Exec("DELETE FROM users WHERE id = $1", id)
	if err != nil {
		return false, err
	}
	rows, _ := res.RowsAffected()
	return rows > 0, nil
}

func (d *Database) UserCount() (int, error) {
	var count int
	err := d.db.QueryRow("SELECT COUNT(*) FROM users").Scan(&count)
	return count, err
}