mirror of
https://github.com/lifegpc/pixiv_downloader.git
synced 2026-07-08 01:32:41 +08:00
GitHub
@@ -4,12 +4,7 @@ A pixiv downloader written in Rust.
|
||||
* Write exif metatata to picture.
|
||||
* Merge ugoira(GIF) pictures to video files.
|
||||
### TODO
|
||||
- [ ] Add support for multiple threads download.
|
||||
- [ ] Batch download.
|
||||
- [ ] Fanbox support.
|
||||
### TODO - CVE
|
||||
- [ ] Remove parse_duration dependency which used in unittest.
|
||||
## Rust features flags
|
||||
See [issues](https://github.com/lifegpc/pixiv_downloader/issues) or [projects](https://github.com/lifegpc/pixiv_downloader/projects)
|
||||
### all
|
||||
Enable all unconflicted features, this will enable [`exif`](#exif) and [`ugoira`](#ugoira).
|
||||
### exif
|
||||
|
||||
193
src/cookies.rs
193
src/cookies.rs
@@ -1,14 +1,20 @@
|
||||
use crate::ext::replace::ReplaceWith;
|
||||
use crate::ext::rw_lock::GetRwLock;
|
||||
use crate::gettext;
|
||||
use chrono::DateTime;
|
||||
use chrono::TimeZone;
|
||||
use chrono::Utc;
|
||||
use reqwest::IntoUrl;
|
||||
use std::collections::HashMap;
|
||||
#[cfg(test)]
|
||||
use std::fs::create_dir;
|
||||
use std::fs::{remove_file, File};
|
||||
use std::io::BufRead;
|
||||
use std::io::BufReader;
|
||||
use std::io::Write;
|
||||
use std::iter::Iterator;
|
||||
use std::path::Path;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
trait ToNetscapeStr {
|
||||
fn to_netscape_str(&self) -> &'static str;
|
||||
@@ -311,16 +317,27 @@ impl CookieJar {
|
||||
self.cookies.clear();
|
||||
}
|
||||
|
||||
pub fn read(&mut self, file_name: &str) -> bool {
|
||||
#[allow(dead_code)]
|
||||
pub fn get<S: AsRef<str> + ?Sized>(&self, name: &S) -> Option<&Cookie> {
|
||||
let name = name.as_ref();
|
||||
for i in self.cookies.iter() {
|
||||
if i._name == name {
|
||||
return Some(i);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
pub fn read<P: AsRef<Path> + ?Sized>(&mut self, file_name: &P) -> bool {
|
||||
self.cookies.clear();
|
||||
let p = Path::new(file_name);
|
||||
let p = file_name.as_ref();
|
||||
if !p.exists() {
|
||||
println!("{} {}", gettext("Can not find file:"), file_name);
|
||||
println!("{} {}", gettext("Can not find file:"), p.display());
|
||||
return false;
|
||||
}
|
||||
let re = File::open(p);
|
||||
if re.is_err() {
|
||||
println!("{} {}", gettext("Can not open file:"), file_name);
|
||||
println!("{} {}", gettext("Can not open file:"), p.display());
|
||||
return false;
|
||||
}
|
||||
let f = re.unwrap();
|
||||
@@ -361,8 +378,8 @@ impl CookieJar {
|
||||
true
|
||||
}
|
||||
|
||||
pub fn save(&mut self, file_name: &str) -> bool {
|
||||
let p = Path::new(file_name);
|
||||
pub fn save<P: AsRef<Path> + ?Sized>(&mut self, file_name: &P) -> bool {
|
||||
let p = file_name.as_ref();
|
||||
self.check_expired();
|
||||
if p.exists() {
|
||||
let re = remove_file(p);
|
||||
@@ -391,3 +408,165 @@ impl CookieJar {
|
||||
self.cookies.iter()
|
||||
}
|
||||
}
|
||||
|
||||
struct CookieJarManager {
|
||||
jars: RwLock<HashMap<PathBuf, (Arc<RwLock<CookieJar>>, usize)>>,
|
||||
}
|
||||
|
||||
impl CookieJarManager {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
jars: RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_cookie_jar<P: AsRef<Path> + ?Sized>(
|
||||
&self,
|
||||
path: &P,
|
||||
) -> Result<Arc<RwLock<CookieJar>>, ()> {
|
||||
let path = path.as_ref().to_owned();
|
||||
let mut jars = self.jars.get_mut();
|
||||
match jars.get_mut(&path) {
|
||||
Some((jar, count)) => {
|
||||
*count += 1;
|
||||
Ok(jar.clone())
|
||||
}
|
||||
None => {
|
||||
let mut jar = CookieJar::new();
|
||||
if jar.read(&path) {
|
||||
let jar = Arc::new(RwLock::new(jar));
|
||||
jars.insert(path, (jar.clone(), 1));
|
||||
Ok(jar)
|
||||
} else {
|
||||
Err(())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn drop_cookie_jar<P: AsRef<Path> + ?Sized>(&self, path: &P) {
|
||||
let path = path.as_ref().to_owned();
|
||||
let mut jars = self.jars.get_mut();
|
||||
match jars.get_mut(&path) {
|
||||
Some((jar, count)) => {
|
||||
*count -= 1;
|
||||
if *count == 0 {
|
||||
if !jar.get_mut().save(&path) {
|
||||
println!(
|
||||
"{} {}",
|
||||
gettext("Warning: Failed to save cookies file:"),
|
||||
path.display()
|
||||
);
|
||||
}
|
||||
jars.remove(&path);
|
||||
}
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
lazy_static! {
|
||||
static ref MANAGER: CookieJarManager = CookieJarManager::new();
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
/// A cookie jar that make sure there are one cookie interface for a cookies file
|
||||
pub struct ManagedCookieJar {
|
||||
pub jar: Arc<RwLock<CookieJar>>,
|
||||
path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl ManagedCookieJar {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
jar: Arc::new(RwLock::new(CookieJar::new())),
|
||||
path: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read<P: AsRef<Path> + ?Sized>(&mut self, path: &P) -> bool {
|
||||
let path = path.as_ref();
|
||||
let path = match path.canonicalize() {
|
||||
Ok(p) => p,
|
||||
Err(_) => path.to_owned(),
|
||||
};
|
||||
match self.path.as_ref() {
|
||||
Some(p) => {
|
||||
if p == &path {
|
||||
return true;
|
||||
}
|
||||
MANAGER.drop_cookie_jar(p);
|
||||
let jar = match MANAGER.get_cookie_jar(&path) {
|
||||
Ok(jar) => jar,
|
||||
Err(()) => return false,
|
||||
};
|
||||
self.jar.replace_with(jar);
|
||||
self.path.replace(path);
|
||||
true
|
||||
}
|
||||
None => {
|
||||
let jar = match MANAGER.get_cookie_jar(&path) {
|
||||
Ok(jar) => jar,
|
||||
Err(()) => return false,
|
||||
};
|
||||
self.jar.replace_with(jar);
|
||||
self.path.replace(path);
|
||||
true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ManagedCookieJar {
|
||||
fn drop(&mut self) {
|
||||
if let Some(path) = self.path.as_ref() {
|
||||
MANAGER.drop_cookie_jar(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_managed_cookie_jar() {
|
||||
let p = Path::new("./test");
|
||||
if !p.exists() {
|
||||
let re = create_dir("./test");
|
||||
assert!(re.is_ok() || p.exists());
|
||||
}
|
||||
{
|
||||
let mut jar = CookieJar::new();
|
||||
jar.add(Cookie::new(
|
||||
"test",
|
||||
"de",
|
||||
"example.com",
|
||||
true,
|
||||
"/",
|
||||
false,
|
||||
None,
|
||||
));
|
||||
jar.save("./test/cookies.txt");
|
||||
}
|
||||
{
|
||||
let mut jar = ManagedCookieJar::new();
|
||||
assert!(jar.read("./test/cookies.txt"));
|
||||
assert_eq!(jar.jar.get_mut().get("test").unwrap()._value, "de");
|
||||
let mut jar2 = ManagedCookieJar::new();
|
||||
assert!(jar2.read("./test/cookies.txt"));
|
||||
jar2.jar.get_mut().add(Cookie::new(
|
||||
"test2",
|
||||
"de",
|
||||
"example.com",
|
||||
true,
|
||||
"/",
|
||||
false,
|
||||
None,
|
||||
));
|
||||
assert_eq!(jar.jar.get_mut().get("test2").unwrap()._value, "de");
|
||||
}
|
||||
{
|
||||
let mut jar = CookieJar::new();
|
||||
assert!(jar.read("./test/cookies.txt"));
|
||||
assert_eq!(jar.get("test").unwrap()._value, "de");
|
||||
assert_eq!(jar.get("test2").unwrap()._value, "de");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -297,20 +297,3 @@ impl PixivWebClient {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PixivWebClient {
|
||||
fn drop(&mut self) {
|
||||
if self.is_inited() {
|
||||
let c = get_helper().cookies();
|
||||
if c.is_some() {
|
||||
if !self.client.save_cookies(c.as_ref().unwrap()) {
|
||||
println!(
|
||||
"{} {}",
|
||||
gettext("Warning: Failed to save cookies file:"),
|
||||
c.as_ref().unwrap()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use crate::cookies::Cookie;
|
||||
use crate::cookies::CookieJar;
|
||||
use crate::cookies::ManagedCookieJar;
|
||||
use crate::ext::atomic::AtomicQuick;
|
||||
use crate::ext::json::ToJson;
|
||||
use crate::ext::rw_lock::GetRwLock;
|
||||
@@ -59,10 +59,10 @@ impl ToHeaders for JsonValue {
|
||||
/// * `c` - Cookies
|
||||
/// * `url` - URL
|
||||
pub fn gen_cookie_header<U: IntoUrl>(c: &WebClient, url: U) -> String {
|
||||
c.get_cookies_as_mut().check_expired();
|
||||
c.get_cookies_as_mut().jar.get_mut().check_expired();
|
||||
let mut s = String::from("");
|
||||
let u = url.as_str();
|
||||
for a in c.get_cookies().iter() {
|
||||
for a in c.get_cookies().jar.get_ref().iter() {
|
||||
if a.matched(u) {
|
||||
if s.len() > 0 {
|
||||
s += " ";
|
||||
@@ -81,7 +81,7 @@ pub struct WebClient {
|
||||
/// HTTP Headers
|
||||
headers: RwLock<HashMap<String, String>>,
|
||||
/// Cookies
|
||||
cookies: RwLock<CookieJar>,
|
||||
cookies: RwLock<ManagedCookieJar>,
|
||||
/// Verbose logging
|
||||
verbose: Arc<AtomicBool>,
|
||||
/// Retry times, 0 means disable, < 0 means always retry
|
||||
@@ -98,18 +98,18 @@ impl WebClient {
|
||||
Self {
|
||||
client,
|
||||
headers: RwLock::new(HashMap::new()),
|
||||
cookies: RwLock::new(CookieJar::new()),
|
||||
cookies: RwLock::new(ManagedCookieJar::new()),
|
||||
verbose: Arc::new(AtomicBool::new(false)),
|
||||
retry: Arc::new(AtomicI64::new(3)),
|
||||
retry_interval: RwLock::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_cookies_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, CookieJar> {
|
||||
pub fn get_cookies_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, ManagedCookieJar> {
|
||||
self.cookies.get_mut()
|
||||
}
|
||||
|
||||
pub fn get_cookies<'a>(&'a self) -> RwLockReadGuard<'a, CookieJar> {
|
||||
pub fn get_cookies<'a>(&'a self) -> RwLockReadGuard<'a, ManagedCookieJar> {
|
||||
self.cookies.get_ref()
|
||||
}
|
||||
|
||||
@@ -153,7 +153,7 @@ impl WebClient {
|
||||
let c = Cookie::from_set_cookie(u.as_str(), val);
|
||||
match c {
|
||||
Some(c) => {
|
||||
self.get_cookies_as_mut().add(c);
|
||||
self.get_cookies_as_mut().jar.get_mut().add(c);
|
||||
}
|
||||
None => {
|
||||
println!("{}", gettext("Failed to parse Set-Cookie header."));
|
||||
@@ -177,19 +177,11 @@ impl WebClient {
|
||||
let mut c = self.get_cookies_as_mut();
|
||||
let r = c.read(file_name);
|
||||
if !r {
|
||||
c.clear();
|
||||
c.jar.get_mut().clear();
|
||||
}
|
||||
r
|
||||
}
|
||||
|
||||
/// Save cookies to file
|
||||
/// * `file_name`: File name
|
||||
///
|
||||
/// returns true if saved successfully.
|
||||
pub fn save_cookies(&self, file_name: &str) -> bool {
|
||||
self.get_cookies_as_mut().save(file_name)
|
||||
}
|
||||
|
||||
/// Set new HTTP header
|
||||
/// * `key` - The key of the new HTTP header
|
||||
/// * `value` - The value of the new HTTP value
|
||||
|
||||
Reference in New Issue
Block a user