本站源代码
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

167 lines
4.4KB

  1. // Copyright 2017 The Gitea 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 models
  5. import (
  6. "crypto/aes"
  7. "crypto/cipher"
  8. "crypto/md5"
  9. "crypto/rand"
  10. "crypto/sha256"
  11. "crypto/subtle"
  12. "encoding/base64"
  13. "errors"
  14. "fmt"
  15. "io"
  16. "code.gitea.io/gitea/modules/generate"
  17. "code.gitea.io/gitea/modules/setting"
  18. "code.gitea.io/gitea/modules/timeutil"
  19. "github.com/pquerna/otp/totp"
  20. "golang.org/x/crypto/pbkdf2"
  21. )
  22. // TwoFactor represents a two-factor authentication token.
  23. type TwoFactor struct {
  24. ID int64 `xorm:"pk autoincr"`
  25. UID int64 `xorm:"UNIQUE"`
  26. Secret string
  27. ScratchSalt string
  28. ScratchHash string
  29. LastUsedPasscode string `xorm:"VARCHAR(10)"`
  30. CreatedUnix timeutil.TimeStamp `xorm:"INDEX created"`
  31. UpdatedUnix timeutil.TimeStamp `xorm:"INDEX updated"`
  32. }
  33. // GenerateScratchToken recreates the scratch token the user is using.
  34. func (t *TwoFactor) GenerateScratchToken() (string, error) {
  35. token, err := generate.GetRandomString(8)
  36. if err != nil {
  37. return "", err
  38. }
  39. t.ScratchSalt, _ = generate.GetRandomString(10)
  40. t.ScratchHash = hashToken(token, t.ScratchSalt)
  41. return token, nil
  42. }
  43. func hashToken(token, salt string) string {
  44. tempHash := pbkdf2.Key([]byte(token), []byte(salt), 10000, 50, sha256.New)
  45. return fmt.Sprintf("%x", tempHash)
  46. }
  47. // VerifyScratchToken verifies if the specified scratch token is valid.
  48. func (t *TwoFactor) VerifyScratchToken(token string) bool {
  49. if len(token) == 0 {
  50. return false
  51. }
  52. tempHash := hashToken(token, t.ScratchSalt)
  53. return subtle.ConstantTimeCompare([]byte(t.ScratchHash), []byte(tempHash)) == 1
  54. }
  55. func (t *TwoFactor) getEncryptionKey() []byte {
  56. k := md5.Sum([]byte(setting.SecretKey))
  57. return k[:]
  58. }
  59. // SetSecret sets the 2FA secret.
  60. func (t *TwoFactor) SetSecret(secret string) error {
  61. secretBytes, err := aesEncrypt(t.getEncryptionKey(), []byte(secret))
  62. if err != nil {
  63. return err
  64. }
  65. t.Secret = base64.StdEncoding.EncodeToString(secretBytes)
  66. return nil
  67. }
  68. // ValidateTOTP validates the provided passcode.
  69. func (t *TwoFactor) ValidateTOTP(passcode string) (bool, error) {
  70. decodedStoredSecret, err := base64.StdEncoding.DecodeString(t.Secret)
  71. if err != nil {
  72. return false, err
  73. }
  74. secret, err := aesDecrypt(t.getEncryptionKey(), decodedStoredSecret)
  75. if err != nil {
  76. return false, err
  77. }
  78. secretStr := string(secret)
  79. return totp.Validate(passcode, secretStr), nil
  80. }
  81. // aesEncrypt encrypts text and given key with AES.
  82. func aesEncrypt(key, text []byte) ([]byte, error) {
  83. block, err := aes.NewCipher(key)
  84. if err != nil {
  85. return nil, err
  86. }
  87. b := base64.StdEncoding.EncodeToString(text)
  88. ciphertext := make([]byte, aes.BlockSize+len(b))
  89. iv := ciphertext[:aes.BlockSize]
  90. if _, err := io.ReadFull(rand.Reader, iv); err != nil {
  91. return nil, err
  92. }
  93. cfb := cipher.NewCFBEncrypter(block, iv)
  94. cfb.XORKeyStream(ciphertext[aes.BlockSize:], []byte(b))
  95. return ciphertext, nil
  96. }
  97. // aesDecrypt decrypts text and given key with AES.
  98. func aesDecrypt(key, text []byte) ([]byte, error) {
  99. block, err := aes.NewCipher(key)
  100. if err != nil {
  101. return nil, err
  102. }
  103. if len(text) < aes.BlockSize {
  104. return nil, errors.New("ciphertext too short")
  105. }
  106. iv := text[:aes.BlockSize]
  107. text = text[aes.BlockSize:]
  108. cfb := cipher.NewCFBDecrypter(block, iv)
  109. cfb.XORKeyStream(text, text)
  110. data, err := base64.StdEncoding.DecodeString(string(text))
  111. if err != nil {
  112. return nil, err
  113. }
  114. return data, nil
  115. }
  116. // NewTwoFactor creates a new two-factor authentication token.
  117. func NewTwoFactor(t *TwoFactor) error {
  118. _, err := x.Insert(t)
  119. return err
  120. }
  121. // UpdateTwoFactor updates a two-factor authentication token.
  122. func UpdateTwoFactor(t *TwoFactor) error {
  123. _, err := x.ID(t.ID).AllCols().Update(t)
  124. return err
  125. }
  126. // GetTwoFactorByUID returns the two-factor authentication token associated with
  127. // the user, if any.
  128. func GetTwoFactorByUID(uid int64) (*TwoFactor, error) {
  129. twofa := &TwoFactor{UID: uid}
  130. has, err := x.Get(twofa)
  131. if err != nil {
  132. return nil, err
  133. } else if !has {
  134. return nil, ErrTwoFactorNotEnrolled{uid}
  135. }
  136. return twofa, nil
  137. }
  138. // DeleteTwoFactorByID deletes two-factor authentication token by given ID.
  139. func DeleteTwoFactorByID(id, userID int64) error {
  140. cnt, err := x.ID(id).Delete(&TwoFactor{
  141. UID: userID,
  142. })
  143. if err != nil {
  144. return err
  145. } else if cnt != 1 {
  146. return ErrTwoFactorNotEnrolled{userID}
  147. }
  148. return nil
  149. }
上海开阖软件有限公司 沪ICP备12045867号-1