8 Commits

Author SHA1 Message Date
3872c1c4ca Updated README 2024-09-06 14:37:26 +02:00
077f1efb6f Add support for UTF-8 w/ BOM 2024-09-06 14:36:46 +02:00
4cbd64f263 Fix goplay 2024-09-06 14:27:09 +02:00
078d0c774b Seperate everything into different files to prevent me from going insane 2024-09-06 13:28:12 +02:00
1b0a8fde36 Sanitize Filename 2024-09-06 11:58:20 +02:00
f014b06f88 Download subtitles from drmd file 2024-09-06 11:29:19 +02:00
6b192aadc9 Compiler flag 2024-09-04 23:18:04 +02:00
6b24eb7e07 Todo 2024-09-04 01:48:43 +02:00
10 changed files with 541 additions and 234 deletions

1
.gitignore vendored Normal file
View File

@ -0,0 +1 @@
config.toml

View File

@ -36,4 +36,10 @@ To process a file directly from the command line:
./drmdtool -f /path/to/file.drmd
```
This will download the file and save it in the base directory specified in the config.
This will download the file and save it in the base directory specified in the config.
## TODO
- ~~Filename Sanitation (Makes new directory on /... oops)~~
- ~~GoPlay Fix~~
- Windows?
- Proper UI?

41
config.go Normal file
View File

@ -0,0 +1,41 @@
package main
import (
"fmt"
"io"
"os"
"github.com/BurntSushi/toml"
)
type Config struct {
BaseDir string
Format string
N_m3u8DLRE struct {
Path string
}
}
var config Config
func loadConfig() {
configFile, err := os.Open("config.toml")
if err != nil {
fmt.Println("Error opening config file:", err)
return
}
defer configFile.Close()
byteValue, _ := io.ReadAll(configFile)
if _, err := toml.Decode(string(byteValue), &config); err != nil {
fmt.Println("Error decoding config file:", err)
return
}
if config.N_m3u8DLRE.Path == "" {
fmt.Println("Error: N_m3u8DL-RE path is not specified in the config file")
return
}
}

171
downloaders.go Normal file
View File

@ -0,0 +1,171 @@
package main
import (
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"os/exec"
"path/filepath"
"strings"
)
func processInputFile(inputFile string) error {
jsonFile, err := os.Open(inputFile)
if err != nil {
return fmt.Errorf("error opening file %s: %v", inputFile, err)
}
defer jsonFile.Close()
byteValue, err := io.ReadAll(jsonFile)
if err != nil {
return fmt.Errorf("error reading file %s: %v", inputFile, err)
}
byteValue = removeBOM(byteValue)
var items Items
err = json.Unmarshal(byteValue, &items)
if err != nil {
return fmt.Errorf("error unmarshaling JSON: %v", err)
}
for i, item := range items.Items {
updateProgress(filepath.Base(inputFile), float64(i)/float64(len(items.Items))*100, item.Filename)
err := downloadFile(item)
if err != nil {
fmt.Printf("Error downloading file: %v\n", err)
}
}
updateProgress(filepath.Base(inputFile), 100, "")
return nil
}
func removeBOM(input []byte) []byte {
if len(input) >= 3 && input[0] == 0xEF && input[1] == 0xBB && input[2] == 0xBF {
return input[3:]
}
return input
}
func downloadFile(item Item) error {
fmt.Println("Downloading:", item.Filename)
mpdPath := item.MPD
if !isValidURL(item.MPD) {
decodedMPD, err := base64.StdEncoding.DecodeString(item.MPD)
if err != nil {
return fmt.Errorf("error decoding base64 MPD: %v", err)
}
tempFile, err := os.CreateTemp("", "temp_mpd_*.mpd")
if err != nil {
return fmt.Errorf("error creating temporary MPD file: %v", err)
}
defer os.Remove(tempFile.Name())
if _, err := tempFile.Write(decodedMPD); err != nil {
return fmt.Errorf("error writing to temporary MPD file: %v", err)
}
if err := tempFile.Close(); err != nil {
return fmt.Errorf("error closing temporary MPD file: %v", err)
}
mpdPath = tempFile.Name()
} else if strings.HasPrefix(item.MPD, "https://pubads.g.doubleclick.net") {
resp, err := http.Get(item.MPD)
if err != nil {
return fmt.Errorf("error downloading MPD: %v", err)
}
defer resp.Body.Close()
mpdContent, err := io.ReadAll(resp.Body)
if err != nil {
return fmt.Errorf("error reading MPD content: %v", err)
}
fixedMPDContent, err := fixGoPlay(string(mpdContent))
if err != nil {
return fmt.Errorf("error fixing MPD content: %v", err)
}
tempFile, err := os.CreateTemp("", "fixed_mpd_*.mpd")
if err != nil {
return fmt.Errorf("error creating temporary MPD file: %v", err)
}
defer os.Remove(tempFile.Name())
if _, err := tempFile.WriteString(fixedMPDContent); err != nil {
return fmt.Errorf("error writing to temporary MPD file: %v", err)
}
if err := tempFile.Close(); err != nil {
return fmt.Errorf("error closing temporary MPD file: %v", err)
}
mpdPath = tempFile.Name()
}
command := getDownloadCommand(item, mpdPath)
if item.Subtitles != "" {
subtitlePaths, err := downloadAndConvertSubtitles(item.Subtitles)
if err != nil {
fmt.Printf("Error processing subtitles: %v\n", err)
} else {
for _, path := range subtitlePaths {
fmt.Println("Adding subtitle:", path)
command += fmt.Sprintf(" --mux-import \"path=%s:lang=nl:name=Nederlands\"", path)
}
}
}
cmd := exec.Command("bash", "-c", command)
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
err := cmd.Run()
if err != nil {
return fmt.Errorf("error executing download command: %v", err)
}
fmt.Println("Download completed successfully")
return nil
}
func getDownloadCommand(item Item, mpdPath string) string {
metadata := parseMetadata(item.Metadata)
keys := getKeys(item.Keys)
command := fmt.Sprintf("%s %s", config.N_m3u8DLRE.Path, mpdPath)
for _, key := range keys {
if key != "" {
command += fmt.Sprintf(" --key %s", key)
}
}
command += " --auto-select"
sanitizedFilename := sanitizeFilename(item.Filename)
filename := fmt.Sprintf("\"%s\"", sanitizedFilename)
command += fmt.Sprintf(" --save-name %s", filename)
command += fmt.Sprintf(" --mux-after-done format=%s", config.Format)
saveDir := config.BaseDir
if metadata.Type == "serie" {
saveDir = filepath.Join(saveDir, "Series", metadata.Title, metadata.Season)
} else {
saveDir = filepath.Join(saveDir, "Movies", metadata.Title)
}
command += fmt.Sprintf(" --save-dir \"%s\"", saveDir)
fmt.Println(command)
return command
}

13
go.mod
View File

@ -2,4 +2,15 @@ module DRMDTool
go 1.23.0
require github.com/BurntSushi/toml v1.4.0
require (
github.com/BurntSushi/toml v1.4.0
github.com/asticode/go-astisub v0.26.2
github.com/beevik/etree v1.4.1
)
require (
github.com/asticode/go-astikit v0.20.0 // indirect
github.com/asticode/go-astits v1.8.0 // indirect
golang.org/x/net v0.0.0-20200904194848-62affa334b73 // indirect
golang.org/x/text v0.3.2 // indirect
)

31
go.sum
View File

@ -1,2 +1,33 @@
github.com/BurntSushi/toml v1.4.0 h1:kuoIxZQy2WRRk1pttg9asf+WVv6tWQuBNVmK8+nqPr0=
github.com/BurntSushi/toml v1.4.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
github.com/asticode/go-astikit v0.20.0 h1:+7N+J4E4lWx2QOkRdOf6DafWJMv6O4RRfgClwQokrH8=
github.com/asticode/go-astikit v0.20.0/go.mod h1:h4ly7idim1tNhaVkdVBeXQZEE3L0xblP7fCWbgwipF0=
github.com/asticode/go-astisub v0.26.2 h1:cdEXcm+SUSmYCEPTQYbbfCECnmQoIFfH6pF8wDJhfVo=
github.com/asticode/go-astisub v0.26.2/go.mod h1:WTkuSzFB+Bp7wezuSf2Oxulj5A8zu2zLRVFf6bIFQK8=
github.com/asticode/go-astits v1.8.0 h1:rf6aiiGn/QhlFjNON1n5plqF3Fs025XLUwiQ0NB6oZg=
github.com/asticode/go-astits v1.8.0/go.mod h1:DkOWmBNQpnr9mv24KfZjq4JawCFX1FCqjLVGvO0DygQ=
github.com/beevik/etree v1.4.1 h1:PmQJDDYahBGNKDcpdX8uPy1xRCwoCGVUiW669MEirVI=
github.com/beevik/etree v1.4.1/go.mod h1:gPNJNaBGVZ9AwsidazFZyygnd+0pAU38N4D+WemwKNs=
github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/pkg/profile v1.4.0/go.mod h1:NWz/XGvpEW1FyYQ7fCx4dqYBLlfTcE+A9FLAkNKqjFE=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk=
github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20200904194848-62affa334b73 h1:MXfv8rhZWmFeqX3GNZRsd6vOLoaCHjYEX3qkRo3YBUA=
golang.org/x/net v0.0.0-20200904194848-62affa334b73/go.mod h1:/O7V0waA8r7cgGh81Ro3o1hOxt32SMVPicZroKQ2sZA=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.2 h1:tW2bmiBqwgJj/UpqtC8EpXEZVYOwU0yG4iWbprSVAcs=
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

96
handlers.go Normal file
View File

@ -0,0 +1,96 @@
package main
import (
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"time"
)
func handleRoot(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/" {
http.NotFound(w, r)
return
}
progressMutex.Lock()
jobs := make(map[string]*ProgressInfo)
for k, v := range progress {
jobs[k] = v
}
progressMutex.Unlock()
err := templates.ExecuteTemplate(w, "index", struct{ Jobs map[string]*ProgressInfo }{jobs})
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
}
}
func handleUpload(w http.ResponseWriter, r *http.Request) {
file, header, err := r.FormFile("file")
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
defer file.Close()
filename := fmt.Sprintf("%d_%s", time.Now().UnixNano(), header.Filename)
filepath := filepath.Join(uploadDir, filename)
newFile, err := os.Create(filepath)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
defer newFile.Close()
_, err = io.Copy(newFile, file)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
go func() {
err := processInputFile(filepath)
if err != nil {
fmt.Printf("Error processing file: %v\n", err)
}
os.Remove(filepath)
}()
http.Redirect(w, r, "/progress?filename="+filename, http.StatusSeeOther)
}
func handleProgress(w http.ResponseWriter, r *http.Request) {
filename := r.URL.Query().Get("filename")
fmt.Printf("Handling progress request for filename: %s\n", filename)
if r.Header.Get("Accept") == "application/json" {
progressInfo := getProgress(filename)
fmt.Printf("Progress info for %s: %+v\n", filename, progressInfo)
if progressInfo == nil {
w.WriteHeader(http.StatusNotFound)
json.NewEncoder(w).Encode(map[string]string{"error": "No progress information found"})
return
}
w.Header().Set("Content-Type", "application/json")
err := json.NewEncoder(w).Encode(progressInfo)
if err != nil {
fmt.Printf("Error encoding progress info: %v\n", err)
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return
}
return
}
err := templates.ExecuteTemplate(w, "progress", struct{ Filename string }{filename})
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
}
}

234
main.go
View File

@ -1,24 +1,15 @@
package main
import (
"encoding/base64"
"encoding/json"
"flag"
"fmt"
"html/template"
"io"
"net/http"
"net/url"
"os"
"os/exec"
"path/filepath"
"strings"
"sync"
"time"
"embed"
"github.com/BurntSushi/toml"
)
type Item struct {
@ -41,15 +32,6 @@ type Metadata struct {
Season string
}
type Config struct {
BaseDir string
Format string
N_m3u8DLRE struct {
Path string
}
}
var config Config
var progressMutex sync.Mutex
var progress = make(map[string]*ProgressInfo)
@ -62,6 +44,7 @@ type ProgressInfo struct {
var templates *template.Template
//go:embed templates
var templateFS embed.FS
func init() {
@ -73,25 +56,7 @@ func init() {
}
func main() {
configFile, err := os.Open("config.toml")
if err != nil {
fmt.Println("Error opening config file:", err)
return
}
defer configFile.Close()
byteValue, _ := io.ReadAll(configFile)
if _, err := toml.Decode(string(byteValue), &config); err != nil {
fmt.Println("Error decoding config file:", err)
return
}
if config.N_m3u8DLRE.Path == "" {
fmt.Println("Error: N_m3u8DL-RE path is not specified in the config file")
return
}
loadConfig()
inputFile := flag.String("f", "", "Path to the input JSON file")
flag.Parse()
@ -111,121 +76,6 @@ func startWebServer() {
http.ListenAndServe(":8080", nil)
}
func handleRoot(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/" {
http.NotFound(w, r)
return
}
progressMutex.Lock()
jobs := make(map[string]*ProgressInfo)
for k, v := range progress {
jobs[k] = v
}
progressMutex.Unlock()
err := templates.ExecuteTemplate(w, "index", struct{ Jobs map[string]*ProgressInfo }{jobs})
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
}
}
func handleUpload(w http.ResponseWriter, r *http.Request) {
file, header, err := r.FormFile("file")
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
defer file.Close()
filename := fmt.Sprintf("%d_%s", time.Now().UnixNano(), header.Filename)
filepath := filepath.Join(uploadDir, filename)
newFile, err := os.Create(filepath)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
defer newFile.Close()
_, err = io.Copy(newFile, file)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
go func() {
err := processInputFile(filepath)
if err != nil {
fmt.Printf("Error processing file: %v\n", err)
}
os.Remove(filepath)
}()
http.Redirect(w, r, "/progress?filename="+filename, http.StatusSeeOther)
}
func handleProgress(w http.ResponseWriter, r *http.Request) {
filename := r.URL.Query().Get("filename")
fmt.Printf("Handling progress request for filename: %s\n", filename)
if r.Header.Get("Accept") == "application/json" {
progressInfo := getProgress(filename)
fmt.Printf("Progress info for %s: %+v\n", filename, progressInfo)
if progressInfo == nil {
w.WriteHeader(http.StatusNotFound)
json.NewEncoder(w).Encode(map[string]string{"error": "No progress information found"})
return
}
w.Header().Set("Content-Type", "application/json")
err := json.NewEncoder(w).Encode(progressInfo)
if err != nil {
fmt.Printf("Error encoding progress info: %v\n", err)
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return
}
return
}
err := templates.ExecuteTemplate(w, "progress", struct{ Filename string }{filename})
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
}
}
func processInputFile(inputFile string) error {
jsonFile, err := os.Open(inputFile)
if err != nil {
return fmt.Errorf("error opening file %s: %v", inputFile, err)
}
defer jsonFile.Close()
byteValue, err := io.ReadAll(jsonFile)
if err != nil {
return fmt.Errorf("error reading file %s: %v", inputFile, err)
}
var items Items
err = json.Unmarshal(byteValue, &items)
if err != nil {
return fmt.Errorf("error unmarshaling JSON: %v", err)
}
for i, item := range items.Items {
updateProgress(filepath.Base(inputFile), float64(i)/float64(len(items.Items))*100, item.Filename)
err := downloadFile(item)
if err != nil {
fmt.Printf("Error downloading file: %v\n", err)
}
}
updateProgress(filepath.Base(inputFile), 100, "")
return nil
}
func updateProgress(filename string, value float64, currentFile string) {
progressMutex.Lock()
defer progressMutex.Unlock()
@ -257,83 +107,3 @@ func parseMetadata(metadata string) Metadata {
Season: "S" + strings.TrimSpace(parts[2]),
}
}
func getDownloadCommand(item Item, mpdPath string) string {
metadata := parseMetadata(item.Metadata)
keys := getKeys(item.Keys)
command := fmt.Sprintf("%s %s", config.N_m3u8DLRE.Path, mpdPath)
for _, key := range keys {
if key != "" {
command += fmt.Sprintf(" --key %s", key)
}
}
command += " --auto-select"
filename := fmt.Sprintf("\"%s\"", item.Filename)
command += fmt.Sprintf(" --save-name %s", filename)
command += fmt.Sprintf(" --mux-after-done format=%s", config.Format)
saveDir := config.BaseDir
if metadata.Type == "serie" {
saveDir = filepath.Join(saveDir, "Series", metadata.Title, metadata.Season)
} else {
saveDir = filepath.Join(saveDir, "Movies", metadata.Title)
}
command += fmt.Sprintf(" --save-dir \"%s\"", saveDir)
fmt.Println(command)
return command
}
func downloadFile(item Item) error {
fmt.Println("Downloading:", item.Filename)
mpdPath := item.MPD
if !isValidURL(item.MPD) {
decodedMPD, err := base64.StdEncoding.DecodeString(item.MPD)
if err != nil {
return fmt.Errorf("error decoding base64 MPD: %v", err)
}
tempFile, err := os.CreateTemp("", "temp_mpd_*.mpd")
if err != nil {
return fmt.Errorf("error creating temporary MPD file: %v", err)
}
defer os.Remove(tempFile.Name())
if _, err := tempFile.Write(decodedMPD); err != nil {
return fmt.Errorf("error writing to temporary MPD file: %v", err)
}
if err := tempFile.Close(); err != nil {
return fmt.Errorf("error closing temporary MPD file: %v", err)
}
mpdPath = tempFile.Name()
}
command := getDownloadCommand(item, mpdPath)
cmd := exec.Command("bash", "-c", command)
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
err := cmd.Run()
if err != nil {
return fmt.Errorf("error executing download command: %v", err)
}
fmt.Println("Download completed successfully")
return nil
}
func isValidURL(toTest string) bool {
_, err := url.ParseRequestURI(toTest)
return err == nil
}

60
subtitles.go Normal file
View File

@ -0,0 +1,60 @@
package main
import (
"fmt"
"io"
"net/http"
"os"
"strings"
"github.com/asticode/go-astisub"
)
func downloadAndConvertSubtitles(subtitlesURLs string) ([]string, error) {
var subtitlePaths []string
urls := strings.Split(subtitlesURLs, ",")
for _, url := range urls {
vttPath, err := downloadSubtitle(url)
if err != nil {
return nil, fmt.Errorf("error downloading subtitle: %v", err)
}
srtPath, err := convertVTTtoSRT(vttPath)
if err != nil {
return nil, fmt.Errorf("error converting subtitle: %v", err)
}
subtitlePaths = append(subtitlePaths, srtPath)
}
return subtitlePaths, nil
}
func downloadSubtitle(url string) (string, error) {
resp, err := http.Get(url)
if err != nil {
return "", err
}
defer resp.Body.Close()
tempFile, err := os.CreateTemp("", "subtitle_*.vtt")
if err != nil {
return "", err
}
defer tempFile.Close()
_, err = io.Copy(tempFile, resp.Body)
if err != nil {
return "", err
}
return tempFile.Name(), nil
}
func convertVTTtoSRT(vttPath string) (string, error) {
srtPath := strings.TrimSuffix(vttPath, ".vtt") + ".srt"
s1, _ := astisub.OpenFile(vttPath)
s1.Write(srtPath)
return srtPath, nil
}

120
utils.go Normal file
View File

@ -0,0 +1,120 @@
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")
}