Create simple torznab api module. Start asyncifying the indexer requests

This commit is contained in:
seanomik 2022-06-17 23:57:27 -04:00
parent 2a225ac2db
commit 6741d35783
11 changed files with 2980 additions and 2033 deletions

4183
Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@ -7,11 +7,23 @@ edition = "2021"
[dependencies]
tokio = { version = "1.19.2", features = ["full"] }
tracing = "0.1.35"
tracing-subscriber = "0.3.11"
futures = "0.3.21"
toml = "0.5.9"
lava_torrent = "0.7.0" # https://docs.rs/lava_torrent/0.7.0/lava_torrent/
torznab = "0.7.2" # https://docs.rs/torznab/0.7.2/torznab/
magnet-url = "2.0.0"
serde_with = "1.14.0"
serde = { version = "1.0", features = ["derive"] }
figment = { version = "0.10", features = ["toml", "env"] }
wild = "2.0.4"
argmap = "1.1.2"
argmap = "1.1.2"
reqwest = {version = "0.11", default_features = false, features = ["gzip", "json", "rustls-tls"]}
urlencoding = "2.1.0"
# Torznab stuff
rss = "2.0.1"
bytes = "1.1.0"
quick-xml = {version = "0.23.0", features = ["serialize"]}

View File

@ -5,6 +5,8 @@ use std::collections::HashMap;
use figment::{Figment, providers::{Format, Toml, Env}};
use figment::value::Value as FigmentValue;
use crate::torznab::TorznabClient;
use super::CliProvider;
#[derive(Deserialize, Serialize)]
@ -26,13 +28,26 @@ pub struct Config {
pub indexers: Vec<Indexer>,
}
#[derive(Deserialize, Serialize)]
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Indexer {
#[serde(skip_deserializing)]
pub name: String,
pub enabled: Option<bool>,
pub url: String,
pub api_key: String,
#[serde(skip)]
pub client: Option<TorznabClient>,
}
impl Indexer {
pub async fn create_client(&mut self) -> Result<&TorznabClient, crate::torznab::ClientError> {
if self.client.is_none() {
self.client = Some(TorznabClient::new(self.name.clone(), &self.url, &self.api_key).await?);
}
Ok(self.client.as_ref().unwrap())
}
}
// Allow dead code for functions. We should probably remove this later on.

View File

