DRMDTool/src/utils.go

package main

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

	"github.com/beevik/etree"
)

type JobInfo struct {
	AbortChan  chan struct{}
	ResumeChan chan struct{}
	Cmd        *exec.Cmd
	TempDir    string

	mu        sync.RWMutex
	paused    bool
	abortOnce sync.Once
}

func NewJobInfo() *JobInfo {
	return &JobInfo{
		AbortChan:  make(chan struct{}),
		ResumeChan: make(chan struct{}, 1),
	}
}

func (j *JobInfo) SetPaused(value bool) {
	j.mu.Lock()
	defer j.mu.Unlock()
	j.paused = value
}

func (j *JobInfo) IsPaused() bool {
	j.mu.RLock()
	defer j.mu.RUnlock()
	return j.paused
}

func (j *JobInfo) SetCmd(cmd *exec.Cmd) {
	j.mu.Lock()
	defer j.mu.Unlock()
	j.Cmd = cmd
}

func (j *JobInfo) SetTempDir(tempDir string) {
	j.mu.Lock()
	defer j.mu.Unlock()
	j.TempDir = tempDir
}

func (j *JobInfo) GetTempDir() string {
	j.mu.RLock()
	defer j.mu.RUnlock()
	return j.TempDir
}

func (j *JobInfo) KillProcess() {
	j.mu.RLock()
	cmd := j.Cmd
	j.mu.RUnlock()

	if cmd != nil && cmd.Process != nil {
		_ = cmd.Process.Kill()
	}
}

func (j *JobInfo) SignalResume() {
	select {
	case j.ResumeChan <- struct{}{}:
	default:
	}
}

func (j *JobInfo) Abort() {
	j.abortOnce.Do(func() {
		close(j.AbortChan)
	})
}

func (j *JobInfo) IsAborted() bool {
	select {
	case <-j.AbortChan:
		return true
	default:
		return false
	}
}

var (
	ErrDownloadPaused  = errors.New("download paused")
	ErrDownloadAborted = errors.New("download aborted")
)

var sanitizeFilenameRegex = regexp.MustCompile(`[<>:"/\\|?*\x00-\x1f]`)

var windowsReservedNames = map[string]struct{}{
	"CON": {}, "PRN": {}, "AUX": {}, "NUL": {},
	"COM1": {}, "COM2": {}, "COM3": {}, "COM4": {}, "COM5": {}, "COM6": {}, "COM7": {}, "COM8": {}, "COM9": {},
	"LPT1": {}, "LPT2": {}, "LPT3": {}, "LPT4": {}, "LPT5": {}, "LPT6": {}, "LPT7": {}, "LPT8": {}, "LPT9": {},
}

func sanitizeFilename(filename string) string {
	filename = sanitizeFilenameRegex.ReplaceAllString(filename, "_")
	filename = strings.Trim(filename, ". ")

	base := filename
	if idx := strings.LastIndex(filename, "."); idx > 0 {
		base = filename[:idx]
	}
	if _, reserved := windowsReservedNames[strings.ToUpper(base)]; reserved {
		filename = "_" + filename
	}

	if filename == "" {
		filename = "_"
	}
	return filename
}

func safeUploadPath(filename string) (string, error) {
	cleanName := strings.TrimSpace(filename)
	if cleanName == "" {
		return "", fmt.Errorf("filename is required")
	}

	baseName := filepath.Base(cleanName)
	if baseName != cleanName {
		return "", fmt.Errorf("invalid filename")
	}

	if strings.Contains(baseName, "..") {
		return "", fmt.Errorf("invalid filename")
	}

	return filepath.Join(uploadDir, baseName), nil
}

func isValidURL(toTest string) bool {
	_, err := url.ParseRequestURI(toTest)
	return err == nil
}

