diff --git a/Cargo.lock b/Cargo.lock index 376ac63..0127582 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1084,6 +1084,12 @@ version = "0.4.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" +[[package]] +name = "md5" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "490cc448043f947bae3cbee9c203358d62dbee0db12107a74be5c30ccfd09771" + [[package]] name = "memchr" version = "2.6.3" @@ -1465,6 +1471,7 @@ dependencies = [ "json", "lazy_static", "link-cplusplus", + "md5", "modular-bitfield", "multipart", "openssl", diff --git a/Cargo.toml b/Cargo.toml index 587d5e7..369266e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,7 @@ int-enum = "0.5" itertools = "0.10" json = "0.12" lazy_static = "1.4" +md5 = "0.7" modular-bitfield = "0.11" multipart = { features = ["server"], git = 'https://github.com/lifegpc/multipart', optional = true, default-features = false } openssl = { version = "0.10", optional = true } diff --git a/Language/pixiv_downloader.pot b/Language/pixiv_downloader.pot index 927b241..a7c368a 100644 --- a/Language/pixiv_downloader.pot +++ b/Language/pixiv_downloader.pot @@ -906,7 +906,7 @@ msgid "Failed to flush file:" msgstr "" #: src/settings_list.rs:24 -msgid "Pixiv's refresh tokens. Used to login." +msgid "Pixiv's refresh token. Used to login." msgstr "" #: src/settings_list.rs:30 diff --git a/Language/pixiv_downloader.zh_CN.po b/Language/pixiv_downloader.zh_CN.po index c88f259..0dd8b54 100644 --- a/Language/pixiv_downloader.zh_CN.po +++ b/Language/pixiv_downloader.zh_CN.po @@ -915,8 +915,8 @@ msgid "Failed to flush file:" msgstr "无法刷新文件缓冲区：" #: src/settings_list.rs:24 -msgid "Pixiv's refresh tokens. Used to login." -msgstr "Pixiv 的 refresh tokens。用于登录。" +msgid "Pixiv's refresh token. Used to login." +msgstr "Pixiv 的 refresh token。用于登录。" #: src/settings_list.rs:30 msgid "Remove the part which after these parttens." diff --git a/proc_macros/proc_macros.rs b/proc_macros/proc_macros.rs index 5f6e887..89163db 100644 --- a/proc_macros/proc_macros.rs +++ b/proc_macros/proc_macros.rs @@ -944,3 +944,44 @@ pub fn http_error(item: TokenStream) -> TokenStream { ); stream.into() } + +struct PrintError { + pub msg: Expr, + pub expr: Expr, + pub re: Option, +} + +impl Parse for PrintError { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + let msg = input.parse()?; + input.parse::()?; + let expr = input.parse()?; + let mut re = None; + match input.parse::() { + Ok(_) => { + re.replace(input.parse()?); + } + Err(_) => {} + } + Ok(Self { msg, expr, re }) + } +} + +#[proc_macro] +pub fn print_error(item: TokenStream) -> TokenStream { + let PrintError { msg, expr, re } = parse_macro_input!(item as PrintError); + let re = match re { + Some(re) => quote!(#re), + None => quote!(None), + }; + let stream = quote!( + match (#expr) { + Ok(re) => re, + Err(e) => { + println!("{}{}", #msg, e); + return #re; + } + } + ); + stream.into() +} diff --git a/src/download.rs b/src/download.rs index 333e378..86b76ce 100644 --- a/src/download.rs +++ b/src/download.rs @@ -8,6 +8,8 @@ use crate::data::fanbox::FanboxData; use crate::data::json::JSONDataFile; #[cfg(feature = "ugoira")] use crate::data::video::get_video_metadata; +#[cfg(feature = "db")] +use crate::db::open_and_init_database; use crate::downloader::Downloader; use crate::downloader::DownloaderHelper; use crate::downloader::DownloaderResult; @@ -23,6 +25,7 @@ use crate::fanbox::post::FanboxPost; use crate::fanbox_api::FanboxClient; use crate::gettext; use crate::opthelper::get_helper; +use crate::pixiv_app::PixivAppClient; use crate::pixiv_link::FanboxPostID; use crate::pixiv_link::PixivID; use crate::pixiv_web::PixivWebClient; @@ -34,6 +37,7 @@ use crate::utils::get_file_name_from_url; use crate::Main; use indicatif::MultiProgress; use json::JsonValue; +use proc_macros::print_error; use reqwest::IntoUrl; use std::fs::{create_dir_all, File}; use std::io::Write; @@ -45,25 +49,23 @@ impl Main { pub async fn download(&mut self) -> i32 { let pw = Arc::new(PixivWebClient::new()); let fc = Arc::new(FanboxClient::new()); + #[cfg(feature = "db")] + let db = Arc::new(print_error!( + gettext("Failed to open database:"), + open_and_init_database(get_helper().db()).await, + 0 + )); + #[cfg(not(feature = "db"))] + let ac = Arc::new(PixivAppClient::new()); + #[cfg(feature = "db")] + let ac = PixivAppClient::with_db(Some(db)); let tasks = TaskManager::new_post(); let download_multiple_posts = get_helper().download_multiple_posts(); for id in self.cmd.as_ref().unwrap().ids.iter() { match id { PixivID::Artwork(id) => { - if !pw.is_inited() { - if !pw.init() { - println!("{}", gettext("Failed to initialize pixiv web api client.")); - return 1; - } - if !pw.check_login().await { - return 1; - } - if !pw.logined() { - println!("{}", gettext("Warning: Web api client not logged in, some future may not work.")); - } - } tasks - .add_task(download_artwork(Arc::clone(&pw), id.clone())) + .add_task(download_artwork(ac.clone(), Arc::clone(&pw), id.clone())) .await; if !download_multiple_posts { tasks.join().await; @@ -249,9 +251,52 @@ pub async fn download_artwork_link( } pub async fn download_artwork( + ac: PixivAppClient, pw: Arc, id: u64, ) -> Result<(), PixivDownloaderError> { + let helper = get_helper(); + let app_ok = helper.refresh_token().is_some(); + if app_ok && helper.use_app_api() { + if let Err(e) = download_artwork_app(ac, pw.clone(), id).await { + println!("{}{}", gettext("Warning: Failed to download artwork with app api, trying to download with web api: "), e); + download_artwork_web(pw.clone(), id).await?; + } + } else if app_ok { + if let Err(_) = download_artwork_web(pw.clone(), id).await { + download_artwork_app(ac, pw.clone(), id).await?; + } + } else { + download_artwork_web(pw, id).await?; + } + Ok(()) +} + +pub async fn download_artwork_app( + ac: PixivAppClient, + pw: Arc, + id: u64, +) -> Result<(), PixivDownloaderError> { + let data = ac.get_illust_details(id).await?; + Ok(()) +} + +pub async fn download_artwork_web( + pw: Arc, + id: u64, +) -> Result<(), PixivDownloaderError> { + if !pw.is_login_checked() { + if !pw.check_login().await { + println!("{}", gettext("Failed to check login status.")); + } else { + if !pw.logined() { + println!( + "{}", + gettext("Warning: Web api client not logged in, some future may not work.") + ); + } + } + } let mut re = None; let pages; let mut ajax_ver = true; diff --git a/src/downloader/downloader.rs b/src/downloader/downloader.rs index 9ebb2ad..760a119 100644 --- a/src/downloader/downloader.rs +++ b/src/downloader/downloader.rs @@ -63,7 +63,6 @@ pub trait SetLen { fn set_len(&mut self, size: u64) -> Result<(), DownloaderError>; } -#[derive(Debug)] /// A file downloader pub struct DownloaderInternal< T: Write + Seek + Send + Sync + ClearFile + GetTargetFileName + SetLen, diff --git a/src/error.rs b/src/error.rs index 8eb87bb..b3f906e 100644 --- a/src/error.rs +++ b/src/error.rs @@ -23,7 +23,6 @@ pub enum PixivDownloaderError { FromUtf8Error(std::string::FromUtf8Error), #[cfg(feature = "server")] ToStrError(http::header::ToStrError), - #[cfg(all(feature = "server", feature = "db_sqlite", test))] JSONError(json::Error), #[cfg(feature = "openssl")] OpenSSLError(openssl::error::ErrorStack), diff --git a/src/main.rs b/src/main.rs index 9dfd041..2c0e2fd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -38,6 +38,7 @@ mod opt; mod opthelper; mod opts; mod parser; +mod pixiv_app; mod pixiv_link; mod pixiv_web; mod push; diff --git a/src/opthelper.rs b/src/opthelper.rs index f8cb790..da17f5e 100644 --- a/src/opthelper.rs +++ b/src/opthelper.rs @@ -294,6 +294,16 @@ impl OptHelper { 65536 } + pub fn refresh_token(&self) -> Option { + if self.opt.get_ref().refresh_token.is_some() { + self.opt.get_ref().refresh_token.clone() + } else if self.settings.get_ref().have_str("refresh-token") { + self.settings.get_ref().get_str("refresh-token") + } else { + None + } + } + pub fn update(&self, opt: CommandOpts, settings: SettingStore) { if settings.have("author-name-filters") { self._author_name_filters.replace_with2( @@ -323,6 +333,17 @@ impl OptHelper { self.settings.replace_with2(settings); } + /// Whether to use Pixiv APP API first. + pub fn use_app_api(&self) -> bool { + if self.opt.get_ref().use_app_api.is_some() { + return self.opt.get_ref().use_app_api.unwrap(); + } + if self.settings.get_ref().have("use-app-api") { + return self.settings.get_ref().get_bool("use-app-api").unwrap(); + } + false + } + pub fn overwrite(&self) -> Option { self.opt.get_ref().overwrite } diff --git a/src/opts.rs b/src/opts.rs index 507b35f..6911e9c 100644 --- a/src/opts.rs +++ b/src/opts.rs @@ -120,6 +120,10 @@ pub struct CommandOpts { #[cfg(feature = "ugoira")] pub ugoira_max_fps: Option, pub fanbox_page_number: Option, + /// Pixiv's refresh token. Used to login. + pub refresh_token: Option, + /// Whether to use Pixiv APP API first. + pub use_app_api: Option, } impl CommandOpts { @@ -163,6 +167,8 @@ impl CommandOpts { #[cfg(feature = "ugoira")] ugoira_max_fps: None, fanbox_page_number: None, + refresh_token: None, + use_app_api: None, } } @@ -591,6 +597,26 @@ pub fn parse_cmd() -> Option { HasArg::Maybe, getopts::Occur::Optional, ); + opts.optopt( + "", + "refresh-token", + gettext("Pixiv's refresh token. Used to login."), + "TOKEN", + ); + opts.opt( + "", + "use-app-api", + format!( + "{} ({} {})", + gettext("Whether to use Pixiv APP API first."), + gettext("Default:"), + "yes" + ) + .as_str(), + "yes/no", + HasArg::Maybe, + getopts::Occur::Optional, + ); let result = match opts.parse(&argv[1..]) { Ok(m) => m, Err(err) => { @@ -952,6 +978,20 @@ pub fn parse_cmd() -> Option { return None; } } + re.as_mut().unwrap().refresh_token = result.opt_str("refresh-token"); + match parse_optional_opt(&result, "use-app-api", true, parse_bool) { + Ok(b) => re.as_mut().unwrap().use_app_api = b, + Err(e) => { + println!( + "{} {}", + gettext("Failed to parse :") + .replace("", "use-app-api") + .as_str(), + e + ); + return None; + } + } re } diff --git a/src/pixiv_app.rs b/src/pixiv_app.rs new file mode 100644 index 0000000..199372a --- /dev/null +++ b/src/pixiv_app.rs @@ -0,0 +1,300 @@ +#[cfg(feature = "db")] +use crate::db::PixivDownloaderDb; +use crate::error::PixivDownloaderError; +use crate::ext::atomic::AtomicQuick; +use crate::ext::replace::ReplaceWith2; +use crate::ext::rw_lock::GetRwLock; +use crate::opthelper::OptHelper; +use crate::webclient::{ReqMiddleware, WebClient}; +use crate::{get_helper, gettext}; +use chrono::{DateTime, Local, SecondsFormat, Utc}; +use json::JsonValue; +use reqwest::{Client, Request, RequestBuilder}; +use std::collections::HashMap; +use std::ops::Deref; +use std::sync::atomic::AtomicBool; +use std::sync::Arc; +use std::sync::RwLock; + +pub struct PixivAppMiddleware { + internal: Arc, + helper: Arc, +} + +impl PixivAppMiddleware { + pub fn new(internal: Arc) -> Self { + Self { + internal, + helper: get_helper(), + } + } +} + +impl ReqMiddleware for PixivAppMiddleware { + fn handle(&self, r: Request, c: Client) -> Result { + let now = Local::now().to_rfc3339_opts(SecondsFormat::Secs, false); + let hash = format!( + "{:x}", + md5::compute(format!( + "{}{}", + now, "28c1fdd170a5204386cb1313c7077b34f83e4aaf4aa829ce78c231e05b0bae2c" + )) + ); + if self.helper.verbose() { + println!("X-Client-Hash: {}", hash); + println!("X-Client-Time: {}", now); + } + let is_app_api = r.url().host_str().unwrap_or("") == "app-api.pixiv.net"; + let mut r = RequestBuilder::from_parts(c, r) + .header("X-Client-Hash", hash) + .header("X-Client-Time", now); + if is_app_api { + if let Some(token) = self.internal.access_token.get_ref().as_ref() { + r = r.header("Authorization", format!("Bearer {}", token)); + } + } + Ok(r.build()?) + } +} + +pub struct PixivAppClientInternal { + client: WebClient, + #[cfg(feature = "db")] + db: Option>>, + access_token: RwLock>, + refresh_token: RwLock>, + /// true if in is initialized + inited: AtomicBool, + access_token_expired: RwLock>>, +} + +impl PixivAppClientInternal { + pub fn new() -> Self { + Self { + client: WebClient::default(), + #[cfg(feature = "db")] + db: None, + access_token: RwLock::new(None), + refresh_token: RwLock::new(None), + inited: AtomicBool::new(false), + access_token_expired: RwLock::new(None), + } + } + + #[cfg(feature = "db")] + pub fn with_db(db: Option>>) -> Self { + Self { + client: WebClient::default(), + db, + access_token: RwLock::new(None), + refresh_token: RwLock::new(None), + inited: AtomicBool::new(false), + access_token_expired: RwLock::new(None), + } + } + + pub fn is_inited(&self) -> bool { + self.inited.qload() + } + + pub async fn init(&self, refresh_token: String) -> bool { + self.refresh_token.replace_with2(Some(refresh_token)); + let helper = get_helper(); + self.client + .set_header("User-Agent", "PixivIOSApp/7.16.16 (iOS 16.6; iPad13,18)"); + self.client.set_header("App-OS", "iOS"); + self.client.set_header("App-OS-Version", "16.6"); + self.client.set_header("App-Version", "7.16.16"); + self.client.set_header("Accept", "*/*"); + self.client.set_header( + "Accept-Language", + &helper.language().unwrap_or(String::from("ja")), + ); + #[cfg(feature = "db")] + { + if let Some(db) = &self.db { + if let Ok(Some(token)) = db.get_config("access_token").await { + if let Ok(Some(expired)) = db.get_config("access_token_expired").await { + if let Ok(expired) = DateTime::parse_from_rfc3339(&expired) { + if Local::now() < expired { + self.access_token.replace_with2(Some(token)); + self.access_token_expired + .replace_with2(Some(expired.with_timezone(&Utc))); + } + }; + } + } + } + } + self.inited.qstore(true); + true + } + + async fn auto_init(&self) { + if !self.is_inited() { + let helper = get_helper(); + let r = self + .init( + helper + .refresh_token() + .expect(gettext("refresh_token not setted.")), + ) + .await; + if !r { + panic!("{}", gettext("Failed to initialize pixiv app api client.")); + } + } + } + + fn get_access_token_expired(&self) -> Option> { + self.access_token_expired.get_ref().as_ref().cloned() + } + + async fn auto_handle(&self) -> Result<(), PixivDownloaderError> { + self.auto_init().await; + if self.access_token.get_ref().is_some() { + if let Some(expired) = self.get_access_token_expired() { + if Local::now() >= expired { + self.access_token.get_mut().take(); + self.access_token_expired.get_mut().take(); + self.auth_token().await?; + } + } + } else { + self.auth_token().await?; + } + Ok(()) + } + + fn get_refresh_token(&self) -> Result { + match self.refresh_token.get_ref().as_ref() { + Some(r) => { + return Ok(r.clone()); + } + None => Err(PixivDownloaderError::from(gettext( + "refresh_token not setted.", + ))), + } + } + + async fn auth_token(&self) -> Result<(), PixivDownloaderError> { + self.auto_init().await; + let mut params: HashMap = HashMap::new(); + params.insert( + String::from("client_id"), + String::from("KzEZED7aC0vird8jWyHM38mXjNTY"), + ); + params.insert(String::from("refresh_token"), self.get_refresh_token()?); + params.insert(String::from("include_policy"), String::from("true")); + params.insert( + String::from("client_secret"), + String::from("W9JZoJe00qPvJsiyCGT3CCtC6ZUtdpKpzMbNlUGP"), + ); + params.insert(String::from("grant_type"), String::from("refresh_token")); + let now = Local::now(); + let re = self + .client + .post( + "https://oauth.secure.pixiv.net/auth/token", + None, + Some(params), + ) + .await + .ok_or(gettext("Failed to get auth token."))?; + let status = re.status(); + if status.is_success() { + let obj = json::parse(&re.text().await?)?; + let access_token = obj["access_token"] + .as_str() + .or_else(|| obj["response"]["access_token"].as_str()) + .ok_or("Access token not found.")?; + let expires_in = obj["expires_in"] + .as_i64() + .or_else(|| obj["response"]["expires_in"].as_i64()) + .ok_or("expires_in not found.")?; + let expired = now + chrono::Duration::seconds(expires_in); + self.access_token + .replace_with2(Some(access_token.to_string())); + self.access_token_expired + .replace_with2(Some(expired.with_timezone(&Utc))); + #[cfg(feature = "db")] + { + if let Some(db) = &self.db { + db.set_config("access_token", access_token).await?; + db.set_config("access_token_expired", &expired.to_rfc3339()) + .await?; + } + } + } else { + if let Ok(obj) = json::parse(&re.text().await?) { + if let Some(err) = obj["error_description"].as_str() { + return Err(PixivDownloaderError::from(err)); + } + } + return Err(PixivDownloaderError::from(status.as_str())); + } + Ok(()) + } + + pub async fn get_illust_details(&self, id: u64) -> Result { + self.auto_handle().await?; + let re = self + .client + .get_with_param( + "https://app-api.pixiv.net/v1/illust/detail", + json::object! {"illust_id": id}, + None, + ) + .await + .ok_or(gettext("Failed to get illust details."))?; + let status = re.status(); + let obj = json::parse(&re.text().await?)?; + if get_helper().verbose() { + println!("{}{}", gettext("Illust details:"), obj.pretty(2).as_str()); + } + Ok(obj) + } +} + +#[derive(Clone)] +/// Pixiv App Client +pub struct PixivAppClient { + /// Internal data + internal: Arc, +} + +impl PixivAppClient { + pub fn new() -> Self { + let r = Self { + internal: Arc::new(PixivAppClientInternal::new()), + }; + r.internal + .client + .add_req_middleware(Box::new(PixivAppMiddleware::new(r.internal.clone()))); + r + } + + #[cfg(feature = "db")] + pub fn with_db(db: Option>>) -> Self { + let r = Self { + internal: Arc::new(PixivAppClientInternal::with_db(db)), + }; + r.internal + .client + .add_req_middleware(Box::new(PixivAppMiddleware::new(r.internal.clone()))); + r + } +} + +impl AsRef for PixivAppClient { + fn as_ref(&self) -> &PixivAppClientInternal { + &self.internal + } +} + +impl Deref for PixivAppClient { + type Target = PixivAppClientInternal; + fn deref(&self) -> &Self::Target { + &self.internal + } +} diff --git a/src/pixiv_web.rs b/src/pixiv_web.rs index 5d7689f..e9d2def 100644 --- a/src/pixiv_web.rs +++ b/src/pixiv_web.rs @@ -18,6 +18,7 @@ pub struct PixivWebClient { data: RwLock>, /// Get basic params params: RwLock>, + login_checked: AtomicBool, } impl PixivWebClient { @@ -27,6 +28,7 @@ impl PixivWebClient { inited: AtomicBool::new(false), data: RwLock::new(None), params: RwLock::new(None), + login_checked: AtomicBool::new(false), } } @@ -71,6 +73,7 @@ impl PixivWebClient { pub async fn check_login(&self) -> bool { self.auto_init(); + self.login_checked.qstore(true); let r = self .client .get_with_param("https://www.pixiv.net/", self.get_params(), None) @@ -296,4 +299,8 @@ impl PixivWebClient { } false } + + pub fn is_login_checked(&self) -> bool { + self.login_checked.qload() + } } diff --git a/src/server/context.rs b/src/server/context.rs index f1c0372..a2e6ff0 100644 --- a/src/server/context.rs +++ b/src/server/context.rs @@ -11,10 +11,11 @@ use futures_util::lock::Mutex; use hyper::{http::response::Builder, Body, Request, Response}; use json::JsonValue; use std::collections::BTreeMap; +use std::sync::Arc; pub struct ServerContext { pub cors: CorsContext, - pub db: Box, + pub db: Arc>, pub rsa_key: Mutex>, } @@ -23,7 +24,7 @@ impl ServerContext { Self { cors: CorsContext::default(), db: match open_and_init_database(get_helper().db()).await { - Ok(db) => db, + Ok(db) => Arc::new(db), Err(e) => panic!("{} {}", gettext("Failed to open database:"), e), }, rsa_key: Mutex::new(None), diff --git a/src/server/unittest/mod.rs b/src/server/unittest/mod.rs index 4298986..8740158 100644 --- a/src/server/unittest/mod.rs +++ b/src/server/unittest/mod.rs @@ -28,15 +28,17 @@ impl UnitTestContext { Self { ctx: Arc::new(ServerContext { cors: CorsContext::new(true, vec![], vec![]), - db: open_and_init_database( - PixivDownloaderDbConfig::new(&json::object! { - "type": "sqlite", - "path": "test/server.db", - }) + db: Arc::new( + open_and_init_database( + PixivDownloaderDbConfig::new(&json::object! { + "type": "sqlite", + "path": "test/server.db", + }) + .unwrap(), + ) + .await .unwrap(), - ) - .await - .unwrap(), + ), rsa_key: Mutex::new(None), }), routes: ServerRoutes::new(), diff --git a/src/settings_list.rs b/src/settings_list.rs index 7feec4f..f4ab9b3 100644 --- a/src/settings_list.rs +++ b/src/settings_list.rs @@ -23,7 +23,7 @@ use std::str::FromStr; pub fn get_settings_list() -> Vec { vec![ - SettingDes::new("refresh_tokens", gettext("Pixiv's refresh tokens. Used to login."), JsonValueType::Str, None).unwrap(), + SettingDes::new("refresh-token", gettext("Pixiv's refresh token. Used to login."), JsonValueType::Str, None).unwrap(), SettingDes::new("cookies", gettext("The location of cookies file. Used for web API."), JsonValueType::Str, None).unwrap(), SettingDes::new("language", gettext("The language of translated tags."), JsonValueType::Str, None).unwrap(), SettingDes::new("retry", gettext("Max retry count if request failed."), JsonValueType::Number, Some(check_i64)).unwrap(), @@ -64,6 +64,7 @@ pub fn get_settings_list() -> Vec { SettingDes::new("fanbox-page-number", gettext("Use page number for pictures' file name in fanbox."), JsonValueType::Boolean, None).unwrap(), #[cfg(feature = "server")] SettingDes::new("cors-allow-all", gettext("Whether to allow all domains to send CORS requests."), JsonValueType::Boolean, None).unwrap(), + SettingDes::new("use-app-api", gettext("Whether to use Pixiv APP API first."), JsonValueType::Boolean, None).unwrap(), ] } diff --git a/src/webclient.rs b/src/webclient.rs index cba5299..a4ca4da 100644 --- a/src/webclient.rs +++ b/src/webclient.rs @@ -1,5 +1,6 @@ use crate::cookies::Cookie; use crate::cookies::ManagedCookieJar; +use crate::error::PixivDownloaderError; use crate::ext::atomic::AtomicQuick; use crate::ext::json::ToJson; use crate::ext::rw_lock::GetRwLock; @@ -7,7 +8,8 @@ use crate::gettext; use crate::list::NonTailList; use crate::opthelper::get_helper; use json::JsonValue; -use reqwest::{Client, ClientBuilder, IntoUrl, RequestBuilder, Response}; +use proc_macros::print_error; +use reqwest::{Client, ClientBuilder, IntoUrl, Request, Response}; use std::collections::HashMap; use std::default::Default; use std::sync::atomic::AtomicBool; @@ -72,7 +74,19 @@ pub fn gen_cookie_header(c: &WebClient, url: U) -> String { s } -#[derive(Debug)] +pub trait ReqMiddleware { + fn handle(&self, r: Request, c: Client) -> Result; +} + +impl ReqMiddleware for Arc +where + T: ReqMiddleware, +{ + fn handle(&self, r: Request, c: Client) -> Result { + self.as_ref().handle(r, c) + } +} + /// A Web Client pub struct WebClient { /// Basic Web Client @@ -87,6 +101,8 @@ pub struct WebClient { retry: Arc, /// Retry interval retry_interval: RwLock>>, + /// Request middlewares + req_middlewares: RwLock>>, } impl WebClient { @@ -101,9 +117,22 @@ impl WebClient { verbose: Arc::new(AtomicBool::new(false)), retry: Arc::new(AtomicI64::new(3)), retry_interval: RwLock::new(None), + req_middlewares: RwLock::new(Vec::new()), } } + fn handle_req_middlewares(&self, r: Request) -> Result { + let mut r = r; + for i in self.req_middlewares.get_ref().iter() { + r = i.handle(r, self.client.clone())?; + } + Ok(r) + } + + pub fn add_req_middleware(&self, m: Box) { + self.req_middlewares.get_mut().push(m); + } + pub fn get_cookies_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, ManagedCookieJar> { self.cookies.get_mut() } @@ -321,16 +350,11 @@ impl WebClient { /// Send GET requests without retry pub async fn _aget2(&self, url: U, headers: H) -> Option { - let r = self._aget(url, headers); - let r = r.send().await; - match r { - Ok(_) => {} - Err(e) => { - println!("{} {}", gettext("Error when request:"), e); - return None; - } - } - let r = r.unwrap(); + let r = print_error!( + gettext("Failed to generate request:"), + self._aget(url, headers) + ); + let r = print_error!(gettext("Error when request:"), self.client.execute(r).await); self.handle_set_cookie(&r); if self.get_verbose() { println!("{}", r.status()); @@ -339,7 +363,11 @@ impl WebClient { } /// Generate a requests - pub fn _aget(&self, url: U, headers: H) -> RequestBuilder { + pub fn _aget( + &self, + url: U, + headers: H, + ) -> Result { let s = url.as_str(); if self.get_verbose() { println!("GET {}", s); @@ -359,7 +387,7 @@ impl WebClient { if c.len() > 0 { r = r.header("Cookie", c.as_str()); } - r + self.handle_req_middlewares(r.build()?) } pub async fn post( @@ -407,16 +435,11 @@ impl WebClient { headers: H, form: Option>, ) -> Option { - let r = self._apost(url, headers, form); - let r = r.send().await; - match r { - Ok(_) => {} - Err(e) => { - println!("{} {}", gettext("Error when request:"), e); - return None; - } - } - let r = r.unwrap(); + let r = print_error!( + gettext("Failed to generate request:"), + self._apost(url, headers, form) + ); + let r = print_error!(gettext("Error when request:"), self.client.execute(r).await); self.handle_set_cookie(&r); if self.get_verbose() { println!("{}", r.status()); @@ -430,7 +453,7 @@ impl WebClient { url: U, headers: H, form: Option>, - ) -> RequestBuilder { + ) -> Result { let s = url.as_str(); if self.get_verbose() { println!("POST {}", s); @@ -456,7 +479,7 @@ impl WebClient { } None => {} } - r + self.handle_req_middlewares(r.build()?) } }