@ -1,13 +1,20 @@
mod config;
mod torznab;
use config::Config;
use tracing::{info, Level, debug};
use std::path::{Path, PathBuf};
use std::error::Error;
use lava_torrent::torrent::v1::Torrent;
use torznab::Client as TorznabClient;
use crate::torznab::{GenericSearchParameters, SearchFunction};
use crate::torznab::search_parameters::{GenericSearchParametersBuilder, MovieSearchParametersBuilder};
use tokio::sync::RwLock;
use std::sync::Arc;
fn read_torrents(path: &Path) -> Result<Vec<PathBuf>, Box<dyn Error>> {
let mut torrents = Vec::new();
@ -28,48 +35,67 @@ fn read_torrents(path: &Path) -> Result<Vec<PathBuf>, Box<dyn Error>> {
return Ok(torrents);
}
fn main() {
#[tokio::main]
async fn main() {
let subscriber = tracing_subscriber::fmt()
.with_max_level(Level::INFO)
.finish();
tracing::subscriber::set_global_default(subscriber).expect("Failed to set global default log subscriber");
// Get config and debug the torrents
let config = Config::new();//.expect("Failed to get config");
println!("Searching torrents in: {}", config.torrents_path_str());
let config = Config::new();
info!("Searching torrents in: {}", config.torrents_path_str());
println!("Searching {} trackers: ", config.indexers.len());
for indexer in config.indexers.iter() {
println!(" {}: {}", indexer.name, indexer.url);
let mut indexers = config.indexers.clone();
// Create torznab clients for each indexer.
for indexer in indexers.iter_mut() {
indexer.create_client().await.unwrap();
}
let torrents = read_torrents(config.torrents_path()).unwrap();
// Log the trackers
info!("Searching {} trackers: ", indexers.len());
for indexer in indexers.iter() {
info!(" {}: {}", indexer.name, indexer.url);
debug!(" Can Search: {:?}", indexer.client.as_ref().unwrap().capabilities.searching_capabilities);
}
for torrent_path in torrents.iter() {
let torrent = Torrent::read_from_file(torrent_path).unwrap();
println!("{}:", torrent.name);
let torrent_files = read_torrents(config.torrents_path()).unwrap();
info!("Found {} torrents", torrent_files.len());
/* for indexer in config.indexers.iter() {
if indexer.enabled {
let client = TorznabClient::new(indexer.url.clone());
let results = client.search(&torrent).unwrap();
println!("{}", results);
}
} */
//TorznabClient
//panic!("rhfhujergfre");
/*if let Some(announce) = torrent.announce {
println!(" Announce: {}", announce);
}
if let Some(announce_list) = torrent.announce_list {
println!(" Announce list:");
for announce in announce_list {
for ann in announce {
println!(" {}", ann);
// Convert the indexers to be async friendly.
let mut indexers = indexers.iter()
.map(|indexer| Arc::new(RwLock::new(indexer.clone())))
.collect::<Vec<_>>();
let mut indexer_handles = vec![];
for torrent_path in torrent_files.iter() {
let torrent = Arc::new(Torrent::read_from_file(torrent_path).unwrap());
info!("{}:", torrent.name);
for indexer in indexers.iter() {
let mut indexer = Arc::clone(indexer);
let torrent = Arc::clone(&torrent);
indexer_handles.push(tokio::spawn(async move {
let lock = indexer.read().await;
match &lock.client {
Some(client) => {
let generic = GenericSearchParametersBuilder::new()
.query(torrent.name.clone())
.build();
client.search(SearchFunction::Search, generic).await.unwrap();
},
None => {
panic!("idfk");
}
}
}
}));
}
println!(" Files:");
if let Some(files) = torrent.files {
for file in files.iter() {
println!(" {}", file.path.to_str().unwrap());
}
} */
}
futures::future::join_all(indexer_handles).await;
}

174
src/torznab/capabilities.rs Normal file
View File

@ -0,0 +1,174 @@
use std::collections::HashMap;
use serde::{Serialize, Deserialize};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SupportedParam {
Query,
Season,
Episode,
IMDB,
TMDB,
TVDB,
}
impl From<String> for SupportedParam {
fn from(s: String) -> Self {
match s.as_str() {
"q" => SupportedParam::Query,
"season" => SupportedParam::Season,
"ep" => SupportedParam::Episode,
"imdbid" => SupportedParam::IMDB,
"tmdbid" => SupportedParam::TMDB,
"tvdbid" => SupportedParam::TVDB,
_ => panic!("Unsupported param: {}", s),
}
}
}
#[derive(Debug, Hash, PartialEq, Eq, Clone, Deserialize)]
pub enum SearchCapability {
#[serde(rename = "search")]
Search,
#[serde(rename = "tv-search")]
TV,
#[serde(rename = "movie-search")]
Movie,
#[serde(rename = "music-search")]
Music,
#[serde(rename = "audio-search")]
Audio,
#[serde(rename = "book-search")]
Book
}
impl From<String> for SearchCapability {
fn from(s: String) -> Self {
match s.as_str() {
"search" => SearchCapability::Search,
"tv-search" => SearchCapability::TV,
"movie-search" => SearchCapability::Movie,
"music-search" => SearchCapability::Music,
"audio-search" => SearchCapability::Audio,
"book" => SearchCapability::Book,
_ => panic!("Unsupported param: {}", s),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SearchingCapabilities {
supported_functions: HashMap<SearchCapability, Vec<SupportedParam>>,
}
impl SearchingCapabilities {
pub fn new(supported_functions: HashMap<SearchCapability, Vec<SupportedParam>>) -> SearchingCapabilities {
SearchingCapabilities {
supported_functions,
}
}
/// Returns true if the search type is supported.
pub fn does_support_search(&self, search_capability: SearchCapability) -> bool {
self.supported_functions.contains_key(&search_capability)
}
/// Returns true if a search type supports a specific parameter.
pub fn does_search_support_param(&self, capability: SearchCapability, param: SupportedParam) -> bool {
match self.supported_functions.get(&capability) {
Some(params) => params.contains(&param),
None => false,
}
}
}
impl Default for SearchingCapabilities {
fn default() -> Self {
SearchingCapabilities {
supported_functions: HashMap::new(),
}
}
}
impl<'de> Deserialize<'de> for SearchingCapabilities {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let raw: HashMap<SearchCapability, HashMap<String, String>> = Deserialize::deserialize(deserializer).unwrap();
let mut functions: HashMap<SearchCapability, Vec<SupportedParam>> = HashMap::new();
for (key, value) in raw.iter() {
let mut supported_params = Vec::new();
let available_str: String = value.get("available").map(String::to_owned).unwrap_or_default();
let available = available_str == "yes";
if available {
if let Some(params) = value.get("supportedParams") {
for param in params.split(',') {
supported_params.push(param.to_string().into());
}
}
functions.insert(key.clone(), supported_params);
}
}
Ok(SearchingCapabilities {
supported_functions: functions,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Category {
#[serde(with = "serde_with::rust::display_fromstr")]
id: u32,
name: String,
#[serde(rename = "subcat")]
sub_categories: Option<Vec<Category>>,
}
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct Categories {
#[serde(rename = "category")]
pub categories: Vec<Category>,
}
impl Categories {
pub fn new(categories: Vec<Category>) -> Self {
Categories { categories }
}
}
impl Default for Categories {
fn default() -> Self {
Categories {
categories: Vec::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct Capabilities {
pub categories: Categories,
#[serde(rename = "searching")]
pub searching_capabilities: SearchingCapabilities,
}
impl Default for Capabilities {
fn default() -> Self {
Capabilities {
categories: Categories::default(),
searching_capabilities: SearchingCapabilities::default(),
}
}
}

103
src/torznab/client.rs Normal file
View File

@ -0,0 +1,103 @@
use super::{Capabilities, TorznabFunction, SearchFunction, GenericSearchParameters, TorrentResult, ClientError};
use bytes::Bytes;
use bytes::Buf;
use rss::Channel;
use tracing::{span, event, debug, info, Level};
#[derive(Debug, Clone)]
pub struct TorznabClient {
http: reqwest::Client,
pub name: String,
pub base_url: String,
api_key: String,
pub capabilities: Capabilities,
pub client_span: tracing::Span,
}
impl TorznabClient {
fn client_span(name: &String) -> tracing::Span {
span!(Level::INFO, "torznab_client", indexer = %name)
}
/// Construct a new client without getting the capabilities
pub fn new_no_capabilities(name: String, base_url: &str, api_key: &str) -> Self {
TorznabClient {
name: name.clone(),
http: reqwest::Client::new(),
base_url: base_url.to_string(),
api_key: api_key.to_string(),
capabilities: Capabilities::default(),
client_span: Self::client_span(&name),
}
}
/// Construct a new client and get the capabilities.
pub async fn new(name: String, base_url: &str, api_key: &str) -> Result<Self, reqwest::Error> {
let mut client = TorznabClient {
name: name.clone(),
http: reqwest::Client::new(),
base_url: base_url.to_string(),
api_key: api_key.to_string(),
capabilities: Capabilities::default(),
client_span: Self::client_span(&name),
};
// Get capabilities and store them in the client before returning
client.store_capabilities().await?;
Ok(client)
}
/// Send a request to the indexer using the query parameters.
async fn request(&self, param_str: String) -> Result<Bytes, reqwest::Error> {
let span = span!(parent: &self.client_span, Level::INFO, "client request");
let _enter = span.enter();
// Construct the url
let url = format!("{}?apikey={}{}", self.base_url, self.api_key, param_str);
debug!("Url: {}", url);
self.http.get(url).send().await?.error_for_status()?.bytes().await
}
/// Request the capabilities of the indexer and return them.
pub async fn request_capabilities(&self) -> Result<Capabilities, reqwest::Error> {
let params = TorznabFunction::Capabilities.to_params();
let res = self.request(params).await?;
let str_res = String::from_utf8(res.as_ref().to_vec()).unwrap(); // TODO Handle
let cap: Capabilities = quick_xml::de::from_str(&str_res).unwrap();
Ok(cap)
}
/// Request and store the capabilities of the indexer in the struct.
pub async fn store_capabilities(&mut self) -> Result<&Capabilities, reqwest::Error> {
self.capabilities = self.request_capabilities().await?;
Ok(&self.capabilities)
}
/// Search for torrents.
pub async fn search(&self, func: SearchFunction, generic_params: GenericSearchParameters) -> Result<(), ClientError> {
let param_str = format!("{}{}", func.to_params(), generic_params.to_params());
let bytes = self.request(param_str).await?;
let reader = bytes.reader();
let channel = Channel::read_from(reader).unwrap(); // TODO: handle
let items = channel.into_items();
let torrents: Vec<TorrentResult> = items.iter()
.map(TorrentResult::from_item)
.collect::<Result<Vec<TorrentResult>, super::ResultError>>()?;
debug!("Found results: {:?}", torrents);
//Torrent::from
Ok(())
}
}

17
src/torznab/error.rs Normal file
View File

@ -0,0 +1,17 @@
#[derive(Debug)]
pub enum ClientError {
HttpError(reqwest::Error),
SearchResultError(super::ResultError)
}
impl From<reqwest::Error> for ClientError {
fn from(e: reqwest::Error) -> Self {
ClientError::HttpError(e)
}
}
impl From<super::ResultError> for ClientError {
fn from(e: super::ResultError) -> Self {
ClientError::SearchResultError(e)
}
}

100
src/torznab/functions.rs Normal file
View File

@ -0,0 +1,100 @@
use super::search_parameters::*;
#[derive(Debug)]
pub enum SearchFunction {
/// Free text search query.
Search,
/// Search query with tv specific query params and filtering.
TVSearch(TVSearchParameters),
/// Search query with movie specific query params and filtering.
MovieSearch(MovieSearchParameters),
// TODO
/// Search query with music specific query params and filtering.
MusicSearch,
/// Search query with book specific query params and filtering.
BookSearch,
}
impl SearchFunction {
pub fn to_function_str(&self) -> &str {
match self {
SearchFunction::Search => "search",
SearchFunction::TVSearch(_) => "tvsearch",
SearchFunction::MovieSearch(_) => "movie",
SearchFunction::MusicSearch => "music",
SearchFunction::BookSearch => "book",
}
}
pub fn to_params(&self) -> String {
let mut params = String::new();
params.push_str(&format!("&t={}", self.to_function_str()));
// Concatenate the params of the search function.
match self {
SearchFunction::Search => {},//params.push_str(&s),
SearchFunction::TVSearch(p) => params.push_str(&p.to_params()),
SearchFunction::MovieSearch(p) => params.push_str(&p.to_params()),
_ => panic!("Not implemented!"), // TODO
}
params
}
}
#[derive(Debug)]
pub enum TorznabFunction {
/// Returns the capabilities of the api.
Capabilities,
SearchFunction(GenericSearchParameters, SearchFunction),
// TODO
/// (newznab) Returns all details about a particular item.
Details,
/// (newznab) Returns an nfo for a particular item.
GetNFO,
/// (newznab) Returns nzb for the specified item.
GetNZB,
/// (newznab) Adds item to the users cart.
CardAdd,
/// (newznab) Removes item from the users cart.
CardDel,
/// (newznab) Returns all comments known about an item.
Comments,
/// (newznab) Adds a comment to an item.
CommentsAdd,
/// (newznab) Register a new user account.
Register,
/// (newznab) Retrieves information about an user account.
User
}
impl TorznabFunction {
pub fn to_function_str(&self) -> &str {
match self {
TorznabFunction::Capabilities => "caps",
TorznabFunction::SearchFunction(_, func) => func.to_function_str(),
_ => panic!("Not implemented! ({:?})", self),
}
}
pub fn to_params(&self) -> String {
let mut params = String::new();
// Concatenate the params of the search function.
match self {
TorznabFunction::SearchFunction(p, func) => {
params.push_str(&p.to_params());
params.push_str(&func.to_params());
},
_ => params.push_str(&format!("&t={}", self.to_function_str())),
}
params
}
}

17
src/torznab/mod.rs Normal file
View File

@ -0,0 +1,17 @@
pub mod search_parameters;
pub use search_parameters::GenericSearchParameters;
pub mod functions;
pub use functions::*;
pub mod capabilities;
pub use capabilities::*;
pub mod error;
pub use error::*;
pub mod client;
pub use client::*;
pub mod torrent_result;
pub use torrent_result::*;

View File

@ -0,0 +1,250 @@
#[derive(Debug)]
pub struct GenericSearchParameters {
/// The string search query.
pub query: Option<String>,
/// Categories to search in
pub categories: Vec<i32>,
/// Extended attribute names that should be included in results.
pub attributes: Vec<String>,
/// Specifies that all extended attributes should be included in the results. Overrules attrs.
pub extended: Option<bool>,
/// Number of items to skip in the result.
pub offset: Option<i32>,
/// Number of results to return. Limited by the limits element in the Capabilities.
pub limit: Option<i32>,
}
impl GenericSearchParameters {
/// Convert the search parameters to a query string.
/// This will be prefixed with "&"
pub fn to_params(&self) -> String {
let mut params = String::new();
if let Some(query) = &self.query {
let encoded = urlencoding::encode(query);
params.push_str(&format!("&q={}", encoded));
}
if self.categories.len() > 0 {
params.push_str(&format!("&cat={}",
self.categories.iter()
.map(|i| i.to_string())
.collect::<Vec<_>>()
.join(",")));
}
if self.attributes.len() > 0 {
params.push_str(&format!("&attrs={}", self.attributes.join(",")));
}
if let Some(extended) = &self.extended {
// Convert the boolean to an integer.
let i = if *extended { 1 } else { 0 };
params.push_str(&format!("&extended={}", i));
}
if let Some(offset) = &self.offset {
params.push_str(&format!("&offset={}", offset));
}
if let Some(limit) = &self.limit {
params.push_str(&format!("&limit={}", limit));
}
params
}
}
pub struct GenericSearchParametersBuilder {
params: GenericSearchParameters,
}
impl GenericSearchParametersBuilder {
pub fn new() -> GenericSearchParametersBuilder {
GenericSearchParametersBuilder {
params: GenericSearchParameters {
query: None,
categories: Vec::new(),
attributes: Vec::new(),
extended: None,
offset: None,
limit: None,
},
}
}
pub fn query(mut self, query: String) -> GenericSearchParametersBuilder {
self.params.query = Some(query);
self
}
pub fn categories(mut self, categories: &[i32]) -> GenericSearchParametersBuilder {
self.params.categories.extend_from_slice(categories);
self
}
pub fn category(mut self, category: i32) -> GenericSearchParametersBuilder {
self.params.categories.push(category);
self
}
pub fn attributes(mut self, attributes: &[String]) -> GenericSearchParametersBuilder {
self.params.attributes.extend_from_slice(attributes);
self
}
pub fn attribute(mut self, attribute: String) -> GenericSearchParametersBuilder {
self.params.attributes.push(attribute);
self
}
pub fn extended(mut self, extended: bool) -> GenericSearchParametersBuilder {
self.params.extended = Some(extended);
self
}
pub fn offset(mut self, offset: i32) -> GenericSearchParametersBuilder {
self.params.offset = Some(offset);
self
}
pub fn limit(mut self, limit: i32) -> GenericSearchParametersBuilder {
self.params.limit = Some(limit);
self
}
pub fn build(self) -> GenericSearchParameters {
self.params
}
}
#[derive(Debug)]
pub struct TVSearchParameters {
// idk what this is tbh
pub rid: Option<u32>,
/// Id of the show on TVDB.
pub tvdb_id: Option<u32>,
/// Id of the show on TVMaze.
pub tvmaze_id: Option<u32>,
/// Season number
pub season: Option<u16>,
/// Episode number
pub episode: Option<u16>,
}
impl TVSearchParameters {
pub fn to_params(&self) -> String {
let mut params = String::new();
if let Some(rid) = &self.rid {
params.push_str(&format!("&rid={}", rid));
}
if let Some(tvdb_id) = &self.tvdb_id {
params.push_str(&format!("&tvdbid={}", tvdb_id));
}
if let Some(tvmaze_id) = &self.tvmaze_id {
params.push_str(&format!("&tvmazeid={}", tvmaze_id));
}
if let Some(season) = &self.season {
params.push_str(&format!("&season={}", season));
}
if let Some(episode) = &self.episode {
params.push_str(&format!("&ep={}", episode));
}
params
}
}
pub struct TVSearchParametersBuilder {
params: TVSearchParameters,
}
impl TVSearchParametersBuilder {
pub fn new() -> TVSearchParametersBuilder {
TVSearchParametersBuilder {
params: TVSearchParameters {
rid: None,
tvdb_id: None,
tvmaze_id: None,
season: None,
episode: None,
},
}
}
pub fn rid(mut self, rid: u32) -> TVSearchParametersBuilder {
self.params.rid = Some(rid);
self
}
pub fn tvdb_id(mut self, tvdb_id: u32) -> TVSearchParametersBuilder {
self.params.tvdb_id = Some(tvdb_id);
self
}
pub fn tvmaze_id(mut self, tvmaze_id: u32) -> TVSearchParametersBuilder {
self.params.tvmaze_id = Some(tvmaze_id);
self
}
pub fn season(mut self, season: u16) -> TVSearchParametersBuilder {
self.params.season = Some(season);
self
}
pub fn episode(mut self, episode: u16) -> TVSearchParametersBuilder {
self.params.episode = Some(episode);
self
}
pub fn build(self) -> TVSearchParameters {
self.params
}
}
#[derive(Debug)]
pub struct MovieSearchParameters {
/// Id of the movie on IMDB.
pub imdb_id: Option<u32>,
}
impl MovieSearchParameters {
pub fn to_params(&self) -> String {
let mut params = String::new();
if let Some(imdb_id) = &self.imdb_id {
params.push_str(&format!("&imdbid={}", imdb_id));
}
params
}
}
pub struct MovieSearchParametersBuilder {
params: MovieSearchParameters,
}
impl MovieSearchParametersBuilder {
pub fn new() -> MovieSearchParametersBuilder {
MovieSearchParametersBuilder {
params: MovieSearchParameters {
imdb_id: None,
},
}
}
pub fn imdb_id(mut self, imdb_id: u32) -> MovieSearchParametersBuilder {
self.params.imdb_id = Some(imdb_id);
self
}
pub fn build(self) -> MovieSearchParameters {
self.params
}
}

View File

@ -0,0 +1,42 @@
use rss::Item;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ResultError {
MissingTitle,
MissingLink,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TorrentResult<'a> {
name: &'a str,
link: &'a str,
/* size: u64,
categories: Vec<u32>, */
}
impl<'a> TorrentResult<'a> {
pub fn from_item(item: &'a Item) -> Result<Self, ResultError> {
let name = item.title().ok_or(ResultError::MissingTitle)?;
let link = item.link().ok_or(ResultError::MissingLink)?;
/* let size = item.enclosure().map(|e| e.length().parse::<u64>());
let categories = item.categories().ok_or(ResultError::MissingTitle)?; */
Ok(TorrentResult {
name,
link,
/* size,
categories, */
})
}
}
/* impl<'a> From<Item> for TorrentResult<'a> {
fn from(item: Item) -> Self {
TorrentResult {
name: item.title().unwrap(),
link: item.link().unwrap(),
size: item.size().unwrap(),
categories: item.categories().unwrap(),
}
}
} */