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

further cleanup

Browse source
This commit is contained in:
divyam234 2023-12-03 03:16:53 +05:30
parent b040eb9e1e
commit bc7b50f4da
11 changed files with 60 additions and 91 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
api/router.go
View file

@ -53,13 +53,12 @@ func InitRouter() *gin.Engine {
{
users.Use(middleware.Authmiddleware)
users.GET("/profile", c.GetProfilePhoto)
users.GET("/stats", c.Stats)
users.GET("/stats", c.GetStats)
users.GET("/bots", c.GetBots)
users.GET("/channels", c.ListChannels)
users.PATCH("/channels", c.UpdateChannel)
users.POST("/bots", c.AddBots)
users.DELETE("/bots", c.RemoveBots)
users.DELETE("/bots/session", c.RevokeBotSession)
}
}

2
go.mod
View file

@ -6,6 +6,7 @@ require (
github.com/coocood/freecache v1.2.4
github.com/divyam234/cors v1.4.2
github.com/gin-gonic/gin v1.9.1
github.com/go-co-op/gocron v1.36.0
github.com/go-jose/go-jose/v3 v3.0.1
github.com/gotd/contrib v0.19.0
github.com/gotd/td v0.89.0
@ -26,7 +27,6 @@ require (
require (
github.com/cespare/xxhash/v2 v2.2.0 // indirect
github.com/chenzhuoyu/iasm v0.9.1 // indirect
github.com/go-co-op/gocron v1.36.0 // indirect
github.com/google/uuid v1.4.0 // indirect
github.com/jackc/puddle/v2 v2.2.1 // indirect
github.com/robfig/cron/v3 v3.0.1 // indirect

11
pkg/controller/controller.go
View file

@ -14,12 +14,9 @@ type Controller struct {
func NewController() *Controller {
return &Controller{
FileService: &services.FileService{Db: database.DB},
UserService: &services.UserService{Db: database.DB},
UploadService: &services.UploadService{Db: database.DB},
AuthService: &services.AuthService{
Db: database.DB,
SessionMaxAge: 30 * 24 * 60 * 60,
SessionCookieName: "user-session"},
FileService: services.NewFileService(database.DB),
UserService: services.NewUserService(database.DB),
UploadService: services.NewUploadService(database.DB),
AuthService: services.NewAuthService(database.DB),
}
}

14
pkg/controller/user.go
View file

@ -7,8 +7,8 @@ import (
"github.com/gin-gonic/gin"
)
func (uc *Controller) Stats(c *gin.Context) {
res, err := uc.UserService.Stats(c)
func (uc *Controller) GetStats(c *gin.Context) {
res, err := uc.UserService.GetStats(c)
if err != nil {
httputil.NewError(c, err.Code, err.Error)
return
@ -67,16 +67,6 @@ func (uc *Controller) RemoveBots(c *gin.Context) {
c.JSON(http.StatusOK, res)
}
func (uc *Controller) RevokeBotSession(c *gin.Context) {
res, err := uc.UserService.RevokeBotSession(c)
if err != nil {
httputil.NewError(c, err.Code, err.Error)
return
}
c.JSON(http.StatusOK, res)
}
func (uc *Controller) GetProfilePhoto(c *gin.Context) {
if c.Query("photo") != "" {
uc.UserService.GetProfilePhoto(c)

5
pkg/database/database.go
View file

@ -20,7 +20,6 @@ import (
//go:embed migrations/*.sql
var embedMigrations embed.FS
var DB *gorm.DB
var BoltDB *bbolt.DB
var KV kv.KV
func InitDB() {
@ -69,14 +68,14 @@ func InitDB() {
}
}()
BoltDB, err = bbolt.Open(filepath.Join(config.ExecDir, "teldrive.db"), 0666, &bbolt.Options{
boltDB, err := bbolt.Open(filepath.Join(config.ExecDir, "teldrive.db"), 0666, &bbolt.Options{
Timeout: time.Second,
NoGrowSync: false,
})
if err != nil {
panic(err)
}
KV, err = kv.New(kv.Options{Bucket: "teldrive", DB: BoltDB})
KV, err = kv.New(kv.Options{Bucket: "teldrive", DB: boltDB})
if err != nil {
panic(err)

32
pkg/services/auth.go
View file

@ -40,27 +40,21 @@ type AuthService struct {
SessionCookieName string
}
type SessionData struct {
Version int
Data session.Data
}
type SocketMessage struct {
AuthType string `json:"authType"`
Message string `json:"message"`
PhoneNo string `json:"phoneNo,omitempty"`
PhoneCodeHash string `json:"phoneCodeHash,omitempty"`
PhoneCode string `json:"phoneCode,omitempty"`
Password string `json:"password,omitempty"`
func NewAuthService(db *gorm.DB) *AuthService {
return &AuthService{
Db: db,
SessionMaxAge: 30 * 24 * 60 * 60,
SessionCookieName: "user-session"}
}
func IP4toInt(IPv4Address net.IP) int64 {
func ip4toInt(IPv4Address net.IP) int64 {
IPv4Int := big.NewInt(0)
IPv4Int.SetBytes(IPv4Address.To4())
return IPv4Int.Int64()
}
func Pack32BinaryIP4(ip4Address string) []byte {
ipv4Decimal := IP4toInt(net.ParseIP(ip4Address))
func pack32BinaryIP4(ip4Address string) []byte {
ipv4Decimal := ip4toInt(net.ParseIP(ip4Address))
buf := new(bytes.Buffer)
binary.Write(buf, binary.BigEndian, uint32(ipv4Decimal))
@ -78,7 +72,7 @@ func generateTgSession(dcID int, authKey []byte, port int) string {
}
dcIDByte := byte(dcID)
serverAddressBytes := Pack32BinaryIP4(dcMaps[dcID])
serverAddressBytes := pack32BinaryIP4(dcMaps[dcID])
portByte := make([]byte, 2)
binary.BigEndian.PutUint16(portByte, uint16(port))
@ -286,7 +280,7 @@ func (as *AuthService) HandleMultipleLogin(c *gin.Context) {
err = tgClient.Run(c, func(ctx context.Context) error {
for {
message := &SocketMessage{}
message := &types.SocketMessage{}
err := conn.ReadJSON(message)
if err != nil {
@ -319,7 +313,7 @@ func (as *AuthService) HandleMultipleLogin(c *gin.Context) {
return
}
res, _ := sessionStorage.LoadSession(c)
sessionData := &SessionData{}
sessionData := &types.SessionData{}
json.Unmarshal(res, sessionData)
session := prepareSession(user, &sessionData.Data)
conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"})
@ -360,7 +354,7 @@ func (as *AuthService) HandleMultipleLogin(c *gin.Context) {
return
}
res, _ := sessionStorage.LoadSession(c)
sessionData := &SessionData{}
sessionData := &types.SessionData{}
json.Unmarshal(res, sessionData)
session := prepareSession(user, &sessionData.Data)
conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"})
@ -385,7 +379,7 @@ func (as *AuthService) HandleMultipleLogin(c *gin.Context) {
return
}
res, _ := sessionStorage.LoadSession(c)
sessionData := &SessionData{}
sessionData := &types.SessionData{}
json.Unmarshal(res, sessionData)
session := prepareSession(user, &sessionData.Data)
conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"})

4
pkg/services/common.go
View file

@ -98,7 +98,7 @@ func getUserAuth(c *gin.Context) (int64, string) {
return userId, jwtUser.TgSession
}
func getBotInfo(ctx context.Context, token string) (*BotInfo, error) {
func getBotInfo(ctx context.Context, token string) (*types.BotInfo, error) {
client, _ := tgc.BotLogin(ctx, token)
var user *tg.User
err := tgc.RunWithAuth(ctx, client, token, func(ctx context.Context) error {
@ -109,7 +109,7 @@ func getBotInfo(ctx context.Context, token string) (*BotInfo, error) {
if err != nil {
return nil, err
}
return &BotInfo{Id: user.ID, UserName: user.Username, Token: token}, nil
return &types.BotInfo{Id: user.ID, UserName: user.Username, Token: token}, nil
}
func getTGMessages(ctx context.Context, client *telegram.Client, parts []schemas.Part, channelId int64, userID string) (*tg.MessagesChannelMessages, error) {

4
pkg/services/file.go
View file

@ -35,6 +35,10 @@ type FileService struct {
Db *gorm.DB
}
func NewFileService(db *gorm.DB) *FileService {
return &FileService{Db: db}
}
func (fs *FileService) CreateFile(c *gin.Context) (*schemas.FileOut, *types.AppError) {
userId, _ := getUserAuth(c)
var fileIn schemas.CreateFile

4
pkg/services/upload.go
View file

@ -28,6 +28,10 @@ type UploadService struct {
Db *gorm.DB
}
func NewUploadService(db *gorm.DB) *UploadService {
return &UploadService{Db: db}
}
func (us *UploadService) GetUploadFileById(c *gin.Context) (*schemas.UploadOut, *types.AppError) {
uploadId := c.Param("id")
parts := []schemas.UploadPartOut{}

51
pkg/services/user.go
View file

@ -1,7 +1,6 @@
package services
import (
"bytes"
"context"
"errors"
"fmt"
@ -11,7 +10,6 @@ import (
"github.com/divyam234/teldrive/internal/cache"
"github.com/divyam234/teldrive/internal/tgc"
"github.com/divyam234/teldrive/pkg/database"
"github.com/divyam234/teldrive/pkg/models"
"github.com/divyam234/teldrive/pkg/schemas"
"github.com/divyam234/teldrive/pkg/types"
@ -20,7 +18,6 @@ import (
"github.com/gotd/td/telegram/query"
"github.com/gotd/td/tg"
"github.com/thoas/go-funk"
"go.etcd.io/bbolt"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
@ -31,11 +28,8 @@ type UserService struct {
Db *gorm.DB
}
type BotInfo struct {
Id int64
UserName string
AccessHash int64
Token string
func NewUserService(db *gorm.DB) *UserService {
return &UserService{Db: db}
}
func (us *UserService) GetProfilePhoto(c *gin.Context) {
@ -80,7 +74,7 @@ func (us *UserService) GetProfilePhoto(c *gin.Context) {
}
}
func (us *UserService) Stats(c *gin.Context) (*schemas.AccountStats, *types.AppError) {
func (us *UserService) GetStats(c *gin.Context) (*schemas.AccountStats, *types.AppError) {
userId, _ := getUserAuth(c)
var res []schemas.AccountStats
if err := us.Db.Raw("select * from teldrive.account_stats(?);", userId).Scan(&res).Error; err != nil {
@ -210,42 +204,9 @@ func (us *UserService) RemoveBots(c *gin.Context) (*schemas.Message, *types.AppE
}
func (us *UserService) RevokeBotSession(c *gin.Context) (*schemas.Message, *types.AppError) {
pattern := []byte("botsession:")
err := database.BoltDB.Update(func(tx *bbolt.Tx) error {
bucket := tx.Bucket([]byte("teldrive"))
if bucket == nil {
return errors.New("bucket not found")
}
c := bucket.Cursor()
for key, _ := c.First(); key != nil; key, _ = c.Next() {
if bytes.HasPrefix(key, pattern) {
if err := c.Delete(); err != nil {
return err
}
}
}
return nil
})
if err != nil {
return nil, &types.AppError{Error: errors.New("failed to revoke session"),
Code: http.StatusInternalServerError}
}
return &schemas.Message{Message: "session revoked"}, nil
}
func (us *UserService) addBots(c context.Context, client *telegram.Client, userId int64, channelId int64, botsTokens []string) (*schemas.Message, *types.AppError) {
botInfo := []BotInfo{}
botInfo := []types.BotInfo{}
var wg sync.WaitGroup
@ -260,7 +221,7 @@ func (us *UserService) addBots(c context.Context, client *telegram.Client, userI
return err
}
botInfoChannel := make(chan *BotInfo, len(botsTokens))
botInfoChannel := make(chan *types.BotInfo, len(botsTokens))
waitChan := make(chan struct{}, 6)
@ -295,7 +256,7 @@ func (us *UserService) addBots(c context.Context, client *telegram.Client, userI
}
if len(botsTokens) == len(botInfo) {
users := funk.Map(botInfo, func(info BotInfo) tg.InputUser {
users := funk.Map(botInfo, func(info types.BotInfo) tg.InputUser {
return tg.InputUser{UserID: info.Id, AccessHash: info.AccessHash}
})
botsToAdd := users.([]tg.InputUser)

21
pkg/types/types.go
View file

@ -2,6 +2,7 @@ package types
import (
"github.com/go-jose/go-jose/v3/jwt"
"github.com/gotd/td/session"
"github.com/gotd/td/tg"
)
@ -25,3 +26,23 @@ type JWTClaims struct {
IsPremium bool `json:"isPremium"`
Hash string `json:"hash"`
}
type SessionData struct {
Version int
Data session.Data
}
type SocketMessage struct {
AuthType string `json:"authType"`
Message string `json:"message"`
PhoneNo string `json:"phoneNo,omitempty"`
PhoneCodeHash string `json:"phoneCodeHash,omitempty"`
PhoneCode string `json:"phoneCode,omitempty"`
Password string `json:"password,omitempty"`
}
type BotInfo struct {
Id int64
UserName string
AccessHash int64
Token string
}