Initial version of epgpruner

Prune EPG file: leave only items that are in the playlist.

1 files changed, 206 insertions, 0 deletions

diff --git a/epgpruner/epgpruner.go b/epgpruner/epgpruner.go
new file mode 100644
index 0000000..b2a4f54
--- /dev/null
+++ b/epgpruner/epgpruner.go
@@ -0,0 +1,206 @@
+package main
+
+import (
+	"compress/gzip"
+	"encoding/xml"
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"os"
+	"regexp"
+	"strings"
+)
+
+var (
+	docType = `<!DOCTYPE tv SYSTEM "xmltv.dtd">`
+
+	epgRE     = regexp.MustCompile(`url-tvg="([a-z0-9-:/.]+)"`)
+	channelRE = regexp.MustCompile(`tvg-id="([a-z0-9-]+)"`)
+)
+
+func usage() {
+	fmt.Fprintf(os.Stderr, "usage: epgpruner URL\n")
+	flag.PrintDefaults()
+	os.Exit(2)
+}
+
+func main() {
+	flag.Usage = usage
+	flag.Parse()
+	if flag.NArg() != 1 {
+		usage()
+		os.Exit(1)
+	}
+	playlistURL := flag.Args()[0]
+	playlist := Playlist{
+		URL: playlistURL,
+	}
+	err := playlist.Fetch()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	playlist.Reduce()
+
+	reducedEpg, err := xml.MarshalIndent(playlist.Epg, "", "    ")
+	if err != nil {
+		log.Fatal(err)
+	}
+	fmt.Println(
+		xml.Header +
+			docType + "\n" +
+			string(reducedEpg),
+	)
+}
+
+type Playlist struct {
+	URL      string
+	EpgURL   string
+	Channels map[string]struct{}
+	Epg      Epg
+}
+
+func (p *Playlist) Fetch() error {
+	compressed := false
+	body, err := fetch(p.URL, compressed)
+	if err != nil {
+		return err
+	}
+	err = p.parse(body)
+	if err != nil {
+		return err
+	}
+
+	compressed = true
+	body, err = fetch(p.EpgURL, compressed)
+	if err != nil {
+		return err
+	}
+
+	err = xml.Unmarshal(body, &p.Epg)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func fetch(url string, compressed bool) ([]byte, error) {
+	resp, err := http.Get(url)
+	if err != nil {
+		return make([]byte, 0), err
+	}
+	if resp.StatusCode != http.StatusOK {
+		return make([]byte, 0),
+			fmt.Errorf("Bad reponse status: %d", resp.StatusCode)
+	}
+	reader := resp.Body
+	defer reader.Close()
+	if compressed {
+		reader, err = gzip.NewReader(reader)
+		if err != nil {
+			return make([]byte, 0), err
+		}
+	}
+	body, err := io.ReadAll(reader)
+	if err != nil {
+		return make([]byte, 0), err
+	}
+	return body, nil
+}
+
+func (p *Playlist) parse(body []byte) error {
+	epgURL := ""
+	channels := make(map[string]struct{})
+	for _, line := range strings.Split(string(body), "\n") {
+		if strings.HasPrefix(line, "#EXTM3U") {
+			m := epgRE.FindStringSubmatch(line)
+			if m == nil {
+				return fmt.Errorf(
+					"Couldn't match EPG url from: `%s`",
+					line,
+				)
+			}
+			epgURL = m[1]
+		} else if strings.HasPrefix(line, "#EXTINF") {
+			m := channelRE.FindStringSubmatch(line)
+			if m == nil {
+				return fmt.Errorf(
+					"Couldn't match channel name from: `%s`",
+					line,
+				)
+			}
+			channel := m[1]
+			channels[channel] = struct{}{}
+		}
+	}
+
+	p.EpgURL = epgURL
+	p.Channels = channels
+
+	return nil
+}
+
+func (p *Playlist) Reduce() {
+	channels := make([]Channel, 0)
+	programmes := make([]Programme, 0)
+
+	for _, channel := range p.Epg.Channels {
+		if _, ok := p.Channels[channel.ID]; !ok {
+			continue
+		}
+		channels = append(channels, channel)
+	}
+
+	for _, programme := range p.Epg.Programmes {
+		if _, ok := p.Channels[programme.Channel]; !ok {
+			continue
+		}
+		programmes = append(programmes, programme)
+	}
+
+	p.Epg.Channels = channels
+	p.Epg.Programmes = programmes
+}
+
+type Epg struct {
+	XMLName    xml.Name    `xml:"tv"`
+	Info       string      `xml:"generator-info-name,attr"`
+	Channels   []Channel   `xml:"channel"`
+	Programmes []Programme `xml:"programme"`
+}
+
+type Channel struct {
+	ID          string      `xml:"id,attr"`
+	DisplayName DisplayName `xml:"display-name"`
+	Icon        Icon        `xml:"icon"`
+}
+
+type DisplayName struct {
+	Lang string `xml:"lang,attr"`
+	Name string `xml:",chardata"`
+}
+
+type Icon struct {
+	Src string `xml:"src,attr"`
+}
+
+type Programme struct {
+	Start   string `xml:"start,attr"`
+	Stop    string `xml:"stop,attr"`
+	Channel string `xml:"channel,attr"`
+	Title   Title  `xml:"title,omitempty"`
+	Desc    Desc   `xml:"desc,omitempty"`
+}
+
+type Title struct {
+	Lang  string `xml:"lang,attr"`
+	Title string `xml:",chardata"`
+}
+
+type Desc struct {
+	Lang string `xml:"lang,attr"`
+	Desc string `xml:",chardata"`
+}