add progress bar encrypt and decrypt command

robatipoor · Mar 12, 2024 · 09cd537 · 09cd537
1 parent 1d734ba
commit 09cd537
Show file tree

Hide file tree

Showing 5 changed files with 160 additions and 11 deletions.
diff --git a/cli/Cargo.toml b/cli/Cargo.toml
@@ -23,3 +23,4 @@ tracing-subscriber = { workspace = true }
 base64 = { workspace = true }
 cuid2 = { workspace = true }
 indicatif = { workspace = true }
+chacha20poly1305 = { workspace = true }
diff --git a/cli/src/args.rs b/cli/src/args.rs
@@ -96,6 +96,8 @@ pub enum SubCommand {
   },
   #[clap(about = "Encrypt a file before uploading to the server")]
   Encrypt {
+    #[clap(default_value_t = false, short, long)]
+    progress_bar: bool,
     #[clap(short, long, value_parser = parse_source_file)]
     source_file: PathBuf,
     #[clap(short, long, value_parser = parse_destination)]
@@ -105,6 +107,8 @@ pub enum SubCommand {
   },
   #[clap(about = "Decrypt a file after downloading from the server")]
   Decrypt {
+    #[clap(default_value_t = false, short, long)]
+    progress_bar: bool,
     #[clap(short, long, value_parser = parse_source_file)]
     source_file: PathBuf,
     #[clap(short, long, value_parser = parse_destination)]

diff --git a/cli/src/command.rs b/cli/src/command.rs
@@ -211,29 +211,51 @@ pub async fn delete(server_addr: String, url_path: FileUrlPath, auth: Option<(St
   }
 }
 
-pub async fn encrypt_file(key_nonce: &KeyNonce, source_file: &Path, mut destination: PathBuf) {
+pub async fn encrypt_file(
+  progress_bar: bool,
+  key_nonce: &KeyNonce,
+  source_file: &Path,
+  mut destination: PathBuf,
+) {
   if destination.is_dir() {
     destination = destination.join(add_extension(
       PathBuf::from(source_file.file_name().unwrap()),
       "bin",
     ));
   }
-  sdk::util::crypto::encrypt_file(key_nonce, source_file, destination)
-    .await
-    .unwrap();
+  if progress_bar {
+    crate::util::crypto::encrypt_file_with_progress_bar(key_nonce, source_file, destination)
+      .await
+      .unwrap();
+  } else {
+    sdk::util::crypto::encrypt_file(key_nonce, source_file, destination)
+      .await
+      .unwrap();
+  }
 }
 
-pub async fn decrypt_file(key_nonce: &KeyNonce, source_file: &Path, mut destination: PathBuf) {
+pub async fn decrypt_file(
+  progress_bar: bool,
+  key_nonce: &KeyNonce,
+  source_file: &Path,
+  mut destination: PathBuf,
+) {
   if destination.is_dir() {
     destination = destination
       .join(rm_extra_extension(PathBuf::from(source_file.file_name().unwrap())).unwrap());
   }
   if source_file == destination {
     panic!("Please specify the valid destination file path.")
   }
-  sdk::util::crypto::decrypt_file(key_nonce, source_file, destination)
-    .await
-    .unwrap();
+  if progress_bar {
+    crate::util::crypto::decrypt_file_with_progress_bar(key_nonce, source_file, destination)
+      .await
+      .unwrap();
+  } else {
+    sdk::util::crypto::decrypt_file(key_nonce, source_file, destination)
+      .await
+      .unwrap();
+  }
 }
 
 fn print_response_err(err: &BodyResponseError) {

diff --git a/cli/src/main.rs b/cli/src/main.rs
@@ -102,24 +102,26 @@ async fn main() {
       command::delete(server_addr, url_path, args.auth).await
     }
     SubCommand::Encrypt {
+      progress_bar,
       source_file,
       destination,
       key_nonce,
     } => {
       if destination.is_file() && destination == source_file {
         panic!("Destination file has an invalid path.")
       }
-      command::encrypt_file(&key_nonce, &source_file, destination).await;
+      command::encrypt_file(progress_bar, &key_nonce, &source_file, destination).await;
     }
     SubCommand::Decrypt {
+      progress_bar,
       source_file,
       destination,
       key_nonce,
     } => {
       if destination.is_file() && destination == source_file {
         panic!("Destination file has an invalid path.")
       }
-      command::decrypt_file(&key_nonce, &source_file, destination).await;
+      command::decrypt_file(progress_bar, &key_nonce, &source_file, destination).await;
     }
   };
 }
diff --git a/cli/src/util/crypto.rs b/cli/src/util/crypto.rs
@@ -1,9 +1,25 @@
+use anyhow::anyhow;
+use chacha20poly1305::{
+  aead::{
+    stream::{DecryptorBE32, EncryptorBE32},
+    KeyInit,
+  },
+  XChaCha20Poly1305,
+};
+use indicatif::ProgressBar;
 use sdk::util::{
-  crypto::{decrypt_file, encrypt_file, KeyNonce},
+  crypto::{decrypt_file, encrypt_file, KeyNonce, DECRYPT_BUFFER_LEN, ENCRYPT_BUFFER_LEN},
   file::{add_extension, add_parent_dir, rm_extra_extension},
   random::generate_random_string_with_prefix,
 };
+
 use std::path::{Path, PathBuf};
+use tokio::{
+  fs::File,
+  io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
+};
+
+use super::progress::progress_bar;
 
 pub async fn encrypt_upload_file(
   key_nonce: &KeyNonce,
@@ -37,6 +53,110 @@ pub async fn decrypt_download_file(
   Ok(())
 }
 
+pub async fn encrypt_file_with_progress_bar(
+  key_nonce: &KeyNonce,
+  plaintext_file: impl AsRef<Path>,
+  destination_file: impl AsRef<Path>,
+) -> anyhow::Result<()> {
+  let reader = File::open(plaintext_file).await?;
+  let writer = File::create(destination_file).await?;
+  let total_size = reader.metadata().await?.len();
+  let pb = progress_bar(total_size)?;
+  encrypt_with_progress_bar(key_nonce, reader, writer, pb).await?;
+  Ok(())
+}
+
+pub async fn decrypt_file_with_progress_bar(
+  key_nonce: &KeyNonce,
+  encrypted_file: impl AsRef<Path>,
+  destination_file: impl AsRef<Path>,
+) -> anyhow::Result<()> {
+  let reader = File::open(encrypted_file).await?;
+  let writer = File::create(destination_file).await?;
+  let total_size = reader.metadata().await?.len();
+  let pb = progress_bar(total_size)?;
+  decrypt_with_progress_bar(key_nonce, reader, writer, pb).await?;
+  Ok(())
+}
+
+pub async fn encrypt_with_progress_bar<R, W>(
+  KeyNonce { key, nonce }: &KeyNonce,
+  mut reader: R,
+  mut writer: W,
+  pb: ProgressBar,
+) -> anyhow::Result<()>
+where
+  R: AsyncRead + Unpin,
+  W: AsyncWrite + Unpin,
+{
+  let mut buffer = [0u8; ENCRYPT_BUFFER_LEN];
+  let mut stream_encryptor =
+    EncryptorBE32::from_aead(XChaCha20Poly1305::new(key), (*nonce).as_ref().into());
+  let mut total_read = 0;
+  loop {
+    let read_count = reader.read(&mut buffer).await?;
+    total_read += read_count;
+    pb.set_position(total_read as u64);
+    if read_count == ENCRYPT_BUFFER_LEN {
+      let ciphertext = stream_encryptor
+        .encrypt_next(buffer.as_slice())
+        .map_err(|err| anyhow!("Encrypting file failed, Error: {err}"))?;
+      writer.write_all(&ciphertext).await?;
+    } else if read_count == 0 {
+      pb.finish_with_message("Encrypt completed successfully.");
+      break;
+    } else {
+      let ciphertext = stream_encryptor
+        .encrypt_last(&buffer[..read_count])
+        .map_err(|err| anyhow!("Encrypting file failed, Error: {err}"))?;
+      writer.write_all(&ciphertext).await?;
+      break;
+    }
+  }
+  writer.flush().await?;
+
+  Ok(())
+}
+
+pub async fn decrypt_with_progress_bar<R, W>(
+  KeyNonce { key, nonce }: &KeyNonce,
+  mut reader: R,
+  mut writer: W,
+  pb: ProgressBar,
+) -> anyhow::Result<()>
+where
+  R: AsyncRead + Unpin,
+  W: AsyncWrite + Unpin,
+{
+  let mut buffer = [0u8; DECRYPT_BUFFER_LEN];
+  let mut stream_decryptor =
+    DecryptorBE32::from_aead(XChaCha20Poly1305::new(key), nonce.as_ref().into());
+  let mut total_read = 0;
+  loop {
+    let read_count = reader.read(&mut buffer).await?;
+    total_read += read_count;
+    pb.set_position(total_read as u64);
+    if read_count == DECRYPT_BUFFER_LEN {
+      let plaintext = stream_decryptor
+        .decrypt_next(buffer.as_slice())
+        .map_err(|err| anyhow!("Decrypting file failed, Error: {err}"))?;
+      writer.write_all(&plaintext).await?;
+    } else if read_count == 0 {
+      pb.finish_with_message("Decrypt completed successfully.");
+      break;
+    } else {
+      let plaintext = stream_decryptor
+        .decrypt_last(&buffer[..read_count])
+        .map_err(|err| anyhow!("Decrypting file failed, Error: {err}"))?;
+      writer.write_all(&plaintext).await?;
+      break;
+    }
+  }
+  writer.flush().await?;
+
+  Ok(())
+}
+
 #[cfg(test)]
 mod tests {