feat: custom filename template (#110)

client/bot/handlers/config.go

@@ -3,6 +3,7 @@ package handlers
 import (
 	"fmt"
 	"strings"
+	"text/template"
 
 	"github.com/celestix/gotgproto/dispatcher"
 	"github.com/celestix/gotgproto/ext"
@@ -101,3 +102,41 @@ func handleConfigFnameSTCallback(ctx *ext.Context, update *ext.Update) error {
 	})
 	return dispatcher.EndGroups
 }
+
+func handleConfigFnameTmpl(ctx *ext.Context, update *ext.Update) error {
+	userID := update.GetUserChat().GetID()
+	user, err := database.GetUserByChatID(ctx, userID)
+	if err != nil {
+		return err
+	}
+	args := strings.Fields(string(update.EffectiveMessage.Text))
+	if len(args) <= 1 {
+		text := `使用该命令设置文件名模板, 示例:
+/fnametmpl 图片_{{.msgid}}_{{.msgdate}}.jpg
+
+可用变量:
+- {{.msgid}}: 消息ID
+- {{.msgtags}}: 消息中的标签, 将以下划线分隔输出
+- {{.msggen}}: 根据消息生成的文件名
+- {{.msgdate}}: 消息日期, 格式 YYYY-MM-DD_HH-MM-SS
+- {{.origname}}: 媒体的原始文件名 (如果有)`
+		if user.FilenameTemplate != "" {
+			text += fmt.Sprintf("\n\n当前模板: %s", user.FilenameTemplate)
+		}
+		text += "\n\n模板仅在文件名策略设置为 '自定义模板' 时生效, 且模板解析错误时会回退到默认文件名"
+		ctx.Reply(update, ext.ReplyTextString(text), nil)
+		return dispatcher.EndGroups
+	}
+	newTmpl := strings.Join(args[1:], " ")
+	_, err = template.New("filename").Parse(newTmpl)
+	if err != nil {
+		ctx.Reply(update, ext.ReplyTextString("无效的模板, 请检查语法\n"+err.Error()), nil)
+		return dispatcher.EndGroups
+	}
+	user.FilenameTemplate = newTmpl
+	if err := database.UpdateUser(ctx, user); err != nil {
+		return err
+	}
+	ctx.Reply(update, ext.ReplyTextString("已更新文件名模板"), nil)
+	return dispatcher.EndGroups
+}

client/bot/handlers/media.go

@@ -14,7 +14,6 @@ import (
 	"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/shortcut"
 	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
 	"github.com/krau/SaveAny-Bot/database"
-	"github.com/krau/SaveAny-Bot/pkg/enums/fnamest"
 	"github.com/krau/SaveAny-Bot/pkg/tcbdata"
 	"github.com/krau/SaveAny-Bot/pkg/tfile"
 	"github.com/krau/SaveAny-Bot/storage"
@@ -33,12 +32,13 @@ func handleMediaMessage(ctx *ext.Context, update *ext.Update) error {
 	if err != nil {
 		return err
 	}
-	tfOpts := make([]tfile.TGFileOption, 0)
+	// tfOpts := make([]tfile.TGFileOption, 0)
-	switch userDB.FilenameStrategy {
+	// switch userDB.FilenameStrategy {
-	case fnamest.Message.String():
+	// case fnamest.Message.String():
-		tfOpts = append(tfOpts, tfile.WithName(tgutil.GenFileNameFromMessage(*message)))
+	// 	tfOpts = append(tfOpts, tfile.WithName(tgutil.GenFileNameFromMessage(*message)))
-	default:
+	// default:
-	}
+	// }
+	tfOpts := mediautil.TfileOptions(ctx, userDB, message)
 	msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, message, tfOpts...)
 	if err != nil {
 		return err
@@ -74,12 +74,13 @@ func handleSilentSaveMedia(ctx *ext.Context, update *ext.Update) error {
 	if err != nil {
 		return err
 	}
-	tfOpts := make([]tfile.TGFileOption, 0)
+	// tfOpts := make([]tfile.TGFileOption, 0)
-	switch userDB.FilenameStrategy {
+	// switch userDB.FilenameStrategy {
-	case fnamest.Message.String():
+	// case fnamest.Message.String():
-		tfOpts = append(tfOpts, tfile.WithName(tgutil.GenFileNameFromMessage(*message)))
+	// 	tfOpts = append(tfOpts, tfile.WithName(tgutil.GenFileNameFromMessage(*message)))
-	default:
+	// default:
-	}
+	// }
+	tfOpts := mediautil.TfileOptions(ctx, userDB, message)
 	msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, message, tfOpts...)
 	if err != nil {
 		return err

client/bot/handlers/register.go

@@ -42,6 +42,7 @@ func Register(disp dispatcher.Dispatcher) {
 	disp.AddHandler(handlers.NewCommand("unwatch", handleUnwatchCmd))
 	disp.AddHandler(handlers.NewCommand("save", handleSilentMode(handleSaveCmd, handleSilentSaveReplied)))
 	disp.AddHandler(handlers.NewCommand("config", handleConfigCmd))
+	disp.AddHandler(handlers.NewCommand("fnametmpl", handleConfigFnameTmpl))
 	disp.AddHandler(handlers.NewCommand("update", handleUpdateCmd))
 	disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("update"), handleUpdateCallback))
 	disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeAdd), handleAddCallback))

