From b52a0c34b854669d60de782cfd08ed44fc21011a Mon Sep 17 00:00:00 2001
From: Vladimir Hodakov <vladimir@hodakov.me>
Date: Fri, 11 May 2018 18:55:37 +0400
Subject: [PATCH] Add SOCKS5 proxy support to bot

In Soviet Russia... wait, wrong time period.
---
 config.yml.example             |  5 +++
 lib/config/config.go           |  9 ++++++
 lib/connections/connections.go | 57 ++++++++++++++++++++++++++++++++--
 3 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/config.yml.example b/config.yml.example
index 274be16..60c312e 100644
--- a/config.yml.example
+++ b/config.yml.example
@@ -8,6 +8,11 @@ database_connection:
     user: "i2_bot"
     password: "i2_bot"
     database: "i2_bot"
+socks_proxy:
+    enabled: true
+    address: "127.0.0.1:1337"
+    username: "proxy"
+    password: "proxy"
 special_chats:
     academy_id: "group_id"
     bastion_id: "group_id"
diff --git a/lib/config/config.go b/lib/config/config.go
index 6c21653..44ba101 100644
--- a/lib/config/config.go
+++ b/lib/config/config.go
@@ -30,6 +30,14 @@ type TelegramConnection struct {
 	ListenAddress string `yaml:"listen_address"`
 }
 
+// ProxySettings handles settings for SOCKS5 proxy in config.yml
+type ProxySettings struct {
+	Enabled  bool   `yaml:"enabled"`
+	Address  string `yaml:"address,omitempty"`
+	Username string `yaml:"username,omitempty"`
+	Password string `yaml:"password,omitempty"`
+}
+
 // SpecialChats handles settings for special chats
 type SpecialChats struct {
 	AcademyID      string `yaml:"academy_id"`
@@ -48,6 +56,7 @@ type LoggingConfig struct {
 type Config struct {
 	Telegram     TelegramConnection `yaml:"telegram_connection"`
 	Database     DatabaseConnection `yaml:"database_connection"`
+	Proxy        ProxySettings      `yaml:"socks_proxy"`
 	SpecialChats SpecialChats       `yaml:"special_chats"`
 	Logs         LoggingConfig      `yaml:"logs"`
 }
diff --git a/lib/connections/connections.go b/lib/connections/connections.go
index e6352ca..430282c 100644
--- a/lib/connections/connections.go
+++ b/lib/connections/connections.go
@@ -8,11 +8,13 @@ import (
 	_ "github.com/go-sql-driver/mysql" // MySQL driver for sqlx
 	"github.com/go-telegram-bot-api/telegram-bot-api"
 	"github.com/jmoiron/sqlx"
+	"golang.org/x/net/proxy"
+	"net/http"
 	"source.wtfteam.pro/i2_bot/i2_bot/lib/config"
 )
 
-// BotInit initializes connection to Telegram
-func BotInit(cfg *config.Config, lg *mogrus.LoggerHandler) *tgbotapi.BotAPI {
+// botInitDirect used when no proxy in config file
+func botInitDirect(cfg *config.Config, lg *mogrus.LoggerHandler) *tgbotapi.BotAPI {
 	bot, err := tgbotapi.NewBotAPI(cfg.Telegram.APIToken)
 	if err != nil {
 		lg.Fatal(err.Error())
@@ -26,6 +28,57 @@ func BotInit(cfg *config.Config, lg *mogrus.LoggerHandler) *tgbotapi.BotAPI {
 	return bot
 }
 
+// botInitWithProxy used when there is proxy in config file
+func botInitWithProxy(cfg *config.Config, lg *mogrus.LoggerHandler) *tgbotapi.BotAPI {
+	proxyAuth := proxy.Auth{}
+	if cfg.Proxy.Username != "" {
+		proxyAuth.User = cfg.Proxy.Username
+		proxyAuth.Password = cfg.Proxy.Password
+	}
+
+	var dialProxy proxy.Dialer
+	var err error
+	if cfg.Proxy.Username != "" {
+		dialProxy, err = proxy.SOCKS5("tcp", cfg.Proxy.Address, &proxyAuth, proxy.Direct)
+		if err != nil {
+			lg.Fatal(err.Error())
+		}
+	} else {
+		dialProxy, err = proxy.SOCKS5("tcp", cfg.Proxy.Address, &proxyAuth, proxy.Direct)
+		if err != nil {
+			lg.Fatal(err.Error())
+		}
+	}
+
+	proxyTransport := &http.Transport{Dial: dialProxy.Dial}
+	proxyClient := http.Client{Transport: proxyTransport}
+
+	bot, err := tgbotapi.NewBotAPIWithClient(cfg.Telegram.APIToken, &proxyClient)
+	if err != nil {
+		lg.Fatal(err.Error())
+	}
+
+	bot.Debug = true
+
+	lg.Info("Bot version: " + config.VERSION)
+	lg.Info("Authorized on account @", bot.Self.UserName)
+
+	return bot
+}
+
+// External functions
+
+// BotInit initializes connection to Telegram
+func BotInit(cfg *config.Config, lg *mogrus.LoggerHandler) *tgbotapi.BotAPI {
+	if cfg.Proxy.Enabled {
+		lg.Info("Using proxy for bot: " + cfg.Proxy.Address)
+		return botInitWithProxy(cfg, lg)
+	}
+
+	lg.Info("Using direct connection to Telegram")
+	return botInitDirect(cfg, lg)
+}
+
 // DBInit initializes database connection
 func DBInit(cfg *config.Config, lg *mogrus.LoggerHandler) *sqlx.DB {
 	database, err := sqlx.Connect("mysql", cfg.Database.User+":"+cfg.Database.Password+"@tcp("+cfg.Database.Host+":"+cfg.Database.Port+")/"+cfg.Database.Database+"?parseTime=true&charset=utf8mb4,utf8")