From 184276f4499e9cbe6eeadc187b81495ac29117ce Mon Sep 17 00:00:00 2001 From: Avril Date: Sun, 12 Jul 2020 19:13:55 +0100 Subject: [PATCH] Initial commit --- .gitignore | 4 ++ Cargo.toml | 16 +++++++ src/args.rs | 118 ++++++++++++++++++++++++++++++++++++++++++++++ src/config.rs | 73 ++++++++++++++++++++++++++++ src/error.rs | 57 ++++++++++++++++++++++ src/loli.rs | 7 +++ src/main.rs | 48 +++++++++++++++++++ src/url.rs | 7 +++ src/work.rs | 7 +++ src/work_async.rs | 88 ++++++++++++++++++++++++++++++++++ 10 files changed, 425 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.toml create mode 100644 src/args.rs create mode 100644 src/config.rs create mode 100644 src/error.rs create mode 100644 src/loli.rs create mode 100644 src/main.rs create mode 100644 src/url.rs create mode 100644 src/work.rs create mode 100644 src/work_async.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0efc41c --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +/target +*~ +Cargo.lock +test* diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..8fe7ce1 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "lolistealer" +version = "0.1.0" +authors = ["Avril "] +edition = "2018" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[features] +async = ["tokio"] + +[dependencies] +termprogress = { path = "../termprogress" } +lazy_static = "1.4" +tokio = {version = "0.2", features= ["full"], optional=true} +reqwest = {version = "0.10", features= ["stream"]} \ No newline at end of file diff --git a/src/args.rs b/src/args.rs new file mode 100644 index 0000000..6ee8015 --- /dev/null +++ b/src/args.rs @@ -0,0 +1,118 @@ +use super::*; + +use lazy_static::lazy_static; +use std::{ + error, + fmt, + path::{ + PathBuf, + Path, + }, +}; + +lazy_static! { + pub static ref PROGRAM_NAME: &'static str = Box::leak(std::env::args().next().unwrap().into_boxed_str()); +} + +/// Print usage then exit with code `1` +pub fn usage() -> ! +{ + println!("Usage: {} [--rating ] []", &PROGRAM_NAME[..]); + println!("Usage: {} --help", &PROGRAM_NAME[..]); + + std::process::exit(1) +} + +/// Parse the system args +pub fn parse_args() -> Result +{ + parse(std::env::args().skip(1)) +} + +fn try_dir(path: impl AsRef) -> Result +{ + let path = path.as_ref(); + if path.is_dir() { + Ok(config::OutputType::Directory(path.to_owned())) + } else if !path.exists() { + Ok(config::OutputType::File(path.to_owned())) + } else { + Err(Error::FileExists(path.to_owned())) + } +} + + +#[derive(Debug)] +pub enum Mode { + Normal(config::Config), + Help, +} + +fn parse(args: I) -> Result +where I: IntoIterator +{ + let mut args = args.into_iter(); + + let mut rating = config::Rating::default(); + let mut paths = Vec::new(); + let mut one = String::default(); + let mut reading = true; + + macro_rules! take_one { + () => { + { + if let Some(new) = args.next() { + one = new; + true + } else { + false + } + } + }; + } + + while let Some(arg) = args.next() + { + if reading { + match arg.to_lowercase().trim() { + "-" => reading = false, + "--help" => return Ok(Mode::Help), + "--rating" if take_one!() => rating = one.parse::()?, + _ => paths.push(try_dir(arg)?), + } + } else { + paths.push(try_dir(arg)?); + } + } + + if paths.len() < 1 { + return Err(Error::NoOutput); + } + + Ok(Mode::Normal(config::Config{rating, output: paths})) +} + +#[derive(Debug)] +pub enum Error { + UnknownRating(String), + FileExists(PathBuf), + NoOutput, + + Internal(Box), + Unknown, +} +impl error::Error for Error{} + +impl fmt::Display for Error +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result + { + match self { + Error::NoOutput => write!(f, "need at least one output, try `{} --help`", &PROGRAM_NAME[..]), + Error::UnknownRating(rating) => write!(f, "{} is not a valid rating", rating), + Error::FileExists(path) => write!(f, "file already exists: {:?}", path), + Error::Internal(bx) => write!(f, "internal error: {}", bx), + _ => write!(f, "unknown error"), + } + } +} diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..1a080fe --- /dev/null +++ b/src/config.rs @@ -0,0 +1,73 @@ +use super::*; + +use std::{ + path::PathBuf, + str, + fmt, +}; + +#[derive(Debug, PartialEq, Eq, Hash)] +pub enum Rating +{ + Safe, + Questionable, + Explicit, +} + +impl Rating +{ + pub fn as_str(&self) -> &str + { + match self { + Rating::Safe => "s", + Rating::Questionable => "q", + Rating::Explicit => "e", + } + } +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub enum OutputType +{ + File(PathBuf), + Directory(PathBuf), +} +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct Config +{ + pub rating: Rating, + pub output: Vec +} + +impl Default for Rating +{ + fn default() -> Self + { + Self::Safe + } +} +impl Default for Config +{ + fn default() -> Self + { + Self { + rating: Rating::default(), + output: Vec::new(), + } + } +} + +impl str::FromStr for Rating +{ + type Err = args::Error; + + fn from_str(s: &str) -> Result + { + Ok(match s.chars().next() { + Some('e') | Some('E') => Self::Explicit, + Some('q') | Some('Q') => Self::Questionable, + Some('s') | Some('S') => Self::Safe, + _ => return Err(args::Error::UnknownRating(s.to_owned())), + }) + } +} diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..62b8bcc --- /dev/null +++ b/src/error.rs @@ -0,0 +1,57 @@ +use std::{ + error, + fmt, + io, +}; + +#[derive(Debug)] +pub enum Error +{ + Unknown, + IO(io::Error), + HTTP(reqwest::Error), + HTTPStatus(reqwest::StatusCode), +} + +impl error::Error for Error +{ + fn source(&self) -> Option<&(dyn error::Error + 'static)> + { + match &self { + Error::IO(io) => Some(io), + _ => None + } + } +} + +impl fmt::Display for Error +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result + { + match self { + Error::IO(io) => write!(f, "io: {}", io), + Error::HTTP(http) => write!(f, "http internal error: {}", http), + Error::HTTPStatus(status) => write!(f, "response returned status code {}", status), + _ => write!(f, "unknown error"), + } + } +} + +impl From for Error +{ + fn from(er: io::Error) -> Self + { + Self::IO(er) + } +} + +impl From for Error +{ + fn from(er: reqwest::Error) -> Self + { + match er.status() { + Some(status) => Self::HTTPStatus(status), + None => Self::HTTP(er), + } + } +} diff --git a/src/loli.rs b/src/loli.rs new file mode 100644 index 0000000..d96b846 --- /dev/null +++ b/src/loli.rs @@ -0,0 +1,7 @@ +use super::*; + +#[derive(Debug)] +pub struct Loli +{ + +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..72f4c54 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,48 @@ +#![allow(dead_code)] + +use termprogress::{ + progress, + ProgressBar, +}; + +mod config; +mod args; +mod url; +mod error; +mod loli; + +#[cfg(feature="async")] +mod work_async; +#[cfg(not(feature="async"))] +mod work; + +pub fn parse_args() -> Result +{ + match args::parse_args()? { + args::Mode::Normal(conf) => Ok(conf), + args::Mode::Help => args::usage(), + } +} + +#[cfg(feature="async")] +#[cfg_attr(feature="async", tokio::main)] +async fn main() -> Result<(), Box> +{ + let conf = match parse_args() { + Ok(v) => v, + Err(e) => { + println!("Failed to parse args: {}", e); + std::process::exit(1) + }, + }; + + work_async::work(conf).await +} + +#[cfg(not(feature="async"))] +fn main() -> Result<(), Box> +{ + let conf = parse_args()?; + + work::work(conf) +} diff --git a/src/url.rs b/src/url.rs new file mode 100644 index 0000000..7c9545a --- /dev/null +++ b/src/url.rs @@ -0,0 +1,7 @@ +use super::*; + +/// Parse to loli url +pub fn parse(rating: &config::Rating) -> String +{ + format!("https://plum.moe/loli?rating={}", rating.as_str()) +} diff --git a/src/work.rs b/src/work.rs new file mode 100644 index 0000000..c17b3d6 --- /dev/null +++ b/src/work.rs @@ -0,0 +1,7 @@ +use super::*; + +pub fn work(conf: config::Config) -> Result<(), Box> +{ + + Ok(()) +} diff --git a/src/work_async.rs b/src/work_async.rs new file mode 100644 index 0000000..6903700 --- /dev/null +++ b/src/work_async.rs @@ -0,0 +1,88 @@ +use super::*; +use std::{ + path::Path, +}; +use tokio::{ + fs::{ + OpenOptions, + }, + prelude::*, + stream::StreamExt, +}; +use termprogress::{ + Display, ProgressBar, Spinner, progress, spinner, +}; + +//TODO: Create a progress task module for atomically printing infos and stuff +//TODO: Create a module for temp files, pass the temp file to `perform` and do the regex fixing after `perform` + +/// Download a loli async +pub async fn perform(url: impl AsRef, path: impl AsRef) -> Result +{ + let url = url.as_ref(); + let path = path.as_ref(); + + let mut progress = spinner::Spin::with_title("Starting request...", Default::default()); + progress.refresh(); + let mut resp = reqwest::get(url).await?; + progress.bump(); + progress.set_title(&format!("Starting download to {:?}...", path)); + let mut file = OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open(path).await?; + + let len = resp.content_length(); + let mut version_inc = || { + progress.bump(); + }; + version_inc(); + + let mut bytes = resp.bytes_stream(); + while let Some(buffer) = bytes.next().await { + file.write(buffer?.as_ref()).await?; + version_inc(); + } + + progress.complete_with("OK"); + + loop{} +} + +pub async fn work(conf: config::Config) -> Result<(), Box> +{ + let rating = conf.rating; + let mut children = Vec::new(); + + for path in conf.output.into_iter() + { + let url = url::parse(&rating); + children.push(tokio::task::spawn(async move { + println!("Starting download ({})...", url); + let path = match path { + config::OutputType::File(file) => file, + config::OutputType::Directory(dir) => { + //TODO: Implement downloading to temp and renaming to hash + unimplemented!(); + }, + }; + match perform(&url, &path).await { + Err(e) => panic!("Failed downloading {} -> {:?}: {}", url, path, e), + Ok(v) => v, + } + })); + } + + for child in children.into_iter() + { + match child.await { + Ok(v) => (), + Err(err) => { + println!("Child failed: {}", err); + }, + } + } + + Ok(()) +}