client/bot/handlers/save.go

@@ -16,6 +16,7 @@ import (
 	"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
 	"github.com/krau/SaveAny-Bot/common/utils/strutil"
 	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
+	"github.com/krau/SaveAny-Bot/database"
 	"github.com/krau/SaveAny-Bot/pkg/tcbdata"
 	"github.com/krau/SaveAny-Bot/pkg/tfile"
 
@@ -33,18 +34,27 @@ func handleSaveCmd(ctx *ext.Context, update *ext.Update) error {
 		ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgSaveHelpText)), nil)
 		return dispatcher.EndGroups
 	}
-	genFilename := func() string {
+	// genFilename := func() string {
-		if len(args) > 1 {
+	// 	if len(args) > 1 {
-			return args[1]
+	// 		return args[1]
-		}
+	// 	}
-		filename := tgutil.GenFileNameFromMessage(*replyTo.Message)
+	// 	filename := tgutil.GenFileNameFromMessage(*replyTo.Message)
-		return filename
+	// 	return filename
-	}()
+	// }()
-	option := tfile.WithNameIfEmpty(genFilename)
+	// option := tfile.WithNameIfEmpty(genFilename)
-	if len(args) > 1 {
+	// if len(args) > 1 {
-		option = tfile.WithName(genFilename)
+	// 	option = tfile.WithName(genFilename)
+	// }
+	userDB, err := database.GetUserByChatID(ctx, update.GetUserChat().GetID())
+	if err != nil {
+		return err
 	}
-	msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, replyTo.Message, option)
+	opts := mediautil.TfileOptions(ctx, userDB, replyTo.Message)
+	if len(args) > 1 {
+		// custom filename via command arg
+		opts = append(opts, tfile.WithName(strings.Join(args[1:], " ")))
+	}
+	msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, replyTo.Message, opts...)
 	if err != nil {
 		return err
 	}
@@ -77,18 +87,27 @@ func handleSilentSaveReplied(ctx *ext.Context, update *ext.Update) error {
 		ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgSaveHelpText)), nil)
 		return dispatcher.EndGroups
 	}
-	genFilename := func() string {
+	// genFilename := func() string {
-		if len(args) > 1 {
+	// 	if len(args) > 1 {
-			return args[1]
+	// 		return args[1]
-		}
+	// 	}
-		filename := tgutil.GenFileNameFromMessage(*replyTo.Message)
+	// 	filename := tgutil.GenFileNameFromMessage(*replyTo.Message)
-		return filename
+	// 	return filename
-	}()
+	// }()
-	option := tfile.WithNameIfEmpty(genFilename)
+	// option := tfile.WithNameIfEmpty(genFilename)
-	if len(args) > 1 {
+	// if len(args) > 1 {
-		option = tfile.WithName(genFilename)
+	// 	option = tfile.WithName(genFilename)
+	// }
+	userDB, err := database.GetUserByChatID(ctx, update.GetUserChat().GetID())
+	if err != nil {
+		return err
 	}
-	msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, replyTo.Message, option)
+	opts := mediautil.TfileOptions(ctx, userDB, replyTo.Message)
+	if len(args) > 1 {
+		// custom filename via command arg
+		opts = append(opts, tfile.WithName(strings.Join(args[1:], " ")))
+	}
+	msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, replyTo.Message, opts...)
 	if err != nil {
 		return err
 	}
