lifegpc/pixiv_downloader
1
0
Fork 0
mirror of https://github.com/lifegpc/pixiv_downloader.git synced 2026-07-08 01:32:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

add support to download image

Browse Source
This commit is contained in:
lifegpc
2022-02-27 17:50:15 +08:00
parent 7326df4aaa
commit 23cbb89d10
19 changed files with 1043 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/cookies.rs
View File

@@ -234,7 +234,7 @@ impl Cookie {
} else {
self._domain.clone()
};
if subdomain && !host.starts_with(&domain) && host != &domain[1..] {
if subdomain && !host.ends_with(&domain) && host != &domain[1..] {
return false;
}
if !subdomain && host != domain {

76
src/data/json.rs Normal file
View File

@@ -0,0 +1,76 @@
use crate::gettext;
use crate::pixiv_link::PixivID;
use crate::pixiv_link::ToPixivID;
use json::JsonValue;
use std::collections::HashMap;
use std::ffi::OsStr;
use std::fs::File;
use std::fs::remove_file;
use std::io::Write;
use std::path::Path;
pub trait ToJson {
fn to_json(&self) -> Option<JsonValue>;
}
/// Store metadata informations in JSON file
pub struct JSONDataFile {
id: PixivID,
maps: HashMap<String, JsonValue>,
}
impl JSONDataFile {
pub fn new<T: ToPixivID>(id: T) -> Option<Self> {
let i = id.to_pixiv_id();
if i.is_none() {
return None;
}
Some(Self {
id: i.unwrap(),
maps: HashMap::new(),
})
}
pub fn save<S: AsRef<OsStr> + ?Sized>(&self, path: &S) -> bool {
let p = Path::new(path);
if p.exists() {
let r = remove_file(p);
if r.is_err() {
println!("{} {}", gettext("Failed to remove file:"), r.unwrap_err());
return false;
}
}
let value = self.to_json();
if value.is_none() {
return false;
}
let value = value.unwrap();
let f = File::create(p);
if f.is_err() {
println!("{} {}", gettext("Failed to create file:"), f.unwrap_err());
return false;
}
let mut f = f.unwrap();
let r = f.write(value.pretty(2).as_bytes());
if r.is_err() {
println!("{} {}", gettext("Failed to write file:"), r.unwrap_err());
return false;
}
true
}
}
impl ToJson for JSONDataFile {
fn to_json(&self) -> Option<JsonValue> {
let mut value = json::object! {};
if value.insert("id", self.id.to_json()).is_err() {
return None;
}
for (k, v) in self.maps.iter() {
if value.insert(k, v.clone()).is_err() {
return None;
}
}
Some(value)
}
}

1
src/data/mod.rs Normal file
View File

@@ -0,0 +1 @@
pub mod json;

88
src/download.rs
View File

@@ -1,6 +1,12 @@
use crate::data::json::JSONDataFile;
use crate::gettext;
use crate::pixiv::PixivWebClient;
use crate::pixiv_link::PixivID;
use crate::pixiv_web::PixivWebClient;
use crate::utils::get_file_name_from_url;
use crate::webclient::WebClient;
use crate::Main;
use json::JsonValue;
use std::path::PathBuf;
impl Main {
pub fn download(&mut self) -> i32 {
@@ -10,6 +16,86 @@ impl Main {
return 1;
}
pw.check_login();
if !pw.logined() {
println!(
"{}",
gettext("Warning: Web api client not logined, some future may not work.")
);
}
for id in self.cmd.as_ref().unwrap().ids.iter() {
match id {
PixivID::Artwork(id) => {
let r = self.download_artwork(&mut pw, id.clone());
if r != 0 {
return r;
}
}
}
}
0
}
pub fn download_artwork(&self, pw: &mut PixivWebClient, id: u64) -> i32 {
let re = pw.get_artwork(id);
if re.is_none() {
return 1;
}
let pages = (&(re.as_ref().unwrap())["illust"][format!("{}", id).as_str()]["sl"]).as_u64();
if pages.is_none() {
println!("{}", gettext("Failed to get page count."));
return 1;
}
let pages = pages.unwrap();
let mut pages_data: Option<JsonValue> = None;
if pages > 1 {
pages_data = pw.get_illust_pages(id);
}
if pages > 1 && pages_data.is_none() {
println!("{}", gettext("Failed to get pages' data."));
return 1;
}
let base = PathBuf::from(".");
let json_file = base.join(format!("{}.json", id));
let json_data = JSONDataFile::new(id).unwrap();
if !json_data.save(&json_file) {
println!("{}", gettext("Failed to save metadata to JSON file."));
return 1;
}
if pages_data.is_some() {
let pages_data = pages_data.as_ref().unwrap();
for page in pages_data.members() {
let link = &page["urls"]["original"];
if !link.is_string() {
println!("{}", gettext("Failed to get original picture's link."));
return 1;
}
let link = link.as_str().unwrap();
let file_name = get_file_name_from_url(link);
if file_name.is_none() {
println!("{} {}", gettext("Failed to get file name from url:"), link);
return 1;
}
let file_name = file_name.unwrap();
let file_name = base.join(file_name);
let r = pw.download_image(link);
if r.is_none() {
println!("{} {}", gettext("Failed to download image:"), link);
return 1;
}
let r = r.unwrap();
let re = WebClient::download_stream(&file_name, r, pw.helper.overwrite());
if re.is_err() {
println!("{} {}", gettext("Failed to download image:"), link);
return 1;
}
println!(
"{} {} -> {}",
gettext("Downloaded image:"),
link,
file_name.to_str().unwrap_or("(null)")
)
}
}
0
}
}

1
src/i18n.rs
View File

@@ -150,6 +150,7 @@ impl I18n {
}
lazy_static! {
#[doc(hidden)]
static ref I18NT: I18n = I18n::new();
}

6
src/main.rs
View File

@@ -1,18 +1,22 @@
extern crate chrono;
extern crate dateparser;
extern crate futures_util;
extern crate json;
#[macro_use]
extern crate lazy_static;
extern crate tokio;
extern crate regex;
extern crate reqwest;
mod cookies;
mod data;
mod download;
mod i18n;
mod opthelper;
mod opts;
mod parser;
mod pixiv;
mod pixiv_link;
mod pixiv_web;
mod settings;
mod settings_list;
mod utils;

13
src/opthelper.rs
View File

@@ -1,13 +1,17 @@
use crate::opts::CommandOpts;
use crate::settings::SettingStore;
/// An sturct to access all available settings/command line switches
#[derive(Clone)]
pub struct OptHelper<'a> {
/// Command Line Options
opt: &'a CommandOpts,
/// Settings
settings: &'a SettingStore,
}
impl<'a> OptHelper<'a> {
/// return cookies location, no any check
pub fn cookies(&self) -> Option<String> {
if self.opt.cookies.is_some() {
self.opt.cookies.clone()
@@ -18,6 +22,7 @@ impl<'a> OptHelper<'a> {
}
}
/// return language
pub fn language(&self) -> Option<String> {
if self.opt.language.is_some() {
self.opt.language.clone()
@@ -31,4 +36,12 @@ impl<'a> OptHelper<'a> {
pub fn new(opt: &'a CommandOpts, settings: &'a SettingStore) -> Self {
Self { opt, settings }
}
pub fn overwrite(&self) -> Option<bool> {
self.opt.overwrite
}
pub fn verbose(&self) -> bool {
self.opt.verbose
}
}

49
src/opts.rs
View File

@@ -1,6 +1,7 @@
extern crate getopts;
use crate::gettext;
use crate::pixiv_link::PixivID;
use crate::utils::check_file_exists;
use crate::utils::get_exe_path_else_current;
use getopts::Options;
@@ -35,8 +36,8 @@ impl PartialEq<ConfigCommand> for &ConfigCommand {
pub struct CommandOpts {
/// Command
pub cmd: Command,
/// URLs
pub urls: Vec<String>,
/// IDs
pub ids: Vec<PixivID>,
/// Config location
pub _config: Option<String>,
/// Config command
@@ -45,17 +46,23 @@ pub struct CommandOpts {
pub cookies: Option<String>,
/// The language of translated tags
pub language: Option<String>,
/// Verbose logging
pub verbose: bool,
/// Whether to overwrite file
pub overwrite: Option<bool>,
}
impl CommandOpts {
pub fn new(cmd: Command) -> Self {
Self {
cmd,
urls: Vec::new(),
ids: Vec::new(),
_config: None,
config_cmd: None,
cookies: None,
language: None,
verbose: false,
overwrite: None,
}
}
@@ -87,7 +94,7 @@ impl CommandOpts {
pub fn print_usage(prog: &str, opts: &Options) {
let brief = format!(
"{}
{} download [options] <url> [<url>] {}
{} download [options] <id/url> [<id/url>] {}
{} config fix [options] {}
{} config help [options] {}",
gettext("Usage:"),
@@ -123,6 +130,9 @@ pub fn parse_cmd() -> Option<CommandOpts> {
gettext("The language of translated tags."),
"LANG",
);
opts.optflag("v", "verbose", gettext("Verbose logging."));
opts.optflag("y", "yes", gettext("Overwrite existing file."));
opts.optflag("n", "no", gettext("Skip overwrite existing file."));
let result = match opts.parse(&argv[1..]) {
Ok(m) => m,
Err(err) => {
@@ -148,16 +158,21 @@ pub fn parse_cmd() -> Option<CommandOpts> {
}
match re.as_ref().unwrap().cmd {
Command::Download => {
let mut urls = Vec::new();
let mut ids = Vec::new();
for url in result.free.iter().skip(1) {
urls.push(url.to_string());
let id = PixivID::parse(url);
if id.is_none() {
println!("{} {}", gettext("Failed to parse ID:"), url);
return None;
}
ids.push(id.unwrap());
}
if urls.is_empty() {
println!("{}", gettext("No URL specified."));
if ids.is_empty() {
println!("{}", gettext("No URL or ID specified."));
print_usage(&argv[0], &opts);
return None;
}
re.as_mut().unwrap().urls = urls;
re.as_mut().unwrap().ids = ids;
}
Command::Config => {
if result.free.len() < 2 {
@@ -189,5 +204,21 @@ pub fn parse_cmd() -> Option<CommandOpts> {
if result.opt_present("language") {
re.as_mut().unwrap().language = Some(result.opt_str("language").unwrap());
}
re.as_mut().unwrap().verbose = result.opt_present("verbose");
let yes = result.opt_present("yes");
let no = result.opt_present("no");
re.as_mut().unwrap().overwrite = if yes && no {
if result.opt_positions("yes").last().unwrap() > result.opt_positions("no").last().unwrap() {
Some(true)
} else {
Some(false)
}
} else if yes {
Some(true)
} else if no {
Some(false)
} else {
None
};
re
}

22
src/parser/mainpage.rs → src/parser/metadata.rs
View File

@@ -2,14 +2,19 @@ use crate::gettext;
use html_parser::Dom;
use html_parser::Node;
use json::JsonValue;
use std::default::Default;
pub struct MainPageParser {
pub struct MetaDataParser {
pub key: String,
pub value: Option<JsonValue>,
}
impl MainPageParser {
pub fn new() -> Self {
impl MetaDataParser {
/// Create an new instance
/// * `key` - meta's name
pub fn new(key: &str) -> Self {
Self {
key: key.to_string(),
value: None,
}
}
@@ -26,13 +31,14 @@ impl MainPageParser {
if name.is_none() {
return false;
}
if name.as_ref().unwrap() != "global-data" {
if name.as_ref().unwrap() != self.key.as_str() {
return false;
}
if e.id.is_none() {
return false;
}
if e.id.as_ref().unwrap() != "meta-global-data" {
let mkey = format!("meta-{}", self.key.as_str());
if e.id.as_ref().unwrap() != mkey.as_str() {
return false;
}
let c = e.attributes.get("content");
@@ -85,3 +91,9 @@ impl MainPageParser {
false
}
}
impl Default for MetaDataParser {
fn default() -> Self {
Self::new("global-data")
}
}

2
src/parser/mod.rs
View File

@@ -1,3 +1,3 @@
extern crate html_parser;
pub mod mainpage;
pub mod metadata;

91
src/pixiv.rs
View File

@@ -1,91 +0,0 @@
use crate::gettext;
use crate::opthelper::OptHelper;
use crate::parser::mainpage::MainPageParser;
use crate::webclient::WebClient;
use crate::Main;
use spin_on::spin_on;
/// A client which use Pixiv's web API
pub struct PixivWebClient<'a> {
client: WebClient,
helper: OptHelper<'a>,
/// true if in is initialized
inited: bool,
}
impl<'a> PixivWebClient<'a> {
pub fn new(m: &'a Main) -> Self {
Self {
client: WebClient::new(),
helper: OptHelper::new(m.cmd.as_ref().unwrap(), m.settings.as_ref().unwrap()),
inited: false,
}
}
pub fn init(&mut self) -> bool {
let c = self.helper.cookies();
if c.is_some() {
if !self.client.read_cookies(c.as_ref().unwrap()) {
return false;
}
}
self.client.set_header("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36");
let l = self.helper.language();
if l.is_some() {
self.client.set_header("Accept-Language", l.as_ref().unwrap());
} else {
self.client.set_header("Accept-Language", "ja");
}
self.inited = true;
true
}
pub fn auto_init(&mut self) {
if !self.inited {
let r = self.init();
if !r {
panic!("{}", gettext("Failed to initialize pixiv web api client."));
}
}
}
pub fn check_login(&mut self) -> bool {
self.auto_init();
let r = self.client.get("https://www.pixiv.net/");
if r.is_none() {
return false;
}
let r = r.unwrap();
let status = r.status();
let code = status.as_u16();
if code >= 400 {
println!("{} {}", gettext("Failed to get main page:"), status);
return false;
}
let data = spin_on(r.text_with_charset("UTF-8"));
if data.is_err() {
println!("{} {}", gettext("Failed to get main page:"), data.unwrap_err());
return false;
}
let data = data.unwrap();
let mut p = MainPageParser::new();
if !p.parse(data.as_str()) {
println!("{}", gettext("Failed to parse main page."));
return false;
}
true
}
}
impl <'a> Drop for PixivWebClient<'a> {
fn drop(&mut self) {
if self.inited {
let c = self.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());
}
}
}
}
}

73
src/pixiv_link.rs Normal file
View File

@@ -0,0 +1,73 @@
use crate::data::json::ToJson;
use json::JsonValue;
use regex::Regex;
lazy_static! {
#[doc(hidden)]
static ref RE: Regex = Regex::new("^(https?://)?(www\\.)?pixiv\\.net/artworks/(?P<id>\\d+)").unwrap();
}
/// Repesent an Pixiv ID
#[derive(Clone, Debug)]
pub enum PixivID {
/// Artwork Id, include illust and manga
Artwork(u64),
}
pub trait ToPixivID {
fn to_pixiv_id(&self) -> Option<PixivID>;
}
impl PixivID {
pub fn parse(s: &str) -> Option<PixivID> {
let p = s.trim();
let num = p.parse::<u64>();
if num.is_ok() {
return Some(PixivID::Artwork(num.unwrap()));
}
let re = RE.captures(s);
if re.is_some() {
let r = re.unwrap().name("id");
if r.is_some() {
let r = r.unwrap().as_str();
let num = r.parse::<u64>();
return Some(PixivID::Artwork(num.unwrap()));
}
}
None
}
}
impl ToJson for PixivID {
fn to_json(&self) -> Option<JsonValue> {
match *self {
PixivID::Artwork(id) => {
Some(json::value!({"type": "artwork", "id": id}))
}
}
}
}
impl ToPixivID for &PixivID {
fn to_pixiv_id(&self) -> Option<PixivID> {
Some((*self).clone())
}
}
impl ToPixivID for &str {
fn to_pixiv_id(&self) -> Option<PixivID> {
PixivID::parse(*self)
}
}
impl ToPixivID for String {
fn to_pixiv_id(&self) -> Option<PixivID> {
PixivID::parse(self)
}
}
impl ToPixivID for u64 {
fn to_pixiv_id(&self) -> Option<PixivID> {
Some(PixivID::Artwork(*self))
}
}

222
src/pixiv_web.rs Normal file
View File

@@ -0,0 +1,222 @@
use crate::gettext;
use crate::opthelper::OptHelper;
use crate::parser::metadata::MetaDataParser;
use crate::webclient::WebClient;
use crate::Main;
use json::JsonValue;
use reqwest::IntoUrl;
use reqwest::Response;
use spin_on::spin_on;
/// A client which use Pixiv's web API
pub struct PixivWebClient<'a> {
client: WebClient,
pub helper: OptHelper<'a>,
/// true if in is initialized
inited: bool,
data: Option<JsonValue>,
}
impl<'a> PixivWebClient<'a> {
pub fn new(m: &'a Main) -> Self {
Self {
client: WebClient::new(),
helper: OptHelper::new(m.cmd.as_ref().unwrap(), m.settings.as_ref().unwrap()),
inited: false,
data: None,
}
}
pub fn init(&mut self) -> bool {
let c = self.helper.cookies();
if c.is_some() {
if !self.client.read_cookies(c.as_ref().unwrap()) {
return false;
}
}
self.client.set_header("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36");
let l = self.helper.language();
if l.is_some() {
self.client.set_header("Accept-Language", l.as_ref().unwrap());
} else {
self.client.set_header("Accept-Language", "ja");
}
self.client.verbose = self.helper.verbose();
self.inited = true;
true
}
pub fn auto_init(&mut self) {
if !self.inited {
let r = self.init();
if !r {
panic!("{}", gettext("Failed to initialize pixiv web api client."));
}
}
}
pub fn check_login(&mut self) -> bool {
self.auto_init();
let r = self.client.get("https://www.pixiv.net/", None);
if r.is_none() {
return false;
}
let r = r.unwrap();
let status = r.status();
let code = status.as_u16();
if code >= 400 {
println!("{} {}", gettext("Failed to get main page:"), status);
return false;
}
let data = spin_on(r.text_with_charset("UTF-8"));
if data.is_err() {
println!("{} {}", gettext("Failed to get main page:"), data.unwrap_err());
return false;
}
let data = data.unwrap();
let mut p = MetaDataParser::default();
if !p.parse(data.as_str()) {
println!("{}", gettext("Failed to parse main page."));
return false;
}
if self.helper.verbose() {
println!("{}\n{}", gettext("Main page's data:"), p.value.as_ref().unwrap().pretty(2).as_str());
}
self.data = Some(p.value.unwrap());
true
}
pub fn deal_json(&mut self, r: Response) -> Option<JsonValue> {
let status = r.status();
let code = status.as_u16();
let is_status_err = code >= 400;
let data = spin_on(r.text_with_charset("UTF-8"));
if data.is_err() {
if is_status_err {
println!("HTTP ERROR {}", status);
}
println!("{} {}", gettext("Network error:"), data.unwrap_err());
return None;
}
let data = data.unwrap();
let re = json::parse(data.as_str());
if re.is_err() {
if is_status_err {
println!("HTTP ERROR {}", status);
} else {
println!("{} {}", gettext("Failed to parse JSON:"), re.unwrap_err());
}
return None;
}
let value = re.unwrap();
let error = (&value["error"]).as_bool();
if error.is_none() {
if is_status_err {
println!("HTTP ERROR {}", status);
}
println!("{}", gettext("Failed to detect error."));
return None;
}
let error = error.unwrap();
if error {
let message = (&value["message"]).as_str();
if is_status_err {
println!("HTTP ERROR {}", status);
}
if message.is_some() {
println!("{}", message.unwrap());
}
return None;
}
let body = &value["body"];
if body.is_empty() || body.is_null() {
return Some(value);
}
Some(body.clone())
}
pub fn download_image<U: IntoUrl>(&mut self, url: U) -> Option<Response> {
self.auto_init();
let r = self.client.get(url, json::object!{"referer": "https://www.pixiv.net/"});
if r.is_none() {
return None;
}
let r = r.unwrap();
let status = r.status();
let code = status.as_u16();
if code >= 400 {
println!("{} {}", gettext("Failed to download image:"), status);
return None;
}
Some(r)
}
pub fn get_artwork(&mut self, id: u64) -> Option<JsonValue> {
self.auto_init();
let r = self.client.get(format!("https://www.pixiv.net/artworks/{}", id), None);
if r.is_none() {
return None;
}
let r = r.unwrap();
let status = r.status();
let code = status.as_u16();
if code >= 400 {
println!("{} {}", gettext("Failed to get artwork page:"), status);
return None;
}
let data = spin_on(r.text_with_charset("UTF-8"));
if data.is_err() {
println!("{} {}", gettext("Failed to get artwork page:"), data.unwrap_err());
return None;
}
let data = data.unwrap();
let mut p = MetaDataParser::new("preload-data");
if !p.parse(data.as_str()) {
println!("{}", gettext("Failed to parse artwork page."));
return None;
}
if self.helper.verbose() {
println!("{} {}", gettext("Artwork's data:"), p.value.as_ref().unwrap().pretty(2));
}
Some(p.value.unwrap())
}
pub fn get_illust_pages(&mut self, id: u64) -> Option<JsonValue> {
self.auto_init();
let r = self.client.get(format!("https://www.pixiv.net/ajax/illust/{}/pages", id), None);
if r.is_none() {
return None;
}
let r = r.unwrap();
let v = self.deal_json(r);
if self.helper.verbose() {
println!("{} {}", gettext("Artwork's page data:"), v.as_ref().unwrap().pretty(2));
}
v
}
pub fn logined(&self) -> bool {
if self.data.is_none() {
return false;
}
let value = self.data.as_ref().unwrap();
let d = &value["userData"];
if d.is_object() {
return true;
}
false
}
}
impl <'a> Drop for PixivWebClient<'a> {
fn drop(&mut self) {
if self.inited {
let c = self.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());
}
}
}
}
}

40
src/utils.rs
View File

@@ -1,3 +1,5 @@
use crate::gettext;
use reqwest::IntoUrl;
use std::env;
use std::path::Path;
use std::path::PathBuf;
@@ -22,3 +24,41 @@ pub fn check_file_exists(path: &str) -> bool {
let p = Path::new(path);
p.exists()
}
pub fn ask_need_overwrite(path: &str) -> bool {
let s = gettext("Do you want to delete file \"<file>\"?").replace("<file>", path);
print!("{}(y/n)", s.as_str());
let mut d = String::from("");
loop {
let re = std::io::stdin().read_line(&mut d);
if re.is_err() {
continue;
}
let d = d.trim().to_lowercase();
if d == "y" {
return true;
} else {
return false;
}
}
}
pub fn get_file_name_from_url<U: IntoUrl>(url: U) -> Option<String> {
let u = url.into_url();
if u.is_err() {
println!("{} {}", gettext("Can not parse URL:"), u.unwrap_err());
return None;
}
let u = u.unwrap();
let path = Path::new(u.path());
let re = path.file_name();
if re.is_none() {
println!("{} {}", gettext("Failed to get file name from path:"), u.path());
return None;
}
let r = re.unwrap().to_str();
if r.is_none() {
return None;
}
Some(String::from(r.unwrap()))
}

175
src/webclient.rs
View File

@@ -3,9 +3,51 @@ extern crate spin_on;
use crate::cookies::Cookie;
use crate::cookies::CookieJar;
use crate::gettext;
use crate::utils::ask_need_overwrite;
use futures_util::StreamExt;
use json::JsonValue;
use reqwest::{Client, IntoUrl, RequestBuilder, Response};
use std::collections::HashMap;
use spin_on::spin_on;
use std::collections::HashMap;
use std::ffi::OsStr;
use std::fs::File;
use std::fs::remove_file;
use std::io::Write;
use std::path::Path;
pub trait ToHeaders {
fn to_headers(&self) -> Option<HashMap<String, String>>;
}
impl ToHeaders for Option<HashMap<String, String>> {
fn to_headers(&self) -> Option<HashMap<String, String>> {
self.clone()
}
}
impl ToHeaders for HashMap<String, String> {
fn to_headers(&self) -> Option<HashMap<String, String>> {
Some(self.clone())
}
}
impl ToHeaders for JsonValue {
fn to_headers(&self) -> Option<HashMap<String, String>> {
if !self.is_object() {
return None;
}
let mut h = HashMap::new();
for (k, v) in self.entries() {
let d = if v.is_string() {
String::from(v.as_str().unwrap())
} else {
v.dump()
};
h.insert(String::from(k), d);
}
Some(h)
}
}
/// Generate `cookie` header for a url
/// * `c` - Cookies
@@ -31,6 +73,7 @@ pub struct WebClient {
/// HTTP Headers
pub headers: HashMap<String, String>,
cookies: CookieJar,
pub verbose: bool,
}
impl WebClient {
@@ -39,6 +82,7 @@ impl WebClient {
client: Client::new(),
headers: HashMap::new(),
cookies: CookieJar::new(),
verbose: false,
}
}
@@ -82,10 +126,81 @@ impl WebClient {
pub fn set_header(&mut self, key: &str, value: &str) -> Option<String> {
self.headers.insert(String::from(key), String::from(value))
}
/// Send GET requests with parameters
/// * `param` - GET parameters. Should be a JSON object/array. If value in map is not a string, will dump it
/// # Examples
/// ```
/// let client = WebClient::new();
/// client.verbose = true;
/// client.get_with_param("https://test.com/a", json::object!{"data": "param1"});
/// client.get_with_param("https://test.com/a", json::object!{"daa": {"ad": "test"}});
/// client.get_with_param("https://test.com/a", json::array![["daa", "param1"]]);
/// ```
/// It will GET `https://test.com/a?data=param1`, `https://test.com/a?daa=%7B%22ad%22%3A%22test%22%7D`, `https://test.com/a?daa=param1`
pub fn get_with_param<U: IntoUrl>(&mut self, url: U, param: JsonValue) -> Option<Response> {
let u = url.into_url();
if u.is_err() {
println!("{} \"{}\"", gettext("Can not parse URL:"), u.unwrap_err());
return None;
}
let mut u = u.unwrap();
if !param.is_object() && !param.is_array() {
println!(
"{} \"{}\"",
gettext("Parameters should be object or array:"),
param
);
return None;
}
{
let mut query = u.query_pairs_mut();
if param.is_object() {
for (k, v) in param.entries() {
let s: String;
if v.is_string() {
s = String::from(v.as_str().unwrap());
} else {
s = v.dump();
}
query.append_pair(k, s.as_str());
}
} else {
for v in param.members() {
if !v.is_object() {
println!("{} \"{}\"", gettext("Parameters should be array:"), v);
return None;
}
if v.len() < 2 {
println!("{} \"{}\"", gettext("Parameters need at least a value:"), v);
return None;
}
let okey = &v[0];
let key: String;
if okey.is_string() {
key = String::from(okey.as_str().unwrap());
} else {
key = okey.dump();
}
let mut mems = v.members();
mems.next();
for val in mems {
let s: String;
if val.is_string() {
s = String::from(val.as_str().unwrap());
} else {
s = val.dump();
}
query.append_pair(key.as_str(), s.as_str());
}
}
}
}
self.get(u.as_str(), None)
}
/// Send GET requests
pub fn get<U: IntoUrl>(&mut self, url: U) -> Option<Response> {
let r = self.aget(url);
pub fn get<U: IntoUrl, H: ToHeaders>(&mut self, url: U, headers: H) -> Option<Response> {
let r = self.aget(url, headers);
let r = r.send();
let r = spin_on(r);
match r {
@@ -97,19 +212,71 @@ impl WebClient {
}
let r = r.unwrap();
self.handle_set_cookie(&r);
if self.verbose {
println!("{}", r.status());
}
Some(r)
}
pub fn aget<U: IntoUrl>(&mut self, url: U) -> RequestBuilder {
pub fn aget<U: IntoUrl, H: ToHeaders>(&mut self, url: U, headers: H) -> RequestBuilder {
let s = url.as_str();
if self.verbose {
println!("GET {}", s);
}
let mut r = self.client.get(s);
for (k, v) in self.headers.iter() {
r = r.header(k, v);
}
let headers = headers.to_headers();
if headers.is_some() {
let h = headers.unwrap();
for (k, v) in h.iter() {
r = r.header(k, v);
}
}
let c = gen_cookie_header(&mut self.cookies, s);
if c.len() > 0 {
r = r.header("Cookie", c.as_str());
}
r
}
pub fn download_stream<S: AsRef<OsStr> + ?Sized>(file_name: &S, r: Response, overwrite: Option<bool>) -> Result<(), ()> {
let p = Path::new(file_name);
if p.exists() {
let overwrite = if overwrite.is_none() {
ask_need_overwrite(p.to_str().unwrap())
} else {
overwrite.unwrap()
};
if !overwrite {
return Ok(());
}
let re = remove_file(p);
if re.is_err() {
println!("{} {}", gettext("Failed to remove file:"), re.unwrap_err());
return Err(());
}
}
let f = File::create(p);
if f.is_err() {
println!("{} {}", gettext("Failed to create file:"), f.unwrap_err());
return Err(());
}
let mut f = f.unwrap();
let mut stream = r.bytes_stream();
while let Some(data) = spin_on(stream.next()) {
if data.is_err() {
println!("{} {}", gettext("Error when downloading file:"), data.unwrap_err());
return Err(());
}
let data = data.unwrap();
let r = f.write(&data);
if r.is_err() {
println!("{} {}", gettext("Failed to write file:"), r.unwrap_err());
return Err(());
}
}
Ok(())
}
}