Support loading Telegram bot token from file

Signed-off-by: Andrey Mishakin <stieroglif@gmail.com>
This commit is contained in:
Andrey Mishakin 2023-01-22 00:43:34 +03:00
parent f59460bfd4
commit 6c9c58015e
5 changed files with 114 additions and 7 deletions

View File

@ -739,6 +739,7 @@ type TelegramConfig struct {
APIUrl *URL `yaml:"api_url" json:"api_url,omitempty"`
BotToken Secret `yaml:"bot_token,omitempty" json:"token,omitempty"`
BotTokenFile string `yaml:"bot_token_file,omitempty" json:"token_file,omitempty"`
ChatID int64 `yaml:"chat_id,omitempty" json:"chat,omitempty"`
Message string `yaml:"message,omitempty" json:"message,omitempty"`
DisableNotifications bool `yaml:"disable_notifications,omitempty" json:"disable_notifications,omitempty"`
@ -752,8 +753,11 @@ func (c *TelegramConfig) UnmarshalYAML(unmarshal func(interface{}) error) error
if err := unmarshal((*plain)(c)); err != nil {
return err
}
if c.BotToken == "" {
return fmt.Errorf("missing bot_token on telegram_config")
if c.BotToken == "" && c.BotTokenFile == "" {
return fmt.Errorf("missing bot_token or bot_token_file on telegram_config")
}
if c.BotToken != "" && c.BotTokenFile != "" {
return fmt.Errorf("at most one of bot_token & bot_token_file must be configured")
}
if c.ChatID == 0 {
return fmt.Errorf("missing chat_id on telegram_config")

View File

@ -958,6 +958,70 @@ http_config:
}
}
func TestTelegramConfiguration(t *testing.T) {
tc := []struct {
name string
in string
expected error
}{
{
name: "with both bot_token & bot_token_file - it fails",
in: `
bot_token: xyz
bot_token_file: /file
`,
expected: errors.New("at most one of bot_token & bot_token_file must be configured"),
},
{
name: "with no bot_token & bot_token_file - it fails",
in: `
bot_token: ''
bot_token_file: ''
`,
expected: errors.New("missing bot_token or bot_token_file on telegram_config"),
},
{
name: "with bot_token and chat_id set - it succeeds",
in: `
bot_token: xyz
chat_id: 123
`,
},
{
name: "with bot_token_file and chat_id set - it succeeds",
in: `
bot_token_file: /file
chat_id: 123
`,
},
{
name: "with no chat_id set - it fails",
in: `
bot_token: xyz
`,
expected: errors.New("missing chat_id on telegram_config"),
},
{
name: "with unknown parse_mode - it fails",
in: `
bot_token: xyz
chat_id: 123
parse_mode: invalid
`,
expected: errors.New("unknown parse_mode on telegram_config, must be Markdown, MarkdownV2, HTML or empty string"),
},
}
for _, tt := range tc {
t.Run(tt.name, func(t *testing.T) {
var cfg TelegramConfig
err := yaml.UnmarshalStrict([]byte(tt.in), &cfg)
require.Equal(t, tt.expected, err)
})
}
}
func newBoolPointer(b bool) *bool {
return &b
}

View File

@ -1050,9 +1050,12 @@ attributes:
# If not specified, default API URL will be used.
[ api_url: <string> | default = global.telegram_api_url ]
# Telegram bot token.
# Telegram bot token. It is mutually exclusive with `bot_token_file`.
[ bot_token: <secret> ]
# Read the Telegram bot token from a file. It is mutually exclusive with `bot_token`.
[ bot_token_file: <filepath> ]
# ID of the chat where to send the messages.
[ chat_id: <int> ]

View File

@ -17,6 +17,8 @@ import (
"context"
"fmt"
"net/http"
"os"
"strings"
"github.com/go-kit/log"
"github.com/go-kit/log/level"
@ -48,7 +50,7 @@ func New(conf *config.TelegramConfig, t *template.Template, l log.Logger, httpOp
return nil, err
}
client, err := createTelegramClient(conf.BotToken, conf.APIUrl.String(), conf.ParseMode, httpclient)
client, err := createTelegramClient(conf.APIUrl.String(), conf.ParseMode, httpclient)
if err != nil {
return nil, err
}
@ -83,6 +85,11 @@ func (n *Notifier) Notify(ctx context.Context, alert ...*types.Alert) (bool, err
level.Warn(n.logger).Log("msg", "Truncated message", "alert", key, "max_runes", maxMessageLenRunes)
}
n.client.Token, err = n.getBotToken()
if err != nil {
return true, err
}
message, err := n.client.Send(telebot.ChatID(n.conf.ChatID), messageText, &telebot.SendOptions{
DisableNotification: n.conf.DisableNotifications,
DisableWebPagePreview: true,
@ -95,10 +102,8 @@ func (n *Notifier) Notify(ctx context.Context, alert ...*types.Alert) (bool, err
return false, nil
}
func createTelegramClient(token config.Secret, apiURL, parseMode string, httpClient *http.Client) (*telebot.Bot, error) {
secret := string(token)
func createTelegramClient(apiURL, parseMode string, httpClient *http.Client) (*telebot.Bot, error) {
bot, err := telebot.NewBot(telebot.Settings{
Token: secret,
URL: apiURL,
ParseMode: parseMode,
Client: httpClient,
@ -110,3 +115,14 @@ func createTelegramClient(token config.Secret, apiURL, parseMode string, httpCli
return bot, nil
}
func (n *Notifier) getBotToken() (string, error) {
if len(n.conf.BotTokenFile) > 0 {
content, err := os.ReadFile(n.conf.BotTokenFile)
if err != nil {
return "", fmt.Errorf("could not read %s: %w", n.conf.BotTokenFile, err)
}
return strings.TrimSpace(string(content)), nil
}
return string(n.conf.BotToken), nil
}

View File

@ -21,6 +21,7 @@ import (
"net/http"
"net/http/httptest"
"net/url"
"os"
"testing"
"time"
@ -84,6 +85,13 @@ func TestTelegramRetry(t *testing.T) {
}
func TestTelegramNotify(t *testing.T) {
token := "secret"
fileWithToken, err := os.CreateTemp("", "telegram-bot-token")
require.NoError(t, err, "creating temp file failed")
_, err = fileWithToken.WriteString(token)
require.NoError(t, err, "writing to temp file failed")
for _, tc := range []struct {
name string
cfg config.TelegramConfig
@ -94,6 +102,7 @@ func TestTelegramNotify(t *testing.T) {
cfg: config.TelegramConfig{
Message: "<code>x < y</code>",
HTTPConfig: &commoncfg.HTTPClientConfig{},
BotToken: config.Secret(token),
},
expText: "<code>x < y</code>",
},
@ -103,13 +112,24 @@ func TestTelegramNotify(t *testing.T) {
ParseMode: "HTML",
Message: "<code>x < y</code>",
HTTPConfig: &commoncfg.HTTPClientConfig{},
BotToken: config.Secret(token),
},
expText: "<code>x &lt; y</code>",
},
{
name: "Bot token from file",
cfg: config.TelegramConfig{
Message: "test",
HTTPConfig: &commoncfg.HTTPClientConfig{},
BotTokenFile: fileWithToken.Name(),
},
expText: "test",
},
} {
t.Run(tc.name, func(t *testing.T) {
var out []byte
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Equal(t, "/bot"+token+"/sendMessage", r.URL.Path)
var err error
out, err = io.ReadAll(r.Body)
require.NoError(t, err)