two_factor.go 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201
  1. // Copyright 2017 The Gogs Authors. All rights reserved.
  2. // Use of this source code is governed by a MIT-style
  3. // license that can be found in the LICENSE file.
  4. package db
  5. import (
  6. "encoding/base64"
  7. "fmt"
  8. "strings"
  9. "time"
  10. "github.com/pquerna/otp/totp"
  11. "github.com/unknwon/com"
  12. log "unknwon.dev/clog/v2"
  13. "xorm.io/xorm"
  14. "gogs.io/gogs/internal/conf"
  15. "gogs.io/gogs/internal/db/errors"
  16. "gogs.io/gogs/internal/tool"
  17. )
  18. // TwoFactor represents a two-factor authentication token.
  19. type TwoFactor struct {
  20. ID int64
  21. UserID int64 `xorm:"UNIQUE"`
  22. Secret string
  23. Created time.Time `xorm:"-" json:"-"`
  24. CreatedUnix int64
  25. }
  26. func (t *TwoFactor) BeforeInsert() {
  27. t.CreatedUnix = time.Now().Unix()
  28. }
  29. func (t *TwoFactor) AfterSet(colName string, _ xorm.Cell) {
  30. switch colName {
  31. case "created_unix":
  32. t.Created = time.Unix(t.CreatedUnix, 0).Local()
  33. }
  34. }
  35. // ValidateTOTP returns true if given passcode is valid for two-factor authentication token.
  36. // It also returns possible validation error.
  37. func (t *TwoFactor) ValidateTOTP(passcode string) (bool, error) {
  38. secret, err := base64.StdEncoding.DecodeString(t.Secret)
  39. if err != nil {
  40. return false, fmt.Errorf("DecodeString: %v", err)
  41. }
  42. decryptSecret, err := com.AESGCMDecrypt(tool.MD5Bytes(conf.Security.SecretKey), secret)
  43. if err != nil {
  44. return false, fmt.Errorf("AESGCMDecrypt: %v", err)
  45. }
  46. return totp.Validate(passcode, string(decryptSecret)), nil
  47. }
  48. // IsUserEnabledTwoFactor returns true if user has enabled two-factor authentication.
  49. func IsUserEnabledTwoFactor(userID int64) bool {
  50. has, err := x.Where("user_id = ?", userID).Get(new(TwoFactor))
  51. if err != nil {
  52. log.Error("IsUserEnabledTwoFactor [user_id: %d]: %v", userID, err)
  53. }
  54. return has
  55. }
  56. func generateRecoveryCodes(userID int64) ([]*TwoFactorRecoveryCode, error) {
  57. recoveryCodes := make([]*TwoFactorRecoveryCode, 10)
  58. for i := 0; i < 10; i++ {
  59. code, err := tool.RandomString(10)
  60. if err != nil {
  61. return nil, fmt.Errorf("RandomString: %v", err)
  62. }
  63. recoveryCodes[i] = &TwoFactorRecoveryCode{
  64. UserID: userID,
  65. Code: strings.ToLower(code[:5] + "-" + code[5:]),
  66. }
  67. }
  68. return recoveryCodes, nil
  69. }
  70. // NewTwoFactor creates a new two-factor authentication token and recovery codes for given user.
  71. func NewTwoFactor(userID int64, secret string) error {
  72. t := &TwoFactor{
  73. UserID: userID,
  74. }
  75. // Encrypt secret
  76. encryptSecret, err := com.AESGCMEncrypt(tool.MD5Bytes(conf.Security.SecretKey), []byte(secret))
  77. if err != nil {
  78. return fmt.Errorf("AESGCMEncrypt: %v", err)
  79. }
  80. t.Secret = base64.StdEncoding.EncodeToString(encryptSecret)
  81. recoveryCodes, err := generateRecoveryCodes(userID)
  82. if err != nil {
  83. return fmt.Errorf("generateRecoveryCodes: %v", err)
  84. }
  85. sess := x.NewSession()
  86. defer sess.Close()
  87. if err = sess.Begin(); err != nil {
  88. return err
  89. }
  90. if _, err = sess.Insert(t); err != nil {
  91. return fmt.Errorf("insert two-factor: %v", err)
  92. } else if _, err = sess.Insert(recoveryCodes); err != nil {
  93. return fmt.Errorf("insert recovery codes: %v", err)
  94. }
  95. return sess.Commit()
  96. }
  97. // GetTwoFactorByUserID returns two-factor authentication token of given user.
  98. func GetTwoFactorByUserID(userID int64) (*TwoFactor, error) {
  99. t := new(TwoFactor)
  100. has, err := x.Where("user_id = ?", userID).Get(t)
  101. if err != nil {
  102. return nil, err
  103. } else if !has {
  104. return nil, errors.TwoFactorNotFound{UserID: userID}
  105. }
  106. return t, nil
  107. }
  108. // DeleteTwoFactor removes two-factor authentication token and recovery codes of given user.
  109. func DeleteTwoFactor(userID int64) (err error) {
  110. sess := x.NewSession()
  111. defer sess.Close()
  112. if err = sess.Begin(); err != nil {
  113. return err
  114. }
  115. if _, err = sess.Where("user_id = ?", userID).Delete(new(TwoFactor)); err != nil {
  116. return fmt.Errorf("delete two-factor: %v", err)
  117. } else if err = deleteRecoveryCodesByUserID(sess, userID); err != nil {
  118. return fmt.Errorf("deleteRecoveryCodesByUserID: %v", err)
  119. }
  120. return sess.Commit()
  121. }
  122. // TwoFactorRecoveryCode represents a two-factor authentication recovery code.
  123. type TwoFactorRecoveryCode struct {
  124. ID int64
  125. UserID int64
  126. Code string `xorm:"VARCHAR(11)"`
  127. IsUsed bool
  128. }
  129. // GetRecoveryCodesByUserID returns all recovery codes of given user.
  130. func GetRecoveryCodesByUserID(userID int64) ([]*TwoFactorRecoveryCode, error) {
  131. recoveryCodes := make([]*TwoFactorRecoveryCode, 0, 10)
  132. return recoveryCodes, x.Where("user_id = ?", userID).Find(&recoveryCodes)
  133. }
  134. func deleteRecoveryCodesByUserID(e Engine, userID int64) error {
  135. _, err := e.Where("user_id = ?", userID).Delete(new(TwoFactorRecoveryCode))
  136. return err
  137. }
  138. // RegenerateRecoveryCodes regenerates new set of recovery codes for given user.
  139. func RegenerateRecoveryCodes(userID int64) error {
  140. recoveryCodes, err := generateRecoveryCodes(userID)
  141. if err != nil {
  142. return fmt.Errorf("generateRecoveryCodes: %v", err)
  143. }
  144. sess := x.NewSession()
  145. defer sess.Close()
  146. if err = sess.Begin(); err != nil {
  147. return err
  148. }
  149. if err = deleteRecoveryCodesByUserID(sess, userID); err != nil {
  150. return fmt.Errorf("deleteRecoveryCodesByUserID: %v", err)
  151. } else if _, err = sess.Insert(recoveryCodes); err != nil {
  152. return fmt.Errorf("insert new recovery codes: %v", err)
  153. }
  154. return sess.Commit()
  155. }
  156. // UseRecoveryCode validates recovery code of given user and marks it is used if valid.
  157. func UseRecoveryCode(userID int64, code string) error {
  158. recoveryCode := new(TwoFactorRecoveryCode)
  159. has, err := x.Where("code = ?", code).And("is_used = ?", false).Get(recoveryCode)
  160. if err != nil {
  161. return fmt.Errorf("get unused code: %v", err)
  162. } else if !has {
  163. return errors.TwoFactorRecoveryCodeNotFound{Code: code}
  164. }
  165. recoveryCode.IsUsed = true
  166. if _, err = x.Id(recoveryCode.ID).Cols("is_used").Update(recoveryCode); err != nil {
  167. return fmt.Errorf("mark code as used: %v", err)
  168. }
  169. return nil
  170. }