bf1942/bili-sync-ai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: 用更 idiomatic 的方式改写一些代码 (#54)

Browse Source 
* refactor: Config 采用 arc_swap 而非锁

* refactor: 改进 config 的检查，及其他一些细微优化

* refactor: 不再拆分 lib.rs 和 main.rs
This commit is contained in:
idlercloud
2024-04-04 18:39:41 +08:00
committed by GitHub
parent 2521fe932b
commit 4ba23ce8fc
12 changed files with 279 additions and 255 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
src/bilibili/client.rs
View File

@@ -1,3 +1,5 @@
use std::sync::Arc;
use anyhow::Result;
use reqwest::{header, Method};
@@ -54,31 +56,30 @@ impl Default for Client {
}
pub struct BiliClient {
credential: Option<Credential>,
client: Client,
}
impl BiliClient {
pub fn new(credential: Option<Credential>) -> Self {
pub fn new() -> Self {
let client = Client::new();
Self { credential, client }
Self { client }
}
pub fn request(&self, method: Method, url: &str) -> reqwest::RequestBuilder {
self.client.request(method, url, self.credential.as_ref())
let credential = CONFIG.credential.load();
self.client.request(method, url, credential.as_deref())
}
pub async fn check_refresh(&mut self) -> Result<()> {
let Some(credential) = self.credential.as_mut() else {
pub async fn check_refresh(&self) -> Result<()> {
let credential = CONFIG.credential.load();
let Some(credential) = credential.as_deref() else {
return Ok(());
};
if !credential.need_refresh(&self.client).await? {
return Ok(());
}
credential.refresh(&self.client).await?;
let mut config = CONFIG.lock().unwrap();
config.credential = Some(credential.clone());
config.save()
let new_credential = credential.refresh(&self.client).await?;
CONFIG.credential.store(Some(Arc::new(new_credential)));
CONFIG.save()
}
}

21
src/bilibili/credential.rs
View File

@@ -12,7 +12,7 @@ use serde::{Deserialize, Serialize};
use super::error::BiliError;
use crate::bilibili::Client;
#[derive(Default, Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Credential {
pub sessdata: String,
pub bili_jct: String,
@@ -32,6 +32,16 @@ impl Credential {
}
}
const fn empty() -> Self {
Self {
sessdata: String::new(),
bili_jct: String::new(),
buvid3: String::new(),
dedeuserid: String::new(),
ac_time_value: String::new(),
}
}
/// 检查凭据是否有效
pub async fn need_refresh(&self, client: &Client) -> Result<bool> {
let res = client
@@ -55,13 +65,12 @@ impl Credential {
res["data"]["refresh"].as_bool().ok_or(anyhow!("check refresh failed"))
}
pub async fn refresh(&mut self, client: &Client) -> Result<()> {
pub async fn refresh(&self, client: &Client) -> Result<Self> {
let correspond_path = Self::get_correspond_path();
let csrf = self.get_refresh_csrf(client, correspond_path).await?;
let new_credential = self.get_new_credential(client, &csrf).await?;
self.confirm_refresh(client, &new_credential).await?;
*self = new_credential;
Ok(())
Ok(new_credential)
}
fn get_correspond_path() -> String {
@@ -125,9 +134,9 @@ JNrRuoEUXpabUzGB8QIDAQAB
bail!(BiliError::RequestFailed(code, msg.to_owned()));
}
let set_cookies = headers.get_all(header::SET_COOKIE);
let mut credential = Credential {
let mut credential = Self {
buvid3: self.buvid3.clone(),
..Default::default()
..Self::empty()
};
let required_cookies = HashSet::from(["SESSDATA", "bili_jct", "DedeUserID"]);
let cookies: Vec<Cookie> = set_cookies

1
src/bilibili/favorite_list.rs
View File

@@ -3,7 +3,6 @@ use async_stream::stream;
use chrono::serde::ts_seconds;
use chrono::{DateTime, Utc};
use futures::Stream;
use log::error;
use serde_json::Value;
use crate::bilibili::error::BiliError;

137
src/config.rs
View File

@@ -1,107 +1,114 @@
use std::borrow::Cow;
use std::collections::HashMap;
use std::path::Path;
use std::sync::Mutex;
use std::path::PathBuf;
use anyhow::{anyhow, Result};
use log::warn;
use anyhow::Result;
use arc_swap::ArcSwapOption;
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};
use crate::bilibili::{Credential, FilterOption};
pub static CONFIG: Lazy<Mutex<Config>> = Lazy::new(|| {
let config = Config::new();
// 保存一次，确保配置文件存在
config.save().unwrap();
pub static CONFIG: Lazy<Config> = Lazy::new(|| {
let config = Config::load().unwrap_or_else(|err| {
warn!("Failed loading config: {err}");
let new_config = Config::new();
// 保存一次，确保配置文件存在
new_config.save().unwrap();
new_config
});
// 检查配置文件内容
config.check();
Mutex::new(Config::new())
config
});
pub static CONFIG_DIR: Lazy<PathBuf> =
Lazy::new(|| dirs::config_dir().expect("No config path found").join("bili-sync"));
#[derive(Serialize, Deserialize)]
pub struct Config {
pub credential: Option<Credential>,
pub credential: ArcSwapOption<Credential>,
pub filter_option: FilterOption,
pub favorite_list: HashMap<String, String>,
pub video_name: String,
pub page_name: String,
pub favorite_list: HashMap<String, PathBuf>,
pub video_name: Cow<'static, str>,
pub page_name: Cow<'static, str>,
pub interval: u64,
pub upper_path: String,
pub upper_path: PathBuf,
}
impl Default for Config {
fn default() -> Self {
Self {
credential: Some(Credential::default()),
filter_option: FilterOption::default(),
favorite_list: HashMap::new(),
video_name: "{{bvid}}".to_string(),
page_name: "{{bvid}}".to_string(),
interval: 1200,
upper_path: dirs::config_dir()
.unwrap()
.join("bili-sync")
.join("upper_face")
.to_str()
.unwrap()
.to_string(),
}
Self::new()
}
}
impl Config {
fn new() -> Self {
Config::load().unwrap_or_default()
Self {
credential: ArcSwapOption::empty(),
filter_option: FilterOption::default(),
favorite_list: HashMap::new(),
video_name: Cow::Borrowed("{{bvid}}"),
page_name: Cow::Borrowed("{{bvid}}"),
interval: 1200,
upper_path: CONFIG_DIR.join("upper_face"),
}
}
/// 简单的预检查
pub fn check(&self) {
assert!(
!self.favorite_list.is_empty(),
"No favorite list found, program won't do anything"
);
for path in self.favorite_list.values() {
assert!(Path::new(path).is_absolute(), "Path in favorite list must be absolute");
let mut ok = true;
if self.favorite_list.is_empty() {
ok = false;
error!("No favorite list found, program won't do anything");
}
assert!(
Path::new(&self.upper_path).is_absolute(),
"Upper face path must be absolute"
);
assert!(!self.video_name.is_empty(), "No video name template found");
assert!(!self.page_name.is_empty(), "No page name template found");
match self.credential {
Some(ref credential) => {
assert!(
!(credential.sessdata.is_empty()
|| credential.bili_jct.is_empty()
|| credential.buvid3.is_empty()
|| credential.dedeuserid.is_empty()
|| credential.ac_time_value.is_empty()),
"Credential is incomplete"
)
for path in self.favorite_list.values() {
if !path.is_absolute() {
ok = false;
error!("Path in favorite list must be absolute: {}", path.display());
}
None => {
warn!("No credential found, can't access high quality video");
}
if !self.upper_path.is_absolute() {
ok = false;
error!("Upper face path must be absolute");
}
if self.video_name.is_empty() {
ok = false;
error!("No video name template found");
}
if self.page_name.is_empty() {
ok = false;
error!("No page name template found");
}
let credential = self.credential.load();
if let Some(credential) = credential.as_deref() {
if credential.sessdata.is_empty()
|| credential.bili_jct.is_empty()
|| credential.buvid3.is_empty()
|| credential.dedeuserid.is_empty()
|| credential.ac_time_value.is_empty()
{
ok = false;
error!("Credential is incomplete");
}
} else {
warn!("No credential found, can't access high quality video");
}
if !ok {
panic!("Config in {} is invalid", CONFIG_DIR.join("config.toml").display());
}
}
fn load() -> Result<Self> {
let config_path = dirs::config_dir()
.ok_or(anyhow!("No config path found"))?
.join("bili-sync")
.join("config.toml");
let config_path = CONFIG_DIR.join("config.toml");
let config_content = std::fs::read_to_string(config_path)?;
Ok(toml::from_str(&config_content)?)
}
pub fn save(&self) -> Result<()> {
let config_path = dirs::config_dir()
.ok_or(anyhow!("No config path found"))?
.join("bili-sync")
.join("config.toml");
if let Some(parent) = config_path.parent() {
std::fs::create_dir_all(parent)?;
}
let config_path = CONFIG_DIR.join("config.toml");
std::fs::create_dir_all(&*CONFIG_DIR)?;
std::fs::write(config_path, toml::to_string_pretty(self)?)?;
Ok(())
}

12
src/core/command.rs
View File

@@ -7,7 +7,6 @@ use entity::{favorite, page, video};
use filenamify::filenamify;
use futures::stream::{FuturesOrdered, FuturesUnordered};
use futures::{pin_mut, Future, StreamExt};
use log::{error, info, warn};
use sea_orm::entity::prelude::*;
use sea_orm::ActiveValue::Set;
use sea_orm::TransactionTrait;
@@ -31,7 +30,7 @@ use crate::error::DownloadAbortError;
pub async fn process_favorite_list(
bili_client: &BiliClient,
fid: &str,
path: &str,
path: &Path,
connection: &DatabaseConnection,
) -> Result<()> {
let favorite_model = refresh_favorite_list(bili_client, fid, path, connection).await?;
@@ -43,7 +42,7 @@ pub async fn process_favorite_list(
pub async fn refresh_favorite_list(
bili_client: &BiliClient,
fid: &str,
path: &str,
path: &Path,
connection: &DatabaseConnection,
) -> Result<favorite::Model> {
let bili_favorite_list = FavoriteList::new(bili_client, fid.to_owned());
@@ -141,11 +140,6 @@ pub async fn download_unprocessed_videos(
for (video_model, _) in &unhandled_videos_pages {
uppers_mutex.insert(video_model.upper_id, (Mutex::new(()), Mutex::new(())));
}
let upper_path = {
let config = CONFIG.lock().unwrap();
config.upper_path.clone()
};
let upper_path = Path::new(&upper_path);
let mut tasks = unhandled_videos_pages
.into_iter()
.map(|(video_model, pages_model)| {
@@ -157,7 +151,7 @@ pub async fn download_unprocessed_videos(
connection,
&semaphore,
&downloader,
upper_path,
&CONFIG.upper_path,
upper_mutex,
)
})

13
src/core/utils.rs
View File

@@ -21,13 +21,10 @@ use crate::config::CONFIG;
pub static TEMPLATE: Lazy<handlebars::Handlebars> = Lazy::new(|| {
let mut handlebars = handlebars::Handlebars::new();
let config = CONFIG.lock().unwrap();
handlebars
.register_template_string("video", config.video_name.clone())
.unwrap();
handlebars
.register_template_string("page", config.page_name.clone())
.register_template_string("video", &CONFIG.video_name)
.unwrap();
handlebars.register_template_string("page", &CONFIG.page_name).unwrap();
handlebars
});
@@ -48,13 +45,13 @@ pub struct NFOSerializer<'a>(pub ModelWrapper<'a>, pub NFOMode);
/// 根据获得的收藏夹信息，插入或更新数据库中的收藏夹，并返回收藏夹对象
pub async fn handle_favorite_info(
info: &FavoriteListInfo,
path: &str,
path: &Path,
connection: &DatabaseConnection,
) -> Result<favorite::Model> {
favorite::Entity::insert(favorite::ActiveModel {
f_id: Set(info.id),
name: Set(info.title.to_string()),
path: Set(path.to_owned()),
name: Set(info.title.clone()),
path: Set(path.to_string_lossy().to_string()),
..Default::default()
})
.on_conflict(

17
src/database.rs
View File

@@ -1,18 +1,13 @@
use anyhow::{anyhow, Result};
use anyhow::Result;
use migration::{Migrator, MigratorTrait};
use sea_orm::{Database, DatabaseConnection};
use tokio::fs;
use crate::config::CONFIG_DIR;
pub async fn database_connection() -> Result<DatabaseConnection> {
let config_dir = dirs::config_dir().ok_or(anyhow!("No config path found"))?;
let target = config_dir.join("bili-sync").join("data.sqlite");
if let Some(parent) = target.parent() {
fs::create_dir_all(parent).await?;
}
Ok(Database::connect(format!(
"sqlite://{}?mode=rwc",
config_dir.join("bili-sync").join("data.sqlite").to_str().unwrap()
))
.await?)
let target = CONFIG_DIR.join("data.sqlite");
fs::create_dir_all(&*CONFIG_DIR).await?;
Ok(Database::connect(format!("sqlite://{}?mode=rwc", target.to_str().unwrap())).await?)
}
pub async fn migrate_database(connection: &DatabaseConnection) -> Result<()> {

6
src/lib.rs
View File

@@ -1,6 +0,0 @@
pub mod bilibili;
pub mod config;
pub mod core;
pub mod database;
pub mod downloader;
pub mod error;

34
src/main.rs
View File

@@ -1,34 +1,44 @@
use bili_sync::bilibili::BiliClient;
use bili_sync::core::command::process_favorite_list;
use bili_sync::database::{database_connection, migrate_database};
use log::error;
#[macro_use]
extern crate log;
mod bilibili;
mod config;
mod core;
mod database;
mod downloader;
mod error;
use once_cell::sync::Lazy;
use self::bilibili::BiliClient;
use self::config::CONFIG;
use self::core::command::process_favorite_list;
use self::database::{database_connection, migrate_database};
#[tokio::main]
async fn main() -> ! {
env_logger::init();
Lazy::force(&CONFIG);
let mut anchor = chrono::Local::now().date_naive();
let (credential, interval, favorites) = {
let config = bili_sync::config::CONFIG.lock().unwrap();
(config.credential.clone(), config.interval, config.favorite_list.clone())
};
let mut bili_client = BiliClient::new(credential);
let bili_client = BiliClient::new();
let connection = database_connection().await.unwrap();
migrate_database(&connection).await.unwrap();
loop {
if anchor != chrono::Local::now().date_naive() {
if let Err(e) = bili_client.check_refresh().await {
error!("Error: {e}");
tokio::time::sleep(std::time::Duration::from_secs(interval)).await;
tokio::time::sleep(std::time::Duration::from_secs(CONFIG.interval)).await;
continue;
}
anchor = chrono::Local::now().date_naive();
}
for (fid, path) in &favorites {
for (fid, path) in &CONFIG.favorite_list {
let res = process_favorite_list(&bili_client, fid, path, &connection).await;
if let Err(e) = res {
error!("Error: {e}");
}
}
tokio::time::sleep(std::time::Duration::from_secs(interval)).await;
tokio::time::sleep(std::time::Duration::from_secs(CONFIG.interval)).await;
}
}