Use buffer writer and remove unless function

README.md

@@ -2,8 +2,7 @@
 
 M3u8 - a mini m3u8 downloader written in Golang for downloading and merging TS(Transport Stream) files.
 
-You only need to specify the flags(`u`, `o`, `c`) to run, downloader will automatically download the m3u8 and parse it, 
-then download and merge all TS files.
+You only need to specify the flags(`u`, `o`, `c`) to run, downloader will automatically download all TS files and consolidate them into a single TS file.
 
 
 ## Features

dl/dowloader.go

@@ -3,7 +3,6 @@ package dl
 import (
 	"bufio"
 	"fmt"
-	"io"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -34,6 +33,7 @@ type Downloader struct {
 	result *parse.Result
 }
 
+// NewTask returns a Task instance
 func NewTask(output string, url string) (*Downloader, error) {
 	result, err := parse.FromURL(url)
 	if err != nil {
@@ -67,6 +67,7 @@ func NewTask(output string, url string) (*Downloader, error) {
 	return d, nil
 }
 
+// Start runs downloader
 func (d *Downloader) Start(concurrency int) error {
 	var wg sync.WaitGroup
 	limitChan := make(chan byte, concurrency)
@@ -130,7 +131,6 @@ func (d *Downloader) download(segIndex int) error {
 			}
 		}
 	}
-
 	// https://en.wikipedia.org/wiki/MPEG_transport_stream
 	// Some TS files do not start with SyncByte 0x47, they can not be played after merging,
 	// Need to remove the bytes before the SyncByte 0x47(71).
@@ -146,14 +146,15 @@ func (d *Downloader) download(segIndex int) error {
 	if _, err := w.Write(bytes); err != nil {
 		return fmt.Errorf("write TS bytes failed: %s", err.Error())
 	}
+	// Release file resource to rename file
 	_ = f.Close()
 	if err = os.Rename(fTemp, fPath); err != nil {
 		return err
 	}
 	// Maybe it will be safer in this way...
 	atomic.AddInt32(&d.finish, 1)
 	tool.DrawProgressBar("Downloading",
-		float32(d.finish)/float32(d.segLen), progressWidth, tsFilename)
+		float32(d.finish)/float32(d.segLen), progressWidth)
 	return nil
 }
 
@@ -187,78 +188,47 @@ func (d *Downloader) back(segIndex int) error {
 
 func (d *Downloader) merge() error {
 	// In fact, the number of downloaded segments should be equal to number of m3u8 segments
+	missingCount := 0
 	for idx := 0; idx < d.segLen; idx++ {
 		tsFilename := tsFilename(idx)
 		f := filepath.Join(d.tsFolder, tsFilename)
 		if _, err := os.Stat(f); err != nil {
-			fmt.Printf("Missing the TS file：%s\n", tsFilename)
+			missingCount++
 		}
 	}
+	if missingCount > 0 {
+		fmt.Printf("Warning: %d TS files missing\n", missingCount)
+	}
+
 	// Create a TS file for merging, all segment files will be written to this file.
 	mFile, err := os.Create(filepath.Join(d.folder, mergeTSFilename))
 	if err != nil {
-		return fmt.Errorf("merge TS files failed：%s", err.Error())
+		return fmt.Errorf("create main TS file failed：%s", err.Error())
 	}
 	//noinspection GoUnhandledErrorResult
 	defer mFile.Close()
-	// Move to EOF
-	ls, err := mFile.Seek(0, io.SeekEnd)
-	if err != nil {
-		return err
-	}
-	fmt.Print("/r")
+
+	writer := bufio.NewWriter(mFile)
 	mergedCount := 0
 	for segIndex := 0; segIndex < len(d.result.M3u8.Segments); segIndex++ {
 		tsFilename := tsFilename(segIndex)
 		bytes, err := ioutil.ReadFile(filepath.Join(d.tsFolder, tsFilename))
-		s, err := mFile.WriteAt(bytes, ls)
+		_, err = writer.Write(bytes)
 		if err != nil {
 			return err
 		}
-		ls += int64(s)
 		mergedCount++
 		tool.DrawProgressBar("Merging",
-			float32(mergedCount)/float32(len(d.result.M3u8.Segments)), progressWidth, tsFilename)
+			float32(mergedCount)/float32(len(d.result.M3u8.Segments)), progressWidth)
 	}
+	_ = writer.Flush()
+
 	fmt.Println()
-	_ = mFile.Sync()
 	// Remove `ts` folder
 	_ = os.RemoveAll(d.tsFolder)
 	return nil
 }
 
-func (d *Downloader) rangeMerge(start, end int) error {
-	if start < 0 || end > d.segLen-1 {
-		return fmt.Errorf("invalid segment index range：%d,%d", start, end)
-	}
-	if end-start < 2 {
-		return nil
-	}
-	mFilename := tsFilename(start)
-	mFile, err := os.Create(filepath.Join(d.tsFolder, mFilename))
-	if err != nil {
-		return fmt.Errorf("merge TS files failed：%s", err.Error())
-	}
-	//noinspection GoUnhandledErrorResult
-	defer mFile.Close()
-	// Move to EOF
-	ls, err := mFile.Seek(0, io.SeekEnd)
-	if err != nil {
-		return err
-	}
-	for idx := start + 1; idx <= end; idx++ {
-		tsFilename := tsFilename(idx)
-		bytes, err := ioutil.ReadFile(filepath.Join(d.tsFolder, tsFilename))
-		s, err := mFile.WriteAt(bytes, ls)
-		if err != nil {
-			return err
-		}
-		ls += int64(s)
-	}
-	_ = mFile.Sync()
-	return nil
-}
-
 func (d *Downloader) tsURL(segIndex int) string {
 	seg := d.result.M3u8.Segments[segIndex]
 	return tool.ResolveURL(d.result.URL, seg.URI)

parse/m3u8.go

@@ -2,8 +2,10 @@
 package parse
 
 import (
+	"bufio"
 	"errors"
 	"fmt"
+	"io"
 	"regexp"
 	"strconv"
 	"strings"
@@ -22,6 +24,7 @@ const (
 	CryptMethodNONE CryptMethod = "NONE"
 )
 
+// regex pattern for extracting `key=value` parameters from a line
 var linePattern = regexp.MustCompile(`([a-zA-Z-]+)=("[^"]+"|[^",]+)`)
 
 type M3u8 struct {
@@ -61,7 +64,13 @@ type Key struct {
 	IV     string
 }
 
-func parseLines(lines []string) (*M3u8, error) {
+func parse(reader io.Reader) (*M3u8, error) {
+	s := bufio.NewScanner(reader)
+	var lines []string
+	for s.Scan() {
+		lines = append(lines, s.Text())
+	}
+
 	var (
 		i     = 0
 		count = len(lines)
@@ -72,6 +81,7 @@ func parseLines(lines []string) (*M3u8, error) {
 		extInf  bool
 		extByte bool
 	)
+
 	for ; i < count; i++ {
 		line := strings.TrimSpace(lines[i])
 		if i == 0 {
@@ -104,7 +114,7 @@ func parseLines(lines []string) (*M3u8, error) {
 				return nil, err
 			}
 		case strings.HasPrefix(line, "#EXT-X-STREAM-INF:"):
-			mp, err := parseStreamInfo(line)
+			mp, err := parseMasterPlaylist(line)
 			if err != nil {
 				return nil, err
 			}
@@ -198,10 +208,11 @@ func parseLines(lines []string) (*M3u8, error) {
 			continue
 		}
 	}
+
 	return m3u8, nil
 }
 
-func parseStreamInfo(line string) (*MasterPlaylist, error) {
+func parseMasterPlaylist(line string) (*MasterPlaylist, error) {
 	params := parseLineParameters(line)
 	if len(params) == 0 {
 		return nil, errors.New("empty parameter")
@@ -230,6 +241,7 @@ func parseStreamInfo(line string) (*MasterPlaylist, error) {
 	return mp, nil
 }
 
+// parseLineParameters extra parameters in string `line`
 func parseLineParameters(line string) map[string]string {
 	r := linePattern.FindAllStringSubmatch(line, -1)
 	params := make(map[string]string)

parse/parser.go

@@ -1,7 +1,6 @@
 package parse
 
 import (
-	"bufio"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -28,12 +27,7 @@ func FromURL(link string) (*Result, error) {
 	}
 	//noinspection GoUnhandledErrorResult
 	defer body.Close()
-	s := bufio.NewScanner(body)
-	var lines []string
-	for s.Scan() {
-		lines = append(lines, s.Text())
-	}
-	m3u8, err := parseLines(lines)
+	m3u8, err := parse(body)
 	if err != nil {
 		return nil, err
 	}

tool/crypt_test.go

@@ -1,7 +1,6 @@
 package tool
 
 import (
-	"fmt"
 	"testing"
 )
 
@@ -11,15 +10,14 @@ func Test_AES128Encrypt_AND_AES128Decrypt(t *testing.T) {
 	iv := "xduio1f8a12348u4"
 	encrypt, err := AES128Encrypt([]byte(expected), []byte(key), []byte(iv))
 	if err != nil {
-		t.Error(err)
+		t.Fatal(err)
 	}
 	decrypt, err := AES128Decrypt(encrypt, []byte(key), []byte(iv))
 	if err != nil {
-		t.Error(err)
+		t.Fatal(err)
 	}
 	de := string(decrypt)
 	if de != expected {
-		t.Errorf("expected: %s, result: %s", expected, de)
+		t.Fatalf("expected: %s, result: %s", expected, de)
 	}
-	fmt.Printf("Decrypt result: %s\n", de)
 }

tool/http_test.go

@@ -1,7 +1,6 @@
 package tool
 
 import (
-	"fmt"
 	"io/ioutil"
 	"testing"
 )
@@ -12,9 +11,8 @@ func TestGet(t *testing.T) {
 		t.Error(err)
 	}
 	defer body.Close()
-	bytes, err := ioutil.ReadAll(body)
+	_, err = ioutil.ReadAll(body)
 	if err != nil {
 		t.Error(err)
 	}
-	fmt.Println(string(bytes))
 }

tool/util_test.go

@@ -1,7 +1,6 @@
 package tool
 
 import (
-	"fmt"
 	"net/url"
 	"testing"
 )
@@ -12,22 +11,22 @@ func TestResolveURL(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
+
 	result := ResolveURL(u, "videos/111111.ts")
 	expected := "http://www.example.com/test/videos/111111.ts"
 	if result != expected {
-		t.Error(fmt.Errorf("wrong URL, expected: %s, result: %s", expected, result))
+		t.Fatalf("wrong URL, expected: %s, result: %s", expected, result)
 	}
-	fmt.Printf("expected: %s, result: %s\n", expected, result)
+
 	result = ResolveURL(u, "/videos/2222222.ts")
 	expected = "http://www.example.com/videos/2222222.ts"
 	if result != expected {
-		t.Error(fmt.Errorf("wrong URL, expected: %s, result: %s", expected, result))
+		t.Fatalf("wrong URL, expected: %s, result: %s", expected, result)
 	}
-	fmt.Printf("expected: %s, result: %s\n", expected, result)
+
 	result = ResolveURL(u, "https://test.com/11111.key")
 	expected = "https://test.com/11111.key"
 	if result != expected {
-		t.Error(fmt.Errorf("wrong URL, expected: %s, result: %s", expected, result))
+		t.Fatalf("wrong URL, expected: %s, result: %s", expected, result)
 	}
-	fmt.Printf("expected: %s, result: %s\n", expected, result)
 }