cpolis/cmd/data/db.go

package data

import (
	"database/sql"
	"fmt"

	"github.com/go-sql-driver/mysql"
	"golang.org/x/crypto/bcrypt"
)

type DB struct {
	*sql.DB
}

func OpenDB(dbName string) (*DB, error) {
	var err error
	db := DB{DB: &sql.DB{}}

	cfg := mysql.NewConfig()
	cfg.DBName = dbName
	cfg.User, cfg.Passwd, err = getCredentials()
	if err != nil {
		return nil, fmt.Errorf("error reading user credentials for DB: %v", err)
	}

	db.DB, err = sql.Open("mysql", cfg.FormatDSN())
	if err != nil {
		return nil, fmt.Errorf("error opening DB: %v", err)
	}
	if err = db.Ping(); err != nil {
		return nil, fmt.Errorf("error pinging DB: %v", err)
	}

	return &db, nil
}

func (db *DB) AddUser(user User, pass string) error {
	hashedPass, err := bcrypt.GenerateFromPassword([]byte(pass), bcrypt.DefaultCost)
	if err != nil {
		return fmt.Errorf("error creating password hash: %v", err)
	}

	query := `
    INSERT INTO users
        (username, password, first_name, last_name, role)
    VALUES
        (?, ?, ?, ?, ?)
    `
	_, err = db.Exec(query, user.UserName, string(hashedPass), user.FirstName, user.LastName, user.Role)
	if err != nil {
		return fmt.Errorf("error inserting user into DB: %v", err)
	}

	return nil
}

func (db *DB) GetID(userName string) (int64, error) {
	var id int64

	query := `
    SELECT id FROM
        users
    WHERE
        username = ?
    `
	row := db.QueryRow(query, userName)
	if err := row.Scan(&id); err != nil {
		return 0, fmt.Errorf("user not in DB: %v", err)
	}

	return id, nil
}

func (db *DB) CheckPassword(id int64, pass string) error {
	var queriedPass string

	query := `
    SELECT password FROM
        users
    WHERE
        id = ?
    `
	row := db.QueryRow(query, id)
	if err := row.Scan(&queriedPass); err != nil {
		return fmt.Errorf("error reading password from DB: %v", err)
	}

	if err := bcrypt.CompareHashAndPassword([]byte(queriedPass), []byte(pass)); err != nil {
		return fmt.Errorf("incorrect password: %v", err)
	}

	return nil
}

func (db *DB) ChangePassword(id int64, oldPass, newPass string) error {
	if err := db.CheckPassword(id, oldPass); err != nil {
		return fmt.Errorf("error checking password: %v", err)
	}

	newHashedPass, err := bcrypt.GenerateFromPassword([]byte(newPass), bcrypt.DefaultCost)
	if err != nil {
		return fmt.Errorf("error creating password hash: %v", err)
	}

	query := `
    UPDATE users SET
        password = ?
    WHERE
        id = ?
    `
	_, err = db.Exec(query, string(newHashedPass), id)
	if err != nil {
		return fmt.Errorf("error updating password in DB: %v", err)
	}

	return nil
}

func (db *DB) CountEntries() (int64, error) {
	var count int64

	query := `SELECT COUNT(*) FROM users`
	row := db.QueryRow(query)
	if err := row.Scan(&count); err != nil {
		return 0, fmt.Errorf("error counting rows in user DB: %v", err)
	}

	return count, nil
}