m3u82json/main.go

package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"

	"github.com/grafov/m3u8"
)

const baseScheme = "http"
const MASTER = "master"
const PLAYLIST = "playlist"

type m3u8Json struct {
	URL      string
	BaseUrl  string
	Type     string
	Manifest interface{}
}

func nilIndex(input []*m3u8.MediaSegment) int {
	output := -1

	for i := range input {
		if input[i] == nil {
			return i
		}
	}
	return output
}

func main() {
	isUrl := false
	var data *bufio.Reader
	var outputJson m3u8Json
	recursive := flag.Bool("r", false, "Recursive download master and playlists")
	abs := flag.Bool("a", false, "make URI absolute")
	baseUrl := flag.String("b", "", "Base URL")
	flag.Parse()
	inputs := flag.Args()
	if len(inputs) != 1 {
		log.Fatal("Only 1 input is allowed")
	}

	input := inputs[0]
	u, err := url.Parse(input)
	if err != nil {
		log.Fatal("not a valid URL")
	}

	if strings.Contains(u.Scheme, baseScheme) {
		isUrl = true
	}

	if !isUrl {
		var f *os.File
		var err error
		if input == "-" {
			data = bufio.NewReader(os.Stdin)
		} else {
			f, err = os.Open(input)
			if err != nil {
				log.Fatal(err)
			}
			data = bufio.NewReader(f)
		}

	} else {
		if *baseUrl == "" {
			*baseUrl = fmt.Sprintf("%v://%v%v", u.Scheme, u.Host, path.Dir(u.EscapedPath()))
		}
		r, err := http.Get(u.String())
		if err != nil {
			log.Fatal(err)
		}
		outputJson.URL = u.String()
		outputJson.BaseUrl = *baseUrl
		data = bufio.NewReader(r.Body)
	}

	p, listType, err := m3u8.DecodeFrom(data, true)
	if err != nil {
		log.Fatal(err)
	}
	switch listType {
	case m3u8.MEDIA:
		outputJson.Type = PLAYLIST
		playlist := p.(*m3u8.MediaPlaylist)
		playlist.Segments = playlist.Segments[:nilIndex(playlist.Segments)]
		for _, entry := range playlist.Segments {
			u, err := url.Parse(entry.URI)
			if err != nil {
				log.Fatal("not a valid URL")
			}
			if *abs && !u.IsAbs() && *baseUrl != "" {
				entry.URI = *baseUrl + "/" + entry.URI
			}
		}
		outputJson.Manifest = playlist
	case m3u8.MASTER:
		outputJson.Type = MASTER
		master := p.(*m3u8.MasterPlaylist)
		for _, entry := range master.Variants {
			var chunckUrl string
			u, err := url.Parse(entry.URI)
			if err != nil {
				log.Fatal("not a valid URL")
			}
			if !u.IsAbs() {
				chunckUrl = *baseUrl + "/" + entry.URI
				if *abs && *baseUrl != "" {
					entry.URI = chunckUrl
				}
			} else {
				chunckUrl = entry.URI
			}
			if isUrl && *recursive {
				clist, err := http.Get(chunckUrl)
				if err != nil {
					log.Fatal(err)
				}
				cdata := bufio.NewReader(clist.Body)
				p, _, _ := m3u8.DecodeFrom(cdata, true)
				playlist := p.(*m3u8.MediaPlaylist)
				playlist.Segments = playlist.Segments[:nilIndex(playlist.Segments)]
				for _, entry := range playlist.Segments {
					u, err := url.Parse(entry.URI)
					if err != nil {
						log.Fatal("not a valid URL")
					}
					if *abs && !u.IsAbs() {
						entry.URI = *baseUrl + "/" + entry.URI
					}
				}
				entry.Chunklist = playlist
			}
		}
		outputJson.Manifest = master
	}

	output, err := json.MarshalIndent(outputJson, "", " ")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s", output)
}