feat: allow setting custom referer

michalczmiel · Jan 23, 2024 · b0812da · b0812da
1 parent be19b2c
commit b0812da
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 5 deletions.
diff --git a/cmd/html.go b/cmd/html.go
@@ -51,7 +51,7 @@ func runHtmlCmd(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	doc, err := internal.GetHtmlDocFromUrl(url, parameters.UserAgent)
+	doc, err := internal.GetHtmlDocFromUrl(url, parameters)
 	if err != nil {
 		return err
 	}

diff --git a/cmd/root.go b/cmd/root.go
@@ -34,11 +34,17 @@ func getRootParameters(cmd *cobra.Command) (*internal.Parameters, error) {
 		return nil, err
 	}
 
+	referer, err := cmd.Flags().GetString("referer")
+	if err != nil {
+		return nil, err
+	}
+
 	parameters := &internal.Parameters{
 		ImageTypes: imageTypesToDownload,
 		Directory:  directory,
 		Concurrent: concurrentWorkersCount,
 		UserAgent:  userAgent,
+		Referer:    referer,
 	}
 
 	return parameters, nil
@@ -62,6 +68,7 @@ func addRootFlags(cmd *cobra.Command) {
 	cmd.Flags().IntP("concurrency", "c", 10, "number of concurrent downloads")
 	cmd.Flags().StringP("dir", "d", internal.DefaultPath, "directory to save images to")
 	cmd.Flags().String("user-agent", "", "custom user agent to use for requests")
+	cmd.Flags().String("referer", "", "custom referer to use for requests")
 }
 
 func init() {

diff --git a/internal/http.go b/internal/http.go
@@ -61,9 +61,14 @@ func request(url, userAgent, referer string) (*http.Response, error) {
 }
 
 func DownloadImageFromUrl(url, filePath string, parameters *Parameters) error {
-	rootUrl := getRootUrl(url)
+	var referer string
+	if parameters.Referer == "" {
+		referer = getRootUrl(url)
+	} else {
+		referer = parameters.Referer
+	}
 
-	response, err := request(url, parameters.UserAgent, rootUrl)
+	response, err := request(url, parameters.UserAgent, referer)
 	if err != nil {
 		return err
 	}
@@ -91,8 +96,15 @@ func DownloadImageFromUrl(url, filePath string, parameters *Parameters) error {
 	return nil
 }
 
-func GetHtmlDocFromUrl(url string, userAgent string) (*html.Node, error) {
-	response, err := request(url, userAgent, DefaultReferer)
+func GetHtmlDocFromUrl(url string, parameters *Parameters) (*html.Node, error) {
+	var referer string
+	if parameters.Referer == "" {
+		referer = DefaultReferer
+	} else {
+		referer = parameters.Referer
+	}
+
+	response, err := request(url, parameters.UserAgent, referer)
 	if err != nil {
 		return nil, err
 	}

diff --git a/internal/images.go b/internal/images.go
@@ -10,6 +10,7 @@ type Parameters struct {
 	Directory  string
 	Concurrent int
 	UserAgent  string
+	Referer    string
 }
 
 func downloadWorker(wg *sync.WaitGroup, parameters *Parameters, linksToProcess <-chan string, results chan<- DownloadResult) {