package services import ( "context" "errors" "fmt" "net/http" "strconv" "sync" "github.com/divyam234/teldrive/internal/cache" "github.com/divyam234/teldrive/internal/config" "github.com/divyam234/teldrive/internal/kv" "github.com/divyam234/teldrive/internal/tgc" "github.com/divyam234/teldrive/pkg/logging" "github.com/divyam234/teldrive/pkg/models" "github.com/divyam234/teldrive/pkg/schemas" "github.com/divyam234/teldrive/pkg/types" "github.com/gotd/td/telegram" "github.com/gotd/td/telegram/message/peer" "github.com/gotd/td/telegram/query" "github.com/gotd/td/tg" "github.com/thoas/go-funk" "go.uber.org/zap" "github.com/gin-gonic/gin" "gorm.io/gorm" "gorm.io/gorm/clause" ) type UserService struct { db *gorm.DB cnf *config.Config kv kv.KV } func NewUserService(db *gorm.DB, cnf *config.Config, kv kv.KV) *UserService { return &UserService{db: db, cnf: cnf, kv: kv} } func (us *UserService) GetProfilePhoto(c *gin.Context) { _, session := GetUserAuth(c) client, err := tgc.AuthClient(c, &us.cnf.TG, session) if err != nil { c.AbortWithError(http.StatusInternalServerError, err) return } err = tgc.RunWithAuth(c, client, "", func(ctx context.Context) error { self, err := client.Self(c) if err != nil { return err } peer := self.AsInputPeer() if self.Photo == nil { return nil } photo, ok := self.Photo.AsNotEmpty() if !ok { return errors.New("profile not found") } location := &tg.InputPeerPhotoFileLocation{Big: false, Peer: peer, PhotoID: photo.PhotoID} buff, err := iterContent(c, client, location) if err != nil { return err } content := buff.Bytes() c.Writer.Header().Set("Content-Type", "image/jpeg") c.Writer.Header().Set("Cache-Control", "public, max-age=86400") c.Writer.Header().Set("Content-Length", strconv.Itoa(len(content))) c.Writer.Header().Set("Content-Disposition", fmt.Sprintf("inline; filename=\"%s\"", "profile.jpeg")) c.Writer.Write(content) return nil }) if err != nil { c.AbortWithError(http.StatusNotFound, err) return } } 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 { return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} } return &res[0], nil } func (us *UserService) GetBots(c *gin.Context) ([]string, *types.AppError) { userID, _ := GetUserAuth(c) var ( channelId int64 err error ) if c.Param("channelId") != "" { channelId, _ = strconv.ParseInt(c.Param("channelId"), 10, 64) } else { channelId, err = GetDefaultChannel(c, us.db, userID) if err != nil { return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} } } tokens, err := getBotsToken(c, us.db, userID, channelId) if err != nil { return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} } return tokens, nil } func (us *UserService) UpdateChannel(c *gin.Context) (*schemas.Message, *types.AppError) { cache := cache.FromContext(c) userId, _ := GetUserAuth(c) var payload schemas.Channel if err := c.ShouldBindJSON(&payload); err != nil { return nil, &types.AppError{Error: err, Code: http.StatusBadRequest} } channel := &models.Channel{ChannelID: payload.ChannelID, ChannelName: payload.ChannelName, UserID: userId, Selected: true} if err := us.db.Clauses(clause.OnConflict{ Columns: []clause.Column{{Name: "channel_id"}}, DoUpdates: clause.Assignments(map[string]interface{}{"selected": true}), }).Create(channel).Error; err != nil { return nil, &types.AppError{Error: errors.New("failed to update channel"), Code: http.StatusInternalServerError} } us.db.Model(&models.Channel{}).Where("channel_id != ?", payload.ChannelID). Where("user_id = ?", userId).Update("selected", false) key := fmt.Sprintf("users:channel:%d", userId) cache.Set(key, payload.ChannelID, 0) return &schemas.Message{Message: "channel updated"}, nil } func (us *UserService) ListChannels(c *gin.Context) (interface{}, *types.AppError) { _, session := GetUserAuth(c) client, _ := tgc.AuthClient(c, &us.cnf.TG, session) channels := make(map[int64]*schemas.Channel) client.Run(c, func(ctx context.Context) error { dialogs, _ := query.GetDialogs(client.API()).BatchSize(100).Collect(ctx) for _, dialog := range dialogs { if !dialog.Deleted() { for _, channel := range dialog.Entities.Channels() { _, exists := channels[channel.ID] if !exists && channel.AdminRights.AddAdmins { channels[channel.ID] = &schemas.Channel{ChannelID: channel.ID, ChannelName: channel.Title} } } } } return nil }) return funk.Values(channels), nil } func (us *UserService) AddBots(c *gin.Context) (*schemas.Message, *types.AppError) { userId, session := GetUserAuth(c) client, _ := tgc.AuthClient(c, &us.cnf.TG, session) var botsTokens []string if err := c.ShouldBindJSON(&botsTokens); err != nil { return nil, &types.AppError{Error: err, Code: http.StatusBadRequest} } if len(botsTokens) == 0 { return &schemas.Message{Message: "no bots to add"}, nil } channelId, err := GetDefaultChannel(c, us.db, userId) if err != nil { return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} } return us.addBots(c, client, userId, channelId, botsTokens) } func (us *UserService) RemoveBots(c *gin.Context) (*schemas.Message, *types.AppError) { cache := cache.FromContext(c) userID, _ := GetUserAuth(c) channelId, err := GetDefaultChannel(c, us.db, userID) if err != nil { return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} } if err := us.db.Where("user_id = ?", userID).Where("channel_id = ?", channelId). Delete(&models.Bot{}).Error; err != nil { return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} } cache.Delete(fmt.Sprintf("users:bots:%d:%d", userID, channelId)) return &schemas.Message{Message: "bots deleted"}, nil } func (us *UserService) addBots(c context.Context, client *telegram.Client, userId int64, channelId int64, botsTokens []string) (*schemas.Message, *types.AppError) { botInfo := []types.BotInfo{} var wg sync.WaitGroup logger := logging.FromContext(c) cache := cache.FromContext(c) err := tgc.RunWithAuth(c, client, "", func(ctx context.Context) error { channel, err := GetChannelById(ctx, client, channelId, strconv.FormatInt(userId, 10)) if err != nil { logger.Error("error", zap.Error(err)) return err } if err != nil { return err } botInfoChannel := make(chan *types.BotInfo, len(botsTokens)) waitChan := make(chan struct{}, 6) for _, token := range botsTokens { waitChan <- struct{}{} wg.Add(1) go func(t string) { info, err := getBotInfo(c, us.kv, &us.cnf.TG, t) if err != nil { return } botPeerClass, err := peer.DefaultResolver(client.API()).ResolveDomain(ctx, info.UserName) if err != nil { logger.Error("error", zap.Error(err)) return } botPeer := botPeerClass.(*tg.InputPeerUser) info.AccessHash = botPeer.AccessHash defer func() { <-waitChan wg.Done() }() if err == nil { botInfoChannel <- info } }(token) } wg.Wait() close(botInfoChannel) for result := range botInfoChannel { botInfo = append(botInfo, *result) } if len(botsTokens) == len(botInfo) { users := funk.Map(botInfo, func(info types.BotInfo) tg.InputUser { return tg.InputUser{UserID: info.Id, AccessHash: info.AccessHash} }) botsToAdd := users.([]tg.InputUser) for _, user := range botsToAdd { payload := &tg.ChannelsEditAdminRequest{ Channel: channel, UserID: tg.InputUserClass(&user), AdminRights: tg.ChatAdminRights{ ChangeInfo: true, PostMessages: true, EditMessages: true, DeleteMessages: true, BanUsers: true, InviteUsers: true, PinMessages: true, ManageCall: true, Other: true, ManageTopics: true, }, Rank: "bot", } _, err := client.API().ChannelsEditAdmin(ctx, payload) if err != nil { logger.Error("error", zap.Error(err)) return err } } } else { return errors.New("failed to fetch bots") } return nil }) if err != nil { return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} } payload := []models.Bot{} for _, info := range botInfo { payload = append(payload, models.Bot{UserID: userId, Token: info.Token, BotID: info.Id, BotUserName: info.UserName, ChannelID: channelId, }) } cache.Delete(fmt.Sprintf("users:bots:%d:%d", userId, channelId)) if err := us.db.Clauses(clause.OnConflict{DoNothing: true}).Create(&payload).Error; err != nil { return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} } return &schemas.Message{Message: "bots added"}, nil }