From e6a47be656cb4b9f78a7cd5b05cdf57249961da6 Mon Sep 17 00:00:00 2001 From: Ben Sarmiento Date: Sat, 25 May 2024 21:09:12 +0200 Subject: [PATCH] Reapply download map and unrestrict map --- go.mod | 1 + go.sum | 2 + internal/dav/infuse.go | 2 +- internal/dav/listing.go | 2 +- internal/http/listing.go | 2 +- internal/torrent/file_types.go | 4 +- internal/torrent/manager.go | 93 ++++++++++++++++++++++++++-------- 7 files changed, 82 insertions(+), 24 deletions(-) diff --git a/go.mod b/go.mod index c46af5d..70dc785 100644 --- a/go.mod +++ b/go.mod @@ -23,6 +23,7 @@ require ( github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 // indirect github.com/modern-go/reflect2 v1.0.2 // indirect gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect + gopkg.in/vansante/go-ffprobe.v2 v2.1.1 // indirect ) require ( diff --git a/go.sum b/go.sum index 3f974a7..3d2f116 100644 --- a/go.sum +++ b/go.sum @@ -52,6 +52,8 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc= gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc= +gopkg.in/vansante/go-ffprobe.v2 v2.1.1 h1:DIh5fMn+tlBvG7pXyUZdemVmLdERnf2xX6XOFF+0BBU= +gopkg.in/vansante/go-ffprobe.v2 v2.1.1/go.mod h1:qF0AlAjk7Nqzqf3y333Ly+KxN3cKF2JqA3JT5ZheUGE= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/dav/infuse.go b/internal/dav/infuse.go index 56a7e85..3539bfe 100644 --- a/internal/dav/infuse.go +++ b/internal/dav/infuse.go @@ -114,7 +114,7 @@ func ServeDownloadsListForInfuse(torMgr *torrent.TorrentManager) ([]byte, error) for _, filename := range filenames { download, ok := torMgr.DownloadMap.Get(filename) - if !ok || strings.HasPrefix(download.Link, "https://real-debrid.com/d/") { + if !ok { continue } buf.WriteString(dav.File(download.Filename, download.Filesize, download.Generated)) diff --git a/internal/dav/listing.go b/internal/dav/listing.go index 68dc585..0c448d9 100644 --- a/internal/dav/listing.go +++ b/internal/dav/listing.go @@ -125,7 +125,7 @@ func ServeDownloadsList(torMgr *torrent.TorrentManager) ([]byte, error) { sort.Strings(filenames) for _, filename := range filenames { download, ok := torMgr.DownloadMap.Get(filename) - if !ok || strings.HasPrefix(download.Link, "https://real-debrid.com/d/") { + if !ok { continue } buf.WriteString(dav.File(download.Filename, download.Filesize, download.Generated)) diff --git a/internal/http/listing.go b/internal/http/listing.go index 2a49afb..98d6b00 100644 --- a/internal/http/listing.go +++ b/internal/http/listing.go @@ -98,7 +98,7 @@ func ServeDownloadsList(torMgr *torrent.TorrentManager) ([]byte, error) { sort.Strings(filenames) for _, filename := range filenames { download, ok := torMgr.DownloadMap.Get(filename) - if !ok || strings.HasPrefix(download.Link, "https://real-debrid.com/d/") { + if !ok { continue } filePath := filepath.Join(config.DOWNLOADS, url.PathEscape(filename)) diff --git a/internal/torrent/file_types.go b/internal/torrent/file_types.go index 89ac9bb..fe5c8d5 100644 --- a/internal/torrent/file_types.go +++ b/internal/torrent/file_types.go @@ -5,6 +5,7 @@ import ( "github.com/debridmediamanager/zurg/pkg/realdebrid" "github.com/looplab/fsm" + "gopkg.in/vansante/go-ffprobe.v2" ) type File struct { @@ -14,7 +15,8 @@ type File struct { State *fsm.FSM `json:"-"` - Rename string `json:"Rename"` + Rename string `json:"Rename"` + MediaInfo *ffprobe.ProbeData `json:"MediaInfo"` } func (f *File) MarshalJSON() ([]byte, error) { diff --git a/internal/torrent/manager.go b/internal/torrent/manager.go index 81d00e2..24e027f 100644 --- a/internal/torrent/manager.go +++ b/internal/torrent/manager.go @@ -1,6 +1,7 @@ package torrent import ( + "context" "io" "os" "path/filepath" @@ -15,6 +16,7 @@ import ( mapset "github.com/deckarep/golang-set/v2" cmap "github.com/orcaman/concurrent-map/v2" "github.com/panjf2000/ants/v2" + "gopkg.in/vansante/go-ffprobe.v2" ) const ( @@ -30,8 +32,9 @@ type TorrentManager struct { log *logutil.Logger repairLog *logutil.Logger - DirectoryMap cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent - DownloadMap cmap.ConcurrentMap[string, *realdebrid.Download] + DirectoryMap cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent + DownloadMap cmap.ConcurrentMap[string, *realdebrid.Download] + UnrestrictMap cmap.ConcurrentMap[string, *realdebrid.Download] RootNode *fs.FileNode @@ -63,8 +66,9 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w log: log, repairLog: repairLog, - DirectoryMap: cmap.New[cmap.ConcurrentMap[string, *Torrent]](), - DownloadMap: cmap.New[*realdebrid.Download](), + DirectoryMap: cmap.New[cmap.ConcurrentMap[string, *Torrent]](), + DownloadMap: cmap.New[*realdebrid.Download](), + UnrestrictMap: cmap.New[*realdebrid.Download](), RootNode: fs.NewFileNode("root", true), @@ -77,34 +81,55 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w t.initializeBins() t.initializeDirectoryMaps() + var wg sync.WaitGroup + wg.Add(2) t.workerPool.Submit(func() { t.refreshTorrents() + wg.Done() t.setNewLatestState(t.getCurrentState()) t.StartRefreshJob() t.StartRepairJob() t.TriggerRepair(nil) }) t.workerPool.Submit(func() { - t.mountDownloads() + t.mountNewDownloads() + wg.Done() t.StartDownloadsJob() }) + t.workerPool.Submit(func() { + wg.Wait() + t.log.Info("Applying media info details to all torrents") + allTorrents, _ := t.DirectoryMap.Get(INT_ALL) + allTorrents.IterCb(func(_ string, torrent *Torrent) { + t.applyMediaInfoDetails(torrent) + t.writeTorrentToFile(torrent) + t.log.Debugf("Applied media info details to torrent %s", torrent.Hash) + }) + }) return t } // proxy function func (t *TorrentManager) UnrestrictLinkUntilOk(link string, checkFirstByte bool) *realdebrid.Download { - if strings.HasPrefix(link, "https://real-debrid.com/d/") && t.DownloadMap.Has(link[0:39]) { - ret, _ := t.DownloadMap.Get(link[0:39]) + isRealDebrid := strings.HasPrefix(link, "https://real-debrid.com/d/") + if isRealDebrid && t.UnrestrictMap.Has(link[0:39]) { + ret, _ := t.UnrestrictMap.Get(link[0:39]) + return ret + } else if !isRealDebrid && t.UnrestrictMap.Has(link) { + ret, _ := t.UnrestrictMap.Get(link) return ret } ret, err := t.api.UnrestrictLink(link, checkFirstByte) - t.DownloadMap.Set(ret.Link[0:39], ret) if err != nil { t.log.Warnf("Cannot unrestrict link %s: %v", link, err) return nil } - t.DownloadMap.Set(ret.Filename, ret) + if isRealDebrid { + t.UnrestrictMap.Set(ret.Link[0:39], ret) + } else { + t.UnrestrictMap.Set(ret.Link, ret) + } return ret } @@ -152,7 +177,7 @@ func (t *TorrentManager) getTorrentFiles(parentDir string) mapset.Set[string] { } func (t *TorrentManager) writeTorrentToFile(torrent *Torrent) { - filePath := "data/" + torrent.Hash + ".zurgtorrent" + filePath := "data/" + t.GetKey(torrent) + ".zurgtorrent" file, err := os.Create(filePath) if err != nil { t.log.Warnf("Cannot create file %s: %v", filePath, err) @@ -176,6 +201,28 @@ func (t *TorrentManager) writeTorrentToFile(torrent *Torrent) { t.log.Debugf("Saved torrent %s (hash=%s) to file", t.GetKey(torrent), torrent.Hash) } +func (t *TorrentManager) applyMediaInfoDetails(torrent *Torrent) { + torrent.SelectedFiles.IterCb(func(_ string, file *File) { + if file.MediaInfo != nil { + return + } + unrestrict := t.UnrestrictFileUntilOk(file, false) + if unrestrict == nil { + file.State.Event(context.Background(), "break_file") + return + } + ctx, cancelFn := context.WithTimeout(context.Background(), 5*time.Second) + defer cancelFn() + data, err := ffprobe.ProbeURL(ctx, unrestrict.Download) + if err != nil { + t.log.Warnf("Cannot probe file %s: %v", file.Path, err) + file.State.Event(context.Background(), "break_file") + return + } + file.MediaInfo = data + }) +} + func (t *TorrentManager) readTorrentFromFile(filePath string) *Torrent { file, err := os.Open(filePath) if err != nil { @@ -268,15 +315,20 @@ func (t *TorrentManager) deleteInfoFile(torrentID string) { /// end info functions -func (t *TorrentManager) mountDownloads() { - t.DownloadMap.Clear() - _ = t.workerPool.Submit(func() { - downloads := t.api.GetDownloads() - for i := range downloads { - idx := i - t.DownloadMap.Set(downloads[idx].Filename, &downloads[idx]) +func (t *TorrentManager) mountNewDownloads() { + downloads := t.api.GetDownloads() + for _, download := range downloads { + isRealDebrid := strings.HasPrefix(download.Link, "https://real-debrid.com/d/") + if isRealDebrid && !t.UnrestrictMap.Has(download.Link[0:39]) { + t.UnrestrictMap.Set(download.Link[0:39], &download) + } else if !isRealDebrid { + if !t.UnrestrictMap.Has(download.Link) { + t.UnrestrictMap.Set(download.Link, &download) + } + filename := filepath.Base(download.Filename) + t.DownloadMap.Set(filename, &download) } - }) + } } func (t *TorrentManager) StartDownloadsJob() { @@ -287,9 +339,10 @@ func (t *TorrentManager) StartDownloadsJob() { for { select { case <-remountTicker.C: - t.mountDownloads() + t.mountNewDownloads() case <-t.RemountTrigger: - t.mountDownloads() + t.DownloadMap.Clear() + t.mountNewDownloads() } } })