go-m3u8/cmd/main.go

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"

	"github.com/urfave/cli/v2"
	gom3u8 "sirherobrine23.com.br/Sirherobrine23/go-m3u8"
)

func ProcessRequestLocal(ctx *cli.Context) error {
	sPath, output, exts := ctx.String("file"), ctx.String("output"), ctx.StringSlice("filter")

	var videos *gom3u8.Playlist
	if sPath[0:4] == "http" {
		var err error
		videos, err = gom3u8.GetParse(sPath)
		if err != nil {
			return err
		}
	} else {
		inputFile, err := os.Open(sPath)
		if err != nil {
			return err
		}
		if videos, err = gom3u8.Parser(inputFile); err != nil {
			return err
		}
	}

	// Filter channels
	videos.FilterChannels(exts...)

	if strings.HasSuffix(output, "m3u") {
		fmt.Printf("Writing %d channels to file\n", len(videos.Channels))
		file, err := os.Create(output)
		if err != nil {
			return err
		}
		defer file.Close()
		videos.WriteTo(file)
		return nil
	}

	// JSON save
	var et *json.Encoder
	if len(output) == 0 || ctx.Bool("print") {
		et = json.NewEncoder(os.Stdout)
	} else {
		fmt.Printf("Writing %d channels to file\n", len(videos.Channels))
		file, err := os.Create(output)
		if err != nil {
			return err
		}
		defer file.Close()
		et = json.NewEncoder(file)
	}
	et.SetIndent("", "  ")
	return et.Encode(videos)
}

func ProcessServer(ctx *cli.Context) error {
	videos, err := gom3u8.GetPlayer(ctx.String("server"), ctx.String("username"), ctx.String("password"))
	if err != nil {
		return err
	}

	// Filter channels
	output, exts := ctx.String("output"), ctx.StringSlice("filter")
	videos.FilterChannels(exts...)

	if strings.HasSuffix(output, "m3u") {
		fmt.Printf("Writing %d channels to file\n", len(videos.Channels))
		file, err := os.Create(output)
		if err != nil {
			return err
		}
		defer file.Close()
		videos.WriteTo(file)
		return nil
	}

	// JSON save
	var et *json.Encoder
	if len(output) == 0 || ctx.Bool("print") {
		et = json.NewEncoder(os.Stdout)
	} else {
		fmt.Printf("Writing %d channels to file\n", len(videos.Channels))
		file, err := os.Create(output)
		if err != nil {
			return err
		}
		defer file.Close()
		et = json.NewEncoder(file)
	}
	et.SetIndent("", "  ")
	return et.Encode(videos)
}

func main() {
	app := cli.NewApp()
	app.Name = "gom3u"
	app.Usage = "Parse and filter m3u8 files"
	app.EnableBashCompletion = true

	app.Commands = []*cli.Command{
		{
			Name:    "local",
			Aliases: []string{"request"},
			Action:  ProcessRequestLocal,
			Flags: []cli.Flag{
				&cli.StringFlag{
					Name:     "file",
					Required: true,
					Aliases: []string{
						"input",
						"i",
					},
					DefaultText: "./iptv.m3u",
				},
				&cli.BoolFlag{
					Name:     "print",
					Aliases:  []string{"p"},
					Required: false,
				},
				&cli.StringFlag{
					Name:     "output",
					Aliases:  []string{"o"},
					Required: false,
				},
				&cli.MultiStringFlag{
					Target: &cli.StringSliceFlag{
						Name:    "filter",
						Aliases: []string{"f"},
					},
					Value: []string{
						".mp4",
						".mkv",
						".avi",
						".mov",
					},
				},
			},
		},
		{
			Name:   "server",
			Action: ProcessServer,
			Flags: []cli.Flag{
				&cli.BoolFlag{
					Name:     "print",
					Aliases:  []string{"p"},
					Required: false,
				},
				&cli.StringFlag{
					Name:     "output",
					Aliases:  []string{"o"},
					Required: false,
				},
				&cli.StringFlag{
					Name:     "server",
					Aliases:  []string{"s"},
					Required: true,
				},
				&cli.StringFlag{
					Name:     "username",
					Aliases:  []string{"u"},
					Required: true,
				},
				&cli.StringFlag{
					Name:     "password",
					Aliases:  []string{"P"},
					Required: true,
				},
				&cli.MultiStringFlag{
					Target: &cli.StringSliceFlag{
						Name:    "filter",
						Aliases: []string{"f"},
					},
					Value: []string{
						".mp4",
						".mkv",
						".avi",
						".mov",
					},
				},
			},
		},
	}

	err := app.Run(os.Args)
	if err != nil {
		// app.Run already exits for errors implementing ErrorCoder,
		// so we only handle generic errors with code 1 here.
		fmt.Fprintf(app.ErrWriter, "Error: %v\n", err)
		os.Exit(1)
	}
}