Adding Sonarr support into the code base.

This commit is contained in:
Matt Burchett 2020-05-17 21:06:52 -05:00
parent d6ed947bb9
commit 4e846885dd
6 changed files with 342 additions and 29 deletions

View File

@ -16,8 +16,10 @@ type Config struct {
} `json:"telegram"` } `json:"telegram"`
Sonarr struct { Sonarr struct {
URL string `json:"url"` URL string `json:"url"`
APIKey string `json:"apiKey"` APIKey string `json:"apiKey"`
SeasonLimit int `json:"seasonLimit"`
ProfileID int `json:"proFileId"`
} `json:"sonarr"` } `json:"sonarr"`
} }

View File

@ -8,8 +8,8 @@ import (
"github.com/yanzay/tbot/v2" "github.com/yanzay/tbot/v2"
) )
// SonarrStatus contains the Sonarr request for system status. // Status contains the Sonarr request for system status.
func SonarrStatus(m *tbot.Message, config config.Config) (string, error) { func Status(m *tbot.Message, config config.Config) (string, error) {
r, err := http.Get(config.Sonarr.URL + "system/status?apikey=" + config.Sonarr.APIKey) r, err := http.Get(config.Sonarr.URL + "system/status?apikey=" + config.Sonarr.APIKey)
if err != nil { if err != nil {
return "Failed to contact Sonarr for data", err return "Failed to contact Sonarr for data", err

View File

@ -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 ""
}

View File

@ -1,31 +1,36 @@
package telegram package telegram
import ( import (
"fmt"
"strconv" "strconv"
"github.com/yanzay/tbot/v2" "github.com/yanzay/tbot/v2"
) )
func (tb *Bot) myID(m *tbot.Message) { 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)) 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) { 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. // adminCheck checks for valid bot admins.
func (tb *Bot) AdminCheck(m *tbot.Message) bool { func (tb *Bot) adminCheck(id int, callback bool) bool {
for _, admin := range tb.Config.Telegram.Admins { for _, admin := range tb.Config.Telegram.Admins {
if m.From.ID == admin { if id == admin {
return true return true
} }
} }
tb.Client.SendMessage(m.Chat.ID, "You are not an authorized admin.")
return false return false
} }

View File

@ -2,22 +2,71 @@ package telegram
import ( import (
"fmt" "fmt"
"strings"
"github.com/mattburchett/go_telegram/pkg/service/sonarr" "github.com/mattburchett/go_telegram/pkg/service/sonarr"
"github.com/yanzay/tbot/v2" "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) { func (tb *Bot) sonarrStatus(m *tbot.Message) {
if tb.AdminCheck(m) { if tb.adminCheck(m.From.ID, false) {
request, err := sonarr.SonarrStatus(m, tb.Config) request, err := sonarr.Status(m, tb.Config)
if err != nil { if err != nil {
tb.Client.SendMessage(m.Chat.ID, fmt.Sprintf("%v: \n %v", request, err)) tb.Client.SendMessage(m.Chat.ID, fmt.Sprintf("%v: \n %v", request, err))
} else { return
tb.Client.SendMessage(m.Chat.ID, "Sonarr Status:")
tb.Client.SendMessage(m.Chat.ID, request)
} }
tb.Client.SendMessage(m.Chat.ID, "Sonarr Status:")
tb.Client.SendMessage(m.Chat.ID, request)
} }
} }

View File

@ -2,6 +2,7 @@ package telegram
import ( import (
"log" "log"
"strings"
"time" "time"
"github.com/mattburchett/go_telegram/pkg/core/config" "github.com/mattburchett/go_telegram/pkg/core/config"
@ -17,6 +18,15 @@ type Bot struct {
CallbackMessageID int 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. // New creates an active telegram bot and loads the handlers.
func (tb *Bot) New(token string) { func (tb *Bot) New(token string) {
tb.Bot = tbot.New(token) tb.Bot = tbot.New(token)
@ -35,7 +45,8 @@ func (tb *Bot) Handler() {
tb.Client.SendMessage(m.Chat.ID, "pong") 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) tb.Bot.HandleMessage("/admin sonarrStatus", tb.sonarrStatus)
// telegram/testhandler.go // telegram/testhandler.go
@ -53,19 +64,20 @@ func (tb *Bot) Handler() {
tb.Bot.HandleCallback(tb.callbackHandler) 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. // callbackHandler handles callbacks.
func (tb *Bot) callbackHandler(cq *tbot.CallbackQuery) { 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) tb.Client.SendMessage(tb.CallbackChatID, cq.Data)
} }
func (tb *Bot) helpHandler(m *tbot.Message) { func (tb *Bot) helpHandler(m *tbot.Message) {