From 173424c10eaa6fa3537ee820c873f3e409a9a786 Mon Sep 17 00:00:00 2001
From: Alex Ferrari <alex@thealexferrari.com>
Date: Wed, 17 Jan 2024 09:57:11 +0100
Subject: [PATCH] Implement passing cookie name (Fixes #4)

---
 README.md   | 30 ++++++++++++++++++++++--------
 cmd/root.go | 50 ++++++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 66 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index 66cb428..7d3d7a0 100644
--- a/README.md
+++ b/README.md
@@ -22,13 +22,14 @@ Available Commands:
   version     Print the version number of sbstck-dl
 
 Flags:
-      --after string    Download posts published after this date (format: YYYY-MM-DD)
-      --before string   Download posts published before this date (format: YYYY-MM-DD)
-  -h, --help            help for sbstck-dl
-  -x, --proxy string    Specify the proxy url
-  -i, --sid string      Specify the substack.sid required to download private newsletters (found in the Substack's cookie)
-  -r, --rate int        Specify the rate of requests per second (default 2)
-  -v, --verbose         Enable verbose output
+      --after string             Download posts published after this date (format: YYYY-MM-DD)
+      --before string            Download posts published before this date (format: YYYY-MM-DD)
+      --cookie_name cookieName   Either substack.sid or connect.sid, based on your cookie (required for private newsletters)
+      --cookie_val string        The substack.sid/connect.sid cookie value (required for private newsletters)
+  -h, --help                     help for sbstck-dl
+  -x, --proxy string             Specify the proxy url
+  -r, --rate int                 Specify the rate of requests per second (default 2)
+  -v, --verbose                  Enable verbose output
 
 Use "sbstck-dl [command] --help" for more information about a command.
 ```
@@ -74,9 +75,22 @@ Global Flags:
   -v, --verbose         Enable verbose output
 ```
 
+### Private Newsletters
+
+In order to download the full text of private newsletters you need to provide the cookie name and value of your session.
+The cookie name is either `substack.sid` or `connect.sid`, based on your cookie.
+To get the cookie value you can use the developer tools of your browser.
+Once you have the cookie name and value, you can pass them to the downloader using the `--cookie_name` and `--cookie_val` flags.
+
+#### Example
+
+```bash
+sbstck-dl download --url https://example.substack.com --cookie_name substack.sid --cookie_val COOKIE_VALUE
+```
+
 ## Thanks
 
-- [wemoveon2](https://github.com/wemoveon2) for the discussion and help implementing the support for private newsletters
+- [wemoveon2](https://github.com/wemoveon2) and [lenzj](https://github.com/lenzj) for the discussion and help implementing the support for private newsletters
 
 ## TODO
 
diff --git a/cmd/root.go b/cmd/root.go
index 29867b6..c2b8569 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"context"
+	"errors"
 	"log"
 	"net/http"
 	"net/url"
@@ -13,13 +14,39 @@ import (
 
 // rootCmd represents the base command when called without any subcommands
 
+type cookieName string
+
+const (
+	substackSid cookieName = "substack.sid"
+	connectSid  cookieName = "connect.sid"
+)
+
+func (c *cookieName) String() string {
+	return string(*c)
+}
+
+func (c *cookieName) Set(val string) error {
+	switch val {
+	case "substack.sid", "connect.sid":
+		*c = cookieName(val)
+	default:
+		return errors.New("invalid cookie name: must be either substack.sid or connect.sid")
+	}
+	return nil
+}
+
+func (c *cookieName) Type() string {
+	return "cookieName"
+}
+
 var (
 	proxyURL       string
 	verbose        bool
 	ratePerSecond  int
 	beforeDate     string
 	afterDate      string
-	substackID     string
+	IdCookieName   cookieName
+	IdCookieVal    string
 	ctx            = context.Background()
 	parsedProxyURL *url.URL
 	fetcher        *lib.Fetcher
@@ -46,10 +73,20 @@ func Execute() {
 	if ratePerSecond == 0 {
 		log.Fatal("rate must be greater than 0")
 	}
-	if substackID != "" {
-		cookie = &http.Cookie{
-			Name:  "substack.sid",
-			Value: substackID,
+	if IdCookieVal != "" && IdCookieName == "" {
+		log.Fatal("You must specify the cookie name when using a cookie value")
+	}
+	if IdCookieVal != "" && IdCookieName != "" {
+		if IdCookieName == substackSid {
+			cookie = &http.Cookie{
+				Name:  "substack.sid",
+				Value: IdCookieVal,
+			}
+		} else if IdCookieName == connectSid {
+			cookie = &http.Cookie{
+				Name:  "connect.sid",
+				Value: IdCookieVal,
+			}
 		}
 	}
 	fetcher = lib.NewFetcher(lib.WithRatePerSecond(ratePerSecond), lib.WithProxyURL(parsedProxyURL), lib.WithCookie(cookie))
@@ -62,7 +99,8 @@ func Execute() {
 
 func init() {
 	rootCmd.PersistentFlags().StringVarP(&proxyURL, "proxy", "x", "", "Specify the proxy url")
-	rootCmd.PersistentFlags().StringVarP(&substackID, "sid", "i", "", "The substack.sid cookie value (required for private newsletters)")
+	rootCmd.PersistentFlags().Var(&IdCookieName, "cookie_name", "Either \"substack.sid\" or \"connect.sid\", based on the cookie you have (required for private newsletters)")
+	rootCmd.PersistentFlags().StringVar(&IdCookieVal, "cookie_val", "", "The substack.sid/connect.sid cookie value (required for private newsletters)")
 	rootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "Enable verbose output")
 	rootCmd.PersistentFlags().IntVarP(&ratePerSecond, "rate", "r", lib.DefaultRatePerSecond, "Specify the rate of requests per second")
 	rootCmd.PersistentFlags().StringVar(&beforeDate, "before", "", "Download posts published before this date (format: YYYY-MM-DD)")