From ae2f2239ef2e133e373324d3672a4ae6036bc5b9 Mon Sep 17 00:00:00 2001 From: lifegpc Date: Mon, 6 Nov 2023 13:11:44 +0000 Subject: [PATCH] Update push module --- src/db/push_task.rs | 9 + src/opthelper.rs | 6 + src/opts.rs | 15 ++ src/pixiv_app.rs | 24 ++ src/pixiv_web.rs | 28 ++ src/pixivapp/illusts.rs | 8 +- src/server/push/task/mod.rs | 12 + src/server/push/task/pixiv_follow.rs | 5 +- src/server/push/task/pixiv_illusts.rs | 360 ++++++++++++++++++++++++++ src/server/service.rs | 5 +- 10 files changed, 468 insertions(+), 4 deletions(-) create mode 100644 src/server/push/task/pixiv_illusts.rs diff --git a/src/db/push_task.rs b/src/db/push_task.rs index f4deb39..311fb8b 100644 --- a/src/db/push_task.rs +++ b/src/db/push_task.rs @@ -26,6 +26,10 @@ fn default_mode() -> PixivMode { PixivMode::All } +fn default_max_len_used() -> usize { + 50 +} + #[derive(Debug, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "camelCase")] pub enum PushTaskPixivAction { @@ -51,6 +55,11 @@ pub enum PushTaskPixivAction { Illusts { /// User ID uid: u64, + #[serde(default = "default_max_len_used")] + /// Maximum count of checked artworks. + /// + /// Only supported when using Pixiv Web API. + max_len_used: usize, }, } diff --git a/src/opthelper.rs b/src/opthelper.rs index 0392629..02adf9c 100644 --- a/src/opthelper.rs +++ b/src/opthelper.rs @@ -617,6 +617,12 @@ impl OptHelper { } 4 } + + #[cfg(feature = "server")] + /// Whether to prevent to run push task. + pub fn disable_push_task(&self) -> bool { + self.opt.get_ref().disable_push_task + } } impl Default for OptHelper { diff --git a/src/opts.rs b/src/opts.rs index 1b92f96..929e240 100644 --- a/src/opts.rs +++ b/src/opts.rs @@ -136,6 +136,9 @@ pub struct CommandOpts { #[cfg(feature = "server")] /// The maximum number of tasks to push to client at the same time. pub push_task_max_push_count: Option, + #[cfg(feature = "server")] + /// Whether to prevent to run push task. + pub disable_push_task: bool, } impl CommandOpts { @@ -187,6 +190,8 @@ impl CommandOpts { push_task_max_count: None, #[cfg(feature = "server")] push_task_max_push_count: None, + #[cfg(feature = "server")] + disable_push_task: false, } } @@ -686,6 +691,12 @@ pub fn parse_cmd() -> Option { gettext("The maximum number of tasks to push to client at the same time."), "COUNT", ); + #[cfg(feature = "server")] + opts.optflag( + "", + "disable-push-task", + gettext("Prevent to run push task."), + ); let result = match opts.parse(&argv[1..]) { Ok(m) => m, Err(err) => { @@ -1115,6 +1126,10 @@ pub fn parse_cmd() -> Option { return None; } } + #[cfg(feature = "server")] + { + re.as_mut().unwrap().disable_push_task = result.opt_present("disable-push-task"); + } re } diff --git a/src/pixiv_app.rs b/src/pixiv_app.rs index 9c97524..cee901d 100644 --- a/src/pixiv_app.rs +++ b/src/pixiv_app.rs @@ -325,6 +325,22 @@ impl PixivAppClientInternal { log::debug!("{}{}", info, obj.pretty(2).as_str()); Ok(obj) } + + pub async fn get_user_illusts(&self, uid: u64) -> Result { + self.auto_handle().await?; + let re = self + .client + .get_with_param( + "https://app-api.pixiv.net/v1/user/illusts", + json::object! {"type": "illust", "user_id": uid, "filter": "for_ios"}, + None, + ) + .await + .ok_or(gettext("Failed to get user's illusts."))?; + let obj = handle_error(re).await?; + log::debug!("{}{}", gettext("User's illusts: "), obj.pretty(2).as_str()); + Ok(obj) + } } #[derive(Clone)] @@ -372,6 +388,14 @@ impl PixivAppClient { let obj = self.internal.get_follow(restrict).await?; PixivAppIllusts::new(self.internal.clone(), obj) } + + pub async fn get_user_illusts( + &self, + uid: u64, + ) -> Result { + let obj = self.internal.get_user_illusts(uid).await?; + PixivAppIllusts::new(self.internal.clone(), obj) + } } impl AsRef for PixivAppClient { diff --git a/src/pixiv_web.rs b/src/pixiv_web.rs index 9f8d628..863436b 100644 --- a/src/pixiv_web.rs +++ b/src/pixiv_web.rs @@ -319,6 +319,34 @@ impl PixivWebClient { v } + /// Get user's works + /// * `uid` - User's id + pub async fn get_user_works(&self, uid: u64) -> Option { + self.auto_init(); + let r = self + .client + .get_with_param( + format!("https://www.pixiv.net/ajax/user/{}/profile/all", uid), + self.get_params(), + None, + ) + .await; + let r = match r { + Some(r) => r, + None => return None, + }; + let v = self.deal_json(r).await; + if v.is_some() { + log::debug!( + target: "pixiv_web", + "{} {}", + gettext("User's works: "), + v.as_ref().unwrap().pretty(2) + ); + } + v + } + pub fn logined(&self) -> bool { let data = self.data.get_ref(); if data.is_none() { diff --git a/src/pixivapp/illusts.rs b/src/pixivapp/illusts.rs index 8144700..297fc87 100644 --- a/src/pixivapp/illusts.rs +++ b/src/pixivapp/illusts.rs @@ -9,6 +9,7 @@ pub struct PixivAppIllusts { pub illusts: Vec, next_url: Option, client: Arc, + pub user: Option, } impl PixivAppIllusts { @@ -49,14 +50,19 @@ impl PixivAppIllusts { for item in oillusts.members() { illusts.push(PixivAppIllust::new(item.clone())); } - let next_url = match value["nextUrl"].as_str() { + let next_url = match value["next_url"].as_str() { Some(next_url) => Some(next_url.to_owned()), None => None, }; + let user = match value["user"].is_null() { + true => None, + false => Some(value["user"].clone()), + }; Ok(Self { illusts, next_url, client, + user, }) } } diff --git a/src/server/push/task/mod.rs b/src/server/push/task/mod.rs index fa1deed..375b71a 100644 --- a/src/server/push/task/mod.rs +++ b/src/server/push/task/mod.rs @@ -1,4 +1,5 @@ pub mod pixiv_follow; +pub mod pixiv_illusts; pub mod pixiv_send_message; use super::super::preclude::*; @@ -66,6 +67,17 @@ pub async fn run_push_task( pixiv_follow::run_push_task(ctx, task.clone(), config, restrict, mode, send_mode) .await } + PushTaskPixivAction::Illusts { uid, max_len_used } => { + pixiv_illusts::run_push_task( + ctx, + task.clone(), + config, + uid.clone(), + max_len_used.clone(), + send_mode, + ) + .await + } _ => Ok(()), }, } diff --git a/src/server/push/task/pixiv_follow.rs b/src/server/push/task/pixiv_follow.rs index dac2a3a..0e39a88 100644 --- a/src/server/push/task/pixiv_follow.rs +++ b/src/server/push/task/pixiv_follow.rs @@ -176,9 +176,10 @@ impl<'a> RunContext<'a> { self.web_illust(i, &data).await?; } } else { - let index = m.to_index(illusts.len()); + let len = illusts.len(); + let index = m.to_index(len); if let Some(index) = index { - self.web_illust(&illusts[index], &data).await?; + self.web_illust(&illusts[len - index - 1], &data).await?; } } } diff --git a/src/server/push/task/pixiv_illusts.rs b/src/server/push/task/pixiv_illusts.rs new file mode 100644 index 0000000..4e51ee1 --- /dev/null +++ b/src/server/push/task/pixiv_illusts.rs @@ -0,0 +1,360 @@ +use super::super::super::preclude::*; +use super::pixiv_send_message::send_message; +use super::TestSendMode; +use crate::db::push_task::{PushTask, PushTaskPixivConfig}; +use crate::ext::atomic::AtomicQuick; +use crate::ext::replace::ReplaceWith2; +use crate::ext::rw_lock::GetRwLock; +use crate::pixivapp::illust::PixivAppIllust; +use crate::task_manager::{MaxCount, TaskManagerWithId}; +use crate::{concat_pixiv_downloader_error, get_helper}; +use futures_util::lock::Mutex; +use std::collections::HashMap; +use std::sync::atomic::AtomicBool; +use std::sync::RwLock; + +struct PixivIllustsData { + web_data: RwLock>, +} + +impl PixivIllustsData { + pub fn new() -> Self { + Self { + web_data: RwLock::new(HashMap::new()), + } + } + + pub fn get_web_data(&self, id: u64) -> Option { + self.web_data.get_ref().get(&id).cloned() + } + + pub fn set_web_data(&self, id: u64, data: JsonValue) { + self.web_data.get_mut().insert(id, data); + } +} + +struct RunContext<'a> { + ctx: Arc, + task: Arc, + config: &'a PushTaskPixivConfig, + uid: u64, + max_len_used: usize, + send_mode: Option<&'a TestSendMode>, + data: Arc, + use_app_api: bool, + use_web_description: bool, + use_webpage: bool, + pushed: RwLock>, + first_run: AtomicBool, + push_manager: TaskManagerWithId>, +} + +impl<'a> RunContext<'a> { + pub fn new( + ctx: Arc, + task: Arc, + config: &'a PushTaskPixivConfig, + uid: u64, + max_len_used: usize, + send_mode: Option<&'a TestSendMode>, + ) -> Self { + Self { + ctx, + task, + config, + uid, + max_len_used, + send_mode, + data: Arc::new(PixivIllustsData::new()), + use_app_api: config.use_app_api.unwrap_or(true), + use_web_description: config.use_web_description.unwrap_or(true), + use_webpage: config.use_webpage.unwrap_or(true), + pushed: RwLock::new(Vec::new()), + first_run: AtomicBool::new(true), + push_manager: TaskManagerWithId::new( + Arc::new(Mutex::new(0)), + MaxCount::new(get_helper().push_task_max_push_count()), + ), + } + } + + pub async fn handle_finished_tasks(&self) -> Result<(), PixivDownloaderError> { + let tasks = self.push_manager.take_finished_tasks(); + let mut error = Ok(()); + for (i, task) in tasks { + let cfg = self + .task + .push_configs + .get(i) + .ok_or("push config not found")?; + let re = task.await; + match re { + Ok(re) => match re { + Ok(_) => { + log::debug!(target: "pixiv_illusts", "Push task success (task id: {}, index: {}).", self.task.id, i); + } + Err(e) => { + if cfg.allow_failed() { + log::warn!(target: "pixiv_illusts", "Push task error (task id: {}, index: {}): {}", self.task.id, i, e); + } else { + log::debug!(target: "pixiv_illusts", "Push task error (task id: {}, index: {}): {}", self.task.id, i, e); + let e: Result<(), _> = Err(e); + concat_pixiv_downloader_error!(error, e); + } + } + }, + Err(e) => { + if cfg.allow_failed() { + log::error!(target: "pixiv_illusts", "Push task join error (task id: {}, index: {}): {}", self.task.id, i, e); + } else { + log::info!(target: "pixiv_illusts", "Push task join error (task id: {}, index: {}): {}", self.task.id, i, e); + let e: Result<(), _> = Err(e); + concat_pixiv_downloader_error!(error, e); + } + } + } + } + error + } + + pub async fn run(&self) -> Result<(), PixivDownloaderError> { + let now = chrono::Utc::now(); + if self.send_mode.is_none() { + match self.ctx.db.get_push_task_data(self.task.id).await? { + Some(data) => match serde_json::from_str(&data) { + Ok(data) => { + self.pushed.replace_with2(data); + } + Err(e) => { + log::warn!(target: "pixiv_illusts", "Failed to parse push task data: {}", e); + log::debug!(target: "pixiv_illusts", "Push task data: {}", data); + } + }, + None => { + self.first_run.qstore(true); + } + } + } + if self.use_app_api { + self.app_run().await?; + } else { + self.web_run().await?; + } + if self.send_mode.is_none() { + let len = self.pushed.get_ref().len(); + if len > self.config.max_len { + self.pushed.get_mut().drain(0..(len - self.config.max_len)); + } + let data = serde_json::to_string(self.pushed.get_ref().as_slice())?; + self.ctx.db.set_push_task_data(self.task.id, &data).await?; + self.ctx + .db + .update_push_task_last_updated(self.task.id, &now) + .await?; + } + Ok(()) + } + + pub async fn web_run(&self) -> Result<(), PixivDownloaderError> { + let pw = self.ctx.pixiv_web_client().await; + let data = pw + .get_user_works(self.uid) + .await + .ok_or("Failed to get user works.")?; + let mut illusts = data["illusts"] + .entries() + .filter_map(|(k, _)| k.parse::().ok()) + .collect::>(); + illusts.sort(); + let illusts = if illusts.len() > self.max_len_used { + &illusts[illusts.len() - self.max_len_used..] + } else { + &illusts + }; + match self.send_mode { + Some(m) => { + if m.is_all() { + for i in illusts { + self.web_illust(i).await?; + } + } else { + let len = illusts.len(); + let index = m.to_index(len); + if let Some(index) = index { + self.web_illust(&illusts[len - index - 1]).await?; + } + } + } + None => { + if self.first_run.qload() { + for i in illusts.iter().rev() { + self.pushed.get_mut().push(i.clone()); + } + } else { + for i in illusts.iter().rev() { + self.web_illust(i).await?; + } + } + } + } + Ok(()) + } + + pub async fn web_illust(&self, id: &u64) -> Result<(), PixivDownloaderError> { + if self.send_mode.is_none() && self.pushed.get_ref().contains(id) { + return Ok(()); + } + let id = id.clone(); + let wdata = match self.data.get_web_data(id) { + Some(d) => d, + None => { + let pw = self.ctx.pixiv_web_client().await; + let wdata = if self.use_webpage { + pw.get_artwork(id).await.ok_or("Failed to get artwork.")?["illust"] + [format!("{}", id)] + .clone() + } else { + pw.get_artwork_ajax(id) + .await + .ok_or("Failed to get artwork ajax.")? + }; + self.data.set_web_data(id, wdata.clone()); + wdata + } + }; + let len = wdata["pageCount"].as_u64().unwrap_or(1); + let pdata = if len != 1 { + let pw = self.ctx.pixiv_web_client().await; + let pdata = pw + .get_illust_pages(id) + .await + .ok_or("Failed to get illust pages.")?; + Some(Arc::new(pdata)) + } else { + None + }; + if self.send_mode.is_none() { + self.pushed.get_mut().push(id); + } + let wdata = Arc::new(wdata); + let mut index = 0; + for i in self.task.push_configs.iter() { + let cfg = Arc::new(i.clone()); + self.push_manager + .add_task( + index, + send_message( + self.ctx.clone(), + None, + Some(wdata.clone()), + pdata.clone(), + None, + None, + cfg, + ), + true, + ) + .await; + index += 1; + } + self.push_manager.join().await; + self.handle_finished_tasks().await?; + Ok(()) + } + + pub async fn app_run(&self) -> Result<(), PixivDownloaderError> { + let app = self.ctx.pixiv_app_client().await; + let app_data = app.get_user_illusts(self.uid).await?; + match self.send_mode { + Some(m) => { + if m.is_all() { + for i in app_data.illusts.iter() { + self.app_illust(i).await?; + } + } else { + let index = m.to_index(app_data.illusts.len()); + if let Some(index) = index { + self.app_illust(&app_data.illusts[index]).await?; + } + } + } + None => { + if self.first_run.qload() { + for i in app_data.illusts.iter().rev() { + if let Some(id) = i.id() { + self.pushed.get_mut().push(id); + } + } + } else { + for i in app_data.illusts.iter().rev() { + self.app_illust(i).await?; + } + } + } + } + Ok(()) + } + + pub async fn app_illust(&self, illust: &PixivAppIllust) -> Result<(), PixivDownloaderError> { + let id = illust.id().ok_or("illust id is none")?; + if self.send_mode.is_none() && self.pushed.get_ref().contains(&id) { + return Ok(()); + } + let data = match self.data.get_web_data(id) { + Some(d) => Some(Arc::new(d)), + None => { + if self.use_web_description && illust.caption_is_empty() { + let pw = self.ctx.pixiv_web_client().await; + match pw.get_artwork_ajax(id).await { + Some(data) => { + self.data.set_web_data(id, data.clone()); + Some(Arc::new(data)) + } + None => None, + } + } else { + None + } + } + }; + if self.send_mode.is_none() { + self.pushed.get_mut().push(id); + } + let illust = Arc::new(illust.clone()); + let mut index = 0; + for i in self.task.push_configs.iter() { + let i = Arc::new(i.clone()); + self.push_manager + .add_task( + index, + send_message( + self.ctx.clone(), + Some(illust.clone()), + data.clone(), + None, + None, + None, + i, + ), + true, + ) + .await; + index += 1; + } + self.push_manager.join().await; + self.handle_finished_tasks().await?; + Ok(()) + } +} + +pub async fn run_push_task( + ctx: Arc, + task: Arc, + config: &PushTaskPixivConfig, + uid: u64, + max_len_used: usize, + send_mode: Option<&TestSendMode>, +) -> Result<(), PixivDownloaderError> { + let ctx = RunContext::new(ctx, task, config, uid, max_len_used, send_mode); + ctx.run().await +} diff --git a/src/server/service.rs b/src/server/service.rs index fbd55e2..b2b8745 100644 --- a/src/server/service.rs +++ b/src/server/service.rs @@ -1,6 +1,7 @@ use super::context::ServerContext; use super::preclude::*; use super::route::ServerRoutes; +use crate::get_helper; use hyper::server::conn::AddrIncoming; use hyper::server::Server; use hyper::service::Service; @@ -98,6 +99,8 @@ pub async fn start_server( let ctx = Arc::new(ServerContext::default().await); let ser = Server::try_bind(addr)?.serve(PixivDownloaderMakeSvc::new(&ctx)); tokio::spawn(super::timer::start_timer(ctx.clone())); - tokio::task::spawn(super::push::task::run_checking(ctx)); + if !get_helper().disable_push_task() { + tokio::spawn(super::push::task::run_checking(ctx)); + } Ok(ser) }