DRMDTool/utils.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"github.com/beevik/etree"
)

var (
	jobsMutex sync.Mutex
	jobs      = make(map[string]chan struct{})
)

func sanitizeFilename(filename string) string {
	filename = regexp.MustCompile(`[<>:"/\\|?*]`).ReplaceAllString(filename, "_")

	filename = strings.Trim(filename, ".")

	return filename
}

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(inputFile string) ([]Item, error) {
	jsonFile, err := os.Open(inputFile)
	if err != nil {
		return nil, fmt.Errorf("error opening file %s: %v", inputFile, err)
	}
	defer jsonFile.Close()

	byteValue, err := io.ReadAll(jsonFile)
	if err != nil {
		return nil, fmt.Errorf("error reading file %s: %v", inputFile, err)
	}

	byteValue = removeBOM(byteValue)

	var items Items
	err = json.Unmarshal(byteValue, &items)
	if err != nil {
		return nil, fmt.Errorf("error unmarshaling JSON: %v", 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 {
	var filtered []Item
	for _, item := range items {
		for _, selected := range selectedItems {
			if item.Filename == selected {
				filtered = append(filtered, item)
				break
			}
		}
	}
	return filtered
}

func processItems(filename string, items []Item) error {
	jobsMutex.Lock()
	abortChan := make(chan struct{})
	jobs[filename] = abortChan
	jobsMutex.Unlock()

	defer func() {
		jobsMutex.Lock()
		delete(jobs, filename)
		jobsMutex.Unlock()
	}()

	for i, item := range items {
		select {
		case <-abortChan:
			updateProgress(filename, 100, "Aborted")
			return fmt.Errorf("download aborted")
		default:
			updateProgress(filename, float64(i)/float64(len(items))*100, item.Filename)
			err := downloadFile(item)
			if err != nil {
				fmt.Printf("Error downloading file: %v\n", err)
			}
		}
	}
	updateProgress(filename, 100, "")
	return nil
}