feat(chat): export messages WIP

app/chat/export.go

@@ -0,0 +1,85 @@
+package chat
+
+import (
+	"context"
+	"fmt"
+	"github.com/fatih/color"
+	"github.com/gotd/contrib/middleware/ratelimit"
+	"github.com/gotd/td/telegram/peers"
+	"github.com/gotd/td/telegram/query"
+	"github.com/gotd/td/tg"
+	"github.com/iyear/tdl/app/internal/tgc"
+	"github.com/iyear/tdl/pkg/prog"
+	"github.com/iyear/tdl/pkg/utils"
+	"github.com/jedib0t/go-pretty/v6/progress"
+	"golang.org/x/time/rate"
+	"os"
+	"time"
+)
+
+const (
+	layout       = "2006-01-02 15:04:05"
+	rateInterval = 475 * time.Millisecond
+	rateBucket   = 2
+)
+
+func Export(ctx context.Context, chat string, from, to int, media bool, output string) error {
+	c, _, err := tgc.NoLogin(ratelimit.New(rate.Every(rateInterval), rateBucket))
+	if err != nil {
+		return err
+	}
+
+	return tgc.RunWithAuth(ctx, c, func(ctx context.Context) error {
+		manager := peers.Options{}.Build(c.API())
+		peer, err := utils.Telegram.GetInputPeer(ctx, manager, chat)
+		if err != nil {
+			return err
+		}
+
+		color.Blue("Indexing... %s/%d: %s ~ %s", peer.VisibleName(), peer.ID(), time.Unix(int64(from), 0).Format(layout), time.Unix(int64(to), 0).Format(layout))
+		color.Cyan("Fetch speed is determined by Telegram server")
+
+		pw := prog.New(progress.FormatNumber)
+		pw.SetUpdateFrequency(200 * time.Millisecond)
+		pw.Style().Visibility.TrackerOverall = false
+		pw.Style().Visibility.ETA = false
+		pw.Style().Visibility.Percentage = false
+
+		tracker := prog.AppendTracker(pw, progress.FormatNumber, fmt.Sprintf("%s-%d", peer.VisibleName(), peer.ID()), 0)
+
+		go pw.Render()
+
+		batchSize := 100
+		count := int64(0)
+		iter := query.Messages(c.API()).GetHistory(peer.InputPeer()).
+			OffsetDate(to).BatchSize(batchSize).Iter()
+
+		f, err := os.Create(output)
+		if err != nil {
+			return err
+		}
+		_ = f
+
+		for iter.Next(ctx) {
+			msg := iter.Value()
+			if msg.Msg.GetDate() < from {
+				break
+			}
+
+			m, ok := msg.Msg.(*tg.Message)
+			if !ok {
+				continue
+			}
+
+			_, ok = m.GetMedia()
+			if media && !ok {
+				continue
+			}
+
+			count++
+			tracker.SetValue(count)
+		}
+
+		return nil
+	})
+}

cmd/chat/chat.go

@@ -10,5 +10,5 @@ var Cmd = &cobra.Command{
 }
 
 func init() {
-	Cmd.AddCommand(cmdList)
+	Cmd.AddCommand(cmdList, cmdExport)
 }

cmd/chat/export.go

@@ -0,0 +1,34 @@
+package chat
+
+import (
+	"github.com/iyear/tdl/app/chat"
+	"github.com/spf13/cobra"
+	"time"
+)
+
+var (
+	_chat    string
+	from, to int
+	media    bool
+	output   string
+)
+
+var cmdExport = &cobra.Command{
+	Use:   "export",
+	Short: "export messages from (protected) chat",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		if to == 0 {
+			to = int(time.Now().Unix())
+		}
+
+		return chat.Export(cmd.Context(), _chat, from, to, media, output)
+	},
+}
+
+func init() {
+	cmdExport.Flags().StringVarP(&_chat, "chat", "c", "", "")
+	cmdExport.Flags().IntVar(&from, "from", 0, "timestamp of the starting message")
+	cmdExport.Flags().IntVar(&to, "to", 0, "timestamp of the ending message, default value is NOW")
+	cmdExport.Flags().BoolVar(&media, "media", false, "only export messages that contains media")
+	cmdExport.Flags().StringVarP(&output, "output", "o", "tdl-export.json", "output JSON file path")
+}

pkg/utils/telegram.go

@@ -52,6 +52,7 @@ func (t telegram) ParseChannelMsgLink(ctx context.Context, manager *peers.Manage
 	return ch, msgid, nil
 }
 
+// GetInputChannel TODO(iyear): use GetInputPeer
 func (t telegram) GetInputChannel(ctx context.Context, manager *peers.Manager, from string) (*tg.InputChannel, error) {
 	id, err := strconv.ParseInt(from, 10, 64)
 	if err != nil {
@@ -77,6 +78,33 @@ func (t telegram) GetInputChannel(ctx context.Context, manager *peers.Manager, f
 	return &tg.InputChannel{ChannelID: ch.Raw().ID, AccessHash: ch.Raw().AccessHash}, nil
 }
 
+func (t telegram) GetInputPeer(ctx context.Context, manager *peers.Manager, from string) (peers.Peer, error) {
+	id, err := strconv.ParseInt(from, 10, 64)
+	if err != nil {
+		// from is username
+		p, err := manager.Resolve(ctx, from)
+		if err != nil {
+			return nil, err
+		}
+
+		return p, nil
+	}
+
+	if p, err := manager.ResolveChannelID(ctx, id); err == nil {
+		return p, nil
+	}
+
+	if p, err := manager.ResolveUserID(ctx, id); err == nil {
+		return p, nil
+	}
+
+	if p, err := manager.ResolveChatID(ctx, id); err == nil {
+		return p, nil
+	}
+
+	return nil, fmt.Errorf("failed to get result from %d", id)
+}
+
 func (t telegram) GetPeerID(peer tg.PeerClass) int64 {
 	switch p := peer.(type) {
 	case *tg.PeerUser: