diff --git a/bot/utils.go b/bot/utils.go index 2e1e9f7..2bfc9b6 100644 --- a/bot/utils.go +++ b/bot/utils.go @@ -4,7 +4,7 @@ import ( "fmt" "regexp" - "github.com/krau/SaveAny-Bot/config" + "github.com/krau/SaveAny-Bot/storage" "github.com/mymmrac/telego" "github.com/mymmrac/telego/telegoutil" ) @@ -20,15 +20,18 @@ func ReplyMessage(replyTo telego.Message, format string, args ...any) (*telego.M })) } +var StorageDisplayNames = map[string]string{ + "all": "全部", + "local": "服务器磁盘", + "alist": "Alist", + "webdav": "WebDAV", +} + func AddTaskReplyMarkup(messageID int) *telego.InlineKeyboardMarkup { storageButtons := make([]telego.InlineKeyboardButton, 0) - if config.Cfg.Storage.Local.Enable { - storageButtons = append(storageButtons, telegoutil.InlineKeyboardButton("服务器磁盘"). - WithCallbackData(fmt.Sprintf("add %d local", messageID))) - } - if config.Cfg.Storage.Alist.Enable { - storageButtons = append(storageButtons, telegoutil.InlineKeyboardButton("Alist"). - WithCallbackData(fmt.Sprintf("add %d alist", messageID))) + for name := range storage.Storages { + storageButtons = append(storageButtons, telegoutil.InlineKeyboardButton(StorageDisplayNames[string(name)]). + WithCallbackData(fmt.Sprintf("add %d %s", messageID, name))) } if len(storageButtons) > 1 { diff --git a/config.example.toml b/config.example.toml index 1bbaf76..6f777e6 100644 --- a/config.example.toml +++ b/config.example.toml @@ -22,3 +22,10 @@ token_exp = 86400 [storage.local] enable = true base_path = "downloads/" + +[storage.webdav] +enable = true +base_path = "/telegram" +username = "admin" +password = "password" +url = "https://alist.com" \ No newline at end of file diff --git a/config/viper.go b/config/viper.go index b0aea59..6b56bbc 100644 --- a/config/viper.go +++ b/config/viper.go @@ -37,8 +37,9 @@ type telegramConfig struct { } type storageConfig struct { - Alist alistConfig `toml:"alist" mapstructure:"alist"` - Local localConfig `toml:"local" mapstructure:"local"` + Alist alistConfig `toml:"alist" mapstructure:"alist"` + Local localConfig `toml:"local" mapstructure:"local"` + Webdav webdavConfig `toml:"webdav" mapstructure:"webdav"` } type alistConfig struct { @@ -55,6 +56,14 @@ type localConfig struct { BasePath string `toml:"base_path" mapstructure:"base_path"` } +type webdavConfig struct { + Enable bool `toml:"enable" mapstructure:"enable"` + URL string `toml:"url" mapstructure:"url"` + Username string `toml:"username" mapstructure:"username"` + Password string `toml:"password" mapstructure:"password"` + BasePath string `toml:"base_path" mapstructure:"base_path"` +} + var Cfg *Config func Init() { diff --git a/go.mod b/go.mod index 725e1ad..6a2cca1 100644 --- a/go.mod +++ b/go.mod @@ -33,6 +33,7 @@ require ( github.com/quic-go/quic-go v0.47.0 // indirect github.com/refraction-networking/utls v1.6.7 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + github.com/studio-b12/gowebdav v0.9.0 // indirect github.com/tcnksm/go-gitconfig v0.1.2 // indirect github.com/ulikunitz/xz v0.5.9 // indirect go.uber.org/mock v0.4.0 // indirect diff --git a/go.sum b/go.sum index 418542a..6899cec 100644 --- a/go.sum +++ b/go.sum @@ -158,6 +158,8 @@ github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/studio-b12/gowebdav v0.9.0 h1:1j1sc9gQnNxbXXM4M/CebPOX4aXYtr7MojAVcN4dHjU= +github.com/studio-b12/gowebdav v0.9.0/go.mod h1:bHA7t77X/QFExdeAnDzK6vKM34kEZAcE1OX4MfiwjkE= github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8= github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU= github.com/tcnksm/go-gitconfig v0.1.2 h1:iiDhRitByXAEyjgBqsKi9QU4o2TNtv9kPP3RgPgXBPw= diff --git a/storage/storage.go b/storage/storage.go index abe3815..8612a48 100644 --- a/storage/storage.go +++ b/storage/storage.go @@ -9,6 +9,7 @@ import ( "github.com/krau/SaveAny-Bot/logger" "github.com/krau/SaveAny-Bot/storage/alist" "github.com/krau/SaveAny-Bot/storage/local" + "github.com/krau/SaveAny-Bot/storage/webdav" "github.com/krau/SaveAny-Bot/types" ) @@ -29,6 +30,10 @@ func Init() { Storages[types.Local] = new(local.Local) Storages[types.Local].Init() } + if config.Cfg.Storage.Webdav.Enable { + Storages[types.Webdav] = new(webdav.Webdav) + Storages[types.Webdav].Init() + } logger.L.Debug("Storage initialized") } diff --git a/storage/webdav/webdav.go b/storage/webdav/webdav.go new file mode 100644 index 0000000..cb87289 --- /dev/null +++ b/storage/webdav/webdav.go @@ -0,0 +1,48 @@ +package webdav + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + + "github.com/krau/SaveAny-Bot/config" + "github.com/krau/SaveAny-Bot/logger" + "github.com/studio-b12/gowebdav" +) + +type Webdav struct{} + +var ( + Client *gowebdav.Client + basePath string +) + +func (w *Webdav) Init() { + webdavConfig := config.Cfg.Storage.Webdav + basePath = strings.TrimSuffix(webdavConfig.BasePath, "/") + Client = gowebdav.NewClient(webdavConfig.URL, webdavConfig.Username, webdavConfig.Password) + if err := Client.Connect(); err != nil { + logger.L.Fatalf("Failed to connect to webdav server: %v", err) + os.Exit(1) + } +} + +func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error { + storagePath = filepath.Join(basePath, storagePath) + if err := Client.MkdirAll(filepath.Dir(storagePath), os.ModePerm); err != nil { + logger.L.Errorf("Failed to create directory %s: %v", filepath.Dir(storagePath), err) + return errors.New("webdav: failed to create directory") + } + fileBytes, err := os.ReadFile(filePath) + if err != nil { + logger.L.Errorf("Failed to read file %s: %v", filePath, err) + return err + } + if err := Client.Write(storagePath, fileBytes, os.ModePerm); err != nil { + logger.L.Errorf("Failed to write file %s: %v", storagePath, err) + return errors.New("webdav: failed to write file") + } + return nil +} diff --git a/types/types.go b/types/types.go index d8e1384..6822866 100644 --- a/types/types.go +++ b/types/types.go @@ -16,10 +16,11 @@ type StorageType string var ( StorageAll StorageType = "all" Local StorageType = "local" + Webdav StorageType = "webdav" Alist StorageType = "alist" ) -var StorageTypes = []StorageType{Local, Alist, StorageAll} +var StorageTypes = []StorageType{Local, Alist, Webdav, StorageAll} type Task struct { Ctx context.Context