Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Allow multiple telegram IDs #52

Merged
merged 13 commits into from
Dec 11, 2022
Merged
15 changes: 15 additions & 0 deletions .github/workflows/test.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
name: Test
on: [push, pull_request]
jobs:
go-test:
runs-on: ubuntu-latest
steps:
- name: Check out source code
uses: actions/checkout@v3
- name: Setup
uses: actions/setup-go@v3
with:
go-version-file: "go.mod"
cache: true
- name: Test
run: go test -v ./...
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ After you download the file, extract it into a folder and open the `env.example`
- `TELEGRAM_ID` (Optional): Your Telegram User ID
- If you set this, only you will be able to interact with the bot.
- To get your ID, message `@userinfobot` on Telegram.
- Multiple IDs can be provided, separated by commas.
- `EDIT_WAIT_SECONDS` (Optional): Amount of seconds to wait between edits
- This is set to `1` by default, but you can increase if you start getting a lot of `Too Many Requests` errors.
- Save the file, and rename it to `.env`.
Expand Down
4 changes: 3 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -5,21 +5,23 @@ go 1.19
require (
github.com/go-telegram-bot-api/telegram-bot-api/v5 v5.5.1
github.com/google/uuid v1.3.0
github.com/joho/godotenv v1.4.0
github.com/launchdarkly/eventsource v1.7.1
github.com/playwright-community/playwright-go v0.2000.1
github.com/spf13/viper v1.14.0
github.com/stretchr/testify v1.8.1
)

