feat: support user-agent customize

bot/bot.go

@@ -3,10 +3,9 @@ package bot
 import (
 	"github.com/indes/flowerss-bot/bot/fsm"
 	"github.com/indes/flowerss-bot/config"
-	"golang.org/x/net/proxy"
+	"github.com/indes/flowerss-bot/util"
 	tb "gopkg.in/tucnak/telebot.v2"
 	"log"
-	"net/http"
 	"time"
 )
 
@@ -23,34 +22,17 @@ func init() {
 		}
 		return true
 	})
-
-	botSettings := tb.Settings{
-		URL:    config.TelegramEndpoint,
-		Token:  config.BotToken,
-		Poller: spamProtected,
-	}
-
-	if config.Socks5 != "" {
-		log.Printf("Bot Token: %s Proxy: %s Endpoint: %s\n", config.BotToken, config.Socks5, config.TelegramEndpoint)
-
-		dialer, err := proxy.SOCKS5("tcp", config.Socks5, nil, proxy.Direct)
-		if err != nil {
-			log.Fatal("Error creating dialer, aborting.")
-		}
-
-		httpTransport := &http.Transport{}
-		httpClient := &http.Client{Transport: httpTransport}
-		httpTransport.Dial = dialer.Dial
-		botSettings.Client = httpClient
-
-	} else {
-		log.Printf("Bot Token: %s Endpoint: %s\n", config.BotToken, config.TelegramEndpoint)
-	}
+	log.Printf("Bot Token: %s Endpoint: %s\n", config.BotToken, config.TelegramEndpoint)
 
 	// create bot
 	var err error
 
-	B, err = tb.NewBot(botSettings)
+	B, err = tb.NewBot(tb.Settings{
+		URL:    config.TelegramEndpoint,
+		Token:  config.BotToken,
+		Poller: spamProtected,
+		Client: util.HttpClient,
+	})
 
 	if err != nil {
 		log.Fatal(err)

config.yml.sample

@@ -2,6 +2,7 @@ bot_token:
 telegraph_token:
 socks5:
 update_interval: 10
+user_agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/51.0.2704.103 Safari/537.36
 
 mysql:
   host:

config/autoload.go

@@ -33,6 +33,7 @@ var (
 	MessageTpl            *template.Template
 	MessageMode           tb.ParseMode
 	TelegramEndpoint      string
+	UserAgent             string
 )
 
 const (
@@ -125,6 +126,7 @@ func init() {
 
 	BotToken = viper.GetString("bot_token")
 	Socks5 = viper.GetString("socks5")
+	UserAgent = viper.GetString("user_agent")
 
 	if viper.IsSet("telegraph_token") {
 		EnableTelegraph = true

model/source.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/SlyMarbo/rss"
 	"github.com/indes/flowerss-bot/config"
+	"github.com/indes/flowerss-bot/util"
 	"github.com/jinzhu/gorm"
 	"io/ioutil"
 	"log"
@@ -46,16 +47,30 @@ func GetSourceByUrl(url string) (*Source, error) {
 }
 
 func fetchFunc(url string) (resp *http.Response, err error) {
-	resp, err = http.Get(url)
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	if config.UserAgent != "" {
+		req.Header.Set("User-Agent", config.UserAgent)
+	} else {
+		req.Header.Set("User-Agent", "flowerss/2.0")
+	}
+
+	resp, err = util.HttpClient.Do(req)
+
 	if err != nil {
 		return nil, err
 	}
+	defer resp.Body.Close()
+
 	var data []byte
 	if data, err = ioutil.ReadAll(resp.Body); err != nil {
-		_ = resp.Body.Close()
+
 		return nil, err
 	}
-	_ = resp.Body.Close()
+
 	resp.Body = ioutil.NopCloser(strings.NewReader(strings.Map(func(r rune) rune {
 		if unicode.IsPrint(r) {
 			return r

util/client.go

@@ -0,0 +1,27 @@
+package util
+
+import (
+	"github.com/indes/flowerss-bot/config"
+	"golang.org/x/net/proxy"
+	"log"
+	"net/http"
+)
+
+var (
+	HttpClient *http.Client
+)
+
+func clientInit() {
+	httpTransport := &http.Transport{}
+	HttpClient = &http.Client{Transport: httpTransport}
+	// set proxy
+	if config.Socks5 != "" {
+		log.Printf("Proxy: %s\n", config.Socks5)
+
+		dialer, err := proxy.SOCKS5("tcp", config.Socks5, nil, proxy.Direct)
+		if err != nil {
+			log.Fatal("Error creating dialer, aborting.")
+		}
+		httpTransport.Dial = dialer.Dial
+	}
+}

util/util.go

@@ -0,0 +1,5 @@
+package util
+
+func init() {
+	clientInit()
+}