func fixGoPlay(mpdContent string) (string, error) {
	doc := etree.NewDocument()
	if err := doc.ReadFromString(mpdContent); err != nil {
		return "", fmt.Errorf("error parsing MPD content: %v", err)
	}

	root := doc.Root()

	// Remove ad periods
	for _, period := range root.SelectElements("Period") {
		if strings.Contains(period.SelectAttrValue("id", ""), "-ad-") {
			root.RemoveChild(period)
		}
	}

	// Find highest bandwidth for video
	highestBandwidth := 0
	for _, adaptationSet := range root.FindElements("//AdaptationSet") {
		if strings.Contains(adaptationSet.SelectAttrValue("mimeType", ""), "video") {
			for _, representation := range adaptationSet.SelectElements("Representation") {
				bandwidth, _ := strconv.Atoi(representation.SelectAttrValue("bandwidth", "0"))
				if bandwidth > highestBandwidth {
					highestBandwidth = bandwidth
				}
			}
		}
	}

	// Remove lower bitrate representations
	for _, adaptationSet := range root.FindElements("//AdaptationSet") {
		if strings.Contains(adaptationSet.SelectAttrValue("mimeType", ""), "video") {
			for _, representation := range adaptationSet.SelectElements("Representation") {
				bandwidth, _ := strconv.Atoi(representation.SelectAttrValue("bandwidth", "0"))
				if bandwidth != highestBandwidth {
					adaptationSet.RemoveChild(representation)
				}
			}
		}
	}

	// Combine periods
	periods := root.SelectElements("Period")
	if len(periods) > 1 {
		firstPeriod := periods[0]
		var newVideoTimeline, newAudioTimeline *etree.Element

		// Find or create SegmentTimeline elements
		for _, adaptationSet := range firstPeriod.SelectElements("AdaptationSet") {
			mimeType := adaptationSet.SelectAttrValue("mimeType", "")
			if strings.Contains(mimeType, "video") && newVideoTimeline == nil {
				newVideoTimeline = findOrCreateSegmentTimeline(adaptationSet)
			} else if strings.Contains(mimeType, "audio") && newAudioTimeline == nil {
				newAudioTimeline = findOrCreateSegmentTimeline(adaptationSet)
			}
		}

		for _, period := range periods[1:] {
			for _, adaptationSet := range period.SelectElements("AdaptationSet") {
				mimeType := adaptationSet.SelectAttrValue("mimeType", "")
				var timeline *etree.Element
				if strings.Contains(mimeType, "video") {
					timeline = newVideoTimeline
				} else if strings.Contains(mimeType, "audio") {
					timeline = newAudioTimeline
				}

				if timeline != nil {
					segmentTimeline := findOrCreateSegmentTimeline(adaptationSet)
					for _, s := range segmentTimeline.SelectElements("S") {
						timeline.AddChild(s.Copy())
					}
				}
			}
			root.RemoveChild(period)
		}
	}

	return doc.WriteToString()
}

func findOrCreateSegmentTimeline(adaptationSet *etree.Element) *etree.Element {
	for _, representation := range adaptationSet.SelectElements("Representation") {
		for _, segmentTemplate := range representation.SelectElements("SegmentTemplate") {
			timeline := segmentTemplate.SelectElement("SegmentTimeline")
			if timeline != nil {
				return timeline
			}
		}
	}

	// If no SegmentTimeline found, create one
	representation := adaptationSet.CreateElement("Representation")
	segmentTemplate := representation.CreateElement("SegmentTemplate")
	return segmentTemplate.CreateElement("SegmentTimeline")
}

func parseInputFile(filename string) ([]Item, error) {
	fileInfo, err := os.Stat(filename)
	if err != nil {
		return nil, err
	}
	if fileInfo.IsDir() {
		return nil, fmt.Errorf("%s is a directory", filename)
	}

	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	byteValue, err := io.ReadAll(file)
	if err != nil {
		return nil, err
	}

	byteValue = removeBOM(byteValue)

	var items Items
	err = json.Unmarshal(byteValue, &items)
	if err != nil {
		return nil, err
	}

	return items.Items, nil
}

func groupItemsBySeason(items []Item) map[string][]Item {
	grouped := make(map[string][]Item)
	for _, item := range items {
		metadata := parseMetadata(item.Metadata)
		if metadata.Type == "serie" {
			key := fmt.Sprintf("%s - %s", metadata.Title, metadata.Season)
			grouped[key] = append(grouped[key], item)
		} else {
			grouped["Movies"] = append(grouped["Movies"], item)
		}
	}
	return grouped
}

func filterSelectedItems(items []Item, selectedItems []string) []Item {
	set := make(map[string]struct{}, len(selectedItems))
	for _, s := range selectedItems {
		set[s] = struct{}{}
	}
	filtered := make([]Item, 0, len(selectedItems))
	for _, item := range items {
		if _, ok := set[item.Filename]; ok {
			filtered = append(filtered, item)
		}
	}
	return filtered
}

