fix links parserHEADmaster

1 files changed, 22 insertions, 3 deletions

diff --git a/main.go b/main.go
index 6c77033..07d550f 100644
--- a/main.go
+++ b/main.go
@@ -8,6 +8,7 @@ import (
 
 	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
 	"github.com/kkdai/youtube/v2"
+	"mvdan.cc/xurls/v2"
 )
 
 func main() {
@@ -30,20 +31,33 @@ func main() {
 		if update.Message != nil { // If we got a message
 
 			message := update.Message
-			link, err := url.Parse(message.Text)
+			xRelaxed := xurls.Relaxed()
+			linkString := xRelaxed.FindString(message.Text)
+			if linkString == "" {
+				msg := tgbotapi.NewMessage(message.Chat.ID, "Сообщение должно содержать ссылку")
+				log.Println(err)
+				msg.ReplyToMessageID = message.MessageID
+				bot.Send(msg)
+				continue
+			}
+			link, err := url.Parse(linkString)
 			if err != nil {
+				msg := tgbotapi.NewMessage(message.Chat.ID, "Сообщение должно содержать корректную ссылку на youtube")
+				log.Println(err)
+				msg.ReplyToMessageID = message.MessageID
+				bot.Send(msg)
 				continue
 			}
 
 			go func(message *tgbotapi.Message) {
-				if link.Host == "youtube.com" && link.Query().Get("v") != "" {
+				if strings.HasSuffix(link.Host, "youtube.com") && link.Query().Get("v") != "" {
 					if err := download(bot, message, link.Query().Get("v")); err != nil {
 						msg := tgbotapi.NewMessage(message.Chat.ID, "Ошибка: "+err.Error())
 						log.Println(err)
 						msg.ReplyToMessageID = message.MessageID
 						bot.Send(msg)
 					}
-				} else if link.Host == "youtu.be" && link.Path != "" {
+				} else if strings.HasSuffix(link.Host, "youtu.be") && link.Path != "" {
 					pathParts := strings.SplitN(link.Path, "?", 2)
 					if err := download(bot, message, pathParts[0]); err != nil {
 						msg := tgbotapi.NewMessage(message.Chat.ID, "Ошибка: "+err.Error())
@@ -51,6 +65,11 @@ func main() {
 						msg.ReplyToMessageID = message.MessageID
 						bot.Send(msg)
 					}
+				} else {
+					msg := tgbotapi.NewMessage(message.Chat.ID, "Сообщение должно содержать корректную ссылку на youtube")
+					log.Println(err)
+					msg.ReplyToMessageID = message.MessageID
+					bot.Send(msg)
 				}
 			}(message)
 		}