Skip to content

Commit 9d23a22

Browse files
authored
feat: Allow multiple telegram IDs (#52)
* ChatGPT as conversation manager * fix: sending of error from SendMessage. Separate bot into its own package * Resolve conflicts - add edit interval to tgbot.New * Add env config, load with viper * Update config.Init - use viper instance rather than global * Allow multiple ids for TELEGRAM_ID env variable * Add tests for multiple ids being provided * Test action * Remove commented tests * Update LoadEnvConfig to work when no file is provided
1 parent 072d00d commit 9d23a22

File tree

9 files changed

+271
-49
lines changed

9 files changed

+271
-49
lines changed

.github/workflows/test.yaml

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
name: Test
2+
on: [push, pull_request]
3+
jobs:
4+
go-test:
5+
runs-on: ubuntu-latest
6+
steps:
7+
- name: Check out source code
8+
uses: actions/checkout@v3
9+
- name: Setup
10+
uses: actions/setup-go@v3
11+
with:
12+
go-version-file: "go.mod"
13+
cache: true
14+
- name: Test
15+
run: go test -v ./...

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ After you download the file, extract it into a folder and open the `env.example`
1818
- `TELEGRAM_ID` (Optional): Your Telegram User ID
1919
- If you set this, only you will be able to interact with the bot.
2020
- To get your ID, message `@userinfobot` on Telegram.
21+
- Multiple IDs can be provided, separated by commas.
2122
- `EDIT_WAIT_SECONDS` (Optional): Amount of seconds to wait between edits
2223
- This is set to `1` by default, but you can increase if you start getting a lot of `Too Many Requests` errors.
2324
- Save the file, and rename it to `.env`.

go.mod

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,23 @@ go 1.19
55
require (
66
github.com/go-telegram-bot-api/telegram-bot-api/v5 v5.5.1
77
github.com/google/uuid v1.3.0
8-
github.com/joho/godotenv v1.4.0
98
github.com/launchdarkly/eventsource v1.7.1
109
github.com/playwright-community/playwright-go v0.2000.1
1110
github.com/spf13/viper v1.14.0
11+
github.com/stretchr/testify v1.8.1
1212
)
1313

1414
require (
1515
github.com/danwakefield/fnmatch v0.0.0-20160403171240-cbb64ac3d964 // indirect
16+
github.com/davecgh/go-spew v1.1.1 // indirect
1617
github.com/fsnotify/fsnotify v1.6.0 // indirect
1718
github.com/go-stack/stack v1.8.1 // indirect
1819
github.com/hashicorp/hcl v1.0.0 // indirect
1920
github.com/magiconair/properties v1.8.6 // indirect
2021
github.com/mitchellh/mapstructure v1.5.0 // indirect
2122
github.com/pelletier/go-toml v1.9.5 // indirect
2223
github.com/pelletier/go-toml/v2 v2.0.5 // indirect
24+
github.com/pmezard/go-difflib v1.0.0 // indirect
2325
github.com/spf13/afero v1.9.2 // indirect
2426
github.com/spf13/cast v1.5.0 // indirect
2527
github.com/spf13/jwalterweatherman v1.1.0 // indirect

go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -132,8 +132,6 @@ github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
132132
github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
133133
github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
134134
github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
135-
github.com/joho/godotenv v1.4.0 h1:3l4+N6zfMWnkbPEXKng2o2/MR5mSwTrBih4ZEkkz1lg=
136-
github.com/joho/godotenv v1.4.0/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
137135
github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
138136
github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk=
139137
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
@@ -176,6 +174,7 @@ github.com/spf13/viper v1.14.0 h1:Rg7d3Lo706X9tHsJMUjdiwMpHB7W8WnSVOssIY+JElU=
176174
github.com/spf13/viper v1.14.0/go.mod h1:WT//axPky3FdvXHzGw33dNdXXXfFQqmEalje+egj8As=
177175
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
178176
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
177+
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
179178
github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
180179
github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
181180
github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
@@ -184,6 +183,7 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
184183
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
185184
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
186185
github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
186+
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
187187
github.com/subosito/gotenv v1.4.1 h1:jyEFiXpy21Wm81FBN71l9VoMMV8H8jG+qIK3GCpY6Qs=
188188
github.com/subosito/gotenv v1.4.1/go.mod h1:ayKnFf/c6rvx/2iiLrJUk1e6plDbT3edrFNGqEflhK0=
189189
github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

main.go

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -5,54 +5,44 @@ import (
55
"log"
66
"os"
77
"os/signal"
8-
"strconv"
98
"syscall"
109
"time"
1110

12-
"github.com/joho/godotenv"
1311
"github.com/m1guelpf/chatgpt-telegram/src/chatgpt"
1412
"github.com/m1guelpf/chatgpt-telegram/src/config"
1513
"github.com/m1guelpf/chatgpt-telegram/src/session"
1614
"github.com/m1guelpf/chatgpt-telegram/src/tgbot"
1715
)
1816

1917
func main() {
20-
config, err := config.Init()
18+
persistentConfig, err := config.LoadOrCreatePersistentConfig()
2119
if err != nil {
2220
log.Fatalf("Couldn't load config: %v", err)
2321
}
2422

25-
if config.OpenAISession == "" {
26-
session, err := session.GetSession()
23+
if persistentConfig.OpenAISession == "" {
24+
token, err := session.GetSession()
2725
if err != nil {
2826
log.Fatalf("Couldn't get OpenAI session: %v", err)
2927
}
3028

31-
err = config.Set("OpenAISession", session)
32-
if err != nil {
29+
if err = persistentConfig.SetSessionToken(token); err != nil {
3330
log.Fatalf("Couldn't save OpenAI session: %v", err)
3431
}
3532
}
3633

37-
chatGPT := chatgpt.Init(config)
34+
chatGPT := chatgpt.Init(persistentConfig)
3835
log.Println("Started ChatGPT")
3936

40-
err = godotenv.Load()
37+
envConfig, err := config.LoadEnvConfig(".env")
4138
if err != nil {
42-
log.Printf("Couldn't load .env file: %v. Using shell exposed env variables...", err)
39+
log.Fatalf("Couldn't load .env config: %v", err)
4340
}
44-
45-
editInterval := 1 * time.Second
46-
if os.Getenv("EDIT_WAIT_SECONDS") != "" {
47-
editSecond, err := strconv.ParseInt(os.Getenv("EDIT_WAIT_SECONDS"), 10, 64)
48-
if err != nil {
49-
log.Printf("Couldn't convert your edit seconds setting into int: %v", err)
50-
editSecond = 1
51-
}
52-
editInterval = time.Duration(editSecond) * time.Second
41+
if err := envConfig.ValidateWithDefaults(); err != nil {
42+
log.Fatalf("Invalid .env config: %v", err)
5343
}
5444

55-
bot, err := tgbot.New(os.Getenv("TELEGRAM_TOKEN"), editInterval)
45+
bot, err := tgbot.New(envConfig.TelegramToken, time.Duration(envConfig.EditWaitSeconds))
5646
if err != nil {
5747
log.Fatalf("Couldn't start Telegram bot: %v", err)
5848
}
@@ -76,10 +66,11 @@ func main() {
7666
updateText = update.Message.Text
7767
updateChatID = update.Message.Chat.ID
7868
updateMessageID = update.Message.MessageID
69+
updateUserID = update.Message.From.ID
7970
)
8071

81-
userId := strconv.FormatInt(update.Message.Chat.ID, 10)
82-
if os.Getenv("TELEGRAM_ID") != "" && userId != os.Getenv("TELEGRAM_ID") {
72+
if len(envConfig.TelegramID) != 0 && !envConfig.HasTelegramID(updateUserID) {
73+
log.Printf("User %d is not allowed to use this bot", updateUserID)
8374
bot.Send(updateChatID, updateMessageID, "You are not authorized to use this bot.")
8475
continue
8576
}

src/chatgpt/chatgpt.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ type ChatResponse struct {
4848
Message string
4949
}
5050

51-
func Init(config config.Config) *ChatGPT {
51+
func Init(config *config.Config) *ChatGPT {
5252
return &ChatGPT{
5353
AccessTokenMap: expirymap.New(),
5454
SessionToken: config.OpenAISession,

src/config/config.go

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -9,46 +9,46 @@ import (
99
)
1010

1111
type Config struct {
12+
v *viper.Viper
13+
1214
OpenAISession string
1315
}
1416

15-
// init tries to read the config from the file, and creates it if it doesn't exist.
16-
func Init() (Config, error) {
17+
// LoadOrCreatePersistentConfig uses the default config directory for the current OS
18+
// to load or create a config file named "chatgpt.json"
19+
func LoadOrCreatePersistentConfig() (*Config, error) {
1720
configPath, err := os.UserConfigDir()
1821
if err != nil {
19-
return Config{}, errors.New(fmt.Sprintf("Couldn't get user config dir: %v", err))
22+
return nil, errors.New(fmt.Sprintf("Couldn't get user config dir: %v", err))
2023
}
21-
viper.SetConfigType("json")
22-
viper.SetConfigName("chatgpt")
23-
viper.AddConfigPath(configPath)
24+
v := viper.New()
25+
v.SetConfigType("json")
26+
v.SetConfigName("chatgpt")
27+
v.AddConfigPath(configPath)
2428

25-
if err := viper.ReadInConfig(); err != nil {
29+
if err := v.ReadInConfig(); err != nil {
2630
if _, ok := err.(viper.ConfigFileNotFoundError); ok {
27-
if err := viper.SafeWriteConfig(); err != nil {
28-
return Config{}, errors.New(fmt.Sprintf("Couldn't create config file: %v", err))
31+
if err := v.SafeWriteConfig(); err != nil {
32+
return nil, errors.New(fmt.Sprintf("Couldn't create config file: %v", err))
2933
}
3034
} else {
31-
return Config{}, errors.New(fmt.Sprintf("Couldn't read config file: %v", err))
35+
return nil, errors.New(fmt.Sprintf("Couldn't read config file: %v", err))
3236
}
3337
}
3438

3539
var cfg Config
36-
err = viper.Unmarshal(&cfg)
40+
err = v.Unmarshal(&cfg)
3741
if err != nil {
38-
return Config{}, errors.New(fmt.Sprintf("Error parsing config: %v", err))
42+
return nil, errors.New(fmt.Sprintf("Error parsing config: %v", err))
3943
}
44+
cfg.v = v
4045

41-
return cfg, nil
46+
return &cfg, nil
4247
}
4348

44-
// key should be part of the Config struct
45-
func (cfg *Config) Set(key string, value interface{}) error {
46-
viper.Set(key, value)
47-
48-
err := viper.Unmarshal(&cfg)
49-
if err != nil {
50-
return errors.New(fmt.Sprintf("Error parsing config: %v", err))
51-
}
52-
53-
return viper.WriteConfig()
49+
func (cfg *Config) SetSessionToken(token string) error {
50+
// key must match the struct field name
51+
cfg.v.Set("OpenAISession", token)
52+
cfg.OpenAISession = token
53+
return cfg.v.WriteConfig()
5454
}

src/config/config_test.go

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
package config
2+
3+
import (
4+
"fmt"
5+
"os"
6+
"testing"
7+
8+
"github.com/stretchr/testify/require"
9+
)
10+
11+
func createFile(name string, content string) (remove func(), err error) {
12+
f, err := os.Create(name)
13+
if err != nil {
14+
return nil, err
15+
}
16+
defer f.Close()
17+
18+
if _, err := f.WriteString(content); err != nil {
19+
return nil, err
20+
}
21+
22+
return func() {
23+
if err := os.Remove(name); err != nil {
24+
panic(fmt.Sprintf("failed to remove file: %s", err))
25+
}
26+
}, nil
27+
}
28+
29+
func setEnvVariables(vals map[string]string) func() {
30+
for k, v := range vals {
31+
os.Setenv(k, v)
32+
}
33+
return func() {
34+
for k := range vals {
35+
os.Unsetenv(k)
36+
}
37+
}
38+
}
39+
40+
func TestLoadEnvConfig(t *testing.T) {
41+
for label, test := range map[string]struct {
42+
fileContent string
43+
envVars map[string]string
44+
want *EnvConfig
45+
}{
46+
"all values empty in file and env": {
47+
fileContent: `TELEGRAM_ID=
48+
TELEGRAM_TOKEN=
49+
EDIT_WAIT_SECONDS=`,
50+
want: &EnvConfig{
51+
TelegramID: []int64{},
52+
TelegramToken: "",
53+
EditWaitSeconds: 0,
54+
},
55+
},
56+
"no file, all values through env": {
57+
envVars: map[string]string{
58+
"TELEGRAM_ID": "123,456",
59+
"TELEGRAM_TOKEN": "token",
60+
"EDIT_WAIT_SECONDS": "10",
61+
},
62+
want: &EnvConfig{
63+
TelegramID: []int64{123, 456},
64+
TelegramToken: "token",
65+
EditWaitSeconds: 10,
66+
},
67+
},
68+
"all values provided in file, single TELEGRAM_ID": {
69+
fileContent: `TELEGRAM_ID=123
70+
TELEGRAM_TOKEN=abc
71+
EDIT_WAIT_SECONDS=10`,
72+
want: &EnvConfig{
73+
TelegramID: []int64{123},
74+
TelegramToken: "abc",
75+
EditWaitSeconds: 10,
76+
},
77+
},
78+
"multiple TELEGRAM_IDs provided in file": {
79+
fileContent: `TELEGRAM_ID=123,456
80+
TELEGRAM_TOKEN=abc
81+
EDIT_WAIT_SECONDS=10`,
82+
envVars: map[string]string{},
83+
want: &EnvConfig{
84+
TelegramID: []int64{123, 456},
85+
TelegramToken: "abc",
86+
EditWaitSeconds: 10,
87+
},
88+
},
89+
"env variables should override file values": {
90+
fileContent: `TELEGRAM_ID=123
91+
TELEGRAM_TOKEN=abc
92+
EDIT_WAIT_SECONDS=10`,
93+
envVars: map[string]string{
94+
"TELEGRAM_ID": "456",
95+
"TELEGRAM_TOKEN": "def",
96+
"EDIT_WAIT_SECONDS": "20",
97+
},
98+
want: &EnvConfig{
99+
TelegramID: []int64{456},
100+
TelegramToken: "def",
101+
EditWaitSeconds: 20,
102+
},
103+
},
104+
"multiple TELEGRAM_IDs provided in env": {
105+
fileContent: `TELEGRAM_ID=123
106+
TELEGRAM_TOKEN=abc
107+
EDIT_WAIT_SECONDS=10`,
108+
envVars: map[string]string{
109+
"TELEGRAM_ID": "456,789",
110+
},
111+
want: &EnvConfig{
112+
TelegramID: []int64{456, 789},
113+
TelegramToken: "abc",
114+
EditWaitSeconds: 10,
115+
},
116+
},
117+
} {
118+
t.Run(label, func(t *testing.T) {
119+
unset := setEnvVariables(test.envVars)
120+
t.Cleanup(unset)
121+
122+
if test.fileContent != "" {
123+
remove, err := createFile("test.env", test.fileContent)
124+
require.NoError(t, err)
125+
t.Cleanup(remove)
126+
}
127+
128+
cfg, err := LoadEnvConfig("test.env")
129+
require.NoError(t, err)
130+
require.Equal(t, test.want, cfg)
131+
})
132+
}
133+
}

0 commit comments

Comments
 (0)