feat(dl): support expose files to HTTP server with serve flag

app/dl/dl.go

@@ -37,6 +37,10 @@ type Options struct {
 
 	// resume opts
 	Continue, Restart bool
+
+	// serve
+	Serve bool
+	Port  int
 }
 
 type parser struct {
@@ -65,6 +69,10 @@ func Run(ctx context.Context, opts *Options) error {
 		logger.From(ctx).Debug("Collect dialogs",
 			zap.Any("dialogs", dialogs))
 
+		if opts.Serve {
+			return serve(ctx, kvd, pool, dialogs, opts.Port, opts.Takeout)
+		}
+
 		iter, err := dliter.New(ctx, &dliter.Options{
 			Pool:     pool,
 			KV:       kvd,

app/dl/serve.go

@@ -0,0 +1,171 @@
+package dl
+
+import (
+	"bytes"
+	"context"
+	_ "embed"
+	"fmt"
+	"html/template"
+	"net/http"
+	"strconv"
+	"sync"
+
+	"github.com/fatih/color"
+	"github.com/go-faster/errors"
+	"github.com/gorilla/mux"
+	"github.com/gotd/contrib/http_io"
+	"github.com/gotd/contrib/partio"
+	"github.com/gotd/contrib/tg_io"
+	"github.com/gotd/td/telegram/peers"
+	"github.com/gotd/td/telegram/query"
+	"github.com/gotd/td/telegram/query/messages"
+	"github.com/gotd/td/tg"
+	"github.com/spf13/viper"
+
+	"github.com/iyear/tdl/app/internal/dliter"
+	"github.com/iyear/tdl/pkg/consts"
+	"github.com/iyear/tdl/pkg/dcpool"
+	"github.com/iyear/tdl/pkg/downloader"
+	"github.com/iyear/tdl/pkg/kv"
+	"github.com/iyear/tdl/pkg/logger"
+	"github.com/iyear/tdl/pkg/storage"
+	"github.com/iyear/tdl/pkg/tmedia"
+	"github.com/iyear/tdl/pkg/utils"
+)
+
+type media struct {
+	*downloader.Item
+	MIME string
+}
+
+//go:embed serve.go.tmpl
+var tmpl string
+
+func serve(ctx context.Context,
+	kvd kv.KV,
+	pool dcpool.Pool,
+	dialogs [][]*dliter.Dialog,
+	port int,
+	takeout bool,
+) error {
+	manager := peers.Options{Storage: storage.NewPeers(kvd)}.Build(pool.Default(ctx))
+
+	router := mux.NewRouter()
+
+	cache := &sync.Map{} // map[string]*media
+	router.Handle("/{peer}/{message:[0-9]+}", handler(func(w http.ResponseWriter, r *http.Request) error {
+		vars := mux.Vars(r)
+		peer := vars["peer"]
+		messageStr := vars["message"]
+
+		var item *media
+		if t, ok := cache.Load(peer + messageStr); ok {
+			item = t.(*media)
+		} else {
+			message, err := strconv.Atoi(messageStr)
+			if err != nil {
+				return errors.Wrap(err, "invalid message id")
+			}
+
+			p, err := utils.Telegram.GetInputPeer(ctx, manager, peer)
+			if err != nil {
+				return errors.Wrap(err, "resolve peer")
+			}
+
+			iter := query.Messages(pool.Default(ctx)).
+				GetHistory(p.InputPeer()).OffsetID(message + 1).
+				BatchSize(1).Iter()
+			if !iter.Next(ctx) {
+				return errors.New("no messages")
+			}
+
+			item, err = convItem(iter.Value())
+			if err != nil {
+				return errors.Wrap(err, "convItem")
+			}
+
+			cache.Store(peer+messageStr, item)
+		}
+
+		api := pool.Client(ctx, item.DC)
+		if takeout {
+			api = pool.Takeout(ctx, item.DC)
+		}
+
+		u := partio.NewStreamer(
+			tg_io.NewDownloader(api).ChunkSource(item.Size, item.InputFileLoc),
+			int64(viper.GetInt(consts.FlagPartSize)))
+
+		w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, item.Name))
+
+		http_io.NewHandler(u, item.Size).
+			WithContentType(item.MIME).
+			WithLog(logger.From(ctx).Named("serve")).
+			ServeHTTP(w, r)
+		return nil
+	}))
+
+	items := make([]string, 0)
+	for _, dialog := range dialogs {
+		for _, d := range dialog {
+			for _, m := range d.Messages {
+				items = append(items, fmt.Sprintf("%d/%d", utils.Telegram.GetInputPeerID(d.Peer), m))
+			}
+		}
+	}
+
+	list := bytes.NewBuffer(nil)
+	err := template.Must(template.New("serve.go.tmpl").Parse(tmpl)).Execute(list, items)
+	if err != nil {
+		return errors.Wrap(err, "execute template")
+	}
+
+	router.Handle("/", handler(func(w http.ResponseWriter, r *http.Request) error {
+		_, err := fmt.Fprint(w, list.String())
+		return err
+	}))
+
+	s := http.Server{
+		Addr:    fmt.Sprintf(":%d", port),
+		Handler: router,
+	}
+
+	go func() {
+		<-ctx.Done()
+		_ = s.Shutdown(ctx)
+	}()
+
+	color.Green("(Beta) Serving on http://localhost:%d", port)
+
+	return s.ListenAndServe()
+}
+
+func handler(h func(w http.ResponseWriter, r *http.Request) error) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := h(w, r); err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+		}
+	})
+}
+
+func convItem(elem messages.Elem) (*media, error) {
+	msg, ok := elem.Msg.(*tg.Message)
+	if !ok {
+		return nil, errors.New("value is not a message")
+	}
+
+	item, ok := tmedia.GetMedia(msg)
+	if !ok {
+		return nil, errors.New("message is not a media")
+	}
+
+	file, ok := elem.File()
+	if !ok {
+		return nil, errors.New("message has no file")
+	}
+
+	return &media{
+		Item: item,
+		MIME: file.MIMEType,
+	}, nil
+}