require (
github.com/danwakefield/fnmatch v0.0.0-20160403171240-cbb64ac3d964 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/fsnotify/fsnotify v1.6.0 // indirect
github.com/go-stack/stack v1.8.1 // indirect
github.com/hashicorp/hcl v1.0.0 // indirect
github.com/magiconair/properties v1.8.6 // indirect
github.com/mitchellh/mapstructure v1.5.0 // indirect
github.com/pelletier/go-toml v1.9.5 // indirect
github.com/pelletier/go-toml/v2 v2.0.5 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/spf13/afero v1.9.2 // indirect
github.com/spf13/cast v1.5.0 // indirect
github.com/spf13/jwalterweatherman v1.1.0 // indirect
Expand Down
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -132,8 +132,6 @@ github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
github.com/joho/godotenv v1.4.0 h1:3l4+N6zfMWnkbPEXKng2o2/MR5mSwTrBih4ZEkkz1lg=
github.com/joho/godotenv v1.4.0/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk=
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
Expand Down Expand Up @@ -176,6 +174,7 @@ github.com/spf13/viper v1.14.0 h1:Rg7d3Lo706X9tHsJMUjdiwMpHB7W8WnSVOssIY+JElU=
github.com/spf13/viper v1.14.0/go.mod h1:WT//axPky3FdvXHzGw33dNdXXXfFQqmEalje+egj8As=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
Expand All @@ -184,6 +183,7 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
github.com/subosito/gotenv v1.4.1 h1:jyEFiXpy21Wm81FBN71l9VoMMV8H8jG+qIK3GCpY6Qs=
github.com/subosito/gotenv v1.4.1/go.mod h1:ayKnFf/c6rvx/2iiLrJUk1e6plDbT3edrFNGqEflhK0=
github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
Expand Down
35 changes: 13 additions & 22 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,54 +5,44 @@ import (
"log"
"os"
"os/signal"
"strconv"
"syscall"
"time"

"github.com/joho/godotenv"
"github.com/m1guelpf/chatgpt-telegram/src/chatgpt"
"github.com/m1guelpf/chatgpt-telegram/src/config"
"github.com/m1guelpf/chatgpt-telegram/src/session"
"github.com/m1guelpf/chatgpt-telegram/src/tgbot"
)

func main() {
config, err := config.Init()
persistentConfig, err := config.LoadOrCreatePersistentConfig()
if err != nil {
log.Fatalf("Couldn't load config: %v", err)
}

if config.OpenAISession == "" {
session, err := session.GetSession()
if persistentConfig.OpenAISession == "" {
token, err := session.GetSession()
if err != nil {
log.Fatalf("Couldn't get OpenAI session: %v", err)
}

err = config.Set("OpenAISession", session)
if err != nil {
if err = persistentConfig.SetSessionToken(token); err != nil {
log.Fatalf("Couldn't save OpenAI session: %v", err)
}
}

chatGPT := chatgpt.Init(config)
chatGPT := chatgpt.Init(persistentConfig)
log.Println("Started ChatGPT")

err = godotenv.Load()
envConfig, err := config.LoadEnvConfig(".env")
if err != nil {
log.Printf("Couldn't load .env file: %v. Using shell exposed env variables...", err)
log.Fatalf("Couldn't load .env config: %v", err)
}

editInterval := 1 * time.Second
if os.Getenv("EDIT_WAIT_SECONDS") != "" {
editSecond, err := strconv.ParseInt(os.Getenv("EDIT_WAIT_SECONDS"), 10, 64)
if err != nil {
log.Printf("Couldn't convert your edit seconds setting into int: %v", err)
editSecond = 1
}
editInterval = time.Duration(editSecond) * time.Second
if err := envConfig.ValidateWithDefaults(); err != nil {
log.Fatalf("Invalid .env config: %v", err)
}

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

userId := strconv.FormatInt(update.Message.Chat.ID, 10)
if os.Getenv("TELEGRAM_ID") != "" && userId != os.Getenv("TELEGRAM_ID") {
if len(envConfig.TelegramID) != 0 && !envConfig.HasTelegramID(updateUserID) {
log.Printf("User %d is not allowed to use this bot", updateUserID)
bot.Send(updateChatID, updateMessageID, "You are not authorized to use this bot.")
continue
}
Expand Down
2 changes: 1 addition & 1 deletion src/chatgpt/chatgpt.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ type ChatResponse struct {
Message string
}

func Init(config config.Config) *ChatGPT {
func Init(config *config.Config) *ChatGPT {
return &ChatGPT{
AccessTokenMap: expirymap.New(),
SessionToken: config.OpenAISession,
Expand Down
46 changes: 23 additions & 23 deletions src/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,46 +9,46 @@ import (
)

type Config struct {
v *viper.Viper

OpenAISession string
}

// init tries to read the config from the file, and creates it if it doesn't exist.
func Init() (Config, error) {
// LoadOrCreatePersistentConfig uses the default config directory for the current OS
// to load or create a config file named "chatgpt.json"
func LoadOrCreatePersistentConfig() (*Config, error) {
configPath, err := os.UserConfigDir()
if err != nil {
return Config{}, errors.New(fmt.Sprintf("Couldn't get user config dir: %v", err))
return nil, errors.New(fmt.Sprintf("Couldn't get user config dir: %v", err))
}
viper.SetConfigType("json")
viper.SetConfigName("chatgpt")
viper.AddConfigPath(configPath)
v := viper.New()
v.SetConfigType("json")
v.SetConfigName("chatgpt")
v.AddConfigPath(configPath)

if err := viper.ReadInConfig(); err != nil {
if err := v.ReadInConfig(); err != nil {
if _, ok := err.(viper.ConfigFileNotFoundError); ok {
if err := viper.SafeWriteConfig(); err != nil {
return Config{}, errors.New(fmt.Sprintf("Couldn't create config file: %v", err))
if err := v.SafeWriteConfig(); err != nil {
return nil, errors.New(fmt.Sprintf("Couldn't create config file: %v", err))
}
} else {
return Config{}, errors.New(fmt.Sprintf("Couldn't read config file: %v", err))
return nil, errors.New(fmt.Sprintf("Couldn't read config file: %v", err))
}
}

var cfg Config
err = viper.Unmarshal(&cfg)
err = v.Unmarshal(&cfg)
if err != nil {
return Config{}, errors.New(fmt.Sprintf("Error parsing config: %v", err))
return nil, errors.New(fmt.Sprintf("Error parsing config: %v", err))
}
cfg.v = v

return cfg, nil
return &cfg, nil
}

// key should be part of the Config struct
func (cfg *Config) Set(key string, value interface{}) error {
viper.Set(key, value)

err := viper.Unmarshal(&cfg)
if err != nil {
return errors.New(fmt.Sprintf("Error parsing config: %v", err))
}

return viper.WriteConfig()
func (cfg *Config) SetSessionToken(token string) error {
// key must match the struct field name
cfg.v.Set("OpenAISession", token)
cfg.OpenAISession = token
return cfg.v.WriteConfig()
}
133 changes: 133 additions & 0 deletions src/config/config_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
package config

import (
"fmt"
"os"
"testing"

"github.com/stretchr/testify/require"
)

func createFile(name string, content string) (remove func(), err error) {
f, err := os.Create(name)
if err != nil {
return nil, err
}
defer f.Close()

if _, err := f.WriteString(content); err != nil {
return nil, err
}

return func() {
if err := os.Remove(name); err != nil {
panic(fmt.Sprintf("failed to remove file: %s", err))
}
}, nil
}

func setEnvVariables(vals map[string]string) func() {
for k, v := range vals {
os.Setenv(k, v)
}
return func() {
for k := range vals {
os.Unsetenv(k)
}
}
}

func TestLoadEnvConfig(t *testing.T) {
for label, test := range map[string]struct {
fileContent string
envVars map[string]string
want *EnvConfig
}{
"all values empty in file and env": {
fileContent: `TELEGRAM_ID=
TELEGRAM_TOKEN=
EDIT_WAIT_SECONDS=`,
want: &EnvConfig{
TelegramID: []int64{},
TelegramToken: "",
EditWaitSeconds: 0,
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should default to 1 right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

defaults to 1 in the Validate method

},
},
"no file, all values through env": {
envVars: map[string]string{
"TELEGRAM_ID": "123,456",
"TELEGRAM_TOKEN": "token",
"EDIT_WAIT_SECONDS": "10",
},
want: &EnvConfig{
TelegramID: []int64{123, 456},
TelegramToken: "token",
EditWaitSeconds: 10,
},
},
"all values provided in file, single TELEGRAM_ID": {
fileContent: `TELEGRAM_ID=123
TELEGRAM_TOKEN=abc
EDIT_WAIT_SECONDS=10`,
want: &EnvConfig{
TelegramID: []int64{123},
TelegramToken: "abc",
EditWaitSeconds: 10,
},
},
"multiple TELEGRAM_IDs provided in file": {
fileContent: `TELEGRAM_ID=123,456
TELEGRAM_TOKEN=abc
EDIT_WAIT_SECONDS=10`,
envVars: map[string]string{},
want: &EnvConfig{
TelegramID: []int64{123, 456},
TelegramToken: "abc",
EditWaitSeconds: 10,
},
},
"env variables should override file values": {
fileContent: `TELEGRAM_ID=123
TELEGRAM_TOKEN=abc
EDIT_WAIT_SECONDS=10`,
envVars: map[string]string{
"TELEGRAM_ID": "456",
"TELEGRAM_TOKEN": "def",
"EDIT_WAIT_SECONDS": "20",
},
want: &EnvConfig{
TelegramID: []int64{456},
TelegramToken: "def",
EditWaitSeconds: 20,
},
},
"multiple TELEGRAM_IDs provided in env": {
fileContent: `TELEGRAM_ID=123
TELEGRAM_TOKEN=abc
EDIT_WAIT_SECONDS=10`,
envVars: map[string]string{
"TELEGRAM_ID": "456,789",
},
want: &EnvConfig{
TelegramID: []int64{456, 789},
TelegramToken: "abc",
EditWaitSeconds: 10,
},
},
} {
t.Run(label, func(t *testing.T) {
unset := setEnvVariables(test.envVars)
t.Cleanup(unset)

if test.fileContent != "" {
remove, err := createFile("test.env", test.fileContent)
require.NoError(t, err)
t.Cleanup(remove)
}

cfg, err := LoadEnvConfig("test.env")
require.NoError(t, err)
require.Equal(t, test.want, cfg)
})
}
}
Loading