diff --git a/pkg/core/config/config.go b/pkg/core/config/config.go index 272e526..addd1ff 100644 --- a/pkg/core/config/config.go +++ b/pkg/core/config/config.go @@ -16,8 +16,10 @@ type Config struct { } `json:"telegram"` Sonarr struct { - URL string `json:"url"` - APIKey string `json:"apiKey"` + URL string `json:"url"` + APIKey string `json:"apiKey"` + SeasonLimit int `json:"seasonLimit"` + ProfileID int `json:"proFileId"` } `json:"sonarr"` } diff --git a/pkg/service/sonarr/admin.go b/pkg/service/sonarr/admin.go index 95293eb..0a20993 100644 --- a/pkg/service/sonarr/admin.go +++ b/pkg/service/sonarr/admin.go @@ -8,8 +8,8 @@ import ( "github.com/yanzay/tbot/v2" ) -// SonarrStatus contains the Sonarr request for system status. -func SonarrStatus(m *tbot.Message, config config.Config) (string, error) { +// Status contains the Sonarr request for system status. +func Status(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 diff --git a/pkg/service/sonarr/sonarr.go b/pkg/service/sonarr/sonarr.go new file mode 100644 index 0000000..2750b58 --- /dev/null +++ b/pkg/service/sonarr/sonarr.go @@ -0,0 +1,245 @@ +package sonarr + +import ( + "bytes" + "encoding/json" + "fmt" + "io/ioutil" + "net/http" + "net/url" + "strings" + + "github.com/mattburchett/go_telegram/pkg/core/config" + "github.com/yanzay/tbot/v2" +) + +// Response holds the information needed for Telegram callback. +type response struct { + Button string `json:"button"` + Callback string `json:"callback"` +} + +type sonarrSearch []struct { + Title string `json:"title"` + SeasonCount int `json:"seasonCount"` + Year int `json:"year"` + TvdbID int `json:"tvdbId"` + Downloaded bool `json:"downloaded"` + QualityProfileID int `json:"qualityProfileId"` + TitleSlug string `json:"titleSlug"` + Images []struct { + CoverType string `json:"coverType"` + URL string `json:"url"` + } `json:"images"` + Seasons []struct { + SeasonNumber int `json:"seasonNumber"` + Monitored bool `json:"monitored"` + } `json:"seasons"` + ProfileID int `json:"profileId"` +} + +type sonarrAdd struct { + Title string `json:"title"` + TvdbID int `json:"tvdbId"` + QualityProfileID int `json:"qualityProfileId"` + TitleSlug string `json:"titleSlug"` + Images []struct { + CoverType string `json:"coverType"` + URL string `json:"url"` + } `json:"images"` + Seasons []struct { + SeasonNumber int `json:"seasonNumber"` + Monitored bool `json:"monitored"` + } `json:"seasons"` + ProfileID int `json:"profileId"` + RootFolderPath string `json:"rootFolderPath"` + AddOptions struct { + IgnoreEpisodesWithFiles bool `json:"ignoreEpisodesWithFiles"` + IgnoreEpisodesWithoutFiles bool `json:"ignoreEpisodesWithoutFiles"` + SearchForMIssingEpisodes bool `json:"searchForMissingEpisodes"` + } `json:"addOptions"` +} + +type RootFolderLookup []struct { + Path string `json:"path"` + FreeSpace int64 `json:"freeSpace"` + TotalSpace int64 `json:"totalSpace"` + UnmappedFolders []struct { + Name string `json:"name"` + Path string `json:"path"` + } `json:"unmappedFolders"` + ID int `json:"id"` +} + +// Search performs the lookup actions within Sonarr. +func Search(m *tbot.Message, config config.Config) ([]response, error) { + + var remote sonarrSearch + var local sonarrSearch + + // Perform series lookup + remoteLookup, err := http.Get(config.Sonarr.URL + + "series/lookup?apikey=" + + config.Sonarr.APIKey + + "&term=" + + url.QueryEscape(strings.TrimPrefix(strings.TrimPrefix(m.Text, "/s"), " "))) + + if err != nil { + return nil, err + } + + // Perform series local database lookup. + localLookup, err := http.Get(config.Sonarr.URL + "series?apikey=" + config.Sonarr.APIKey) + if err != nil { + return nil, err + } + + // Read remote and local Data + remoteData, err := ioutil.ReadAll(remoteLookup.Body) + if err != nil { + return nil, err + } + + localData, err := ioutil.ReadAll(localLookup.Body) + if err != nil { + return nil, err + } + + // Unmarshal remote and local JSON + remoteJSON := json.Unmarshal(remoteData, &remote) + if remoteJSON != nil { + return nil, err + } + + localJSON := json.Unmarshal(localData, &local) + if localJSON != nil { + return nil, err + } + + // Check for downloaded items. + for r := range remote { + for l := range local { + if remote[r].TvdbID == local[l].TvdbID { + remote[r].Downloaded = true + } + } + } + + // Form URLs and return to Telegram + responseData := []response{} + for k := range remote { + responseData = append(responseData, + response{ + fmt.Sprintf("%v%v%v (%v) - %v Seasons", + seasonHandler(remote[k].SeasonCount, config.Sonarr.SeasonLimit), + downloadedHandler(remote[k].Downloaded), + remote[k].Title, + remote[k].Year, + remote[k].SeasonCount), + fmt.Sprintf("%d%s%s", remote[k].TvdbID, + strings.TrimSuffix(downloadedHandler(remote[k].Downloaded), " "), + strings.TrimSuffix(seasonHandler(remote[k].SeasonCount, config.Sonarr.SeasonLimit), " ")), + }) + } + + return responseData, err + +} + +// Add will take the callback data and add the show to Sonarr. +func Add(callback string, config config.Config) string { + // Separate the TVDB ID + tvdbid := strings.TrimPrefix(strings.TrimSuffix(strings.TrimSuffix(callback, "+"), "*"), "tv_") + + // Look it up, to gather the information needed and the title. + seriesLookup, err := http.Get(config.Sonarr.URL + "/series/lookup?apikey=" + config.Sonarr.APIKey + "&term=tvdb:" + tvdbid) + if err != nil { + return err.Error() + } + + seriesData, err := ioutil.ReadAll(seriesLookup.Body) + if err != nil { + return err.Error() + } + + series := sonarrSearch{} + seriesJSON := json.Unmarshal(seriesData, &series) + if seriesJSON != nil { + return err.Error() + } + + // If the "downloaded" asterisk is already in place, go ahead and return. + if strings.Contains(callback, "*") { + return fmt.Sprintf("%v has already been requested for download.", series[0].Title) + } + + // Gather the root folder location. + rootFolderLookup, err := http.Get(config.Sonarr.URL + "/rootfolder?apikey=" + config.Sonarr.APIKey) + if err != nil { + return err.Error() + } + rootFolderData, err := ioutil.ReadAll(rootFolderLookup.Body) + if err != nil { + return err.Error() + } + + rootFolder := RootFolderLookup{} + rootFolderJSON := json.Unmarshal(rootFolderData, &rootFolder) + if rootFolderJSON != nil { + return err.Error() + } + + // Form the JSON needed for adding to Sonarr. + seriesAdd := sonarrAdd{ + TvdbID: series[0].TvdbID, + Title: series[0].Title, + // QualityProfileID: series[0].QualityProfileID, + TitleSlug: series[0].TitleSlug, + Images: series[0].Images, + Seasons: series[0].Seasons, + RootFolderPath: rootFolder[0].Path, + ProfileID: config.Sonarr.ProfileID, + } + seriesAdd.AddOptions.IgnoreEpisodesWithFiles = false + seriesAdd.AddOptions.IgnoreEpisodesWithoutFiles = false + seriesAdd.AddOptions.SearchForMIssingEpisodes = true + + // Post it to Sonarr to be added. + seriesAddJSON, err := json.Marshal(seriesAdd) + if err != nil { + return err.Error() + } + seriesAddReq, err := http.Post(config.Sonarr.URL+"/series?apikey="+config.Sonarr.APIKey, "application/json", bytes.NewBuffer(seriesAddJSON)) + if err != nil { + return err.Error() + } + + if seriesAddReq.StatusCode != 201 { + return "There was an error processing this request." + } + return fmt.Sprintf("%s has been queued for download.", series[0].Title) + +} + +// downloadHandler returns the proper string that should be shown in the Telegram response. +func downloadedHandler(downloaded bool) string { + if downloaded { + return "* " + } + + return "" +} + +// seasonHandler returns the proper string that should be shown in the Telegram response. +// it is also used to prevent non-admins from downloading shows with a high number of seasons. +func seasonHandler(seasonCount int, seasonLimit int) string { + if seasonLimit == 0 { + return "" + } + + if seasonCount >= seasonLimit { + return "+ " + } + + return "" +} diff --git a/pkg/service/telegram/admin.go b/pkg/service/telegram/admin.go index 5f8e4aa..03f8876 100644 --- a/pkg/service/telegram/admin.go +++ b/pkg/service/telegram/admin.go @@ -1,31 +1,36 @@ package telegram import ( - "fmt" "strconv" "github.com/yanzay/tbot/v2" ) func (tb *Bot) myID(m *tbot.Message) { - if tb.AdminCheck(m) { + if tb.adminCheck(m.From.ID, false) { tb.Client.SendMessage(m.Chat.ID, strconv.Itoa(m.From.ID)) - fmt.Println(m.From.ID) + return } + + tb.Client.SendMessage(m.Chat.ID, "You are not an authorized admin.") } func (tb *Bot) chatID(m *tbot.Message) { - tb.Client.SendMessage(m.Chat.ID, m.Chat.ID) + if tb.adminCheck(m.From.ID, false) { + tb.Client.SendMessage(m.Chat.ID, m.Chat.ID) + return + } + + tb.Client.SendMessage(m.Chat.ID, "You are not an authorized admin.") } -// AdminCheck checks for valid bot admins. -func (tb *Bot) AdminCheck(m *tbot.Message) bool { +// adminCheck checks for valid bot admins. +func (tb *Bot) adminCheck(id int, callback bool) bool { for _, admin := range tb.Config.Telegram.Admins { - if m.From.ID == admin { + if 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 index 3487bf5..81f77f0 100644 --- a/pkg/service/telegram/sonarr.go +++ b/pkg/service/telegram/sonarr.go @@ -2,22 +2,71 @@ package telegram import ( "fmt" + "strings" "github.com/mattburchett/go_telegram/pkg/service/sonarr" "github.com/yanzay/tbot/v2" ) +// Sonarr Search +func (tb *Bot) sonarrSearch(m *tbot.Message) { + 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.") + return + } + + request, err := sonarr.Search(m, tb.Config) + if err != nil { + tb.Client.SendMessage(m.Chat.ID, err.Error()) + return + } + + inlineResponse := make([][]tbot.InlineKeyboardButton, 0) + for _, i := range request { + inlineResponse = append(inlineResponse, []tbot.InlineKeyboardButton{{ + Text: i.Button, + CallbackData: "tv_" + i.Callback, + }}) + } + + 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 + +} + +// sonarrAdd will perform the add requests to Sonarr. +func (tb *Bot) sonarrAdd(cq *tbot.CallbackQuery) { + if strings.Contains(cq.Data, "+") { + if tb.adminCheck(cq.From.ID, true) { + tb.Client.SendMessage(tb.CallbackChatID, sonarr.Add(cq.Data, tb.Config)) + } else { + tb.Client.AnswerCallbackQuery(cq.ID, tbot.OptText("This request is over the season limit.")) + } + return + } + + tb.Client.SendMessage(tb.CallbackChatID, sonarr.Add(cq.Data, tb.Config)) + +} + +// Admin Functions + +// sonarrStatus queries Sonarr for it's system status information. func (tb *Bot) sonarrStatus(m *tbot.Message) { - if tb.AdminCheck(m) { - request, err := sonarr.SonarrStatus(m, tb.Config) + if tb.adminCheck(m.From.ID, false) { + request, err := sonarr.Status(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) + return } + + 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 index 2d9605e..725ae33 100644 --- a/pkg/service/telegram/telegram.go +++ b/pkg/service/telegram/telegram.go @@ -2,6 +2,7 @@ package telegram import ( "log" + "strings" "time" "github.com/mattburchett/go_telegram/pkg/core/config" @@ -17,6 +18,15 @@ type Bot struct { CallbackMessageID int } +// 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)) + } +} + // New creates an active telegram bot and loads the handlers. func (tb *Bot) New(token string) { tb.Bot = tbot.New(token) @@ -35,7 +45,8 @@ func (tb *Bot) Handler() { tb.Client.SendMessage(m.Chat.ID, "pong") }) - // sonarr/admin.go + // telegram/sonar.go + tb.Bot.HandleMessage("/s", tb.sonarrSearch) tb.Bot.HandleMessage("/admin sonarrStatus", tb.sonarrStatus) // telegram/testhandler.go @@ -53,19 +64,20 @@ func (tb *Bot) 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.") + 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) {