@@ -126,7 +145,7 @@ func handleBatchSave(ctx *ext.Context, update *ext.Update, args []string) error
 		return dispatcher.EndGroups
 	}
 
-	// TODO: generator istead of get all messages
+	// [TODO]: generator istead of get all messages
 	msgs, err := tgutil.GetMessagesRange(ctx, chatID, int(startID), int(endID))
 	if err != nil {
 		ctx.Reply(update, ext.ReplyTextString("获取消息失败: "+err.Error()), nil)

client/bot/handlers/utils/mediautil/media.go

@@ -1,6 +1,20 @@
 package mediautil
 
-import "github.com/gotd/td/tg"
+import (
+	"context"
+	"fmt"
+	"strings"
+	"text/template"
+	"time"
+
+	"github.com/charmbracelet/log"
+	"github.com/gotd/td/tg"
+	"github.com/krau/SaveAny-Bot/common/utils/strutil"
+	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
+	"github.com/krau/SaveAny-Bot/database"
+	"github.com/krau/SaveAny-Bot/pkg/enums/fnamest"
+	"github.com/krau/SaveAny-Bot/pkg/tfile"
+)
 
 func IsSupported(media tg.MessageMediaClass) bool {
 	switch media.(type) {
@@ -10,3 +24,83 @@ func IsSupported(media tg.MessageMediaClass) bool {
 		return false
 	}
 }
+
+type FilenameTemplateData struct {
+	MsgID    string `json:"msgid,omitempty"`
+	MsgTags  string `json:"msgtags,omitempty"`
+	MsgGen   string `json:"msggen,omitempty"`
+	MsgDate  string `json:"msgdate,omitempty"`
+	OrigName string `json:"origname,omitempty"`
+}
+
+func (f FilenameTemplateData) ToMap() map[string]string {
+	return map[string]string{
+		"msgid":    f.MsgID,
+		"msgtags":  f.MsgTags,
+		"msggen":   f.MsgGen,
+		"msgdate":  f.MsgDate,
+		"origname": f.OrigName,
+	}
+}
+
+func TfileOptions(ctx context.Context, user *database.User, message *tg.Message) []tfile.TGFileOption {
+	opts := make([]tfile.TGFileOption, 0)
+	var fnameOpt tfile.TGFileOption
+	switch user.FilenameStrategy {
+	case fnamest.Message.String():
+		fnameOpt = tfile.WithName(tgutil.GenFileNameFromMessage(*message))
+	case fnamest.Template.String():
+		if user.FilenameTemplate == "" {
+			log.FromContext(ctx).Warnf("empty filename template")
+			fnameOpt = tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*message))
+			break
+		}
+		tmpl, err := template.New("filename").Parse(user.FilenameTemplate)
+		if err != nil {
+			log.FromContext(ctx).Errorf("failed to parse filename template: %s", err)
+			fnameOpt = tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*message))
+			break
+		}
+		data := FilenameTemplateData{
+			MsgID: func() string {
+				id := message.GetID()
+				if id == 0 {
+					return ""
+				}
+				return fmt.Sprintf("%d", id)
+			}(),
+			MsgTags: func() string {
+				tags := strutil.ExtractTagsFromText(message.GetMessage())
+				if len(tags) == 0 {
+					return ""
+				}
+				return strings.Join(tags, "_")
+			}(),
+			MsgGen: tgutil.GenFileNameFromMessage(*message),
+			OrigName: func() string {
+				f, _ := tgutil.GetMediaFileName(message.Media)
+				return f
+			}(),
+			MsgDate: func() string {
+				date := message.GetDate()
+				if date == 0 {
+					return ""
+				}
+				t := time.Unix(int64(date), 0)
+				return t.Format("2006-01-02_15-04-05")
+			}(),
+		}.ToMap()
+		var sb strings.Builder
+		err = tmpl.Execute(&sb, data)
+		if err != nil {
+			log.FromContext(ctx).Errorf("failed to execute filename template: %s", err)
+			fnameOpt = tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*message))
+			break
+		}
+		fnameOpt = tfile.WithName(sb.String())
+	default:
+		fnameOpt = tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*message))
+	}
+	opts = append(opts, fnameOpt, tfile.WithMessage(message))
+	return opts
+}

client/bot/handlers/utils/shortcut/message.go

