rakuten-drive-cui/src/main.rs

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<String>,
    },
    #[clap(about = "Upload file")]
    Upload {
        file: PathBuf,

        /// Parent folder path
        #[clap(short, long)]
        prefix: Option<String>,

        /// Upload folder recursively
        #[clap(short, long)]
        recursive: bool,

        /// Send fake file size to server (byte)
        #[clap(short, long)]
        fake_size: Option<u64>,

        /// 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<String>,

        /// Send fake file size to server (byte)
        #[clap(short, long)]
        fake_size: Option<u64>,

        /// Do not check file existence
        #[clap(long)]
        force: bool,

        /// Set division size in bytes
        #[clap(long)]
        length: Option<u64>,
    },
    #[clap(about = "Download file")]
    Download {
        path: String,

        /// Parent folder path
        #[clap(long)]
        prefix: Option<String>,
    },
    #[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<String>,
    },
    #[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::<TargetFile>::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::<Vec<&str>>()
            //     [..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::<TargetFile>::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::<Vec<_>>();

            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(())
}