Reapply download map and unrestrict map

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 (

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=

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))

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))

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))

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 {
@@ -15,6 +16,7 @@ type File struct {
 	State *fsm.FSM `json:"-"`
 
 	Rename    string             `json:"Rename"`
+	MediaInfo *ffprobe.ProbeData `json:"MediaInfo"`
 }
 
 func (f *File) MarshalJSON() ([]byte, error) {

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 (
@@ -32,6 +34,7 @@ type TorrentManager struct {
 
 	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
 
@@ -65,6 +68,7 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w
 
 		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]) {
+	isRealDebrid := strings.HasPrefix(link, "https://real-debrid.com/d/")
-		ret, _ := t.DownloadMap.Get(link[0:39])
+	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() {
+func (t *TorrentManager) mountNewDownloads() {
-	t.DownloadMap.Clear()
-	_ = t.workerPool.Submit(func() {
 	downloads := t.api.GetDownloads()
-		for i := range downloads {
+	for _, download := range downloads {
-			idx := i
+		isRealDebrid := strings.HasPrefix(download.Link, "https://real-debrid.com/d/")
-			t.DownloadMap.Set(downloads[idx].Filename, &downloads[idx])
+		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()
 			}
 		}
 	})