func sortItems(items []Item) {
	sort.Slice(items, func(i, j int) bool {
		iMeta := parseMetadata(items[i].Metadata)
		jMeta := parseMetadata(items[j].Metadata)

		if iMeta.Title != jMeta.Title {
			return iMeta.Title < jMeta.Title
		}

		iSeason := extractNumber(iMeta.Season)
		jSeason := extractNumber(jMeta.Season)

		if iSeason != jSeason {
			return iSeason < jSeason
		}

		iEpisode := extractEpisodeNumber(items[i].Filename)
		jEpisode := extractEpisodeNumber(items[j].Filename)

		return iEpisode < jEpisode
	})
}

func extractNumber(s string) int {
	num, _ := strconv.Atoi(strings.TrimLeft(s, "S"))
	return num
}

var episodeNumberRegex = regexp.MustCompile(`(?i)S\d+E(\d+)`)

func extractEpisodeNumber(filename string) int {
	match := episodeNumberRegex.FindStringSubmatch(filename)
	if len(match) < 2 {
		return 0
	}
	num, _ := strconv.Atoi(match[1])
	return num
}

func processItems(filename string, items []Item) error {
	jobInfo := NewJobInfo()
	setJob(filename, jobInfo)

	defer func() {
		removeJob(filename)

		tempDir := jobInfo.GetTempDir()
		if tempDir != "" {
			_ = os.RemoveAll(tempDir)
		}
	}()

	sortItems(items)

	for i := 0; i < len(items); i++ {
		select {
		case <-jobInfo.AbortChan:
			updateProgress(filename, 100, "Aborted", "aborted")
			logger.LogJobState(filename, "aborted")
			return ErrDownloadAborted
		default:
			if jobInfo.IsPaused() {
				select {
				case <-jobInfo.ResumeChan:
					jobInfo.SetPaused(false)
					logger.LogJobState(filename, "resumed")
				case <-jobInfo.AbortChan:
					updateProgress(filename, 100, "Aborted", "aborted")
					logger.LogJobState(filename, "aborted")
					return ErrDownloadAborted
				}
			}
			updateProgress(filename, float64(i)/float64(len(items))*100, items[i].Filename, "running")
			err := downloadFile(filename, items[i], jobInfo)
			if err != nil {
				if errors.Is(err, ErrDownloadPaused) {
					logger.LogJobState(filename, "paused")
					if remErr := removeCompletedEpisodes(filename, items[:i]); remErr != nil {
						logger.LogError("Process Items", fmt.Sprintf("Error updating partial progress file: %v", remErr))
					}
					i--
					continue
				}
				if errors.Is(err, ErrDownloadAborted) {
					updateProgress(filename, 100, "Aborted", "aborted")
					logger.LogJobState(filename, "aborted")
					return ErrDownloadAborted
				}

				updateProgress(filename, float64(i)/float64(len(items))*100, items[i].Filename, "failed")
				logger.LogError("Process Items", fmt.Sprintf("Error downloading item %s: %v", items[i].Filename, err))
				return fmt.Errorf("error downloading %s: %w", items[i].Filename, err)
			}
		}
	}
	updateProgress(filename, 100, "", "completed")
	logger.LogJobState(filename, "completed successfully")
	return nil
}

func removeCompletedEpisodes(filename string, completedItems []Item) error {
	inputFile := filepath.Join(uploadDir, filename)
	items, err := parseInputFile(inputFile)
	if err != nil {
		return fmt.Errorf("error parsing input file: %v", err)
	}

	remainingItems := make([]Item, 0)
	for _, item := range items {
		if !isItemCompleted(item, completedItems) || isLastCompletedItem(item, completedItems) {
			remainingItems = append(remainingItems, item)
		}
	}

	updatedItems := Items{Items: remainingItems}
	jsonData, err := json.MarshalIndent(updatedItems, "", "  ")
	if err != nil {
		return fmt.Errorf("error marshaling updated items: %v", err)
	}

	err = os.WriteFile(inputFile, jsonData, 0644)
	if err != nil {
		return fmt.Errorf("error writing updated DRMD file: %v", err)
	}

	return nil
}

func isItemCompleted(item Item, completedItems []Item) bool {
	for _, completedItem := range completedItems {
		if item.Filename == completedItem.Filename {
			return true
		}
	}
	return false
}

func isLastCompletedItem(item Item, completedItems []Item) bool {
	if len(completedItems) == 0 {
		return false
	}
	return item.Filename == completedItems[len(completedItems)-1].Filename
}