305 lines
10 KiB
Rust
305 lines
10 KiB
Rust
mod parse;
|
|
mod crypto;
|
|
mod upload;
|
|
mod thumbnailing;
|
|
mod protobuf;
|
|
mod errors;
|
|
|
|
use std::{borrow::Cow, io::stdout};
|
|
use std::io::{Read, Write, self};
|
|
use std::fs::{File, self, OpenOptions};
|
|
use std::path::Path;
|
|
|
|
use clap::Parser;
|
|
use errors::AviaryUploadError;
|
|
use itertools::{Itertools, Either};
|
|
use parse::{CreateArgs, Command, DownloadArgs};
|
|
use ::protobuf::Message;
|
|
use sanitise_file_name::sanitise_with_options;
|
|
|
|
use crate::errors::AviaryDownloadError;
|
|
use crate::protobuf::image::Format;
|
|
|
|
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::parse();
|
|
let server_no_trailing_slash = args.server.trim_end_matches('/');
|
|
let full_server = if server_no_trailing_slash.starts_with("http") {
|
|
Cow::Borrowed(server_no_trailing_slash)
|
|
} else {
|
|
Cow::Owned(format!("https://{}", server_no_trailing_slash))
|
|
};
|
|
match args.command {
|
|
Command::Create(create_args) => create(&*full_server, create_args),
|
|
Command::Download(download_args) =>
|
|
if let Err(e) = download(&*full_server, download_args) {
|
|
println!("{e}")
|
|
},
|
|
}
|
|
}
|
|
|
|
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| AviaryUploadError::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 {
|
|
AviaryUploadError::FileDNE(path) => Either::Left(path),
|
|
AviaryUploadError::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 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| AviaryUploadError::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 (full, w, h, thumbnail, blurhash, format) = thumbnailing::thumbnail(&raw_dat)
|
|
.map_err(|_| AviaryUploadError::ImageFormatError(path.to_owned()))?;
|
|
Ok((full.map(Either::Right).unwrap_or(Either::Left(raw_dat)), w, h, thumbnail, blurhash, format))
|
|
}))
|
|
.inspect(|r| if r.is_ok() { print!("\x1b[32mDone!\n\x1b[37m├─\x1b[0m Encrypting... ")})
|
|
.inspect(|_| drop(stdout().flush()))
|
|
.map_ok(|(raw_dat, w, h, thumbnail, blurhash, format)| {
|
|
let key = crypto::make_key();
|
|
(
|
|
key,
|
|
w, h,
|
|
crypto::encrypt(&key, &crypto::NONCE_A, &raw_dat),
|
|
crypto::encrypt(&key, &crypto::NONCE_B, &thumbnail),
|
|
blurhash,
|
|
format
|
|
)
|
|
})
|
|
.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, w, h, full_img, thumb, blurhash, format)|
|
|
upload::put_data(&agent, server, &thumb)
|
|
.and_then(|thumb_url|
|
|
upload::put_data(&agent, server, &full_img)
|
|
.map(|full_url| (key, w, h, full_url, thumb_url, blurhash, format)))
|
|
.map_err(AviaryUploadError::from_upload_error)
|
|
))
|
|
.map(|r| r.and_then(|(key, w, h, full_url, thumb_url, blurhash, format)| {
|
|
let full_trimmed = trim_url(server, &full_url);
|
|
let thmb_trimmed = trim_url(server, &thumb_url);
|
|
if let (Some(full_url), Some(thmb_url)) = (full_trimmed, thmb_trimmed) {
|
|
Ok((key, w, h, full_url.to_owned(), thmb_url.to_owned(), blurhash, format))
|
|
} else {
|
|
Err(AviaryUploadError::ServerError(format!("Received bad response from server: {}", full_url)))
|
|
}
|
|
}))
|
|
.inspect(|r| if r.is_ok() { println!("\x1b[32mDone!\n")})
|
|
.map_ok(|(key, width, height, full_url, thumb_url, blurhash, format)| protobuf::image::Image {
|
|
key: key.into(),
|
|
format: format.into(),
|
|
width, height, full_url, thumb_url, blurhash,
|
|
special_fields: Default::default()
|
|
})
|
|
.collect::<Result<Vec<_>, _>>()
|
|
.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,
|
|
&crypto::NONCE_A,
|
|
&index.write_to_bytes()
|
|
.expect("Error serializing protocol buffers")
|
|
);
|
|
let encoded_key = base64::encode(index_key);
|
|
print!("\x1b[0mUploading index... ");
|
|
upload::put_data(&agent, server, &encrypted_index)
|
|
.map_err(|e| AviaryUploadError::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),
|
|
}
|
|
}
|
|
}
|
|
|
|
const SANITIZATION_OPTIONS: sanitise_file_name::Options<fn(char) -> Option<char>> = sanitise_file_name::Options {
|
|
length_limit: 127,
|
|
reserve_extra: 0,
|
|
extension_cleverness: false,
|
|
most_fs_safe: true,
|
|
windows_safe: true,
|
|
url_safe: false,
|
|
normalise_whitespace: false,
|
|
trim_spaces_and_full_stops: true,
|
|
trim_more_punctuation: false,
|
|
remove_control_characters: true,
|
|
remove_reordering_characters: false,
|
|
replace_with: |_| Some('_'),
|
|
collapse_replacements: true,
|
|
six_measures_of_barley: "Unnamed-Album"
|
|
};
|
|
|
|
fn download(server: &str, args: DownloadArgs) -> Result<(), AviaryDownloadError> {
|
|
let mut download_buffer = Vec::with_capacity(5_000_000);
|
|
let mut decrypt_buffer = Vec::with_capacity(5_000_000);
|
|
let index_url = format!("{}/{}.bin", server, args.id);
|
|
let key: [u8; 32] = base64::decode(args.key)
|
|
.map_err(|_| AviaryDownloadError::MalformedKey(true))?
|
|
.try_into()
|
|
.map_err(|_| AviaryDownloadError::MalformedKey(false))?;
|
|
let download_agent = upload::get_agent();
|
|
|
|
print!("Downloading index... ");
|
|
let encrypted_index = upload::get_data(&download_agent, &index_url, 500_000, &mut download_buffer)
|
|
.map_err(AviaryDownloadError::from_download_error(true))?;
|
|
let serialized_index = crypto::decrypt(&key, &crypto::NONCE_A, &encrypted_index, &mut decrypt_buffer)
|
|
.ok_or(AviaryDownloadError::KeyMismatch)?;
|
|
let index = protobuf::index::Index::parse_from_bytes(&serialized_index)
|
|
.map_err(|_| AviaryDownloadError::CorruptIndex("Index was not a valid protobuf structure"))?;
|
|
let image_count = index.images.len();
|
|
println!(
|
|
"\x1b[32mDone!\x1b[0m\n\nGallery Name: {}\n Image Count: {}\n\n(desc not implemented)\n",
|
|
index.title.as_ref().map(String::as_str).unwrap_or("<not set>"),
|
|
image_count
|
|
);
|
|
|
|
let dest_dir: Cow<Path> = args.output.map(Cow::Owned)
|
|
.unwrap_or_else(||
|
|
index.title.map(|title|
|
|
Cow::Owned(sanitise_with_options(&title, &SANITIZATION_OPTIONS).into()))
|
|
.unwrap_or(Cow::Borrowed("Unnamed-Album".as_ref())));
|
|
|
|
fs::create_dir_all(&dest_dir).map_err(|e| match e.kind() {
|
|
io::ErrorKind::PermissionDenied =>
|
|
AviaryDownloadError::PermissionDenied(dest_dir.as_ref().to_owned()),
|
|
_ => AviaryDownloadError::FilesystemError(e, dest_dir.as_ref().to_owned()),
|
|
})?;
|
|
for (indx, image) in index.images.into_iter().enumerate() {
|
|
|
|
let extension = image.format.enum_value().map(|f| match f {
|
|
Format::PNG => "png",
|
|
Format::WEBP => "webp",
|
|
Format::AVIF => "avif",
|
|
Format::JPG => "jpg",
|
|
Format::GIF => "gif",
|
|
}).unwrap_or("unk");
|
|
|
|
// Print a little blurhash
|
|
const BH_WIDTH: u32 = 12;
|
|
const BH_HEIGHT: u32 = 5;
|
|
let blur = blurhash::decode(&image.blurhash, BH_WIDTH, BH_HEIGHT, 1.0);
|
|
let mut blur = blur.iter();
|
|
for y in 0..BH_HEIGHT {
|
|
for _ in 0..BH_WIDTH {
|
|
let r = blur.next().unwrap();
|
|
let g = blur.next().unwrap();
|
|
let b = blur.next().unwrap();
|
|
let _ = blur.next().unwrap();
|
|
print!("\x1b[38;2;{r};{g};{b}m█");
|
|
}
|
|
match y {
|
|
1 => println!("\x1b[0m Image {}/{}", indx + 1, image_count),
|
|
2 => println!("\x1b[0m Format: {}", extension),
|
|
_ => println!(""),
|
|
}
|
|
}
|
|
print!("\x1b[37m├─\x1b[0m Reserving file... ");
|
|
stdout().flush().unwrap();
|
|
let path = dest_dir.join(format!("{indx:03}.{extension}"));
|
|
let mut dest_file = OpenOptions::new()
|
|
.write(true)
|
|
.truncate(true)
|
|
.create(true)
|
|
.create_new(!args.force)
|
|
.open(&path)
|
|
.map_err(|e| match e.kind() {
|
|
io::ErrorKind::PermissionDenied =>
|
|
AviaryDownloadError::PermissionDenied(path.clone()),
|
|
io::ErrorKind::AlreadyExists =>
|
|
AviaryDownloadError::NotOverwriting(path.clone()),
|
|
io::ErrorKind::Interrupted =>
|
|
AviaryDownloadError::Cancelled,
|
|
_ => AviaryDownloadError::FilesystemError(e, path.clone())
|
|
})?;
|
|
let key = image.key.try_into()
|
|
.map_err(|_|
|
|
AviaryDownloadError::CorruptIndex("Key of invalid length specified in index")
|
|
)?;
|
|
print!("\x1b[32mDone!\n\x1b[37m├─\x1b[0m Dowloading... ");
|
|
stdout().flush().unwrap();
|
|
let encrypted_thumbnail = upload::get_data(
|
|
&download_agent,
|
|
&format!("{}/{}.bin", server, image.full_url),
|
|
7680_000,
|
|
&mut download_buffer
|
|
).map_err(AviaryDownloadError::from_download_error(false))?;
|
|
print!("\x1b[32mDone!\n\x1b[37m├─\x1b[0m Decrypting... ");
|
|
stdout().flush().unwrap();
|
|
let decrypted_thumb = crypto::decrypt(
|
|
&key,
|
|
&crypto::NONCE_B,
|
|
encrypted_thumbnail,
|
|
&mut decrypt_buffer
|
|
).ok_or(AviaryDownloadError::ExpiredImage)?;
|
|
print!("\x1b[32mDone!\n\x1b[37m└─\x1b[0m Saving... ");
|
|
stdout().flush().unwrap();
|
|
dest_file.write_all(&decrypted_thumb)
|
|
.map_err(|e| match e.kind() {
|
|
io::ErrorKind::PermissionDenied =>
|
|
AviaryDownloadError::PermissionDenied(path.clone()),
|
|
io::ErrorKind::Interrupted =>
|
|
AviaryDownloadError::Cancelled,
|
|
_ => AviaryDownloadError::FilesystemError(e, path.clone())
|
|
})?;
|
|
println!("\x1b[32mDone!\n");
|
|
}
|
|
|
|
Ok(())
|
|
}
|