am-downloader/main.go

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/studio-b12/gowebdav"
	"gorm.io/gorm"
)

type Task struct {
	ID         uint      `gorm:"primaryKey" json:"id"`
	Source     string    `json:"source"`
	RID        string    `json:"rid"`
	Title      string    `json:"title"`
	Artist     string    `json:"artist"`
	Album      string    `json:"album"`
	Pic        string    `json:"pic"`
	Quality    string    `json:"quality"`
	Format     string    `json:"format"`
	Status     string    `json:"status"`
	Error      string    `json:"error,omitempty"`
	LocalPath  string    `json:"localPath,omitempty"`
	RemotePath string    `json:"remotePath,omitempty"`
	CreatedAt  time.Time `json:"createdAt"`
	UpdatedAt  time.Time `json:"updatedAt"`
}

type Setting struct {
	ID    uint   `gorm:"primaryKey"`
	Key   string `gorm:"uniqueIndex"`
	Value string
}
type SongRequest struct {
	Source  string `json:"source"`
	RID     string `json:"rid"`
	Title   string `json:"title"`
	Artist  string `json:"artist"`
	Album   string `json:"album"`
	Pic     string `json:"pic"`
	Quality string `json:"quality"`
}
type Config struct {
	AMAPI, DownloadDir, WebDAVURL, WebDAVUser, WebDAVPass, WebDAVPath, DefaultQuality string
	WorkerCount                                                                       int
}
type App struct {
	db         *gorm.DB
	cfg        Config
	queue      chan uint
	wg         sync.WaitGroup
	httpClient *http.Client
}

type playInfo struct {
	URL     string            `json:"url"`
	Format  string            `json:"format"`
	Headers map[string]string `json:"headers"`
}

func main() {
	dbPath := getenv("DB_PATH", "/data/downloader.db")
	_ = os.MkdirAll(filepath.Dir(dbPath), 0755)
	db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
	must(err)
	must(db.AutoMigrate(&Task{}, &Setting{}))
	a := &App{db: db, cfg: Config{AMAPI: getenv("AM_API_URL", "http://127.0.0.1:9530/api"), DownloadDir: getenv("DOWNLOAD_DIR", "/downloads"), DefaultQuality: getenv("DEFAULT_QUALITY", "standard"), WorkerCount: atoi(getenv("WORKERS", "2"), 2)}, queue: make(chan uint, 100), httpClient: &http.Client{Timeout: 2 * time.Minute}}
	_ = os.MkdirAll(a.cfg.DownloadDir, 0755)
	a.loadSettings()
	a.startWorkers()
	r := gin.Default()
	r.StaticFile("/", "/app/index.html")
	api := r.Group("/api")
	api.GET("/tasks", a.listTasks)
	api.POST("/tasks", a.addTask)
	api.POST("/tasks/:id/retry", a.retryTask)
	api.DELETE("/tasks/:id", a.deleteTask)
	api.GET("/settings", a.getSettings)
	api.PUT("/settings", a.putSettings)
	api.POST("/webdav/test", a.testWebDAV)
	must(r.Run(":" + getenv("PORT", "9540")))
}

