From 7c6b9c4a744b6aeeecaf90e1b6540586fdfb58b7 Mon Sep 17 00:00:00 2001 From: Matt Burchett Date: Sun, 17 May 2020 01:25:29 -0500 Subject: [PATCH] Bigger and better rewrite. --- cmd/main.go | 57 +++---------------------- pkg/{ => core}/config/config.go | 14 +++++-- pkg/service/sonarr/admin.go | 24 +++++++++++ pkg/service/telegram/admin.go | 31 ++++++++++++++ pkg/service/telegram/sonarr.go | 23 ++++++++++ pkg/service/telegram/telegram.go | 65 +++++++++++++++++++++++++++++ pkg/service/telegram/testhandler.go | 25 +++++++++++ 7 files changed, 184 insertions(+), 55 deletions(-) rename pkg/{ => core}/config/config.go (80%) create mode 100644 pkg/service/sonarr/admin.go create mode 100644 pkg/service/telegram/admin.go create mode 100644 pkg/service/telegram/sonarr.go create mode 100644 pkg/service/telegram/telegram.go create mode 100644 pkg/service/telegram/testhandler.go diff --git a/cmd/main.go b/cmd/main.go index 172078a..23e9317 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -2,63 +2,18 @@ package main import ( "log" - "strings" - "github.com/mattburchett/go_telegram/pkg/config" - "github.com/yanzay/tbot/v2" + "github.com/mattburchett/go_telegram/pkg/core/config" + "github.com/mattburchett/go_telegram/pkg/service/telegram" ) -type application struct { - client *tbot.Client - callbackChatID string - callbackMessageID int -} - func main() { - cfg, err := config.GetConfig("config.json") + conf, err := config.GetConfig("config.json") if err != nil { log.Fatal("Failed to read JSON.") } - app := &application{} - - bot := tbot.New(cfg.TelegramToken) - app.client = bot.Client() - c := bot.Client() - bot.HandleMessage("/ping", func(m *tbot.Message) { - c.SendMessage(m.Chat.ID, "pong") - }) - - bot.HandleMessage("/test", app.testHandler) - bot.HandleCallback(app.callbackHandler) - - err = bot.Start() - if err != nil { - log.Fatal(err) - } - -} - -func (a *application) testHandler(m *tbot.Message) { - buttons := make([]string, 0) - buttons = append(buttons, "ping", "is", "stupid") - - inline2 := make([][]tbot.InlineKeyboardButton, 0) - - for _, i := range buttons { - inline2 = append(inline2, []tbot.InlineKeyboardButton{{ - Text: i, - CallbackData: i, - }}) - } - - msg, _ := a.client.SendMessage(m.Chat.ID, "Inline test. "+strings.TrimPrefix(m.Text, "/test "), tbot.OptInlineKeyboardMarkup(&tbot.InlineKeyboardMarkup{InlineKeyboard: inline2})) - a.callbackMessageID = msg.MessageID - a.callbackChatID = m.Chat.ID - -} - -func (a *application) callbackHandler(cq *tbot.CallbackQuery) { - a.client.EditMessageText(a.callbackChatID, a.callbackMessageID, "Callback received.") - a.client.SendMessage(a.callbackChatID, cq.Data) + tgBot := telegram.Bot{} + tgBot.Config = conf + tgBot.New(conf.Telegram.Token) } diff --git a/pkg/config/config.go b/pkg/core/config/config.go similarity index 80% rename from pkg/config/config.go rename to pkg/core/config/config.go index 78d38f2..272e526 100644 --- a/pkg/config/config.go +++ b/pkg/core/config/config.go @@ -9,10 +9,16 @@ import ( // Config - This struct will hold configuration components. type Config struct { - TelegramToken string `json:"telegramToken"` - TelegramChatID string `json:"telegramChatID"` - SonarrAPIURL string `json:"sonarrAPIURL"` - SonarrAPIKey string `json:"sonarrAPIKey"` + Telegram struct { + Token string `json:"token"` + ChatID string `json:"chatID"` + Admins []int `json:"admins"` + } `json:"telegram"` + + Sonarr struct { + URL string `json:"url"` + APIKey string `json:"apiKey"` + } `json:"sonarr"` } //GetConfig gets the configuration values for the api using the file in the supplied configPath. diff --git a/pkg/service/sonarr/admin.go b/pkg/service/sonarr/admin.go new file mode 100644 index 0000000..95293eb --- /dev/null +++ b/pkg/service/sonarr/admin.go @@ -0,0 +1,24 @@ +package sonarr + +import ( + "io/ioutil" + "net/http" + + "github.com/mattburchett/go_telegram/pkg/core/config" + "github.com/yanzay/tbot/v2" +) + +// SonarrStatus contains the Sonarr request for system status. +func SonarrStatus(m *tbot.Message, config config.Config) (string, error) { + r, err := http.Get(config.Sonarr.URL + "system/status?apikey=" + config.Sonarr.APIKey) + if err != nil { + return "Failed to contact Sonarr for data", err + } + + rd, err := ioutil.ReadAll(r.Body) + if err != nil { + return "Failed to read Sonarr status data.", err + } + + return string(rd), err +} diff --git a/pkg/service/telegram/admin.go b/pkg/service/telegram/admin.go new file mode 100644 index 0000000..5f8e4aa --- /dev/null +++ b/pkg/service/telegram/admin.go @@ -0,0 +1,31 @@ +package telegram + +import ( + "fmt" + "strconv" + + "github.com/yanzay/tbot/v2" +) + +func (tb *Bot) myID(m *tbot.Message) { + if tb.AdminCheck(m) { + tb.Client.SendMessage(m.Chat.ID, strconv.Itoa(m.From.ID)) + fmt.Println(m.From.ID) + } +} + +func (tb *Bot) chatID(m *tbot.Message) { + tb.Client.SendMessage(m.Chat.ID, m.Chat.ID) +} + +// AdminCheck checks for valid bot admins. +func (tb *Bot) AdminCheck(m *tbot.Message) bool { + for _, admin := range tb.Config.Telegram.Admins { + if m.From.ID == admin { + return true + } + + } + tb.Client.SendMessage(m.Chat.ID, "You are not an authorized admin.") + return false +} diff --git a/pkg/service/telegram/sonarr.go b/pkg/service/telegram/sonarr.go new file mode 100644 index 0000000..3487bf5 --- /dev/null +++ b/pkg/service/telegram/sonarr.go @@ -0,0 +1,23 @@ +package telegram + +import ( + "fmt" + + "github.com/mattburchett/go_telegram/pkg/service/sonarr" + + "github.com/yanzay/tbot/v2" +) + +func (tb *Bot) sonarrStatus(m *tbot.Message) { + if tb.AdminCheck(m) { + request, err := sonarr.SonarrStatus(m, tb.Config) + + if err != nil { + tb.Client.SendMessage(m.Chat.ID, fmt.Sprintf("%v: \n %v", request, err)) + } else { + tb.Client.SendMessage(m.Chat.ID, "Sonarr Status:") + tb.Client.SendMessage(m.Chat.ID, request) + } + } + +} diff --git a/pkg/service/telegram/telegram.go b/pkg/service/telegram/telegram.go new file mode 100644 index 0000000..fa84fb9 --- /dev/null +++ b/pkg/service/telegram/telegram.go @@ -0,0 +1,65 @@ +package telegram + +import ( + "log" + "time" + + "github.com/mattburchett/go_telegram/pkg/core/config" + "github.com/yanzay/tbot/v2" +) + +// Bot contains all the necessary bot and callback information. +type Bot struct { + Client *tbot.Client + Config config.Config + Bot *tbot.Server + CallbackChatID string + CallbackMessageID int +} + +// New creates an active telegram bot and loads the handlers. +func (tb *Bot) New(token string) { + tb.Bot = tbot.New(token) + tb.Bot.Use(stat) + tb.Client = tb.Bot.Client() + tb.Handler() + 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") + }) + + // sonarr/admin.go + 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) + + // Callback Handler + tb.Bot.HandleCallback(tb.callbackHandler) +} + +// Stat middleware. +func stat(h tbot.UpdateHandler) tbot.UpdateHandler { + return func(u *tbot.Update) { + start := time.Now() + h(u) + log.Printf("Handle time: %v", time.Now().Sub(start)) + } +} + +// callbackHandler handles callbacks. +func (tb *Bot) callbackHandler(cq *tbot.CallbackQuery) { + tb.Client.EditMessageText(tb.CallbackChatID, tb.CallbackMessageID, "Callback received.") + tb.Client.SendMessage(tb.CallbackChatID, cq.Data) +} diff --git a/pkg/service/telegram/testhandler.go b/pkg/service/telegram/testhandler.go new file mode 100644 index 0000000..c312c0f --- /dev/null +++ b/pkg/service/telegram/testhandler.go @@ -0,0 +1,25 @@ +package telegram + +import ( + "strings" + + "github.com/yanzay/tbot/v2" +) + +func (tb *Bot) testHandler(m *tbot.Message) { + buttons := make([]string, 0) + buttons = append(buttons, "ping", "is", "stupid") + + inline2 := make([][]tbot.InlineKeyboardButton, 0) + + for _, i := range buttons { + inline2 = append(inline2, []tbot.InlineKeyboardButton{{ + Text: i, + CallbackData: i, + }}) + } + + msg, _ := tb.Client.SendMessage(m.Chat.ID, "Inline test. "+strings.TrimPrefix(m.Text, "/test "), tbot.OptInlineKeyboardMarkup(&tbot.InlineKeyboardMarkup{InlineKeyboard: inline2})) + tb.CallbackMessageID = msg.MessageID + tb.CallbackChatID = m.Chat.ID +}