tgdrive/teldrive
1
1
Fork 0
mirror of https://github.com/tgdrive/teldrive.git synced 2025-10-08 14:27:46 +08:00
Code Issues Projects Releases Packages Wiki Activity Actions

added 2fa support

Browse source
This commit is contained in:
divyam234 2023-08-22 20:30:14 +05:30
parent 9ba5c8ceb9
commit a4a8ceca29
7 changed files with 74 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
models/file.model.go
View file

@ -15,7 +15,7 @@ type File struct {
Size int64 `gorm:"type:bigint"`
Starred *bool `gorm:"default:false"`
Depth *int `gorm:"type:integer"`
UserID int `gorm:"type:bigint;not null"`
UserID int64 `gorm:"type:bigint;not null"`
Status string `gorm:"type:text"`
ParentID string `gorm:"type:text;index"`
Parts *Parts `gorm:"type:jsonb"`

2
models/user.model.go
View file

@ -5,7 +5,7 @@ import (
)
type User struct {
UserId int `gorm:"type:bigint;primaryKey"`
UserId int64 `gorm:"type:bigint;primaryKey"`
Name string `gorm:"type:text"`
UserName string `gorm:"type:text"`
IsPremium bool `gorm:"type:bool"`

2
routes/auth.go
View file

@ -38,7 +38,7 @@ func addAuthRoutes(rg *gin.RouterGroup) {
})
r.GET("/ws", authService.HandleQrCodeLogin)
r.GET("/ws", authService.HandleMultipleLogin)
r.GET("/session", func(c *gin.Context) {

4
schemas/file.schema.go
View file

@ -27,7 +27,7 @@ type FileQuery struct {
ParentID string `form:"parentId" mapstructure:"parent_id,omitempty"`
UpdatedAt *time.Time `form:"updatedAt" mapstructure:"updated_at,omitempty"`
Status string `mapstructure:"status"`
UserId int `mapstructure:"user_id"`
UserID int64 `mapstructure:"user_id"`
}
type FileIn struct {
@ -41,7 +41,7 @@ type FileIn struct {
Starred *bool `json:"starred" mapstructure:"starred,omitempty"`
Depth *int `json:"depth,omitempty" mapstructure:"depth,omitempty"`
Status string `mapstructure:"status,omitempty"`
UserID int `json:"userId" mapstructure:"user_id,omitempty"`
UserID int64 `json:"userId" mapstructure:"user_id,omitempty"`
ParentID string `json:"parentId" mapstructure:"parent_id,omitempty"`
}

87
services/auth.service.go
View file

@ -27,6 +27,7 @@ import (
tgauth "github.com/gotd/td/telegram/auth"
"github.com/gotd/td/telegram/auth/qrlogin"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
"gorm.io/gorm"
)
@ -46,6 +47,7 @@ type SocketMessage struct {
PhoneNo string `json:"phoneNo,omitempty"`
PhoneCodeHash string `json:"phoneCodeHash,omitempty"`
PhoneCode string `json:"phoneCode,omitempty"`
Password string `json:"password,omitempty"`
}
func IP4toInt(IPv4Address net.IP) int64 {
@ -109,7 +111,7 @@ func (as *AuthService) LogIn(c *gin.Context) (*schemas.Message, *types.AppError)
now := time.Now().UTC()
jwtClaims := &types.JWTClaims{Claims: jwt.Claims{
Subject: strconv.Itoa(session.UserID),
Subject: strconv.FormatInt(session.UserID, 10),
IssuedAt: jwt.NewNumericDate(now),
Expiry: jwt.NewNumericDate(now.Add(time.Duration(as.SessionMaxAge) * time.Second)),
}, TgSession: session.Sesssion,
@ -219,7 +221,7 @@ func prepareSession(user *tg.User, data *session.Data) *types.TgSession {
sessionString := generateTgSession(data.DC, data.AuthKey, 443)
session := &types.TgSession{
Sesssion: sessionString,
UserID: int(user.ID),
UserID: user.ID,
Bot: user.Bot,
UserName: user.Username,
Name: fmt.Sprintf("%s %s", user.FirstName, user.LastName),
@ -228,7 +230,7 @@ func prepareSession(user *tg.User, data *session.Data) *types.TgSession {
return session
}
func (as *AuthService) HandleQrCodeLogin(c *gin.Context) {
func (as *AuthService) HandleMultipleLogin(c *gin.Context) {
upgrader := websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
@ -255,11 +257,19 @@ func (as *AuthService) HandleQrCodeLogin(c *gin.Context) {
conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": map[string]string{"token": token.URL()}})
return nil
})
if tgerr.Is(err, "SESSION_PASSWORD_NEEDED") {
conn.WriteJSON(map[string]interface{}{"type": "auth", "message": "2FA required"})
return
}
if err != nil {
conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()})
return
}
user, ok := authorization.User.AsNotEmpty()
if !ok {
conn.WriteJSON(map[string]interface{}{"type": "error", "message": errors.New("auth failed")})
return
}
res, _ := sessionStorage.LoadSession(c)
@ -270,27 +280,60 @@ func (as *AuthService) HandleQrCodeLogin(c *gin.Context) {
}()
}
if message.AuthType == "phone" && message.Message == "sendcode" {
res, err := tgClient.Auth().SendCode(c, message.PhoneNo, tgauth.SendCodeOptions{})
if err != nil {
return
}
code := res.(*tg.AuthSentCode)
conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": map[string]string{"phoneCodeHash": code.PhoneCodeHash}})
go func() {
res, err := tgClient.Auth().SendCode(c, message.PhoneNo, tgauth.SendCodeOptions{})
if err != nil {
conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()})
return
}
code := res.(*tg.AuthSentCode)
conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": map[string]string{"phoneCodeHash": code.PhoneCodeHash}})
}()
}
if message.AuthType == "phone" && message.Message == "signin" {
auth, err := tgClient.Auth().SignIn(c, message.PhoneNo, message.PhoneCode, message.PhoneCodeHash)
if err != nil {
return
}
user, ok := auth.User.AsNotEmpty()
if !ok {
return
}
res, _ := sessionStorage.LoadSession(c)
sessionData := &SessionData{}
json.Unmarshal(res, sessionData)
session := prepareSession(user, &sessionData.Data)
conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"})
go func() {
auth, err := tgClient.Auth().SignIn(c, message.PhoneNo, message.PhoneCode, message.PhoneCodeHash)
if errors.Is(err, tgauth.ErrPasswordAuthNeeded) {
conn.WriteJSON(map[string]interface{}{"type": "auth", "message": "2FA required"})
return
}
if err != nil {
conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()})
return
}
user, ok := auth.User.AsNotEmpty()
if !ok {
conn.WriteJSON(map[string]interface{}{"type": "error", "message": errors.New("auth failed")})
return
}
res, _ := sessionStorage.LoadSession(c)
sessionData := &SessionData{}
json.Unmarshal(res, sessionData)
session := prepareSession(user, &sessionData.Data)
conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"})
}()
}
if message.AuthType == "2fa" && message.Password != "" {
go func() {
auth, err := tgClient.Auth().Password(c, message.Password)
if err != nil {
conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()})
return
}
user, ok := auth.User.AsNotEmpty()
if !ok {
conn.WriteJSON(map[string]interface{}{"type": "error", "message": errors.New("auth failed")})
return
}
res, _ := sessionStorage.LoadSession(c)
sessionData := &SessionData{}
json.Unmarshal(res, sessionData)
session := prepareSession(user, &sessionData.Data)
conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"})
}()
}
if err != nil {
log.Println(err)

6
services/file.service.go
View file

@ -33,10 +33,10 @@ type FileService struct {
ChannelID int64
}
func getAuthUserId(c *gin.Context) int {
func getAuthUserId(c *gin.Context) int64 {
val, _ := c.Get("jwtUser")
jwtUser := val.(*types.JWTClaims)
userId, _ := strconv.Atoi(jwtUser.Subject)
userId, _ := strconv.ParseInt(jwtUser.Subject, 10, 64)
return userId
}
@ -160,7 +160,7 @@ func (fs *FileService) ListFiles(c *gin.Context) (*schemas.FileResponse, *types.
var fileQuery schemas.FileQuery
fileQuery.Op = "list"
fileQuery.Status = "active"
fileQuery.UserId = userId
fileQuery.UserID = userId
if err := c.ShouldBindQuery(&fileQuery); err != nil {
return nil, &types.AppError{Error: errors.New(""), Code: http.StatusBadRequest}
}

2
types/main.go
View file

@ -28,7 +28,7 @@ type JWTClaims struct {
type TgSession struct {
Sesssion string `json:"session"`
UserID int `json:"userId"`
UserID int64 `json:"userId"`
Bot bool `json:"bot"`
UserName string `json:"userName"`
Name string `json:"name"`