func (a *App) loadSettings() {
	var ss []Setting
	a.db.Find(&ss)
	for _, s := range ss {
		switch s.Key {
		case "am_api":
			a.cfg.AMAPI = s.Value
		case "download_dir":
			a.cfg.DownloadDir = s.Value
		case "webdav_url":
			a.cfg.WebDAVURL = s.Value
		case "webdav_user":
			a.cfg.WebDAVUser = s.Value
		case "webdav_pass":
			a.cfg.WebDAVPass = s.Value
		case "webdav_path":
			a.cfg.WebDAVPath = s.Value
		case "quality":
			a.cfg.DefaultQuality = s.Value
		}
	}
}
func (a *App) save(k, v string) {
	var s Setting
	if a.db.Where("key = ?", k).First(&s).Error != nil {
		a.db.Create(&Setting{Key: k, Value: v})
	} else {
		s.Value = v
		a.db.Save(&s)
	}
}
func (a *App) startWorkers() {
	for i := 0; i < a.cfg.WorkerCount; i++ {
		a.wg.Add(1)
		go func() {
			defer a.wg.Done()
			for id := range a.queue {
				a.download(id)
			}
		}()
	}
}
func (a *App) addTask(c *gin.Context) {
	var in SongRequest
	if c.BindJSON(&in) != nil || in.RID == "" || in.Title == "" {
		c.JSON(400, gin.H{"error": "source/rid/title required"})
		return
	}
	if in.Quality == "" {
		in.Quality = a.cfg.DefaultQuality
	}
	t := Task{Source: in.Source, RID: in.RID, Title: in.Title, Artist: in.Artist, Album: in.Album, Pic: in.Pic, Quality: in.Quality, Status: "pending"}
	if err := a.db.Create(&t).Error; err != nil {
		c.JSON(500, gin.H{"error": err.Error()})
		return
	}
	a.queue <- t.ID
	c.JSON(202, t)
}
func (a *App) listTasks(c *gin.Context) {
	var ts []Task
	a.db.Order("id desc").Limit(100).Find(&ts)
	c.JSON(200, gin.H{"data": ts})
}
func (a *App) retryTask(c *gin.Context) {
	id, _ := strconv.ParseUint(c.Param("id"), 10, 64)
	var t Task
	if a.db.First(&t, id).Error != nil {
		c.Status(404)
		return
	}
	a.db.Model(&t).Updates(map[string]any{"status": "pending", "error": ""})
	a.queue <- t.ID
	c.JSON(202, t)
}
func (a *App) deleteTask(c *gin.Context) {
	id, _ := strconv.ParseUint(c.Param("id"), 10, 64)
	a.db.Delete(&Task{}, id)
	c.Status(204)
}

func (a *App) download(id uint) {
	var t Task
	if a.db.First(&t, id).Error != nil {
		return
	}
	a.db.Model(&t).Updates(map[string]any{"status": "downloading", "error": ""})
	info, err := a.resolve(t)
	if err != nil {
		a.fail(&t, err)
		return
	}
	ext := info.Format
	if ext == "" {
		ext = "mp3"
	}
	filename := safe(t.Artist) + " - " + safe(t.Title) + "." + safe(ext)
	local := filepath.Join(a.cfg.DownloadDir, filename)
	t.Format = ext
	t.LocalPath = local
	if _, err = os.Stat(local); err == nil {
		a.upload(&t)
		return
	}
	req, err := http.NewRequest("GET", info.URL, nil)
	if err != nil {
		a.fail(&t, err)
		return
	}
	for k, v := range info.Headers {
		req.Header.Set(k, v)
	}
	resp, err := a.httpClient.Do(req)
	if err != nil {
		a.fail(&t, err)
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		a.fail(&t, fmt.Errorf("media http %d", resp.StatusCode))
		return
	}
	f, err := os.Create(local)
	if err != nil {
		a.fail(&t, err)
		return
	}
	_, err = io.Copy(f, resp.Body)
	cerr := f.Close()
	if err != nil {
		a.fail(&t, err)
		return
	}
	if cerr != nil {
		a.fail(&t, cerr)
		return
	}
	a.upload(&t)
}

func (a *App) resolve(t Task) (playInfo, error) {
	u, _ := url.Parse(strings.TrimRight(a.cfg.AMAPI, "/") + "/song/play-url")
	q := u.Query()
	q.Set("source", t.Source)
	q.Set("rid", t.RID)
	q.Set("level", t.Quality)
	u.RawQuery = q.Encode()
	resp, err := a.httpClient.Get(u.String())
	if err != nil {
		return playInfo{}, err
	}
	defer resp.Body.Close()
	var x struct {
		Code    int      `json:"code"`
		Data    playInfo `json:"data"`
		Message string   `json:"message"`
	}
	if err = json.NewDecoder(resp.Body).Decode(&x); err != nil {
		return playInfo{}, err
	}
	if x.Code != 0 || x.Data.URL == "" {
		return playInfo{}, errors.New(x.Message)
	}
	return x.Data, nil
}