app/dl/serve.go.tmpl

@@ -0,0 +1,39 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <title>tdl serve(beta)</title>
+    <style>
+        body {
+            display: flex;
+            justify-content: center;
+            align-items: center;
+            height: 100vh;
+            margin: 0;
+        }
+
+        .file-list {
+            text-align: center;
+        }
+
+        ul {
+            list-style: none;
+            padding: 0;
+        }
+
+        li {
+            margin: 10px 0;
+        }
+    </style>
+</head>
+<body>
+    <div class="file-list">
+        <h1>Files</h1>
+        <h2>You can use sniffer to download all files</h2>
+        <ul>
+            {{range .}}
+                <li><a href="{{.}}">{{.}}</a></li>
+            {{end}}
+        </ul>
+    </div>
+</body>
+</html>

cmd/dl.go

@@ -65,6 +65,10 @@ func NewDownload() *cobra.Command {
 	cmd.Flags().BoolVar(&opts.Continue, _continue, false, "continue the last download directly")
 	cmd.Flags().BoolVar(&opts.Restart, restart, false, "restart the last download directly")
 
+	// serve flags
+	cmd.Flags().BoolVar(&opts.Serve, "serve", false, "serve the media files as a http server instead of downloading them with built-in downloader")
+	cmd.Flags().IntVar(&opts.Port, "port", 8080, "http server port")
+
 	_ = viper.BindPFlag(consts.FlagDlTemplate, cmd.Flags().Lookup(consts.FlagDlTemplate))
 
 	// completion and validation

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/go-faster/jx v1.1.0
 	github.com/go-playground/validator/v10 v10.15.4
 	github.com/google/uuid v1.3.1
+	github.com/gorilla/mux v1.8.0
 	github.com/gotd/contrib v0.19.0
 	github.com/gotd/td v0.88.0
 	github.com/iancoleman/strcase v0.3.0

go.sum

@@ -194,6 +194,8 @@ github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/googleapis/google-cloud-go-testing v0.0.0-20200911160855-bcd43fbb19e8/go.mod h1:dvDLG8qkwmyD9a/MJJN3XJcT3xFxOKAvTZGvuZmac9g=
+github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=
+github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
 github.com/gorilla/websocket v1.4.1 h1:q7AeDBpnBk8AogcD4DSag/Ukw/KV+YhzLj2bP5HvKCM=
 github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/gotd/contrib v0.19.0 h1:O6GvMrRVeFslIHLUcpaHVzcl9/5PcgR2jQTIIeTyds0=