From 94cbb251ac31f56317a1f7849a52be07a91481f5 Mon Sep 17 00:00:00 2001 From: zhangwm404 <442798+zhangwm404@users.noreply.github.com> Date: Thu, 8 Dec 2022 11:46:04 +0800 Subject: [PATCH] support socks5 proxy --- env.example | 2 ++ go.mod | 1 + go.sum | 2 ++ main.go | 30 +++++++++++++++++++++++++++--- src/session/session.go | 13 ++++++++++++- 5 files changed, 44 insertions(+), 4 deletions(-) diff --git a/env.example b/env.example index 47324e9..7b47211 100644 --- a/env.example +++ b/env.example @@ -1,2 +1,4 @@ TELEGRAM_ID= TELEGRAM_TOKEN= +# Set it if you need surf the Internet through a proxy; DO NOT prefix with `socks5://` or `socks5h://` +# SOCKS5_PROXY=ip:port diff --git a/go.mod b/go.mod index 586700e..987812f 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/launchdarkly/eventsource v1.7.1 github.com/playwright-community/playwright-go v0.2000.1 github.com/spf13/viper v1.14.0 + golang.org/x/net v0.0.0-20221014081412-f15817d10f9b ) require ( diff --git a/go.sum b/go.sum index f45b846..137539d 100644 --- a/go.sum +++ b/go.sum @@ -267,6 +267,8 @@ golang.org/x/net v0.0.0-20201031054903-ff519b6c9102/go.mod h1:sp8m0HH+o8qH0wwXwY golang.org/x/net v0.0.0-20201209123823-ac852fbbde11/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20201224014010-6772e930b67b/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= +golang.org/x/net v0.0.0-20221014081412-f15817d10f9b h1:tvrvnPFcdzp294diPnrdZZZ8XUt2Tyj7svb7X52iDuU= +golang.org/x/net v0.0.0-20221014081412-f15817d10f9b/go.mod h1:YDH+HFinaLZZlnHAfSS6ZXJJ9M9t4Dl22yv3iI2vPwk= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= diff --git a/main.go b/main.go index 2431a64..722d098 100644 --- a/main.go +++ b/main.go @@ -3,9 +3,11 @@ package main import ( "fmt" "log" + "net/http" "os" "os/signal" "strconv" + "strings" "syscall" "time" @@ -16,6 +18,7 @@ import ( "github.com/m1guelpf/chatgpt-telegram/src/markdown" "github.com/m1guelpf/chatgpt-telegram/src/ratelimit" "github.com/m1guelpf/chatgpt-telegram/src/session" + "golang.org/x/net/proxy" ) type Conversation struct { @@ -49,9 +52,30 @@ func main() { log.Fatalf("Couldn't load .env file: %v", err) } - bot, err := tgbotapi.NewBotAPI(os.Getenv("TELEGRAM_TOKEN")) - if err != nil { - log.Fatalf("Couldn't start Telegram bot: %v", err) + var bot *tgbotapi.BotAPI + socks5Proxy := strings.TrimSpace(os.Getenv("SOCKS5_PROXY")) + if socks5Proxy != "" { + socks5, err := proxy.SOCKS5("tcp", socks5Proxy, nil, proxy.Direct) + if err != nil { + log.Fatalf("Couldn't set socks5 proxy: %v", err) + } + client := &http.Client{ + Transport: &http.Transport{Dial: socks5.Dial}, + } + bot, err = tgbotapi.NewBotAPIWithClient( + os.Getenv("TELEGRAM_TOKEN"), + tgbotapi.APIEndpoint, + client, + ) + if err != nil { + log.Fatalf("Couldn't start Telegram bot with socks5 proxy: %v", err) + } + } else { + var err error + bot, err = tgbotapi.NewBotAPI(os.Getenv("TELEGRAM_TOKEN")) + if err != nil { + log.Fatalf("Couldn't start Telegram bot: %v", err) + } } c := make(chan os.Signal, 2) diff --git a/src/session/session.go b/src/session/session.go index 505a6c8..a34522c 100644 --- a/src/session/session.go +++ b/src/session/session.go @@ -4,6 +4,8 @@ import ( "errors" "fmt" "log" + "os" + "strings" "sync" "time" @@ -73,7 +75,16 @@ func GetSession() (string, error) { } func launchBrowser(pw *playwright.Playwright, url string, headless bool) (playwright.BrowserContext, playwright.Page, error) { - browser, err := pw.Chromium.LaunchPersistentContext("/tmp/chatgpt", playwright.BrowserTypeLaunchPersistentContextOptions{Headless: playwright.Bool(headless)}) + opts := playwright.BrowserTypeLaunchPersistentContextOptions{ + Headless: playwright.Bool(headless), + } + socks5Proxy := strings.TrimSpace(os.Getenv("SOCKS5_PROXY")) + if socks5Proxy != "" { + opts.Proxy = &playwright.BrowserTypeLaunchPersistentContextOptionsProxy{ + Server: playwright.String("socks5://" + socks5Proxy), + } + } + browser, err := pw.Chromium.LaunchPersistentContext("/tmp/chatgpt", opts) if err != nil { return nil, nil, errors.New(fmt.Sprintf("Couldn't launch headless browser: %v", err)) }