From 79493d2e1350ea0050589580cf4289f95fdc5a90 Mon Sep 17 00:00:00 2001
From: p418 <prince418@gmail.com>
Date: Mon, 12 Oct 2020 03:40:14 +0200
Subject: [PATCH] Add m3u group filtering

---
 cmd/root.go          |  6 ++++++
 pkg/config/config.go |  1 +
 pkg/server/server.go | 15 +++++++++++++++
 3 files changed, 22 insertions(+)

diff --git a/cmd/root.go b/cmd/root.go
index 914cda7c..007735fe 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -68,6 +68,11 @@ var rootCmd = &cobra.Command{
 			}
 		}
 
+		filters := map[string](bool){}
+		for _, groupName := range viper.GetStringSlice("filter-groups") {
+			filters[groupName] = true
+		}
+
 		conf := &config.ProxyConfig{
 			HostConfig: &config.HostConfiguration{
 				Hostname: viper.GetString("hostname"),
@@ -83,6 +88,7 @@ var rootCmd = &cobra.Command{
 			HTTPS:              viper.GetBool("https"),
 			M3UFileName:        viper.GetString("m3u-file-name"),
 			CustomEndpoint:     viper.GetString("custom-endpoint"),
+			FilterGroups:       filters,
 		}
 
 		server, err := server.NewServer(conf)
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 66b5c1c8..b94f83c0 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -53,4 +53,5 @@ type ProxyConfig struct {
 	RemoteURL          *url.URL
 	HTTPS              bool
 	User, Password     CredentialString
+	FilterGroups       map[string](bool)
 }
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 23db2c5f..c16eac4c 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -96,9 +96,24 @@ func (c *Config) playlistInitialization() error {
 // MarshallInto a *bufio.Writer a Playlist.
 func (c *Config) marshallInto(into *os.File, xtream bool) error {
 	into.WriteString("#EXTM3U\n") // nolint: errcheck
+
+TRACKS_LOOP:
 	for _, track := range c.playlist.Tracks {
+
+		// Groups filtering
+		if len(c.FilterGroups) > 0 {
+			for i := range track.Tags {
+				name := track.Tags[i].Name
+				value := track.Tags[i].Value
+				if name == "group-title" && !c.FilterGroups[value] {
+					continue TRACKS_LOOP
+				}
+			}
+		}
+
 		into.WriteString("#EXTINF:")                       // nolint: errcheck
 		into.WriteString(fmt.Sprintf("%d ", track.Length)) // nolint: errcheck
+
 		for i := range track.Tags {
 			if i == len(track.Tags)-1 {
 				into.WriteString(fmt.Sprintf("%s=%q", track.Tags[i].Name, track.Tags[i].Value)) // nolint: errcheck