From 598e460e6cd14dee602300e778f96f6a558ddfb7 Mon Sep 17 00:00:00 2001
From: bytedream <bytedream@protonmail.com>
Date: Sat, 16 Apr 2022 01:07:56 +0200
Subject: [PATCH] Add custom useragent for cli request

---
 Makefile                        | 13 +++++-------
 cmd/crunchyroll-go/cmd/root.go  | 13 ++++++++++--
 cmd/crunchyroll-go/cmd/utils.go | 37 +++++++++++++++++++++++++++------
 3 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/Makefile b/Makefile
index 8a21b56..8f42f36 100644
--- a/Makefile
+++ b/Makefile
@@ -6,8 +6,7 @@ DESTDIR=
 PREFIX=/usr
 
 build:
-		cd cmd/crunchyroll-go && go build -o $(BINARY_NAME)
-		mv cmd/crunchyroll-go/$(BINARY_NAME) .
+		go build -ldflags "-X 'github.com/ByteDream/crunchyroll-go/cmd/crunchyroll-go/cmd.Version=$(VERSION)'" -o $(BINARY_NAME) cmd/crunchyroll-go/main.go
 
 clean:
 		rm -f $(BINARY_NAME) $(VERSION_BINARY_NAME)_*
@@ -25,10 +24,8 @@ uninstall:
 		rm -f $(DESTDIR)$(PREFIX)/share/licenses/crunchyroll-go/LICENSE
 
 release:
-		cd cmd/crunchyroll-go && CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o $(VERSION_BINARY_NAME)_linux
-		cd cmd/crunchyroll-go && CGO_ENABLED=0 GOOS=windows GOARCH=amd64 go build -o $(VERSION_BINARY_NAME)_windows.exe
-		cd cmd/crunchyroll-go && CGO_ENABLED=0 GOOS=darwin GOARCH=amd64 go build -o $(VERSION_BINARY_NAME)_darwin
+		CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags "-X 'github.com/ByteDream/crunchyroll-go/cmd/crunchyroll-go/cmd.Version=$(VERSION)'" -o $(VERSION_BINARY_NAME)_linux cmd/crunchyroll-go/main.go
+		CGO_ENABLED=0 GOOS=windows GOARCH=amd64 go build -ldflags "-X 'github.com/ByteDream/crunchyroll-go/cmd/crunchyroll-go/cmd.Version=$(VERSION)'" -o $(VERSION_BINARY_NAME)_windows.exe cmd/crunchyroll-go/main.go
+		CGO_ENABLED=0 GOOS=darwin GOARCH=amd64 go build -ldflags "-X 'github.com/ByteDream/crunchyroll-go/cmd/crunchyroll-go/cmd.Version=$(VERSION)'" -o $(VERSION_BINARY_NAME)_darwin cmd/crunchyroll-go/main.go
 
-		strip cmd/crunchyroll-go/$(VERSION_BINARY_NAME)_linux
-
-		mv cmd/crunchyroll-go/$(VERSION_BINARY_NAME)_* .
+		strip $(VERSION_BINARY_NAME)_linux
diff --git a/cmd/crunchyroll-go/cmd/root.go b/cmd/crunchyroll-go/cmd/root.go
index 25c8270..c06df04 100644
--- a/cmd/crunchyroll-go/cmd/root.go
+++ b/cmd/crunchyroll-go/cmd/root.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 	"github.com/ByteDream/crunchyroll-go"
 	"github.com/spf13/cobra"
 	"net/http"
@@ -10,6 +11,8 @@ import (
 	"strings"
 )
 
+var Version = "development"
+
 var (
 	client  *http.Client
 	crunchy *crunchyroll.Crunchyroll
@@ -17,7 +20,10 @@ var (
 
 	quietFlag   bool
 	verboseFlag bool
-	proxyFlag   string
+
+	proxyFlag string
+
+	useragentFlag string
 )
 
 var rootCmd = &cobra.Command{
@@ -36,7 +42,7 @@ var rootCmd = &cobra.Command{
 
 		out.DebugLog.Printf("Executing `%s` command with %d arg(s)\n", cmd.Name(), len(args))
 
-		client, err = createOrDefaultClient(proxyFlag)
+		client, err = createOrDefaultClient(proxyFlag, useragentFlag)
 		return
 	},
 }
@@ -44,7 +50,10 @@ var rootCmd = &cobra.Command{
 func init() {
 	rootCmd.PersistentFlags().BoolVarP(&quietFlag, "quiet", "q", false, "Disable all output")
 	rootCmd.PersistentFlags().BoolVarP(&verboseFlag, "verbose", "v", false, "Adds debug messages to the normal output")
+
 	rootCmd.PersistentFlags().StringVarP(&proxyFlag, "proxy", "p", "", "Proxy to use")
+
+	rootCmd.PersistentFlags().StringVar(&useragentFlag, "useragent", fmt.Sprintf("crunchyroll-go/%s", Version), "Useragent to do all request with")
 }
 
 func Execute() {
diff --git a/cmd/crunchyroll-go/cmd/utils.go b/cmd/crunchyroll-go/cmd/utils.go
index 42fe7de..7589010 100644
--- a/cmd/crunchyroll-go/cmd/utils.go
+++ b/cmd/crunchyroll-go/cmd/utils.go
@@ -71,7 +71,23 @@ func allLocalesAsStrings() (locales []string) {
 	return
 }
 
-func createOrDefaultClient(proxy string) (*http.Client, error) {
+type headerRoundTripper struct {
+	http.RoundTripper
+	header map[string]string
+}
+
+func (rht headerRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	resp, err := rht.RoundTripper.RoundTrip(r)
+	if err != nil {
+		return nil, err
+	}
+	for k, v := range rht.header {
+		resp.Header.Set(k, v)
+	}
+	return resp, nil
+}
+
+func createOrDefaultClient(proxy, useragent string) (*http.Client, error) {
 	if proxy == "" {
 		return http.DefaultClient, nil
 	} else {
@@ -80,12 +96,21 @@ func createOrDefaultClient(proxy string) (*http.Client, error) {
 		if err != nil {
 			return nil, err
 		}
+
+		var rt http.RoundTripper = &http.Transport{
+			DisableCompression: true,
+			Proxy:              http.ProxyURL(proxyURL),
+		}
+		if useragent != "" {
+			rt = headerRoundTripper{
+				RoundTripper: rt,
+				header:       map[string]string{"User-Agent": useragent},
+			}
+		}
+
 		client := &http.Client{
-			Transport: &http.Transport{
-				DisableCompression: true,
-				Proxy:              http.ProxyURL(proxyURL),
-			},
-			Timeout: 30 * time.Second,
+			Transport: rt,
+			Timeout:   30 * time.Second,
 		}
 		return client, nil
 	}