gowiki/pkg/db/user.go

package db

import (
	"github.com/boltdb/bolt"
	"golang.org/x/crypto/bcrypt"
)

type User struct {
	Username string
	Password string
	Admin    bool
	Locked   int
}

var (
	errUserNotFound = dbErr("user not found")
	errUserExists   = dbErr("user already exist")
)

func (bs *BoltStore) GetUsers() (users []User, err error) {
	users, err = bs.DumpUsers()
	for i := range users {
		users[i].Password = ""
	}
	return
}

func (bs *BoltStore) GetUserWithoutPassword(username string) (u User, err error) {
	u, err = bs.GetUser(username)
	u.Password = ""
	return
}

func (bs *BoltStore) GetUser(username string) (u User, err error) {
	err = bs.db.View(func(tx *bolt.Tx) error {
		v := tx.Bucket([]byte(userTable)).Get([]byte(username))
		if v == nil {
			return errUserNotFound
		}
		return bs.Unmarshal(v, &u)
	})
	return
}

func (bs *BoltStore) UserIsAdmin(username string) (admin bool) {
	bs.db.View(func(tx *bolt.Tx) error {
		u, _ := bs.GetUser(username)
		admin = u.Admin
		return nil
	})
	return
}

func (bs *BoltStore) UserExists(username string) error {
	return bs.db.View(func(tx *bolt.Tx) error {
		if tx.Bucket([]byte(userTable)).Get([]byte(username)) == nil {
			return errUserNoFound
		}
		return nil
	})
}

func (bs *BoltStore) DumpUsers() (users []User, err error) {
	var buf [][]byte
	err = bs.db.View(func(tx *bolt.Tx) error {
		return tx.Bucket([]byte(userTable)).ForEach(func(k, v []byte) error {
			buf = append(buf, copyBuf(v))
			return nil
		})
	})
	if err != nil {
		return
	}
	for _, b := range buf {
		u := new(User)
		if err = bs.Unmarshal(b, u); err != nil {
			return
		}
		users = append(users, *u)
	}
	return
}

func (bs *BoltStore) RestoreUsers(users []User) error {
	return bs.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(userTable))
		err := b.ForEach(func(k, v []byte) error {
			return b.Delete(k)
		})
		if err != nil {
			return err
		}
		for _, u := range users {
			v, err := bs.Marshal(u)
			if err != nil {
				return err
			}
			err = b.Put([]byte(u.Username), v)
			if err != nil {
				return err
			}
		}
		return nil
	})
}

func (bs *BoltStore) CreateUser(username, password string, admin bool) error {
	if err := bs.UserExists(username); err == nil {
		return errUserExists
	}
	return bs.db.Update(func(tx *bolt.Tx) error {
		hash, err := bcrypt.GenerateFromPassword(validatePassword(password), bcryptCost)
		v, err := bs.Marshal(&User{
			Username: username,
			Password: string(hash),
			Admin:    admin,
		})
		if err != nil {
			return err
		}
		return tx.Bucket([]byte(userTable)).Put([]byte(username), v)
	})
}

func (bs *BoltStore) Login(username, password string) error {
	var e error
	err := bs.db.Update(func(tx *bolt.Tx) error {
		us, err := bs.GetUser(username)
		if err != nil {
			return err
		}
		if us.Locked == 3 {
			return errors.New("db: user locked")
		}
		b := tx.Bucket([]byte(userTable))
		if err := bcrypt.CompareHashAndPassword([]byte(us.Password), validatePassword(password)); err != nil {
			us.Locked++
			v, err := bs.Marshal(us)
			if err != nil {
				return err
			}
			if err = b.Put([]byte(username), v); err != nil {
				return err
			}
			e = errors.New("db: wrong login")
			return nil
		}
		us.Locked = 0
		v, err := bs.Marshal(us)
		if err != nil {
			return err
		}
		return b.Put([]byte(username), v)
	})
	if e != nil {
		return e
	}
	return err
}

func (bs *BoltStore) UnlockUser(username string) error {
	return bs.db.Update(func(tx *bolt.Tx) error {
		u, err := bs.GetUser(username)
		if err != nil {
			return err
		}
		if u.Locked < 3 {
			return nil
		}
		u.Locked = 0
		v, err := bs.Marshal(u)
		if err != nil {
			return err
		}
		return tx.Bucket([]byte(userTable)).Put([]byte(username), v)
	})
}

func (bs *BoltStore) UpdateUserPassword(username, password string) error {
	return bs.db.Update(func(tx *bolt.Tx) error {
		if password != "" {
			us, err := bs.GetUser(username)
			if err != nil {
				return err
			}
			hash, err := bcrypt.GenerateFromPassword(validatePassword(password), bcryptCost)
			if err != nil {
				return err
			}
			us.Password = string(hash)
			v, err := bs.Marshal(us)
			if err != nil {
				return err
			}
			return tx.Bucket([]byte(userTable)).Put([]byte(username), v)
		}
		return nil
	})
}

func (bs *BoltStore) AdminUpdateUser(username, password string, admin bool) error {
	if username == defUser && !admin {
		return errors.New("user '" + defUser + "' cannot lose admin privileges")
	}
	return bs.db.Update(func(tx *bolt.Tx) error {
		us, err := bs.GetUser(username)
		if err != nil {
			return err
		}
		us.Admin = admin
		if password != "" {
			hash, err := bcrypt.GenerateFromPassword(validatePassword(password), bcryptCost)
			if err != nil {
				return err
			}
			us.Password = string(hash)
		}
		v, err := bs.Marshal(us)
		if err != nil {
			return err
		}
		return tx.Bucket([]byte(userTable)).Put([]byte(username), v)
	})
}

func (bs *BoltStore) DeleteUser(username string) error {
	if username == defUser {
		return errors.New("user '" + defUser + "' cannot be deleted")
	}
	if err := bs.UserExists(username); err != nil {
		return err
	}
	return bs.db.Update(func(tx *bolt.Tx) error {
		return tx.Bucket([]byte(userTable)).Delete([]byte(username))
	})
}