From f6ecf916b740ee9eeb475859ccaf5179337a278a Mon Sep 17 00:00:00 2001
From: xbol0 <xbolo@duck.com>
Date: Tue, 4 Oct 2022 17:02:24 +0800
Subject: [PATCH] Add -F flag to try use remote filename

---
 cmd/main.go      |  1 +
 download.go      | 17 +++++++++++++++++
 download_test.go | 18 ++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/cmd/main.go b/cmd/main.go
index 962407b..45e98a4 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -110,6 +110,7 @@ func init() {
 	downloadCmd.Flags().BoolVarP(&gRunReq.Insecure, "Insecure", "", false, "是否允许不安全的请求")
 	downloadCmd.Flags().StringVarP(&gRunReq.SaveDir, "SaveDir", "d", "", "文件保存路径(默认为当前路径)")
 	downloadCmd.Flags().StringVarP(&gRunReq.FileName, "FileName", "f", "", "文件名")
+	downloadCmd.Flags().BoolVarP(&gRunReq.RemoteName, "RemoteFileName", "F", false, "尝试自动获取文件名")
 	downloadCmd.Flags().IntVarP(&gRunReq.SkipTsCountFromHead, "SkipTsCountFromHead", "", 0, "跳过前面几个ts")
 	downloadCmd.Flags().StringVarP(&gRunReq.SetProxy, "SetProxy", "", "", "代理设置, http://127.0.0.1:8080 socks5://127.0.0.1:1089")
 	downloadCmd.Flags().BoolVarP(&gRunReq.SkipRemoveTs, "SkipRemoveTs", "", false, "不删除下载的ts文件")
diff --git a/download.go b/download.go
index 31600fa..28cc88d 100644
--- a/download.go
+++ b/download.go
@@ -15,6 +15,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"path"
 	"path/filepath"
 	"regexp"
 	"strconv"
@@ -85,6 +86,7 @@ type RunDownload_Req struct {
 	Insecure            bool   // "是否允许不安全的请求(默认为false)"
 	SaveDir             string // "文件保存路径(默认为当前路径)"
 	FileName            string // 文件名
+	RemoteName          bool   // 尝试自动获取文件名
 	SkipTsCountFromHead int    // 跳过前面几个ts
 	SetProxy            string
 	HeaderMap           map[string][]string
@@ -116,6 +118,9 @@ func (this *downloadEnv) RunDownload(req RunDownload_Req) (resp RunDownload_Resp
 			return resp
 		}
 	}
+	if req.RemoteName && req.FileName == "" {
+		req.FileName = getFileName(req.M3u8Url)
+	}
 	if req.FileName == "" {
 		req.FileName = "all"
 	}
@@ -671,3 +676,15 @@ func (this *downloadEnv) GetIsCancel() bool {
 		return false
 	}
 }
+
+func getFileName(u string) string {
+	url, err := url.Parse(u)
+	if err != nil {
+		return ""
+	}
+	name := path.Base(url.Path)
+	if strings.HasSuffix(name, ".m3u8") {
+		return name[:len(name)-5]
+	}
+	return name
+}
diff --git a/download_test.go b/download_test.go
index 5a432a1..c36decb 100644
--- a/download_test.go
+++ b/download_test.go
@@ -100,3 +100,21 @@ func TestFull(t *testing.T) {
 		panic("state error")
 	}
 }
+
+func TestGetFileName(t *testing.T) {
+	u1 := "https://example.com/video.m3u8"
+	u2 := "https://example.com/video.m3u8?query=1"
+	u3 := "https://example.com/video-name"
+
+	if getFileName(u1) != "video" {
+		t.Fail()
+	}
+
+	if getFileName(u2) != "video" {
+		t.Fail()
+	}
+
+	if getFileName(u3) != "video-name" {
+		t.Fail()
+	}
+}