Sync from /srv/compose/unified-media-manager
This commit is contained in:
Christopher Mayor
426
internal/service/search.go
Normal file
426
internal/service/search.go
Normal file
@@ -0,0 +1,426 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user