@@ -21,7 +21,6 @@ import (
 	"github.com/krau/SaveAny-Bot/common/utils/tphutil"
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/krau/SaveAny-Bot/database"
-	"github.com/krau/SaveAny-Bot/pkg/enums/fnamest"
 	"github.com/krau/SaveAny-Bot/pkg/telegraph"
 	"github.com/krau/SaveAny-Bot/pkg/tfile"
 )
@@ -42,15 +41,15 @@ func GetFileFromMessageWithReply(ctx *ext.Context, update *ext.Update, message *
 		logger.Errorf("Failed to reply: %s", err)
 		return nil, nil, dispatcher.EndGroups
 	}
-	options := []tfile.TGFileOption{
+	// options := []tfile.TGFileOption{
-		tfile.WithMessage(message),
+	// 	tfile.WithMessage(message),
-	}
+	// }
-	if len(tfileopts) > 0 {
+	// if len(tfileopts) > 0 {
-		options = append(options, tfileopts...)
+	// 	options = append(options, tfileopts...)
-	} else {
+	// } else {
-		options = append(options, tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*message)))
+	// 	options = append(options, tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*message)))
-	}
+	// }
-	file, err = tfile.FromMediaMessage(media, ctx.Raw, message, options...)
+	file, err = tfile.FromMediaMessage(media, ctx.Raw, message, tfileopts...)
 	if err != nil {
 		logger.Errorf("Failed to get file from media: %s", err)
 		ctx.Reply(update, ext.ReplyTextString("获取文件失败: "+err.Error()), nil)
@@ -100,14 +99,15 @@ func GetFilesFromUpdateLinkMessageWithReplyEdit(ctx *ext.Context, update *ext.Up
 			logger.Debugf("message %d has no media", msg.GetID())
 			return
 		}
-		var opt tfile.TGFileOption
+		// var opt tfile.TGFileOption
-		switch user.FilenameStrategy {
+		// switch user.FilenameStrategy {
-		case fnamest.Message.String():
+		// case fnamest.Message.String():
-			opt = tfile.WithName(tgutil.GenFileNameFromMessage(*msg))
+		// 	opt = tfile.WithName(tgutil.GenFileNameFromMessage(*msg))
-		default:
+		// default:
-			opt = tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*msg))
+		// 	opt = tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*msg))
-		}
+		// }
-		file, err := tfile.FromMediaMessage(media, client, msg, opt)
+		opts := mediautil.TfileOptions(ctx, user, msg)
+		file, err := tfile.FromMediaMessage(media, client, msg, opts...)
 		if err != nil {
 			logger.Errorf("failed to create file from media: %s", err)
 			return

database/model.go

@@ -14,6 +14,7 @@ type User struct {
 	Rules            []Rule
 	WatchChats       []WatchChat
 	FilenameStrategy string
+	FilenameTemplate string
 }
 
 type WatchChat struct {

pkg/enums/fnamest/filename_srategy.go

@@ -4,11 +4,12 @@ package fnamest
 
 // FnameST
 /* ENUM(
-default, message
+default, message, template
 ) */
 type FnameST string
 
 var FnameSTDisplay = map[FnameST]string{
-	Default: "默认",
+	Default:  "默认",
-	Message: "优先从消息生成",
+	Message:  "优先从消息生成",
+	Template: "自定义模板",
 }

pkg/enums/fnamest/filename_srategy_enum.go

@@ -16,6 +16,8 @@ const (
 	Default FnameST = "default"
 	// Message is a FnameST of type message.
 	Message FnameST = "message"
+	// Template is a FnameST of type template.
+	Template FnameST = "template"
 )
 
 var ErrInvalidFnameST = fmt.Errorf("not a valid FnameST, try [%s]", strings.Join(_FnameSTNames, ", "))
@@ -23,6 +25,7 @@ var ErrInvalidFnameST = fmt.Errorf("not a valid FnameST, try [%s]", strings.Join
 var _FnameSTNames = []string{
 	string(Default),
 	string(Message),
+	string(Template),
 }
 
 // FnameSTNames returns a list of possible string values of FnameST.
@@ -37,6 +40,7 @@ func FnameSTValues() []FnameST {
 	return []FnameST{
 		Default,
 		Message,
+		Template,
 	}
 }
 
@@ -53,8 +57,9 @@ func (x FnameST) IsValid() bool {
 }
 
 var _FnameSTValue = map[string]FnameST{
-	"default": Default,
+	"default":  Default,
-	"message": Message,
+	"message":  Message,
+	"template": Template,
 }
 
 // ParseFnameST attempts to convert a string to a FnameST.