427 lines
11 KiB
Go
427 lines
11 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"encoding/xml"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/url"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/TopherMayor/unified-media-manager/internal/cardigann"
|
|
"golang.org/x/sync/errgroup"
|
|
)
|
|
|
|
type rssFeed struct {
|
|
XMLName xml.Name `xml:"rss"`
|
|
Channel rssChannel `xml:"channel"`
|
|
}
|
|
|
|
type rssChannel struct {
|
|
Items []rssItem `xml:"item"`
|
|
}
|
|
|
|
type rssItem struct {
|
|
Title string `xml:"title"`
|
|
GUID string `xml:"guid"`
|
|
Link string `xml:"link"`
|
|
PubDate string `xml:"pubDate"`
|
|
Description string `xml:"description"`
|
|
Enclosure *rssEnclosure `xml:"enclosure"`
|
|
Attrs []rssAttr `xml:"attr"`
|
|
Size string `xml:"size"`
|
|
}
|
|
|
|
type rssEnclosure struct {
|
|
URL string `xml:"url,attr"`
|
|
Length string `xml:"length,attr"`
|
|
Type string `xml:"type,attr"`
|
|
}
|
|
|
|
type rssAttr struct {
|
|
Name string `xml:"name,attr"`
|
|
Value string `xml:"value,attr"`
|
|
}
|
|
|
|
type SearchResult struct {
|
|
Title string `json:"title"`
|
|
GUID string `json:"guid"`
|
|
Link string `json:"link"`
|
|
Size int64 `json:"size"`
|
|
PubDate string `json:"pub_date"`
|
|
IndexerName string `json:"indexer_name"`
|
|
IndexerPriority int `json:"indexer_priority"`
|
|
Quality ReleaseInfo `json:"quality"`
|
|
QualityTier *QualityTier `json:"quality_tier"`
|
|
Seeders int `json:"seeders"`
|
|
Peers int `json:"peers"`
|
|
Category string `json:"category"`
|
|
DownloadURL string `json:"download_url"`
|
|
SourceIndexers []string `json:"source_indexers"`
|
|
}
|
|
|
|
type SearchRequest struct {
|
|
Query string `json:"query"`
|
|
MediaType string `json:"media_type"`
|
|
IndexerIDs []int64 `json:"indexer_ids,omitempty"`
|
|
}
|
|
|
|
type SearchService struct {
|
|
indexerSvc *IndexerService
|
|
parser *ReleaseParser
|
|
cardigannEngine *cardigann.CardigannEngine
|
|
httpClient *http.Client
|
|
}
|
|
|
|
func NewSearchService(indexerSvc *IndexerService, parser *ReleaseParser, cardigannEngine *cardigann.CardigannEngine) *SearchService {
|
|
transport := &http.Transport{
|
|
MaxIdleConns: 20,
|
|
MaxIdleConnsPerHost: 5,
|
|
IdleConnTimeout: 90 * time.Second,
|
|
DisableKeepAlives: false,
|
|
}
|
|
return &SearchService{
|
|
indexerSvc: indexerSvc,
|
|
parser: parser,
|
|
cardigannEngine: cardigannEngine,
|
|
httpClient: &http.Client{Timeout: 30 * time.Second, Transport: transport},
|
|
}
|
|
}
|
|
|
|
func (s *SearchService) Search(ctx context.Context, req SearchRequest) ([]SearchResult, error) {
|
|
indexers, err := s.indexerSvc.ListEnabled(ctx)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("get enabled indexers: %w", err)
|
|
}
|
|
|
|
if len(indexers) == 0 {
|
|
return nil, fmt.Errorf("no enabled indexers available")
|
|
}
|
|
|
|
if len(req.IndexerIDs) > 0 {
|
|
idSet := make(map[int64]bool, len(req.IndexerIDs))
|
|
for _, id := range req.IndexerIDs {
|
|
idSet[id] = true
|
|
}
|
|
var filtered []Indexer
|
|
for _, idx := range indexers {
|
|
if idSet[idx.ID] {
|
|
filtered = append(filtered, idx)
|
|
}
|
|
}
|
|
indexers = filtered
|
|
}
|
|
|
|
if len(indexers) == 0 {
|
|
return nil, fmt.Errorf("no matching indexers found")
|
|
}
|
|
|
|
category := MediaTypeToCategory(req.MediaType)
|
|
|
|
searchCtx, searchCancel := context.WithTimeout(ctx, 30*time.Second)
|
|
defer searchCancel()
|
|
|
|
var mu sync.Mutex
|
|
var allResults []SearchResult
|
|
|
|
g, gCtx := errgroup.WithContext(searchCtx)
|
|
g.SetLimit(10)
|
|
|
|
for _, idx := range indexers {
|
|
idx := idx
|
|
g.Go(func() error {
|
|
results, err := s.searchIndexer(gCtx, idx, req.Query, category)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
|
|
mu.Lock()
|
|
allResults = append(allResults, results...)
|
|
mu.Unlock()
|
|
|
|
return nil
|
|
})
|
|
}
|
|
|
|
_ = g.Wait()
|
|
|
|
merged := s.mergeAndDedup(allResults)
|
|
|
|
sort.Slice(merged, func(i, j int) bool {
|
|
ti := 0
|
|
tj := 0
|
|
if merged[i].QualityTier != nil {
|
|
ti = merged[i].QualityTier.Rank
|
|
}
|
|
if merged[j].QualityTier != nil {
|
|
tj = merged[j].QualityTier.Rank
|
|
}
|
|
if ti != tj {
|
|
return ti > tj
|
|
}
|
|
return merged[i].Size < merged[j].Size
|
|
})
|
|
|
|
return merged, nil
|
|
}
|
|
|
|
func (s *SearchService) searchIndexer(ctx context.Context, idx Indexer, query string, category string) ([]SearchResult, error) {
|
|
// Dispatch to Cardigann engine for cardigann implementation
|
|
if idx.Implementation == "cardigann" {
|
|
return s.searchCardigannIndexer(ctx, idx, query)
|
|
}
|
|
|
|
searchURL := fmt.Sprintf("%s/api?t=search&q=%s", idx.URL, url.QueryEscape(query))
|
|
|
|
apiKey := ""
|
|
if idx.APIKey != nil {
|
|
apiKey = *idx.APIKey
|
|
}
|
|
if apiKey != "" {
|
|
searchURL += "&apikey=" + apiKey
|
|
}
|
|
|
|
if category != "" {
|
|
searchURL += "&cat=" + category
|
|
}
|
|
|
|
if idx.Implementation == "torznab" {
|
|
searchURL += "&extended=1"
|
|
}
|
|
|
|
searchURL += "&offset=0&limit=100"
|
|
|
|
reqCtx, reqCancel := context.WithTimeout(ctx, 15*time.Second)
|
|
defer reqCancel()
|
|
|
|
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodGet, searchURL, nil)
|
|
if err != nil {
|
|
go s.indexerSvc.RecordFailure(context.Background(), idx.ID)
|
|
return nil, nil
|
|
}
|
|
|
|
client := s.httpClient
|
|
resp, err := client.Do(httpReq)
|
|
if err != nil {
|
|
go s.indexerSvc.RecordFailure(context.Background(), idx.ID)
|
|
return nil, nil
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode >= 400 {
|
|
go s.indexerSvc.RecordFailure(context.Background(), idx.ID)
|
|
return nil, nil
|
|
}
|
|
|
|
body, err := io.ReadAll(io.LimitReader(resp.Body, 10*1024*1024))
|
|
if err != nil {
|
|
go s.indexerSvc.RecordFailure(context.Background(), idx.ID)
|
|
return nil, nil
|
|
}
|
|
|
|
var feed rssFeed
|
|
if err := xml.Unmarshal(body, &feed); err != nil {
|
|
go s.indexerSvc.RecordFailure(context.Background(), idx.ID)
|
|
slog.Error("failed to parse indexer XML response", "indexer", idx.Name, "error", err)
|
|
return nil, nil
|
|
}
|
|
|
|
go s.indexerSvc.RecordSuccess(context.Background(), idx.ID)
|
|
|
|
var results []SearchResult
|
|
for _, item := range feed.Channel.Items {
|
|
result := s.parseItem(item, idx)
|
|
if result.DownloadURL != "" {
|
|
results = append(results, result)
|
|
}
|
|
}
|
|
|
|
return results, nil
|
|
}
|
|
|
|
func (s *SearchService) parseItem(item rssItem, idx Indexer) SearchResult {
|
|
result := SearchResult{
|
|
Title: item.Title,
|
|
GUID: item.GUID,
|
|
Link: item.Link,
|
|
PubDate: item.PubDate,
|
|
IndexerName: idx.Name,
|
|
IndexerPriority: idx.Priority,
|
|
SourceIndexers: []string{idx.Name},
|
|
}
|
|
|
|
if item.Enclosure != nil && item.Enclosure.URL != "" {
|
|
result.DownloadURL = item.Enclosure.URL
|
|
if item.Enclosure.Length != "" {
|
|
if size, err := strconv.ParseInt(item.Enclosure.Length, 10, 64); err == nil {
|
|
result.Size = size
|
|
}
|
|
}
|
|
} else if item.Link != "" {
|
|
result.DownloadURL = item.Link
|
|
}
|
|
|
|
if result.Size == 0 && item.Size != "" {
|
|
if size, err := strconv.ParseInt(item.Size, 10, 64); err == nil {
|
|
result.Size = size
|
|
}
|
|
}
|
|
|
|
for _, attr := range item.Attrs {
|
|
switch attr.Name {
|
|
case "size":
|
|
if result.Size == 0 {
|
|
if size, err := strconv.ParseInt(attr.Value, 10, 64); err == nil {
|
|
result.Size = size
|
|
}
|
|
}
|
|
case "seeders":
|
|
if v, err := strconv.Atoi(attr.Value); err == nil {
|
|
result.Seeders = v
|
|
}
|
|
case "peers":
|
|
if v, err := strconv.Atoi(attr.Value); err == nil {
|
|
result.Peers = v
|
|
}
|
|
case "category":
|
|
result.Category = attr.Value
|
|
}
|
|
}
|
|
|
|
quality := s.parser.Parse(item.Title)
|
|
result.Quality = quality
|
|
result.QualityTier = s.parser.MatchQuality(quality)
|
|
|
|
return result
|
|
}
|
|
|
|
func (s *SearchService) searchCardigannIndexer(ctx context.Context, idx Indexer, query string) ([]SearchResult, error) {
|
|
cfg, err := s.indexerSvc.GetCardigannConfig(idx.Settings)
|
|
if err != nil {
|
|
slog.Error("failed to get cardigann config", "indexer", idx.Name, "error", err)
|
|
go s.indexerSvc.RecordFailure(context.Background(), idx.ID)
|
|
return nil, nil
|
|
}
|
|
|
|
def, err := cardigann.ParseDefinition([]byte(cfg.YAML))
|
|
if err != nil {
|
|
slog.Error("failed to parse cardigann YAML", "indexer", idx.Name, "error", err)
|
|
go s.indexerSvc.RecordFailure(context.Background(), idx.ID)
|
|
return nil, nil
|
|
}
|
|
|
|
results, err := s.cardigannEngine.Search(ctx, def, cfg.Config, cardigann.SearchQuery{
|
|
Keywords: query,
|
|
})
|
|
if err != nil {
|
|
slog.Error("cardigann search failed", "indexer", idx.Name, "error", err)
|
|
go s.indexerSvc.RecordFailure(context.Background(), idx.ID)
|
|
return nil, nil
|
|
}
|
|
|
|
go s.indexerSvc.RecordSuccess(context.Background(), idx.ID)
|
|
|
|
var searchResults []SearchResult
|
|
for _, cr := range results {
|
|
result := SearchResult{
|
|
Title: cr.Title,
|
|
GUID: cr.GUID,
|
|
DownloadURL: cr.DownloadURL,
|
|
Size: cr.Size,
|
|
PubDate: cr.PubDate,
|
|
IndexerName: idx.Name,
|
|
IndexerPriority: idx.Priority,
|
|
SourceIndexers: []string{idx.Name},
|
|
Seeders: cr.Seeders,
|
|
Peers: cr.Peers,
|
|
Category: cr.Category,
|
|
}
|
|
result.Quality = s.parser.Parse(cr.Title)
|
|
result.QualityTier = s.parser.MatchQuality(result.Quality)
|
|
searchResults = append(searchResults, result)
|
|
}
|
|
|
|
return searchResults, nil
|
|
}
|
|
|
|
func (s *SearchService) mergeAndDedup(results []SearchResult) []SearchResult {
|
|
seen := make(map[string]int)
|
|
var merged []SearchResult
|
|
|
|
for _, r := range results {
|
|
guid := strings.ToLower(r.GUID)
|
|
if guid == "" {
|
|
guid = strings.ToLower(r.DownloadURL)
|
|
}
|
|
if guid == "" {
|
|
merged = append(merged, r)
|
|
continue
|
|
}
|
|
|
|
if existingIdx, ok := seen[guid]; ok {
|
|
existing := &merged[existingIdx]
|
|
if r.IndexerPriority < existing.IndexerPriority {
|
|
existing.SourceIndexers = append(existing.SourceIndexers, r.IndexerName)
|
|
} else {
|
|
newSources := make([]string, 0, len(existing.SourceIndexers)+1)
|
|
newSources = append(newSources, r.IndexerName)
|
|
newSources = append(newSources, existing.SourceIndexers...)
|
|
r.SourceIndexers = newSources
|
|
merged[existingIdx] = r
|
|
}
|
|
} else {
|
|
seen[guid] = len(merged)
|
|
merged = append(merged, r)
|
|
}
|
|
}
|
|
|
|
return merged
|
|
}
|
|
|
|
type GrabRequest struct {
|
|
DownloadURL string `json:"download_url"`
|
|
Title string `json:"title"`
|
|
MediaType string `json:"media_type"`
|
|
Quality ReleaseInfo `json:"quality"`
|
|
IndexerName string `json:"indexer_name"`
|
|
MediaID int64 `json:"media_id"`
|
|
}
|
|
|
|
type GrabResult struct {
|
|
QueueID int64 `json:"queue_id"`
|
|
DownloadID string `json:"download_id"`
|
|
ClientName string `json:"client_name"`
|
|
Protocol string `json:"protocol"`
|
|
}
|
|
|
|
func (s *SearchService) Grab(ctx context.Context, req GrabRequest, downloadClientSvc *DownloadClientService) (*GrabResult, error) {
|
|
protocol := "torrent"
|
|
if strings.HasPrefix(req.DownloadURL, "magnet:?") {
|
|
protocol = "torrent"
|
|
} else if strings.HasSuffix(strings.ToLower(req.DownloadURL), ".nzb") {
|
|
protocol = "nzb"
|
|
}
|
|
|
|
client, cfg, err := downloadClientSvc.GetClient(ctx, protocol)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("get download client: %w", err)
|
|
}
|
|
|
|
downloadID, err := client.Add(ctx, req.DownloadURL, "umm")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("add download: %w", err)
|
|
}
|
|
|
|
return &GrabResult{
|
|
DownloadID: downloadID,
|
|
ClientName: cfg.Name,
|
|
Protocol: cfg.Protocol,
|
|
}, nil
|
|
}
|