Add token management

This commit is contained in:
Ben Adrian Sarmiento
2024-06-28 13:19:09 +02:00
parent 962845fb81
commit c3aea427d0
9 changed files with 100 additions and 114 deletions

View File

@@ -11,18 +11,19 @@ import (
"github.com/debridmediamanager/zurg/internal/config"
zurghttp "github.com/debridmediamanager/zurg/pkg/http"
"github.com/debridmediamanager/zurg/pkg/logutil"
"github.com/debridmediamanager/zurg/pkg/utils"
cmap "github.com/orcaman/concurrent-map/v2"
"github.com/panjf2000/ants/v2"
)
type RealDebrid struct {
torrentsCache []Torrent
UnrestrictMap cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Download]]
TokenManager *DownloadTokenManager
torrentsCache []Torrent
verifiedLinks cmap.ConcurrentMap[string, int64]
apiClient *zurghttp.HTTPClient
unrestrictClient *zurghttp.HTTPClient
downloadClient *zurghttp.HTTPClient
tokenManager *DownloadTokenManager
workerPool *ants.Pool
cfg config.ConfigInterface
log *logutil.Logger
@@ -36,13 +37,13 @@ func NewRealDebrid(apiClient, unrestrictClient, downloadClient *zurghttp.HTTPCli
}
rd := &RealDebrid{
torrentsCache: []Torrent{},
UnrestrictMap: cmap.New[cmap.ConcurrentMap[string, *Download]](),
TokenManager: NewDownloadTokenManager(downloadTokens),
torrentsCache: []Torrent{},
verifiedLinks: cmap.New[int64](),
apiClient: apiClient,
unrestrictClient: unrestrictClient,
downloadClient: downloadClient,
tokenManager: NewDownloadTokenManager(downloadTokens),
workerPool: workerPool,
cfg: cfg,
log: log,
@@ -59,76 +60,59 @@ func NewRealDebrid(apiClient, unrestrictClient, downloadClient *zurghttp.HTTPCli
return rd
}
// currently unused
func (rd *RealDebrid) UnrestrictCheck(link string) (*Download, error) {
data := url.Values{}
data.Set("link", link)
requestBody := strings.NewReader(data.Encode())
req, err := http.NewRequest(http.MethodPost, "https://api.real-debrid.com/rest/1.0/unrestrict/check", requestBody)
if err != nil {
rd.log.Errorf("Error when creating a unrestrict check request: %v", err)
return nil, err
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, err := rd.unrestrictClient.Do(req)
if err != nil {
rd.log.Errorf("Error when executing the unrestrict check request: %v", err)
return nil, err
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
rd.log.Errorf("Error when reading the body of unrestrict check response: %v", err)
return nil, err
}
var response Download
err = json.Unmarshal(body, &response)
if err != nil {
rd.log.Errorf("Error when decoding unrestrict check JSON: %v", err)
return nil, err
}
rd.log.Debugf("Link %s is streamable? %v", response.Streamable)
return &response, nil
}
func (rd *RealDebrid) UnrestrictLink(link string) (*Download, error) {
func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {
for {
token, err := rd.tokenManager.GetCurrentToken()
token, err := rd.TokenManager.GetCurrentToken()
if err != nil {
// when all tokens are expired
return nil, err
}
download, err := rd.UnrestrictLinkWithToken(token, link)
if dlErr, ok := err.(*zurghttp.DownloadErrorResponse); ok && dlErr.Message == "bytes_limit_reached" {
rd.tokenManager.SetCurrentTokenExpired()
// check if the link is already unrestricted
tokenMap, ok := rd.UnrestrictMap.Get(token)
if ok && tokenMap.Has(link) {
download, _ := tokenMap.Get(link)
// check if the link is in the verified links cache
if expiry, ok := rd.verifiedLinks.Get(download.ID); ok && expiry > time.Now().Unix() {
return download, nil
}
err := rd.downloadClient.VerifyLink(download.Download)
if utils.IsBWLimitExceeded(err) {
rd.TokenManager.SetTokenAsExpired(token)
continue
}
if err != nil {
return nil, err
}
rd.verifiedLinks.Set(download.ID, time.Now().Unix()+60*60*24)
return download, nil
}
download, err := rd.UnrestrictLink(link)
if err != nil {
return nil, err
}
tokenMap.Set(link, download)
err = rd.downloadClient.VerifyLink(download.Download)
if utils.IsBWLimitExceeded(err) {
rd.TokenManager.SetTokenAsExpired(token)
continue
}
if err != nil {
return nil, err
}
rd.verifiedLinks.Set(download.ID, time.Now().Unix()+60*60*24)
return download, err
}
}
func (rd *RealDebrid) UnrestrictLinkWithToken(token, link string) (*Download, error) {
// check if the link is already unrestricted
if tokenMap, ok := rd.UnrestrictMap.Get(token); ok {
if d, ok := tokenMap.Get(link); ok {
// check if the link is in the verified links cache
if expiry, ok := rd.verifiedLinks.Get(d.Download); ok && expiry > time.Now().Unix() {
return d, nil
}
err := rd.downloadClient.VerifyLink(d.Download)
if err != nil {
return nil, err
}
rd.verifiedLinks.Set(d.Download, time.Now().Unix()+60*60*24)
return d, nil
}
}
func (rd *RealDebrid) UnrestrictLink(link string) (*Download, error) {
data := url.Values{}
if strings.HasPrefix(link, "https://real-debrid.com/d/") {
// set link to max 39 chars (26 + 13)
@@ -145,7 +129,6 @@ func (rd *RealDebrid) UnrestrictLinkWithToken(token, link string) (*Download, er
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
rd.unrestrictClient.SetToken(token)
// at this point, any errors mean that the link has expired and we need to repair it
resp, err := rd.unrestrictClient.Do(req)
if err != nil {
@@ -167,16 +150,6 @@ func (rd *RealDebrid) UnrestrictLinkWithToken(token, link string) (*Download, er
return nil, fmt.Errorf("undecodable response: %v", err)
}
tokenMap, _ := rd.UnrestrictMap.Get(token)
tokenMap.Set(link, &response)
err = rd.downloadClient.VerifyLink(response.Download)
if err != nil {
return nil, err
}
rd.verifiedLinks.Set(response.Download, time.Now().Unix()+60*60*24)
// rd.log.Debugf("Unrestricted link %s into %s", link, response.Download)
return &response, nil
}
@@ -428,6 +401,6 @@ func (rd *RealDebrid) AvailabilityCheck(hashes []string) (AvailabilityResponse,
return response, nil
}
func (rd *RealDebrid) GetToken() (string, error) {
return rd.tokenManager.GetCurrentToken()
func (rd *RealDebrid) DownloadFile(req *http.Request) (*http.Response, error) {
return rd.downloadClient.Do(req)
}