From 5d4f13827c602832af3036152f018c6753ce094d Mon Sep 17 00:00:00 2001 From: Matt Burchett Date: Sun, 24 May 2020 22:29:08 -0500 Subject: [PATCH] Adding the start to couchpotato and some middleware --- pkg/core/config/config.go | 14 ++++-- pkg/service/couchpotato/admin.go | 1 + pkg/service/couchpotato/couchpotato.go | 11 ++++ pkg/service/sonarr/sonarr.go | 2 +- pkg/service/telegram/handler.go | 70 ++++++++++++++++++++++++++ pkg/service/telegram/sonarr.go | 13 +++++ pkg/service/telegram/telegram.go | 48 ------------------ 7 files changed, 106 insertions(+), 53 deletions(-) create mode 100644 pkg/service/couchpotato/admin.go create mode 100644 pkg/service/couchpotato/couchpotato.go create mode 100644 pkg/service/telegram/handler.go diff --git a/pkg/core/config/config.go b/pkg/core/config/config.go index addd1ff..28bc58c 100644 --- a/pkg/core/config/config.go +++ b/pkg/core/config/config.go @@ -10,17 +10,23 @@ import ( // Config - This struct will hold configuration components. type Config struct { Telegram struct { - Token string `json:"token"` - ChatID string `json:"chatID"` - Admins []int `json:"admins"` + Token string `json:"token"` + ChatID string `json:"chatID"` + Admins []int `json:"admins"` + AuthorizedChats []string `json:"authorizedChats"` } `json:"telegram"` Sonarr struct { URL string `json:"url"` APIKey string `json:"apiKey"` SeasonLimit int `json:"seasonLimit"` - ProfileID int `json:"proFileId"` + ProfileID int `json:"profileId"` } `json:"sonarr"` + CouchPotato struct { + URL string `json:"url"` + APIKey string `json:"apiKey"` + ProfileID string `json:"profileId` + } `json:"couchpotato"` } //GetConfig gets the configuration values for the api using the file in the supplied configPath. diff --git a/pkg/service/couchpotato/admin.go b/pkg/service/couchpotato/admin.go new file mode 100644 index 0000000..c7b3122 --- /dev/null +++ b/pkg/service/couchpotato/admin.go @@ -0,0 +1 @@ +package couchpotato diff --git a/pkg/service/couchpotato/couchpotato.go b/pkg/service/couchpotato/couchpotato.go new file mode 100644 index 0000000..a93b970 --- /dev/null +++ b/pkg/service/couchpotato/couchpotato.go @@ -0,0 +1,11 @@ +package couchpotato + +import ( + "github.com/mattburchett/go_telegram/pkg/core/config" + "github.com/yanzay/tbot/v2" +) + +// Search performs the lookup actions within CouchPotato +func Search(m *tbot.Message, config config.Config) ([]response, error) { + +} diff --git a/pkg/service/sonarr/sonarr.go b/pkg/service/sonarr/sonarr.go index 2da07c3..13443f4 100644 --- a/pkg/service/sonarr/sonarr.go +++ b/pkg/service/sonarr/sonarr.go @@ -178,7 +178,7 @@ func Add(callback string, config config.Config) string { if err != nil { return err.Error() } - rootFolderData, err := ioutil.ReadAll(rootFolderLookuQuery.Body) + rootFolderData, err := ioutil.ReadAll(rootFolderLookupQuery.Body) if err != nil { return err.Error() } diff --git a/pkg/service/telegram/handler.go b/pkg/service/telegram/handler.go new file mode 100644 index 0000000..3139b9a --- /dev/null +++ b/pkg/service/telegram/handler.go @@ -0,0 +1,70 @@ +package telegram + +import ( + "strings" + + "github.com/yanzay/tbot/v2" +) + +// Handler creates the active Telegram handlers. +func (tb *Bot) Handler() { + + // Bot Healthcheck + tb.Bot.HandleMessage("/ping", func(m *tbot.Message) { + tb.Client.SendMessage(m.Chat.ID, "pong") + }) + + // telegram/sonar.go + tb.Bot.HandleMessage("/s", tb.sonarrSearch) + tb.Bot.HandleMessage("/admin sonarrStatus", tb.sonarrStatus) + + // telegram/testhandler.go + tb.Bot.HandleMessage("/test", tb.testHandler) + + // telegram/admin.go + tb.Bot.HandleMessage("/admin myID", tb.myID) + tb.Bot.HandleMessage("/admin chatID", tb.chatID) + + // Help + tb.Bot.HandleMessage("/help$", tb.helpHandler) + tb.Bot.HandleMessage("/h$", tb.helpHandler) + + // Callback Handler + tb.Bot.HandleCallback(tb.callbackHandler) +} + +// callbackHandler handles callbacks. +func (tb *Bot) callbackHandler(cq *tbot.CallbackQuery) { + go func() { + tb.Client.AnswerCallbackQuery(cq.ID, tbot.OptText("Request received.")) + tb.Client.DeleteMessage(tb.CallbackChatID, tb.CallbackMessageID) + }() + + if strings.Contains(cq.Data, "tv_") { + tb.sonarrAdd(cq) + return + } + + tb.Client.SendMessage(tb.CallbackChatID, cq.Data) + +} + +func (tb *Bot) helpHandler(m *tbot.Message) { + if !tb.whitelistHandler(m) { + return + } + + tb.Client.SendMessage(m.Chat.ID, "USAGE:\n\n/movie or /m \n/show or /s \n\nEXAMPLES:\n\n/s The Walking Dead\n/m Avatar") +} + +func (tb *Bot) whitelistHandler(m *tbot.Message) bool { + for _, id := range tb.Config.Telegram.AuthorizedChats { + if id == m.Chat.ID { + return true + } + } + + tb.Client.SendMessage(m.Chat.ID, "This bot is not authorized for use in this chat.") + return false + +} diff --git a/pkg/service/telegram/sonarr.go b/pkg/service/telegram/sonarr.go index 81f77f0..2eb4696 100644 --- a/pkg/service/telegram/sonarr.go +++ b/pkg/service/telegram/sonarr.go @@ -11,6 +11,10 @@ import ( // Sonarr Search func (tb *Bot) sonarrSearch(m *tbot.Message) { + if !tb.whitelistHandler(m) { + return + } + text := strings.TrimPrefix(strings.TrimPrefix(m.Text, "/s"), " ") if len(text) == 0 { tb.Client.SendMessage(m.Chat.ID, "You must specify a show. Type /help for help.") @@ -31,6 +35,11 @@ func (tb *Bot) sonarrSearch(m *tbot.Message) { }}) } + if len(request) == 0 { + tb.Client.SendMessage(m.Chat.ID, "No results found, try harder.") + return + } + response, _ := tb.Client.SendMessage(m.Chat.ID, "Please select the show you would like to download.", tbot.OptInlineKeyboardMarkup(&tbot.InlineKeyboardMarkup{InlineKeyboard: inlineResponse})) tb.CallbackMessageID = response.MessageID tb.CallbackChatID = m.Chat.ID @@ -56,6 +65,10 @@ func (tb *Bot) sonarrAdd(cq *tbot.CallbackQuery) { // sonarrStatus queries Sonarr for it's system status information. func (tb *Bot) sonarrStatus(m *tbot.Message) { + if !tb.whitelistHandler(m) { + return + } + if tb.adminCheck(m.From.ID, false) { request, err := sonarr.Status(m, tb.Config) diff --git a/pkg/service/telegram/telegram.go b/pkg/service/telegram/telegram.go index 725ae33..805169d 100644 --- a/pkg/service/telegram/telegram.go +++ b/pkg/service/telegram/telegram.go @@ -2,7 +2,6 @@ package telegram import ( "log" - "strings" "time" "github.com/mattburchett/go_telegram/pkg/core/config" @@ -36,50 +35,3 @@ func (tb *Bot) New(token string) { tb.Bot.Start() } - -// Handler creates the active Telegram handlers. -func (tb *Bot) Handler() { - - // Bot Healthcheck - tb.Bot.HandleMessage("/ping", func(m *tbot.Message) { - tb.Client.SendMessage(m.Chat.ID, "pong") - }) - - // telegram/sonar.go - tb.Bot.HandleMessage("/s", tb.sonarrSearch) - tb.Bot.HandleMessage("/admin sonarrStatus", tb.sonarrStatus) - - // telegram/testhandler.go - tb.Bot.HandleMessage("/test", tb.testHandler) - - // telegram/admin.go - tb.Bot.HandleMessage("/admin myID", tb.myID) - tb.Bot.HandleMessage("/admin chatID", tb.chatID) - - // Help - tb.Bot.HandleMessage("/help$", tb.helpHandler) - tb.Bot.HandleMessage("/h$", tb.helpHandler) - - // Callback Handler - tb.Bot.HandleCallback(tb.callbackHandler) -} - -// callbackHandler handles callbacks. -func (tb *Bot) callbackHandler(cq *tbot.CallbackQuery) { - go func() { - tb.Client.AnswerCallbackQuery(cq.ID, tbot.OptText("Request received.")) - tb.Client.DeleteMessage(tb.CallbackChatID, tb.CallbackMessageID) - }() - - if strings.Contains(cq.Data, "tv_") { - tb.sonarrAdd(cq) - return - } - - tb.Client.SendMessage(tb.CallbackChatID, cq.Data) - -} - -func (tb *Bot) helpHandler(m *tbot.Message) { - tb.Client.SendMessage(m.Chat.ID, "USAGE:\n\n/movie or /m \n/show or /s \n\nEXAMPLES:\n\n/s The Walking Dead\n/m Avatar") -}