mod parse; mod crypto; mod upload; mod thumbnailing; mod protobuf; mod errors; use std::borrow::Cow; use std::io::Read; use std::fs::File; use std::path::Path; use errors::AviaryError; use itertools::{Itertools, Either}; use parse::{CreateArgs, Command, DownloadArgs}; use ::protobuf::Message; fn trim_url<'a>(base_url: &str, url: &'a str) -> Option<&'a str> { if url.starts_with(base_url) { let shortened = url .trim() .trim_start_matches(base_url) .trim_matches('/') .trim_end_matches(".bin"); if shortened.len() > 50 { None } else { Some(shortened) } } else { None } } fn main() { let args: parse::Args = argh::from_env(); match args.command { Command::Create(create_args) => create(&args.server, create_args), Command::Download(download_args) => download(&args.server, download_args), } } fn create(server: &str, args: CreateArgs) { print!("Checking files..."); let (files, errors): (Vec<_>, Vec<_>) = args.images.iter() .map(Path::new) .map(|path| File::open(path) .map(|file| (path, file)) .map_err(|e| AviaryError::from_open_error(e.kind(), &path)) ) .partition_result(); if !errors.is_empty() { println!(" \x1b[31mError!\x1b[0m"); let (nonexistant, noread): (Vec<_>, Vec<_>) = errors.iter() .partition_map(|e| match e { AviaryError::FileDNE(path) => Either::Left(path), AviaryError::ReadPermissionDenied(path) => Either::Right(path), other => panic!("This error should not be possible! {other:?}") }); if !nonexistant.is_empty() { println!("\nWe didn't see any files at the following locations:"); nonexistant.iter().for_each(|path| println!("\x1b[37m- \x1b[31m{}", path.display())); } if !noread.is_empty() { println!("\x1b[0m\nWe found these files, but didn't have permission to open them:"); noread.iter().for_each(|path| println!("\x1b[37m- \x1b[31m{}", path.display())); } } else { println!(" \x1b[32mDone!\n"); let agent = upload::get_agent(); let full_server = if server.starts_with("http") { Cow::Borrowed(server) } else { Cow::Owned(format!("https://{}", server)) }; let index_url = files.into_iter() .inspect(|(path, _)| print!("\x1b[36m{}\x1b[0m\n\x1b[37m├─\x1b[0m Reading... ", path.display())) .inspect(|_| drop(stdout().flush())) .map(|(path, mut file)| (|| { let mut buff = Vec::with_capacity(file.metadata()?.len() as usize); file.read_to_end(&mut buff)?; Ok((path, buff)) })().map_err(|e| AviaryError::StreamReadError(path.to_owned(), e)) ) .inspect(|r| if r.is_ok() { print!("\x1b[32mDone!\n\x1b[37m├─\x1b[0m Thumbnailing... ") }) .inspect(|_| drop(stdout().flush())) .map(|r| r.and_then(|(path, raw_dat)| { let (thumbnail, blurhash) = thumbnailing::thumbnail(&raw_dat) .map_err(|_| AviaryError::ImageFormatError(path.to_owned()))?; Ok((raw_dat, thumbnail, blurhash)) })) .inspect(|r| if r.is_ok() { print!("\x1b[32mDone!\n\x1b[37m├─\x1b[0m Encrypting... ")}) .inspect(|_| drop(stdout().flush())) .map_ok(|(raw_dat, thumbnail, blurhash)| { let key = crypto::make_key(); ( key, crypto::encrypt(&key, &raw_dat), crypto::encrypt(&key, &thumbnail), blurhash ) }) .inspect(|r| if r.is_ok() { print!("\x1b[32mDone!\n\x1b[37m└─\x1b[0m Uploading... ")}) .inspect(|_| drop(stdout().flush())) .map(|r| r.and_then(|(key, full_img, thumb, blurhash)| upload::put_data(&agent, &*full_server, &thumb) .and_then(|thumb_url| upload::put_data(&agent, &*full_server, &full_img) .map(|full_url| (key, full_url, thumb_url, blurhash))) .map_err(AviaryError::from_upload_error) )) .map(|r| r.and_then(|(key, full_url, thumb_url, blurhash)| { let full_trimmed = trim_url(&*full_server, &full_url); let thmb_trimmed = trim_url(&*full_server, &thumb_url); if let (Some(full_url), Some(thmb_url)) = (full_trimmed, thmb_trimmed) { Ok((key, full_url.to_owned(), thmb_url.to_owned(), blurhash)) } else { Err(AviaryError::ServerError(format!("Received bad response from server: {}", full_url))) } })) .inspect(|r| if r.is_ok() { println!("\x1b[32mDone!\n")}) .map_ok(|(key, full_url, thumb_url, blurhash)| protobuf::image::Image { key: key.into(), full_url, thumb_url, blurhash, special_fields: Default::default() }) .collect::, _>>() .and_then(|image_info|{ let index = protobuf::index::Index { images: image_info, title: args.title, desc: None, special_fields: Default::default() }; let index_key = crypto::make_key(); let encrypted_index = crypto::encrypt( &index_key, &index.write_to_bytes() .expect("Error serializing protocol buffers") ); let encoded_key = base64::encode(index_key); print!("\x1b[0mUploading index... "); upload::put_data(&agent, &*full_server, &encrypted_index) .map_err(|e| AviaryError::from_upload_error(e)) .map(|url| format!("{}#{}", url.trim().trim_end_matches(".bin"), &encoded_key)) }); match index_url { Ok(url) => println!("\x1b[32mDone!\n\n\x1b[34mYour gallery is: \x1b[1;0m{}", url), Err(e) => print!("{}", e), } } } fn download(server: &str, args: DownloadArgs) { drop(server); drop(args); todo!("Unimplemented... yet") }