use std::{ cmp::{max, min}, collections::{HashMap, HashSet}, io::{stdout, Write}, os::unix::fs::FileTypeExt, path::{Path, PathBuf}, str::FromStr, sync::Arc, }; use aws_config::{environment::region, BehaviorVersion, Region, SdkConfig}; use aws_sdk_s3::{ config::Credentials, operation::upload_part, primitives::ByteStream, types::CompletedPart, }; use aws_smithy_runtime::client::http::hyper_014::HyperClientBuilder; use aws_smithy_types::byte_stream::Length; use clap::{Parser, Subcommand}; use constants::REFRESH_TOKEN; use human_bytes::human_bytes; use indicatif::{ProgressBar, ProgressState, ProgressStyle}; use rakuten_drive_cui::{RakutenDriveClient, TargetFile, UploadFile}; use tokio::{ fs::File, io::{AsyncReadExt, BufReader}, sync::Mutex, }; use types::response::ListFilesResponseFile; use util::*; mod client; mod constants; mod types; mod util; #[derive(Parser, Debug)] #[command(version, about, long_about=None)] struct Args { #[command(subcommand)] command: Commands, } #[derive(Subcommand, Debug)] enum Commands { #[clap(about = "List files")] List { /// Parent folder path #[clap(short, long)] prefix: Option, }, #[clap(about = "Upload file")] Upload { file: PathBuf, /// Parent folder path #[clap(short, long)] prefix: Option, /// Upload folder recursively #[clap(short, long)] recursive: bool, /// Send fake file size to server (byte) #[clap(short, long)] fake_size: Option, /// Do not check file existence #[clap(long)] force: bool, /// Stream Upload #[clap(short, long)] stream: bool, }, #[clap(about = "Upload ultra big file")] UploadBig { file: PathBuf, /// Parent folder path #[clap(short, long)] prefix: Option, /// Send fake file size to server (byte) #[clap(short, long)] fake_size: Option, /// Do not check file existence #[clap(long)] force: bool, /// Set division size in bytes #[clap(long)] length: Option, }, #[clap(about = "Download file")] Download { path: String, /// Parent folder path #[clap(long)] prefix: Option, }, #[clap(about = "Move file")] Move { // Source file path path: String, // Destination folder path dest: String, }, #[clap(about = "Delete file")] Delete { path: String, /// Delete folder recursively #[clap(long)] recursive: bool, }, #[clap(about = "Make directory")] Mkdir { name: String, /// Path to create directory #[clap(long)] path: Option, }, #[clap(about = "Copy file")] Copy { /// Source file path src: String, /// Destination file directory dest: String, }, #[clap(about = "Rename file")] Rename { /// Target file path path: String, /// New file name name: String, }, #[clap(about = "Print file detail")] Info { path: String, }, Auth {}, } #[tokio::main] async fn main() -> anyhow::Result<()> { let args = Args::parse(); let client = Arc::new(RakutenDriveClient::try_new(REFRESH_TOKEN.to_string()).await?); match args.command { Commands::List { prefix } => { let res = client.list(prefix.as_deref()).await.unwrap(); res.file.iter().for_each(|f| { let dir_str = if f.is_folder { "d" } else { " " }; println!( "{}\t{}\t{}\t{}", dir_str, human_bytes(f.size as f64), f.last_modified, f.path, ) }); } Commands::Upload { file, prefix, recursive, fake_size, force, stream, } => { // is folder if file.is_dir() && !recursive { println!("Use --recursive option for folder upload"); return Err(anyhow::anyhow!("Use --recursive option for folder upload")); } if stream && recursive { println!("Can't use Stream Upload and Recursive Upload both."); return Err(anyhow::anyhow!( "Can't use Stream Upload and Recursive Upload both." )); } if let Some(prefix) = prefix.as_ref() { if !prefix.ends_with('/') { println!("Prefix must end with /"); return Err(anyhow::anyhow!("Prefix must end with /")); } } println!("is_dir: {}", file.is_dir()); if file.is_dir() { println!("name: {:?}", file.file_name().unwrap().to_str().unwrap()); println!("prefix: {:?}", prefix.as_deref()); client .mkdir( file.file_name().unwrap().to_str().unwrap(), prefix.as_deref(), ) .await?; } let prefix = if file.is_dir() && prefix.is_none() { Some(file.file_name().unwrap().to_str().unwrap().to_string() + "/") } else { prefix }; let mut files = Vec::::new(); // file check if !file.exists() { println!("File not found: {:?}", file); return Err(anyhow::anyhow!("File not found: {:?}", file)); } let target_file = TargetFile { file: file.clone(), path: file.file_name().unwrap().to_str().unwrap().to_string(), }; let path = match prefix.clone() { Some(p) => p + target_file.clone().path.as_str(), None => target_file.clone().path, }; if cfg!(windows) { // replase \ to / files.iter_mut().for_each(|f| { f.path = f.path.replace('\\', "/"); }); } if !force { println!("Checking file existence..."); println!("Checking: {:?}", path); let res = client.info(path.as_str()).await; if res.is_ok() { println!("File already exists."); return Err(anyhow::anyhow!("File already exists.")); } else { println!("{:?}", res.err().unwrap()); } } let file_size = target_file.file.metadata().unwrap().len(); // let file_data = File::open(target_file.file.clone()).await.unwrap(); let pb = ProgressBar::new(file_size); pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})") .unwrap() .with_key("eta", |state: &ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) .progress_chars("#>-")); // let file_name = target_file.file.file_name().unwrap().to_str().unwrap(); // let file_dir = target_file.path.split('/').collect::>() // [..target_file.path.split('/').count() - 1] // .join("/") // + "/"; // println!("file.path: {:?}", file_name); // println!("prefix: {:?}", file_dir); client .upload_from_path( target_file.clone(), None, prefix.as_deref(), fake_size, None, None, Some(pb), ) .await .unwrap(); for i in 0..5 { println!("Checking: {:?}", path); let res = client.info(path.as_str()).await; if res.is_ok() { println!("File exists."); break; } else { println!("{:?}", res.err().unwrap()); if i > 5 { return Err(anyhow::anyhow!("File not exists.")); } tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; continue; } } } Commands::UploadBig { file, prefix, fake_size, force, length, } => { // is folder if file.is_dir() { println!("Can't folder upload"); return Err(anyhow::anyhow!("Can't folder upload")); } if let Some(prefix) = prefix.as_ref() { if !prefix.ends_with('/') { println!("Prefix must end with /"); return Err(anyhow::anyhow!("Prefix must end with /")); } } let mut files = Vec::::new(); // file check if !file.exists() { println!("File not found: {:?}", file); return Err(anyhow::anyhow!("File not found: {:?}", file)); } let target_file = TargetFile { file: file.clone(), path: file.file_name().unwrap().to_str().unwrap().to_string(), }; let path = match prefix.clone() { Some(p) => p + target_file.clone().path.as_str(), None => target_file.clone().path, }; if cfg!(windows) { // replase \ to / files.iter_mut().for_each(|f| { f.path = f.path.replace('\\', "/"); }); } let is_blockdevice = std::fs::metadata(target_file.file.as_path()) .unwrap() .file_type() .is_block_device(); let file_size = if is_blockdevice { let sectors: u64 = std::fs::read_to_string(format!( "/sys/block/{}/size", target_file.file.file_name().unwrap().to_str().unwrap() )) .unwrap() .trim() .parse() .unwrap(); sectors.checked_mul(512).unwrap() } else { target_file.file.metadata().unwrap().len() }; let file_chunk_size = length.unwrap_or(1 * 1024 * 1024 * 1024); // 16GB let mut chunk_count = file_size.checked_div(file_chunk_size).unwrap(); let mut size_of_last_chunk = file_size % file_chunk_size; if size_of_last_chunk == 0 { size_of_last_chunk = file_chunk_size; chunk_count -= 1; } for file_chunk_index in 0..chunk_count + 1 { // if !force { // println!("Checking file existence..."); // println!("Checking: {:?}", path); // let res = client.info(path.as_str()).await; // if res.is_ok() { // println!("File already exists."); // return Err(anyhow::anyhow!("File already exists.")); // } else { // println!("{:?}", res.err().unwrap()); // } // } let this_chunk_size = if chunk_count == file_chunk_index { size_of_last_chunk } else { file_chunk_size }; let offset = file_chunk_index * file_chunk_size; let file_name = target_file .file .file_name() .unwrap() .to_str() .unwrap() .to_string() + "." + (file_chunk_index + 1).to_string().as_str(); let file_path = match prefix.clone() { Some(p) => p + file_name.as_str(), None => file_name.clone(), }; // let file_data = File::open(target_file.file.clone()).await.unwrap(); loop { println!( "Uploading {}/{} {}...", file_chunk_index + 1, chunk_count + 1, file_name ); let pb = ProgressBar::new(this_chunk_size); pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})") .unwrap() .with_key("eta", |state: &ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) .progress_chars("#>-")); match client .upload_from_path( target_file.clone(), Some(file_name.as_str()), prefix.as_deref(), fake_size, Some(offset), Some(this_chunk_size), Some(pb), ) .await { Ok(_) => {} Err(e) => { println!("ERROR: {:#?}", e); continue; } }; for i in 0..20 { println!("Checking: {:?}", file_path); let res = client.info(file_path.as_str()).await; if res.is_ok() { println!("File exists."); break; } else { println!("{:?}", res.err().unwrap()); if i > 20 { return Err(anyhow::anyhow!("File not exists.")); } tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; continue; } } break; } } } Commands::Download { path, prefix } => { client .download(path.as_str(), prefix.as_deref()) .await .unwrap(); } Commands::Move { path, dest } => { client.move_file(&path, &dest).await.unwrap(); } Commands::Delete { path, recursive } => { client.delete(&path, &recursive).await.unwrap(); } Commands::Mkdir { name, path } => { client.mkdir(&name, path.as_deref()).await.unwrap(); } Commands::Copy { src, dest } => { client.copy(&src, &dest).await.unwrap(); } Commands::Rename { path, name } => { client.rename(&path, &name).await.unwrap(); } Commands::Info { path } => { client.info(&path).await.unwrap(); } Commands::Auth {} => { println!("Click the link below to authorize the app:\n"); let link = "https://login.account.rakuten.com/sso/authorize?response_type=code&client_id=rakuten_drive_web&redirect_uri=https://www.rakuten-drive.com/oauth-callback&scope=openid+profile+email&prompt=login&ui_locales=en"; println!("{}\n", link); println!("Paste the URL you were redirected to:"); let mut auth_url = String::new(); std::io::stdin().read_line(&mut auth_url).unwrap(); let auth_url = url::Url::parse(auth_url.trim())?; let params = auth_url.query_pairs().collect::>(); let rid_code = params .iter() .find(|(key, _)| key == "code") .map(|(_, value)| value.to_string()) .ok_or_else(|| anyhow::anyhow!("Code not found in URL"))?; let rid_token_auth_res = client::rid_token_auth(rid_code.as_str()).await?; let token_verify_res = client::get_refresh_token(&rid_token_auth_res.custom_token).await?; println!("Refresh token: {}", token_verify_res.refresh_token); } } Ok(()) }