From 5fcf5b914afb595fbaf6b222895cb5e741f53d04 Mon Sep 17 00:00:00 2001
From: iawia002 <z2d@jifangcheng.com>
Date: Thu, 12 Jul 2018 12:23:46 +0800
Subject: [PATCH] downloader: download captions

---
 README.md                       |  3 ++-
 config/config.go                |  2 ++
 downloader/downloader.go        | 43 ++++++++++++++++++++++++++-------
 extractors/bilibili/bilibili.go |  7 +++++-
 main.go                         |  1 +
 request/request.go              |  4 +--
 6 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index 399449344..24305e340 100644
--- a/README.md
+++ b/README.md
@@ -400,7 +400,8 @@ $ annie -j https://www.bilibili.com/video/av20203945
 ```console
 $ annie -h
 
-Usage of annie:
+Usage of ./annie:
+  -C	Download captions
   -F string
     	URLs file
   -O string
diff --git a/config/config.go b/config/config.go
index f52d8fbee..6f68bab63 100644
--- a/config/config.go
+++ b/config/config.go
@@ -35,6 +35,8 @@ var (
 	PlaylistEnd int
 	// PlaylistItems Playlist video items to download. Separated by commas like: 1,5,6
 	PlaylistItems string
+	// Caption download captions
+	Caption bool
 )
 
 // FakeHeaders fake http headers
diff --git a/downloader/downloader.go b/downloader/downloader.go
index 345ecc791..0fc61387a 100644
--- a/downloader/downloader.go
+++ b/downloader/downloader.go
@@ -46,6 +46,14 @@ type VideoData struct {
 	Formats map[string]FormatData
 }
 
+func progressBar(size int64) *pb.ProgressBar {
+	bar := pb.New64(size).SetUnits(pb.U_BYTES).SetRefreshRate(time.Millisecond * 10)
+	bar.ShowSpeed = true
+	bar.ShowFinalTime = true
+	bar.SetMaxWidth(1000)
+	return bar
+}
+
 func (data *FormatData) calculateTotalSize() {
 	var size int64
 	for _, urlData := range data.URLs {
@@ -54,15 +62,35 @@ func (data *FormatData) calculateTotalSize() {
 	data.Size = size
 }
 
-// urlSave save url file
-func (data FormatData) urlSave(
+// Caption download danmaku, subtitles, etc
+func Caption(url, refer, fileName, ext string) {
+	if !config.Caption || config.InfoOnly {
+		return
+	}
+	fmt.Println("\nDownloading captions...")
+	body := request.Get(url, refer, nil)
+	filePath := utils.FilePath(fileName, ext, false)
+	file, fileError := os.Create(filePath)
+	if fileError != nil {
+		log.Fatal(fileError)
+	}
+	defer file.Close()
+	file.WriteString(body)
+}
+
+// Save save url file
+func Save(
 	urlData URLData, refer, fileName string, bar *pb.ProgressBar,
 ) {
 	filePath := utils.FilePath(fileName, urlData.Ext, false)
 	fileSize, exists := utils.FileSize(filePath)
+	if bar == nil {
+		bar = progressBar(urlData.Size)
+		bar.Start()
+	}
 	// Skip segment file
 	// TODO: Live video URLs will not return the size
-	if fileSize == urlData.Size {
+	if exists && fileSize == urlData.Size {
 		fmt.Printf("%s: file already exists, skipping\n", filePath)
 		bar.Add64(fileSize)
 		return
@@ -201,14 +229,11 @@ func (v VideoData) Download(refer string) {
 		fmt.Printf("%s: file already exists, skipping\n", mergedFilePath)
 		return
 	}
-	bar := pb.New64(data.Size).SetUnits(pb.U_BYTES).SetRefreshRate(time.Millisecond * 10)
-	bar.ShowSpeed = true
-	bar.ShowFinalTime = true
-	bar.SetMaxWidth(1000)
+	bar := progressBar(data.Size)
 	bar.Start()
 	if len(data.URLs) == 1 {
 		// only one fragment
-		data.urlSave(data.URLs[0], refer, title, bar)
+		Save(data.URLs[0], refer, title, bar)
 		bar.Finish()
 		return
 	}
@@ -223,7 +248,7 @@ func (v VideoData) Download(refer string) {
 		wgp.Add()
 		go func(url URLData, refer, fileName string, bar *pb.ProgressBar) {
 			defer wgp.Done()
-			data.urlSave(url, refer, fileName, bar)
+			Save(url, refer, fileName, bar)
 		}(url, refer, partFileName, bar)
 
 	}
diff --git a/extractors/bilibili/bilibili.go b/extractors/bilibili/bilibili.go
index 3d79ca1a3..729eadd20 100644
--- a/extractors/bilibili/bilibili.go
+++ b/extractors/bilibili/bilibili.go
@@ -249,12 +249,17 @@ func bilibiliDownload(url string, options bilibiliOptions) downloader.VideoData
 		}
 		title = tempTitle
 	}
+	title = utils.FileName(title)
 	extractedData := downloader.VideoData{
 		Site:    "哔哩哔哩 bilibili.com",
-		Title:   utils.FileName(title),
+		Title:   title,
 		Type:    "video",
 		Formats: format,
 	}
 	extractedData.Download(url)
+	downloader.Caption(
+		fmt.Sprintf("https://comment.bilibili.com/%s.xml", cid),
+		url, title, "xml",
+	)
 	return extractedData
 }
diff --git a/main.go b/main.go
index 6a05b1e86..95b6f7682 100644
--- a/main.go
+++ b/main.go
@@ -36,6 +36,7 @@ func init() {
 		&config.PlaylistItems, "items", "",
 		"Playlist video items to download. Separated by commas like: 1,5,6",
 	)
+	flag.BoolVar(&config.Caption, "C", false, "Download captions")
 }
 
 func download(videoURL string) {
diff --git a/request/request.go b/request/request.go
index 9a5f4db01..6c93f1b29 100644
--- a/request/request.go
+++ b/request/request.go
@@ -1,8 +1,8 @@
 package request
 
 import (
+	"compress/flate"
 	"compress/gzip"
-	"compress/zlib"
 	"crypto/tls"
 	"fmt"
 	"io"
@@ -133,7 +133,7 @@ func Get(url, refer string, headers map[string]string) string {
 	case "gzip":
 		reader, _ = gzip.NewReader(res.Body)
 	case "deflate":
-		reader, _ = zlib.NewReader(res.Body)
+		reader = flate.NewReader(res.Body)
 	default:
 		reader = res.Body
 	}