diff --git a/.env.example b/.env.example index cf4b8e9..f9b4728 100644 --- a/.env.example +++ b/.env.example @@ -4,3 +4,4 @@ ADDR=:8080 # Telegram client config BOT_TOKEN= +USERNAME_LIMITS= \ No newline at end of file diff --git a/README.md b/README.md index 6b1cfbc..96fa7cf 100644 --- a/README.md +++ b/README.md @@ -126,6 +126,7 @@ If you add a caption to an image it will be used as the question. Env variables: - `BOT_TOKEN` - The token of the bot you created in Telegram - `AI_SERVICE` - AI service (the server-bot's) address (Default: `http://127.0.0.1:8080`) +- `USERNAME_LIMITS` - A comma separated list of usernames that are allowed to use the bot (Default: `""`, means everybody can use it) #### Facebook messenger client diff --git a/cmd/client-telegram/main.go b/cmd/client-telegram/main.go index 5994643..f22be01 100644 --- a/cmd/client-telegram/main.go +++ b/cmd/client-telegram/main.go @@ -8,6 +8,7 @@ import ( "net/http" "os" "os/signal" + "strings" "github.com/go-telegram/bot" "github.com/go-telegram/bot/models" @@ -29,7 +30,14 @@ func main() { aiSrv = "http://127.0.0.1:8080" } - l := New(aiSrv) + usernameLimits := make([]string, 0) + if usernameLimitsEnv := os.Getenv("USERNAME_LIMITS"); usernameLimitsEnv != "" { + for _, u := range strings.Split(usernameLimitsEnv, ",") { + usernameLimits = append(usernameLimits, strings.TrimSpace(u)) + } + } + + l := New(aiSrv, usernameLimits) ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt) defer cancel() @@ -49,17 +57,39 @@ func main() { } type Logic struct { - httpB *httpBotter.Logic + httpB *httpBotter.Logic + userLimits []string } -func New(baseURL string) *Logic { +func New(baseURL string, userLimit []string) *Logic { return &Logic{ - httpB: httpBotter.New(baseURL), + httpB: httpBotter.New(baseURL), + userLimits: userLimit, } } // Handler . func (l *Logic) Handler(ctx context.Context, b *bot.Bot, update *models.Update) { + // If we have any limits set, check them + if len(l.userLimits) > 0 { + found := false + for _, u := range l.userLimits { + if update.Message.From.Username == u { + found = true + break + } + } + + if !found { + _, _ = b.SendMessage(ctx, &bot.SendMessageParams{ + ChatID: update.Message.Chat.ID, + Text: "🙅You are not allowed to use this bot.", + }) + + return + } + } + var payload []byte msg := update.Message.Text