diff --git a/src/db/mod.rs b/src/db/mod.rs index 8b7ddc0..175fe14 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -20,7 +20,28 @@ pub use token::Token; pub use traits::PixivDownloaderDb; #[cfg(feature = "server")] pub use user::User; -pub type PixivDownloaderDbError = anyhow::Error; + +#[derive(Debug, derive_more::Display)] +pub struct PixivDownloaderDbError { + e: anyhow::Error, +} + +impl PixivDownloaderDbError { + pub fn msg(msg: S) -> Self { + Self { + e: anyhow::Error::msg(msg), + } + } +} + +impl From for PixivDownloaderDbError +where + T: Into, +{ + fn from(e: T) -> Self { + Self { e: e.into() } + } +} use crate::gettext; diff --git a/src/db/sqlite/db.rs b/src/db/sqlite/db.rs index 55e076a..21c5a80 100644 --- a/src/db/sqlite/db.rs +++ b/src/db/sqlite/db.rs @@ -22,6 +22,10 @@ background INT, comment TEXT, webpage TEXT );"; +const CONFIG_TABLE: &'static str = "CREATE TABLE config ( +key TEXT PRIMARY KEY, +value TEXT +);"; const FILES_TABLE: &'static str = "CREATE TABLE files ( id INTEGER PRIMARY KEY AUTOINCREMENT, path TEXT, @@ -78,7 +82,7 @@ v3 INT, v4 INT, PRIMARY KEY (id) );"; -const VERSION: [u8; 4] = [1, 0, 0, 5]; +const VERSION: [u8; 4] = [1, 0, 0, 6]; pub struct PixivDownloaderSqlite { db: Mutex, @@ -208,6 +212,9 @@ impl PixivDownloaderSqlite { tx.execute("ALTER TABLE pixiv_artworks ADD is_nsfw BOOLEAN;", [])?; tx.execute("ALTER TABLE pixiv_artworks ADD lock INT;", [])?; } + if db_version < [1, 0, 0, 6] { + tx.execute(CONFIG_TABLE, [])?; + } self._write_version(&tx)?; tx.commit()?; } @@ -252,6 +259,9 @@ impl PixivDownloaderSqlite { if !tables.contains_key("users") { t.execute(USERS_TABLE, [])?; } + if !tables.contains_key("config") { + t.execute(CONFIG_TABLE, [])?; + } t.commit()?; Ok(()) } @@ -294,6 +304,15 @@ impl PixivDownloaderSqlite { Ok(tables) } + async fn get_config(&self, key: &str) -> Result, SqliteError> { + let con = self.db.lock().await; + Ok(con + .query_row("SELECT value FROM config WHERE key = ?;", [key], |row| { + row.get(0) + }) + .optional()?) + } + async fn get_pixiv_artwork(&self, id: u64) -> Result, SqliteError> { let con = self.db.lock().await; Ok(con @@ -452,6 +471,14 @@ impl PixivDownloaderSqlite { Ok(ts.execute("DELETE FROM token WHERE expired_at < ?;", [now])?) } + fn _set_config(ts: &Transaction, key: &str, value: &str) -> Result<(), SqliteError> { + ts.execute( + "INSERT OR REPLACE INTO config (key, value) VALUES (?, ?);", + (key, value), + )?; + Ok(()) + } + #[cfg(feature = "server")] async fn _set_user( &self, @@ -739,6 +766,10 @@ impl PixivDownloaderDb for PixivDownloaderSqlite { Ok(()) } + async fn get_config(&self, key: &str) -> Result, PixivDownloaderDbError> { + Ok(self.get_config(key).await?) + } + async fn get_pixiv_artwork( &self, id: u64, @@ -798,6 +829,14 @@ impl PixivDownloaderDb for PixivDownloaderSqlite { Ok(size) } + async fn set_config(&self, key: &str, value: &str) -> Result<(), PixivDownloaderDbError> { + let mut db = self.db.lock().await; + let mut tx = db.transaction()?; + Self::_set_config(&mut tx, key, value)?; + tx.commit()?; + Ok(()) + } + #[cfg(feature = "server")] async fn set_user( &self, diff --git a/src/db/traits.rs b/src/db/traits.rs index a5fd789..1b4acf8 100644 --- a/src/db/traits.rs +++ b/src/db/traits.rs @@ -95,6 +95,9 @@ pub trait PixivDownloaderDb { id: u64, expired_at: &DateTime, ) -> Result<(), PixivDownloaderDbError>; + /// Get a config from database + /// * `key` - The config key + async fn get_config(&self, key: &str) -> Result, PixivDownloaderDbError>; /// Get an artwork from database /// * `id` - The artwork ID async fn get_pixiv_artwork( @@ -102,6 +105,21 @@ pub trait PixivDownloaderDb { id: u64, ) -> Result, PixivDownloaderDbError>; #[cfg(feature = "server")] + /// Get proxy pixiv secrets + async fn get_proxy_pixiv_secrets(&self) -> Result { + let key = "proxy_pixiv_secrets"; + match self.get_config(key).await? { + Some(v) => Ok(v), + None => { + let mut buf = [0; 32]; + openssl::rand::rand_bytes(&mut buf)?; + let v = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &buf); + self.set_config(key, &v).await?; + Ok(v) + } + } + } + #[cfg(feature = "server")] /// Get token by ID /// * `id` - The token ID async fn get_token(&self, id: u64) -> Result, PixivDownloaderDbError>; @@ -140,6 +158,10 @@ pub trait PixivDownloaderDb { /// Remove all expired tokens /// Return the number of removed tokens async fn revoke_expired_tokens(&self) -> Result; + /// Set a config + /// * `key` - The config key + /// * `value` - The config value + async fn set_config(&self, key: &str, value: &str) -> Result<(), PixivDownloaderDbError>; #[cfg(feature = "server")] /// Set a user's information by ID /// * `id`: The user's ID diff --git a/src/server/context.rs b/src/server/context.rs index 9a4afed..f1c0372 100644 --- a/src/server/context.rs +++ b/src/server/context.rs @@ -137,6 +137,60 @@ impl ServerContext { .ok_or(gettext("No corresponding user was found."))?) } + pub async fn verify_secrets( + &self, + req: &Request, + params: &RequestParams, + secrets: String, + time_needed: bool, + ) -> Result<(), PixivDownloaderError> { + let mut sign = None; + match req.headers().get("X-SIGN") { + Some(v) => { + sign.replace(v.to_str()?.to_owned()); + } + None => match params.get("sign") { + Some(v) => { + sign.replace(v.to_owned()); + } + None => {} + }, + } + let sign = match sign { + Some(sign) => sign, + None => return Err(PixivDownloaderError::from(gettext("Sign not found."))), + }; + if time_needed { + let time = params + .get_u64_mult(&["time", "t"])? + .ok_or(gettext("Time not found."))?; + let now = chrono::Utc::now().timestamp() as u64; + if time < now - 300 || time > now + 300 { + return Err(PixivDownloaderError::from(gettext("Time out of range."))); + } + } + let mut par = BTreeMap::new(); + for (k, v) in params.params.iter() { + if k == "sign" { + continue; + } + par.insert(k, v); + } + let mut sha = openssl::sha::Sha512::new(); + sha.update(secrets.as_bytes()); + for (k, v) in par { + for v in v { + sha.update(k.as_bytes()); + sha.update(v.as_bytes()); + } + } + let sha = hex::encode(sha.finish()); + if sign != sha { + return Err(PixivDownloaderError::from(gettext("Sign not match."))); + } + Ok(()) + } + pub async fn verify( &self, req: &Request, diff --git a/src/server/proxy/pixiv.rs b/src/server/proxy/pixiv.rs index 10b0519..e4fd2b0 100644 --- a/src/server/proxy/pixiv.rs +++ b/src/server/proxy/pixiv.rs @@ -29,7 +29,11 @@ impl ResponseFor>> for ProxyPixivContext { OPTIONS ); let params = req.get_params().await?; - let _ = http_error!(401, self.ctx.verify(&req, ¶ms).await); + let secrets = self.ctx.db.get_proxy_pixiv_secrets().await?; + http_error!( + 401, + self.ctx.verify_secrets(&req, ¶ms, secrets, false).await + ); let url = http_error!(params.get("url").ok_or("Url is required.")); let uri = http_error!(Uri::try_from(url)); let host = uri.host().ok_or("Host is needed.")?;