DRMDTool/utils.go

package main

import (
	"fmt"
	"net/url"
	"regexp"
	"strconv"
	"strings"

	"github.com/beevik/etree"
)

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