From c63ce8da0aa172abd1fcc4628c793d3fdd83a050 Mon Sep 17 00:00:00 2001 From: Ben Adrian Sarmiento Date: Wed, 5 Jun 2024 19:19:03 +0200 Subject: [PATCH] Add media info filters --- internal/config/types.go | 8 +- internal/config/v1.go | 200 +++++++++++++++++++++++++++++++++++- internal/config/v1types.go | 14 +++ internal/torrent/bins.go | 17 +-- internal/torrent/manager.go | 20 ++-- internal/torrent/refresh.go | 61 +++++------ internal/torrent/repair.go | 3 +- 7 files changed, 265 insertions(+), 58 deletions(-) diff --git a/internal/config/types.go b/internal/config/types.go index c3bce18..cb9f9ef 100644 --- a/internal/config/types.go +++ b/internal/config/types.go @@ -1,6 +1,10 @@ package config -import "os" +import ( + "os" + + "gopkg.in/vansante/go-ffprobe.v2" +) type ConfigInterface interface { GetConfig() ZurgConfig @@ -16,7 +20,7 @@ type ConfigInterface interface { GetPassword() string GetProxy() string GetDirectories() []string - MeetsConditions(directory, torrentName string, torrentSize int64, torrentIDs, fileNames []string, fileSizes []int64) bool + MeetsConditions(directory, torrentName string, torrentSize int64, torrentIDs, fileNames []string, fileSizes []int64, mediaInfos []*ffprobe.ProbeData) bool GetOnLibraryUpdate() string GetNetworkBufferSize() int EnableRetainFolderNameExtension() bool diff --git a/internal/config/v1.go b/internal/config/v1.go index a308032..5c7cd0f 100644 --- a/internal/config/v1.go +++ b/internal/config/v1.go @@ -9,6 +9,7 @@ import ( "github.com/debridmediamanager/zurg/pkg/logutil" "github.com/debridmediamanager/zurg/pkg/utils" + "gopkg.in/vansante/go-ffprobe.v2" "gopkg.in/yaml.v3" ) @@ -96,7 +97,7 @@ func (z *ZurgConfigV1) GetGroupMap() map[string][]string { return result } -func (z *ZurgConfigV1) MeetsConditions(directory, torrentName string, torrentSize int64, torrentIDs, fileNames []string, fileSizes []int64) bool { +func (z *ZurgConfigV1) MeetsConditions(directory, torrentName string, torrentSize int64, torrentIDs, fileNames []string, fileSizes []int64, mediaInfos []*ffprobe.ProbeData) bool { if directory == ALL_TORRENTS { return true } @@ -104,14 +105,14 @@ func (z *ZurgConfigV1) MeetsConditions(directory, torrentName string, torrentSiz return false } for _, filter := range z.Directories[directory].Filters { - if z.matchFilter(torrentName, torrentSize, torrentIDs, fileNames, fileSizes, filter) { + if z.matchFilter(torrentName, torrentSize, torrentIDs, fileNames, fileSizes, mediaInfos, filter) { return true } } return false } -func (z *ZurgConfigV1) matchFilter(torrentName string, torrentSize int64, torrentIDs, fileNames []string, fileSizes []int64, filter *FilterConditionsV1) bool { +func (z *ZurgConfigV1) matchFilter(torrentName string, torrentSize int64, torrentIDs, fileNames []string, fileSizes []int64, mediaInfos []*ffprobe.ProbeData, filter *FilterConditionsV1) bool { if filter.ID != "" { for _, torrentID := range torrentIDs { if torrentID == filter.ID { @@ -160,7 +161,7 @@ func (z *ZurgConfigV1) matchFilter(torrentName string, torrentSize int64, torren if len(filter.And) > 0 { andResult := true for _, andFilter := range filter.And { - andResult = andResult && z.matchFilter(torrentName, torrentSize, torrentIDs, fileNames, fileSizes, andFilter) + andResult = andResult && z.matchFilter(torrentName, torrentSize, torrentIDs, fileNames, fileSizes, mediaInfos, andFilter) if !andResult { return false } @@ -169,7 +170,7 @@ func (z *ZurgConfigV1) matchFilter(torrentName string, torrentSize int64, torren } if len(filter.Or) > 0 { for _, orFilter := range filter.Or { - if z.matchFilter(torrentName, torrentSize, torrentIDs, fileNames, fileSizes, orFilter) { + if z.matchFilter(torrentName, torrentSize, torrentIDs, fileNames, fileSizes, mediaInfos, orFilter) { return true } } @@ -292,6 +293,195 @@ func (z *ZurgConfigV1) matchFilter(torrentName string, torrentSize int64, torren } } } + + // media info filters + if filter.MediaInfoResolution != "" { + for _, mediaInfo := range mediaInfos { + for _, stream := range mediaInfo.Streams { + if stream.CodecType != "video" || stream.Level <= 0 { + continue + } else if (stream.Width >= 7680 || stream.Height >= 4320) && filter.MediaInfoResolution == "8k" { + return true + } else if (stream.Width >= 3840 || stream.Height >= 2160) && filter.MediaInfoResolution == "4k" { + return true + } else if (stream.Width >= 1920 || stream.Height >= 1080) && filter.MediaInfoResolution == "1080p" { + return true + } else if (stream.Width >= 1280 || stream.Height >= 720) && filter.MediaInfoResolution == "720p" { + return true + } else if (stream.Width >= 854 || stream.Height >= 480) && filter.MediaInfoResolution == "480p" { + return true + } else if (stream.Width >= 640 || stream.Height >= 360) && filter.MediaInfoResolution == "360p" { + return true + } else if (stream.Width >= 426 || stream.Height >= 240) && filter.MediaInfoResolution == "240p" { + return true + } else if (stream.Width >= 256 || stream.Height >= 144) && filter.MediaInfoResolution == "144p" { + return true + } + } + } + return false + } + if filter.MediaInfoBitRateGreaterThanOrEqual > 0 { + for _, mediaInfo := range mediaInfos { + bitrate, err := strconv.ParseInt(mediaInfo.Format.BitRate, 10, 64) + if err != nil { + continue + } + if bitrate >= filter.MediaInfoBitRateGreaterThanOrEqual { + return true + } + } + return false + } + if filter.MediaInfoBitRateLessThanOrEqual > 0 { + for _, mediaInfo := range mediaInfos { + bitrate, err := strconv.ParseInt(mediaInfo.Format.BitRate, 10, 64) + if err != nil { + continue + } + if bitrate <= filter.MediaInfoBitRateLessThanOrEqual { + return true + } + } + return false + } + if filter.MediaInfoVideoBitRateGreaterThanOrEqual > 0 { + for _, mediaInfo := range mediaInfos { + for _, stream := range mediaInfo.Streams { + if stream.CodecType != "video" || stream.Level <= 0 { + continue + } + bitrate, err := strconv.ParseInt(stream.BitRate, 10, 64) + if err != nil { + continue + } + if bitrate >= filter.MediaInfoVideoBitRateGreaterThanOrEqual { + return true + } + } + } + return false + } + if filter.MediaInfoVideoBitRateLessThanOrEqual > 0 { + for _, mediaInfo := range mediaInfos { + for _, stream := range mediaInfo.Streams { + if stream.CodecType != "video" || stream.Level <= 0 { + continue + } + bitrate, err := strconv.ParseInt(stream.BitRate, 10, 64) + if err != nil { + continue + } + if bitrate <= filter.MediaInfoVideoBitRateLessThanOrEqual { + return true + } + } + } + return false + } + if filter.MediaInfoAudioBitRateGreaterThanOrEqual > 0 { + for _, mediaInfo := range mediaInfos { + for _, stream := range mediaInfo.Streams { + if stream.CodecType != "audio" { + continue + } + bitrate, err := strconv.ParseInt(stream.BitRate, 10, 64) + if err != nil { + continue + } + if bitrate >= filter.MediaInfoAudioBitRateGreaterThanOrEqual { + return true + } + } + } + return false + } + if filter.MediaInfoAudioBitRateLessThanOrEqual > 0 { + for _, mediaInfo := range mediaInfos { + for _, stream := range mediaInfo.Streams { + if stream.CodecType != "audio" { + continue + } + bitrate, err := strconv.ParseInt(stream.BitRate, 10, 64) + if err != nil { + continue + } + if bitrate <= filter.MediaInfoAudioBitRateLessThanOrEqual { + return true + } + } + } + return false + } + if filter.MediaInfoDurationGreaterThanOrEqual > 0 { + for _, mediaInfo := range mediaInfos { + if int64(mediaInfo.Format.DurationSeconds) >= filter.MediaInfoDurationGreaterThanOrEqual { + return true + } + } + return false + } + if filter.MediaInfoDurationLessThanOrEqual > 0 { + for _, mediaInfo := range mediaInfos { + if int64(mediaInfo.Format.DurationSeconds) <= filter.MediaInfoDurationLessThanOrEqual { + return true + } + } + return false + } + if filter.MediaInfoWithAudioLanguage != "" { + for _, mediaInfo := range mediaInfos { + for _, stream := range mediaInfo.Streams { + if stream.CodecType != "audio" { + continue + } + if strings.EqualFold(stream.Tags.Language, filter.MediaInfoWithAudioLanguage) { + return true + } + } + } + return false + } + if filter.MediaInfoWithoutAudioLanguage != "" { + for _, mediaInfo := range mediaInfos { + for _, stream := range mediaInfo.Streams { + if stream.CodecType != "audio" { + continue + } + if strings.EqualFold(stream.Tags.Language, filter.MediaInfoWithoutAudioLanguage) { + return false + } + } + } + return true + } + if filter.MediaInfoWithSubtitleLanguage != "" { + for _, mediaInfo := range mediaInfos { + for _, stream := range mediaInfo.Streams { + if stream.CodecType != "subtitle" { + continue + } + if strings.EqualFold(stream.Tags.Language, filter.MediaInfoWithSubtitleLanguage) { + return true + } + } + } + return false + } + if filter.MediaInfoWithoutSubtitleLanguage != "" { + for _, mediaInfo := range mediaInfos { + for _, stream := range mediaInfo.Streams { + if stream.CodecType != "subtitle" { + continue + } + if strings.EqualFold(stream.Tags.Language, filter.MediaInfoWithoutSubtitleLanguage) { + return false + } + } + } + return true + } + return false } diff --git a/internal/config/v1types.go b/internal/config/v1types.go index b05287d..e275c5d 100644 --- a/internal/config/v1types.go +++ b/internal/config/v1types.go @@ -42,4 +42,18 @@ type FilterConditionsV1 struct { HasEpisodes bool `yaml:"has_episodes"` IsMusic bool `yaml:"is_music"` + + MediaInfoResolution string `yaml:"media_info_resolution"` // possible values: 8k 4k 1080p 720p 480p 360p 240p 144p + MediaInfoBitRateGreaterThanOrEqual int64 `yaml:"media_info_bit_rate_gte"` // bytes per second + MediaInfoBitRateLessThanOrEqual int64 `yaml:"media_info_bit_rate_lte"` // bytes per second + MediaInfoVideoBitRateGreaterThanOrEqual int64 `yaml:"media_info_video_bit_rate_gte"` // bytes per second + MediaInfoVideoBitRateLessThanOrEqual int64 `yaml:"media_info_video_bit_rate_lte"` // bytes per second + MediaInfoAudioBitRateGreaterThanOrEqual int64 `yaml:"media_info_audio_bit_rate_gte"` // bytes per second + MediaInfoAudioBitRateLessThanOrEqual int64 `yaml:"media_info_audio_bit_rate_lte"` // bytes per second + MediaInfoDurationGreaterThanOrEqual int64 `yaml:"media_info_duration_gte"` // seconds + MediaInfoDurationLessThanOrEqual int64 `yaml:"media_info_duration_lte"` // seconds + MediaInfoWithAudioLanguage string `yaml:"media_info_with_audio_language"` // 3 char language code + MediaInfoWithoutAudioLanguage string `yaml:"media_info_without_audio_language"` // 3 char language code + MediaInfoWithSubtitleLanguage string `yaml:"media_info_with_subtitle_language"` // 3 char language code + MediaInfoWithoutSubtitleLanguage string `yaml:"media_info_without_subtitle_language"` // 3 char language code } diff --git a/internal/torrent/bins.go b/internal/torrent/bins.go index 4f4cab9..c47f795 100644 --- a/internal/torrent/bins.go +++ b/internal/torrent/bins.go @@ -136,22 +136,23 @@ func (t *TorrentManager) binOnceDoneErrorCheck(torrentId, status string) bool { // binOnceDone checks if the torrent is in the OnceDoneBin and deletes it if it is. // returns true if the torrent was in the bin and was deleted, false otherwise -func (t *TorrentManager) binOnceDone(torrentId string) bool { - if t.OnceDoneBin.Contains(torrentId) { - if err := t.api.DeleteTorrent(torrentId); err != nil { - t.repairLog.Warnf("Failed to delete torrent %s: %v", torrentId, err) +func (t *TorrentManager) binOnceDone(completedTorrentId string) bool { + if t.OnceDoneBin.Contains(completedTorrentId) { + if err := t.api.DeleteTorrent(completedTorrentId); err != nil { + t.repairLog.Warnf("Failed to delete torrent %s: %v", completedTorrentId, err) } - t.deleteInfoFile(torrentId) - t.OnceDoneBin.Remove(torrentId) + t.deleteInfoFile(completedTorrentId) + t.OnceDoneBin.Remove(completedTorrentId) t.persistBins() return true } // special case: yyy-xxx means if yyy is done, delete xxx - specialCase := fmt.Sprintf("%s-", torrentId) + specialCase := fmt.Sprintf("%s-", completedTorrentId) if !t.OnceDoneBin.Contains(specialCase) { return false } + t.OnceDoneBin.Remove(specialCase) t.OnceDoneBin.Clone().Each(func(entry string) bool { if strings.Contains(entry, specialCase) { idToDelete := strings.Split(entry, "-")[1] @@ -162,7 +163,7 @@ func (t *TorrentManager) binOnceDone(torrentId string) bool { } return false }) - t.deleteInfoFile(torrentId) + t.deleteInfoFile(completedTorrentId) t.OnceDoneBin.Remove(specialCase) t.persistBins() return true diff --git a/internal/torrent/manager.go b/internal/torrent/manager.go index f9fb19c..3e51018 100644 --- a/internal/torrent/manager.go +++ b/internal/torrent/manager.go @@ -14,6 +14,7 @@ import ( "github.com/debridmediamanager/zurg/internal/rar" "github.com/debridmediamanager/zurg/pkg/logutil" "github.com/debridmediamanager/zurg/pkg/realdebrid" + "github.com/debridmediamanager/zurg/pkg/utils" mapset "github.com/deckarep/golang-set/v2" cmap "github.com/orcaman/concurrent-map/v2" "github.com/panjf2000/ants/v2" @@ -107,10 +108,7 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, r if torrent != nil { accessKey := t.GetKey(torrent) allTorrents.Set(accessKey, torrent) - t.assignDirectory(torrent, func(directory string) { - listing, _ := t.DirectoryMap.Get(directory) - listing.Set(accessKey, torrent) - }) + t.assignDirectory(torrent, false) } return false }) @@ -230,13 +228,16 @@ func (t *TorrentManager) writeTorrentToFile(torrent *Torrent) { } func (t *TorrentManager) applyMediaInfoDetails(torrent *Torrent) { + changesApplied := false torrent.SelectedFiles.IterCb(func(_ string, file *File) { - if file.MediaInfo != nil || file.State.Is("broken_file") { + isPlayable := utils.IsPlayable(file.Path) || t.IsPlayable(file.Path) + if file.MediaInfo != nil || file.State.Is("broken_file") || !isPlayable { return } - unrestrict := t.UnrestrictFileUntilOk(file, false) + unrestrict := t.UnrestrictFileUntilOk(file, true) if unrestrict == nil { file.State.Event(context.Background(), "break_file") + changesApplied = true return } ctx, cancelFn := context.WithTimeout(context.Background(), 5*time.Second) @@ -244,11 +245,15 @@ func (t *TorrentManager) applyMediaInfoDetails(torrent *Torrent) { 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 + changesApplied = true }) + if changesApplied { + t.assignDirectory(torrent, true) + t.writeTorrentToFile(torrent) + } } func (t *TorrentManager) readTorrentFromFile(filePath string) *Torrent { @@ -436,7 +441,6 @@ func (t *TorrentManager) analyzeAllTorrents() { idx := 0 allTorrents.IterCb(func(_ string, torrent *Torrent) { t.applyMediaInfoDetails(torrent) - t.writeTorrentToFile(torrent) idx++ t.log.Debugf("Applied media info details to torrent %s (%d/%d)", t.GetKey(torrent), idx, totalCount) }) diff --git a/internal/torrent/refresh.go b/internal/torrent/refresh.go index 0d50a4b..246651d 100644 --- a/internal/torrent/refresh.go +++ b/internal/torrent/refresh.go @@ -13,25 +13,24 @@ import ( "github.com/debridmediamanager/zurg/pkg/utils" mapset "github.com/deckarep/golang-set/v2" cmap "github.com/orcaman/concurrent-map/v2" + "gopkg.in/vansante/go-ffprobe.v2" ) func inProgressStatus(status string) bool { return status == "downloading" || status == "uploading" || status == "queued" || status == "compressing" } -func (t *TorrentManager) refreshTorrents() []string { +func (t *TorrentManager) refreshTorrents() { t.inProgressHashes = mapset.NewSet[string]() instances, _, err := t.api.GetTorrents(false) if err != nil { t.log.Warnf("Cannot get torrents: %v", err) - return nil + return } var wg sync.WaitGroup var mergeChan = make(chan *Torrent, len(instances)) - updatedPaths := mapset.NewSet[string]() - allTorrents, _ := t.DirectoryMap.Get(INT_ALL) freshIDs := mapset.NewSet[string]() @@ -62,15 +61,8 @@ func (t *TorrentManager) refreshTorrents() []string { mainTorrent, exists := allTorrents.Get(accessKey) if !exists { allTorrents.Set(accessKey, torrent) - - t.assignDirectory(torrent, func(directory string) { - listing, _ := t.DirectoryMap.Get(directory) - listing.Set(accessKey, torrent) - - updatedPaths.Add(fmt.Sprintf("%s/%s", directory, accessKey)) - }) - t.writeTorrentToFile(torrent) + t.assignDirectory(torrent, true) } else if !mainTorrent.DownloadedIDs.Contains(tInfo.ID) { forMerging = torrent } @@ -102,20 +94,7 @@ func (t *TorrentManager) refreshTorrents() []string { mainTorrent := t.mergeTorrents(existing, torrent) allTorrents.Set(accessKey, mainTorrent) t.writeTorrentToFile(mainTorrent) - - t.DirectoryMap.IterCb(func(directory string, torrents cmap.ConcurrentMap[string, *Torrent]) { - if strings.HasPrefix(directory, "int__") { - return - } - torrents.Remove(accessKey) - }) - - t.assignDirectory(mainTorrent, func(directory string) { - listing, _ := t.DirectoryMap.Get(directory) - listing.Set(accessKey, mainTorrent) - - updatedPaths.Add(fmt.Sprintf("%s/%s", directory, accessKey)) - }) + t.assignDirectory(mainTorrent, true) } // removed torrents @@ -163,8 +142,6 @@ func (t *TorrentManager) refreshTorrents() []string { } }) }) - - return updatedPaths.ToSlice() } // StartRefreshJob periodically refreshes the torrents @@ -183,10 +160,8 @@ func (t *TorrentManager) StartRefreshJob() { } t.setNewLatestState(checksum) - updatedPaths := t.refreshTorrents() + t.refreshTorrents() t.log.Info("Finished refreshing torrents") - - t.TriggerHookOnLibraryUpdate(updatedPaths) case <-t.RefreshKillSwitch: t.log.Info("Stopping periodic refresh job") return @@ -364,15 +339,28 @@ func (t *TorrentManager) mergeTorrents(existing, toMerge *Torrent) *Torrent { return mergedTorrent } -func (t *TorrentManager) assignDirectory(tor *Torrent, cb func(string)) { +func (t *TorrentManager) assignDirectory(tor *Torrent, triggerHook bool) { + accessKey := t.GetKey(tor) + + t.DirectoryMap.IterCb(func(directory string, torrents cmap.ConcurrentMap[string, *Torrent]) { + if strings.HasPrefix(directory, "int__") { + return + } + torrents.Remove(accessKey) + }) + torrentIDs := tor.DownloadedIDs.ToSlice() // get filenames needed for directory conditions var filenames []string var fileSizes []int64 + var mediaInfos []*ffprobe.ProbeData unplayable := true tor.SelectedFiles.IterCb(func(key string, file *File) { filenames = append(filenames, filepath.Base(file.Path)) fileSizes = append(fileSizes, file.Bytes) + if file.MediaInfo != nil { + mediaInfos = append(mediaInfos, file.MediaInfo) + } if utils.IsPlayable(file.Path) || t.IsPlayable(file.Path) { unplayable = false } @@ -389,8 +377,13 @@ func (t *TorrentManager) assignDirectory(tor *Torrent, cb func(string)) { configV1 := t.Config.(*config.ZurgConfigV1) for _, directories := range configV1.GetGroupMap() { for _, directory := range directories { - if t.Config.MeetsConditions(directory, t.GetKey(tor), tor.ComputeTotalSize(), torrentIDs, filenames, fileSizes) { - cb(directory) + if t.Config.MeetsConditions(directory, t.GetKey(tor), tor.ComputeTotalSize(), torrentIDs, filenames, fileSizes, mediaInfos) { + listing, _ := t.DirectoryMap.Get(directory) + listing.Set(accessKey, tor) + + if triggerHook { + t.TriggerHookOnLibraryUpdate([]string{fmt.Sprintf("%s/%s", directory, accessKey)}) + } break } } diff --git a/internal/torrent/repair.go b/internal/torrent/repair.go index 0b8b3fb..44a583d 100644 --- a/internal/torrent/repair.go +++ b/internal/torrent/repair.go @@ -119,6 +119,7 @@ func (t *TorrentManager) repairAll(torrent *Torrent) { } }) if brokenFileCount > 0 { + t.repairLog.Debugf("Torrent %s has %d broken file(s), adding to repair list", t.GetKey(torrent), brokenFileCount) toRepair.Add(torrent) return } @@ -184,7 +185,7 @@ func (t *TorrentManager) repair(torrent *Torrent) { } allPlayable = false - if t.Config.GetRarAction() == "extract" { + if t.Config.GetRarAction() == "extract" && file.ID != 0 { info, _ := t.redownloadTorrent(torrent, []string{fmt.Sprintf("%d", file.ID)}) if info != nil { t.setToBinOnceDone(info.ID)