From 3481319e1ce6b23afdc7f1daac1f23a65c4e07d3 Mon Sep 17 00:00:00 2001
From: saintliao <saintliao@gmail.com>
Date: Mon, 20 Mar 2023 19:39:43 +0800
Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E=20--cookies-from-browser=20?=
 =?UTF-8?q?=E7=9A=84=E8=AA=AA=E6=98=8E=E5=85=A7=E5=AE=B9=20=E5=9C=A8=20--c?=
 =?UTF-8?q?ookies-from-browser=20=E5=8A=A0=E5=85=A5=E5=8F=AF=E6=8C=87?=
 =?UTF-8?q?=E5=AE=9A=E8=AE=80=E5=85=A5=20profile=20=E7=9A=84=20cookie?=
 =?UTF-8?q?=EF=BC=8C=E4=BB=BF=20yt-dlp=20=E6=A0=BC=E5=BC=8F=E8=A3=BD?=
 =?UTF-8?q?=E4=BD=9C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cookie.go | 34 ++++++++++++++++++++++------------
 main.go   | 21 +++++++++++++++++++--
 2 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/cookie.go b/cookie.go
index a67607f..6a57472 100644
--- a/cookie.go
+++ b/cookie.go
@@ -37,8 +37,11 @@ var (
 	basePath   string
 )
 
-func getBrowserCookiePaths(browserName string) ([]string, error) {
+func getBrowserCookiePaths(browserName, profileName string) ([]string, error) {
 	browserName = strings.ToLower(browserName)
+	if profileName == "" {
+		profileName = "Default"
+	}
 
 	var cookieFile string
 
@@ -90,7 +93,7 @@ func getBrowserCookiePaths(browserName string) ([]string, error) {
 	}
 
 	profileBasePath := filepath.Join(basePath, profileDir)
-	profiles, err := getProfiles(profileBasePath, browserName)
+	profiles, err := getProfiles(profileBasePath, browserName, profileName)
 	if err != nil {
 		return nil, err
 	}
@@ -116,7 +119,7 @@ func getBrowserCookiePaths(browserName string) ([]string, error) {
 	return cookiePaths, nil
 }
 
-func getProfiles(profileBasePath string, browserName string) ([]string, error) {
+func getProfiles(profileBasePath, browserName, profileName string) ([]string, error) {
 	file, err := os.Open(profileBasePath)
 	if err != nil {
 		return nil, err
@@ -135,13 +138,12 @@ func getProfiles(profileBasePath string, browserName string) ([]string, error) {
 				profiles = append(profiles, entry.Name())
 			} else if browserName != "firefox" {
 				// if entry.Name() == "Default" || strings.HasPrefix(entry.Name(), "Profile ") {
-				if entry.Name() == "Default" {
+				if strings.Contains(entry.Name(), profileName) {
 					profiles = append(profiles, entry.Name())
 				}
 			}
 		}
 	}
-
 	return profiles, nil
 }
 
@@ -304,18 +306,25 @@ func decryptChromiumValue(encryptedValue, key []byte) (string, error) {
 	return "", errors.New("unknown platform")
 }
 
-func (di *DownloadInfo) GetCookieFromBrowser(browser string) (*cookiejar.Jar, error) {
+func (di *DownloadInfo) GetCookieFromBrowser(browser, profile string) (*cookiejar.Jar, int, error) {
+	count := 0
+
 	jar, err := cookiejar.New(&cookiejar.Options{
 		PublicSuffixList: publicsuffix.List,
 	})
 	if err != nil {
-		return nil, err
+		return nil, count, err
 	}
 
-	cookiePaths, err := getBrowserCookiePaths(browser)
+	cookiePaths, err := getBrowserCookiePaths(browser, profile)
 	if err != nil {
 		LogError("Error getting browser cookie paths: %s", err.Error())
-		return nil, err
+		return nil, count, err
+	}
+
+	if len(cookiePaths) == 0 {
+		LogError("No cookie paths found for %s", browser)
+		return nil, count, errors.New("no cookie paths found")
 	}
 
 	var chromiumKey []byte
@@ -323,10 +332,9 @@ func (di *DownloadInfo) GetCookieFromBrowser(browser string) (*cookiejar.Jar, er
 		chromiumKey, err = getChromiumKey(browser)
 		if err != nil {
 			LogError("Error getting Chromium key: %v\n", err)
-			return nil, err
+			return nil, count, err
 		}
 	}
-
 	cookieMap := make(map[string][]*http.Cookie)
 	for _, cookiePath := range cookiePaths {
 		cookies, err := getCookies(cookiePath, browser, chromiumKey)
@@ -340,6 +348,7 @@ func (di *DownloadInfo) GetCookieFromBrowser(browser string) (*cookiejar.Jar, er
 				cookieMap[cookie.Domain] = make([]*http.Cookie, 0)
 			}
 			cookieMap[cookie.Domain] = append(cookieMap[cookie.Domain], cookie)
+			count++
 		}
 	}
 
@@ -354,5 +363,6 @@ func (di *DownloadInfo) GetCookieFromBrowser(browser string) (*cookiejar.Jar, er
 			}
 		}
 	}
-	return jar, nil
+
+	return jar, count, nil
 }
diff --git a/main.go b/main.go
index 10328f8..e31a645 100644
--- a/main.go
+++ b/main.go
@@ -80,6 +80,15 @@ Options:
 		the script to access members-only content if you are a member
 		for the given stream's user. Must be netscape cookie format.
 
+	--cookies-from-browser BROWSER[:PROFILE]
+		Attempt to automatically load cookies from your browser. Currently 
+		supported browsers are: firefox, chrome, brave, edge. PROFILE is
+		optional and only applies to chromium core browser. If not provided, 
+		the default	profile will be used. For example, to load cookies from
+		Firefox's default profile, use '--cookie-from-browser firefox', And 
+		wanna load cookies from Chrome's 'Profile 1', use 
+		'--cookies-from-browser chrome:"Profile 1"'.
+
 	--debug
 		Print a lot of extra information.
 
@@ -583,14 +592,22 @@ func run() int {
 	}
 
 	if client.Jar == nil && len(cookieFromBrowser) > 0 {
-		cjar, err := info.GetCookieFromBrowser(cookieFromBrowser)
+		arr := strings.Split(cookieFromBrowser, ":")
+		profile := ""
+		if len(arr) == 2 {
+			cookieFromBrowser = arr[0]
+			profile = arr[1]
+		} else {
+			cookieFromBrowser = arr[0]
+		}
+		cjar, cookieSize, err := info.GetCookieFromBrowser(cookieFromBrowser, profile)
 		if err != nil {
 			LogError("Failed to load cookies from browser: %s", err)
 			return 1
 		}
 
 		client.Jar = cjar
-		LogGeneral("Loaded [%s] cookies.", cookieFromBrowser)
+		LogGeneral("The [%s] browser has (%d) cookies loaded.", cookieFromBrowser, cookieSize)
 	}
 
 	if !info.GVideoDDL && !info.GetVideoInfo() {