Support cookie for downloading from restricted sources

README.md

@@ -9,6 +9,7 @@ You only need to specify the flags(`u`, `o`, `c`) to run, downloader will automa
 ## Features
 
 - Download and parse M3U8（VOD）
+- Pass an auth cookie to the server
 - Retry on download TS failure
 - Parse Master playlist
 - Decrypt TS
@@ -19,7 +20,7 @@ You only need to specify the flags(`u`, `o`, `c`) to run, downloader will automa
 ### source
 
 ```bash
-go run main.go -u=http://example.com/index.m3u8 -o=/data/example
+go run main.go -u=http://example.com/index.m3u8 -o=/data/example -cookie="CloudFront-Policy=...; CloudFront-Signature=...; CloudFront-Key-Pair-Id=..."
 ```
 
 ### binary:

dl/dowloader.go

@@ -4,6 +4,9 @@ import (
 	"bufio"
 	"fmt"
 	"io/ioutil"
+	"net/http"
+	"net/http/cookiejar"
+	"net/url"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -12,6 +15,8 @@ import (
 
 	"github.com/oopsguy/m3u8/parse"
 	"github.com/oopsguy/m3u8/tool"
+
+	"golang.org/x/net/publicsuffix"
 )
 
 const (
@@ -31,11 +36,16 @@ type Downloader struct {
 	segLen   int
 
 	result *parse.Result
+	jar    http.CookieJar
 }
 
 // NewTask returns a Task instance
-func NewTask(output string, url string) (*Downloader, error) {
-	result, err := parse.FromURL(url)
+func NewTask(output string, url string, cookie string) (*Downloader, error) {
+	jar, err := initCookieJar(url, cookie)
+	if err != nil {
+		return nil, fmt.Errorf("init cookie jar failed: %s", err.Error())
+	}
+	result, err := parse.FromURL(url, jar)
 	if err != nil {
 		return nil, err
 	}
@@ -61,6 +71,7 @@ func NewTask(output string, url string) (*Downloader, error) {
 		folder:   folder,
 		tsFolder: tsFolder,
 		result:   result,
+		jar:      jar,
 	}
 	d.segLen = len(result.M3u8.Segments)
 	d.queue = genSlice(d.segLen)
@@ -104,7 +115,7 @@ func (d *Downloader) Start(concurrency int) error {
 func (d *Downloader) download(segIndex int) error {
 	tsFilename := tsFilename(segIndex)
 	tsUrl := d.tsURL(segIndex)
-	b, e := tool.Get(tsUrl)
+	b, e := tool.Get(tsUrl, d.jar)
 	if e != nil {
 		return fmt.Errorf("request %s, %s", tsUrl, e.Error())
 	}
@@ -252,3 +263,20 @@ func genSlice(len int) []int {
 	}
 	return s
 }
+
+func initCookieJar(address string, cookie string) (http.CookieJar, error) {
+	rootUrl, err := url.Parse(address)
+	if err != nil {
+		return nil, fmt.Errorf("parseurl failed: %s", err.Error())
+	}
+	rootUrl.Path = "/"
+	jar, err := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
+	if err != nil {
+		return nil, fmt.Errorf("create cookie jar failed: %s", err.Error())
+	}
+	header := http.Header{}
+	header.Add("Cookie", cookie)
+	request := http.Request{Header: header}
+	jar.SetCookies(rootUrl, request.Cookies())
+	return jar, nil
+}

main.go

@@ -10,12 +10,14 @@ import (
 
 var (
 	url      string
+	cookie   string
 	output   string
 	chanSize int
 )
 
 func init() {
 	flag.StringVar(&url, "u", "", "M3U8 URL, required")
+	flag.StringVar(&cookie, "cookie", "", "Cookie string to use")
 	flag.IntVar(&chanSize, "c", 25, "Maximum number of occurrences")
 	flag.StringVar(&output, "o", "", "Output folder, required")
 }
@@ -37,7 +39,7 @@ func main() {
 	if chanSize <= 0 {
 		panic("parameter 'c' must be greater than 0")
 	}
-	downloader, err := dl.NewTask(output, url)
+	downloader, err := dl.NewTask(output, url, cookie)
 	if err != nil {
 		panic(err)
 	}

parse/parser.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"net/http"
 	"net/url"
 
 	"github.com/oopsguy/m3u8/tool"
@@ -15,13 +16,13 @@ type Result struct {
 	Keys map[int]string
 }
 
-func FromURL(link string) (*Result, error) {
+func FromURL(link string, jar http.CookieJar) (*Result, error) {
 	u, err := url.Parse(link)
 	if err != nil {
 		return nil, err
 	}
 	link = u.String()
-	body, err := tool.Get(link)
+	body, err := tool.Get(link, jar)
 	if err != nil {
 		return nil, fmt.Errorf("request m3u8 URL failed: %s", err.Error())
 	}
@@ -33,7 +34,7 @@ func FromURL(link string) (*Result, error) {
 	}
 	if len(m3u8.MasterPlaylist) != 0 {
 		sf := m3u8.MasterPlaylist[0]
-		return FromURL(tool.ResolveURL(u, sf.URI))
+		return FromURL(tool.ResolveURL(u, sf.URI), jar)
 	}
 	if len(m3u8.Segments) == 0 {
 		return nil, errors.New("can not found any TS file description")
@@ -52,7 +53,7 @@ func FromURL(link string) (*Result, error) {
 			// Request URL to extract decryption key
 			keyURL := key.URI
 			keyURL = tool.ResolveURL(u, keyURL)
-			resp, err := tool.Get(keyURL)
+			resp, err := tool.Get(keyURL, jar)
 			if err != nil {
 				return nil, fmt.Errorf("extract key failed: %s", err.Error())
 			}

tool/http.go

@@ -7,9 +7,10 @@ import (
 	"time"
 )
 
-func Get(url string) (io.ReadCloser, error) {
+func Get(url string, jar http.CookieJar) (io.ReadCloser, error) {
 	c := http.Client{
 		Timeout: time.Duration(60) * time.Second,
+		Jar: jar,
 	}
 	resp, err := c.Get(url)
 	if err != nil {