func (a *App) upload(t *Task) {
	if a.cfg.WebDAVURL == "" {
		t.Status = "done"
		a.db.Save(t)
		return
	}
	a.db.Model(t).Updates(map[string]any{"status": "uploading"})
	remote := strings.TrimRight(a.cfg.WebDAVPath, "/") + "/" + safe(t.Artist) + " - " + safe(t.Title) + "." + safe(t.Format)
	t.RemotePath = remote
	client := gowebdav.NewClient(a.cfg.WebDAVURL, a.cfg.WebDAVUser, a.cfg.WebDAVPass)
	client.SetTimeout(10 * time.Minute)
	if err := client.MkdirAll(strings.TrimSuffix(remote, "/"), 0755); err != nil {
		a.fail(t, err)
		return
	}
	f, err := os.Open(t.LocalPath)
	if err != nil {
		a.fail(t, err)
		return
	}
	defer f.Close()
	if err = client.WriteStream(remote, f, 0644); err != nil {
		a.fail(t, err)
		return
	}
	t.Status = "done"
	a.db.Save(t)
}
func (a *App) fail(t *Task, err error) { t.Status = "error"; t.Error = err.Error(); a.db.Save(t) }

func (a *App) getSettings(c *gin.Context) {
	c.JSON(200, gin.H{"amApi": a.cfg.AMAPI, "downloadDir": a.cfg.DownloadDir, "webdavUrl": a.cfg.WebDAVURL, "webdavUser": a.cfg.WebDAVUser, "webdavPass": a.cfg.WebDAVPass, "webdavPath": a.cfg.WebDAVPath, "quality": a.cfg.DefaultQuality})
}
func (a *App) putSettings(c *gin.Context) {
	var v map[string]string
	if c.BindJSON(&v) != nil {
		c.Status(400)
		return
	}
	if x := v["amApi"]; x != "" {
		a.cfg.AMAPI = x
		a.save("am_api", x)
	}
	if x := v["downloadDir"]; x != "" {
		a.cfg.DownloadDir = x
		a.save("download_dir", x)
		os.MkdirAll(x, 0755)
	}
	a.cfg.WebDAVURL = v["webdavUrl"]
	a.cfg.WebDAVUser = v["webdavUser"]
	a.cfg.WebDAVPass = v["webdavPass"]
	a.cfg.WebDAVPath = strings.Trim(v["webdavPath"], "/")
	a.cfg.DefaultQuality = v["quality"]
	for k, x := range map[string]string{"webdav_url": a.cfg.WebDAVURL, "webdav_user": a.cfg.WebDAVUser, "webdav_pass": a.cfg.WebDAVPass, "webdav_path": a.cfg.WebDAVPath, "quality": a.cfg.DefaultQuality} {
		a.save(k, x)
	}
	c.JSON(200, gin.H{"ok": true})
}
func (a *App) testWebDAV(c *gin.Context) {
	var v struct{ URL, User, Pass, Path string }
	if c.BindJSON(&v) != nil {
		c.Status(400)
		return
	}
	client := gowebdav.NewClient(v.URL, v.User, v.Pass)
	client.SetTimeout(20 * time.Second)
	if err := client.Connect(); err != nil {
		c.JSON(502, gin.H{"ok": false, "error": err.Error()})
		return
	}
	c.JSON(200, gin.H{"ok": true, "path": v.Path})
}

var bad = regexp.MustCompile(`[/\\:*?"<>|]+`)

func safe(s string) string {
	s = strings.TrimSpace(s)
	s = bad.ReplaceAllString(s, "_")
	s = strings.Trim(s, ". ")
	if s == "" {
		return "Unknown"
	}
	return s
}
func getenv(k, d string) string {
	if v := os.Getenv(k); v != "" {
		return v
	}
	return d
}
func atoi(s string, d int) int {
	v, e := strconv.Atoi(s)
	if e != nil || v < 1 {
		return d
	}
	return v
}
func must(e error) {
	if e != nil {
		panic(e)
	}
}