ezbookkeeping/pkg/services/twofactor_authorizations.go

package services

import (
	"time"

	"github.com/pquerna/otp"
	"github.com/pquerna/otp/totp"
	"xorm.io/xorm"

	"github.com/mayswind/lab/pkg/datastore"
	"github.com/mayswind/lab/pkg/errs"
	"github.com/mayswind/lab/pkg/models"
	"github.com/mayswind/lab/pkg/settings"
	"github.com/mayswind/lab/pkg/utils"
	"github.com/mayswind/lab/pkg/uuid"
)

const (
	TWOFACTOR_PERIOD               uint = 30 // seconds
	TWOFACTOR_SECRET_SIZE          uint = 20 // bytes
	TWOFACTOR_RECOVERY_CODE_COUNT  int  = 10
	TWOFACTOR_RECOVERY_CODE_LENGTH int  = 10 // bytes
)

type TwoFactorAuthorizationService struct {
	ServiceUsingDB
	ServiceUsingConfig
	ServiceUsingUuid
}

var (
	TwoFactorAuthorizations = &TwoFactorAuthorizationService{
		ServiceUsingDB: ServiceUsingDB{
			container: datastore.Container,
		},
		ServiceUsingConfig: ServiceUsingConfig{
			container: settings.Container,
		},
		ServiceUsingUuid: ServiceUsingUuid{
			container: uuid.Container,
		},
	}
)

func (s *TwoFactorAuthorizationService) GetUserTwoFactorSettingByUid(uid int64) (*models.TwoFactor, error) {
	if uid <= 0 {
		return nil, errs.ErrUserIdInvalid
	}

	twoFactor := &models.TwoFactor{}
	has, err := s.UserDB().Where("uid=?", uid).Get(twoFactor)

	if err != nil {
		return nil, err
	} else if !has {
		return nil, errs.ErrTwoFactorIsNotEnabled
	}

	twoFactor.Secret, err = utils.DecryptSecret(twoFactor.Secret, s.CurrentConfig().SecretKey)

	if err != nil {
		return nil, err
	}

	return twoFactor, nil
}

func (s *TwoFactorAuthorizationService) GenerateTwoFactorSecret(user *models.User) (*otp.Key, error) {
	if user == nil {
		return nil, errs.ErrUserNotFound
	}

	key, err := totp.Generate(totp.GenerateOpts{
		Issuer:      s.CurrentConfig().AppName,
		AccountName: user.Username,
		Period:      TWOFACTOR_PERIOD,
		SecretSize:  TWOFACTOR_SECRET_SIZE,
	})

	return key, err
}

func (s *TwoFactorAuthorizationService) CreateTwoFactorSetting(twoFactor *models.TwoFactor) error {
	if twoFactor.Uid <= 0 {
		return errs.ErrUserIdInvalid
	}

	var err error
	twoFactor.Secret, err = utils.EncryptSecret(twoFactor.Secret, s.CurrentConfig().SecretKey)

	if err != nil {
		return err
	}

	twoFactor.CreatedUnixTime = time.Now().Unix()

	return s.UserDB().DoTransaction(func(sess *xorm.Session) error {
		_, err := sess.Insert(twoFactor)
		return err
	})
}

func (s *TwoFactorAuthorizationService) DeleteTwoFactorSetting(uid int64) error {
	if uid <= 0 {
		return errs.ErrUserIdInvalid
	}

	return s.UserDB().DoTransaction(func(sess *xorm.Session) error {
		deletedRows, err := sess.Where("uid=?", uid).Delete(&models.TwoFactor{})

		if err != nil {
			return err
		} else if deletedRows < 1 {
			return errs.ErrTwoFactorIsNotEnabled
		}

		return nil
	})
}

func (s *TwoFactorAuthorizationService) ExistsTwoFactorSetting(uid int64) (bool, error) {
	if uid <= 0 {
		return false, errs.ErrUserIdInvalid
	}

	return s.UserDB().Cols("uid").Where("uid=?", uid).Exist(&models.TwoFactor{})
}

func (s *TwoFactorAuthorizationService) GetAndUseUserTwoFactorRecoveryCode(uid int64, recoveryCode string, salt string) error {
	if uid <= 0 {
		return errs.ErrUserIdInvalid
	}

	recoveryCode = utils.EncodePassword(recoveryCode, salt)
	exists, err := s.UserDB().Cols("uid", "recovery_code").Where("uid=? AND recovery_code=? AND used=?", uid, recoveryCode, false).Exist(&models.TwoFactorRecoveryCode{})

	if err != nil {
		return err
	} else if !exists {
		return errs.ErrTwoFactorRecoveryCodeNotExist
	}

	return s.UserDB().DoTransaction(func(sess *xorm.Session) error {
		_, err := sess.Cols("used", "used_unix_time").Where("uid=? AND recovery_code=?", uid, recoveryCode).Update(&models.TwoFactorRecoveryCode{Used: true, UsedUnixTime: time.Now().Unix()})
		return err
	})
}

func (s *TwoFactorAuthorizationService) GenerateTwoFactorRecoveryCodes() ([]string, error) {
	recoveryCodes := make([]string, TWOFACTOR_RECOVERY_CODE_COUNT)

	for i := 0; i < TWOFACTOR_RECOVERY_CODE_COUNT; i++ {
		recoveryCode, err := utils.GetRandomNumberOrLetter(TWOFACTOR_RECOVERY_CODE_LENGTH)

		if err != nil {
			return nil, err
		}

		recoveryCodes[i] = recoveryCode[:5] + "-" + recoveryCode[5:]
	}

	return recoveryCodes, nil
}

func (s *TwoFactorAuthorizationService) CreateTwoFactorRecoveryCodes(uid int64, recoveryCodes []string, salt string) error {
	twoFactorRecoveryCodes := make([]*models.TwoFactorRecoveryCode, len(recoveryCodes))

	for i := 0; i < len(recoveryCodes); i++ {
		twoFactorRecoveryCodes[i] = &models.TwoFactorRecoveryCode{
			Uid:             uid,
			Used:            false,
			RecoveryCode:    utils.EncodePassword(recoveryCodes[i], salt),
			CreatedUnixTime: time.Now().Unix(),
		}
	}

	return s.UserDB().DoTransaction(func(sess *xorm.Session) error {
		_, err := sess.Where("uid=?", uid).Delete(&models.TwoFactorRecoveryCode{})

		if err != nil {
			return err
		}

		for i := 0; i < len(twoFactorRecoveryCodes); i++ {
			twoFactorRecoveryCode := twoFactorRecoveryCodes[i]
			_, err := sess.Insert(twoFactorRecoveryCode)

			if err != nil {
				return err
			}
		}

		return nil
	})
}

func (s *TwoFactorAuthorizationService) DeleteTwoFactorRecoveryCodes(uid int64) error {
	if uid <= 0 {
		return errs.ErrUserIdInvalid
	}

	return s.UserDB().DoTransaction(func(sess *xorm.Session) error {
		_, err := sess.Where("uid=?", uid).Delete(&models.TwoFactorRecoveryCode{})
		return err
	})
}