diff --git a/bot/bot.go b/bot/bot.go index c12c070..f741a21 100644 --- a/bot/bot.go +++ b/bot/bot.go @@ -6,7 +6,6 @@ import ( th "github.com/mymmrac/telego/telegohandler" tu "github.com/mymmrac/telego/telegoutil" "log/slog" - "net/url" "strings" "telegram-ollama-reply-bot/extractor" "telegram-ollama-reply-bot/llm" @@ -153,9 +152,8 @@ func (b *Bot) summarizeHandler(bot *telego.Bot, update telego.Update) { return } - _, err := url.ParseRequestURI(args[1]) - if err != nil { - slog.Error("Provided URL is not valid", "url", args[1]) + if !isValidAndAllowedUrl(args[1]) { + slog.Error("Provided text is not a valid URL", "text", args[1]) _, _ = b.api.SendMessage(b.reply(update.Message, tu.Message( chatID, @@ -295,29 +293,3 @@ func (b *Bot) createLlmRequestContext(update telego.Update) llm.RequestContext { func (b *Bot) escapeMarkdownV1Symbols(input string) string { return b.markdownV1Replacer.Replace(input) } - -func (b *Bot) reply(originalMessage *telego.Message, newMessage *telego.SendMessageParams) *telego.SendMessageParams { - return newMessage.WithReplyParameters(&telego.ReplyParameters{ - MessageID: originalMessage.MessageID, - }) -} - -func (b *Bot) sendTyping(chatId telego.ChatID) { - slog.Debug("Setting 'typing' chat action") - - err := b.api.SendChatAction(tu.ChatAction(chatId, "typing")) - if err != nil { - slog.Error("Cannot set chat action", "error", err) - } -} - -func (b *Bot) trySendReplyError(message *telego.Message) { - if message == nil { - return - } - - _, _ = b.api.SendMessage(b.reply(message, tu.Message( - tu.ID(message.Chat.ID), - "Error occurred while trying to send reply.", - ))) -} diff --git a/bot/helpers.go b/bot/helpers.go new file mode 100644 index 0000000..3c20598 --- /dev/null +++ b/bot/helpers.go @@ -0,0 +1,57 @@ +package bot + +import ( + "github.com/mymmrac/telego" + "github.com/mymmrac/telego/telegoutil" + "log/slog" + "net/url" + "slices" + "strings" +) + +var ( + allowedUrlSchemes = []string{"http", "https"} +) + +func (b *Bot) reply(originalMessage *telego.Message, newMessage *telego.SendMessageParams) *telego.SendMessageParams { + return newMessage.WithReplyParameters(&telego.ReplyParameters{ + MessageID: originalMessage.MessageID, + }) +} + +func (b *Bot) sendTyping(chatId telego.ChatID) { + slog.Debug("Setting 'typing' chat action") + + err := b.api.SendChatAction(telegoutil.ChatAction(chatId, "typing")) + if err != nil { + slog.Error("Cannot set chat action", "error", err) + } +} + +func (b *Bot) trySendReplyError(message *telego.Message) { + if message == nil { + return + } + + _, _ = b.api.SendMessage(b.reply(message, telegoutil.Message( + telegoutil.ID(message.Chat.ID), + "Error occurred while trying to send reply.", + ))) +} + +func isValidAndAllowedUrl(text string) bool { + u, err := url.ParseRequestURI(text) + if err != nil { + slog.Debug("Provided text is not an URL", "text", text) + + return false + } + + if !slices.Contains(allowedUrlSchemes, strings.ToLower(u.Scheme)) { + slog.Debug("Provided URL has disallowed scheme", "scheme", u.Scheme, "allowed-schemes", allowedUrlSchemes) + + return false + } + + return true +}