Refactor: Modularize main logic and simplify argument handling

src/crypto.rs

@@ -0,0 +1,72 @@
+use colored::Colorize;
+use orion::{aead, kdf};
+use std::process::exit;
+
+use crate::utils::into_match;
+
+pub fn encrypt(plaintext: String, secret_key: orion::kdf::SecretKey) -> Vec<u8> {
+    let plaintext = plaintext.into_bytes();
+    match aead::seal(&secret_key, &plaintext) {
+        Ok(temp) => {
+            println!("{} Data was encrypted", "Success:".bright_green());
+            temp
+        }
+        Err(_) => {
+            eprintln!("{} Could not encrypt the data", "Error:".bright_red());
+            exit(1);
+        }
+    }
+}
+
+pub fn decrypt(ciphertext: Vec<u8>, secret_key: orion::kdf::SecretKey) -> String {
+    let plaintext = match aead::open(&secret_key, &ciphertext) {
+        Ok(temp) => {
+            println!("{} Data was decrypted", "Success:".bright_green());
+            temp
+        }
+        Err(_) => {
+            eprintln!(
+                "{} Failed to decrypt the file, please check the password",
+                "Error:".bright_red()
+            );
+            exit(1);
+        }
+    };
+    match String::from_utf8(plaintext) {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("{} Could not convert to String", "Error:".bright_red());
+            exit(1);
+        }
+    }
+}
+
+pub fn get_secret_key(salt_bytes: [u8; 32], password: String) -> orion::kdf::SecretKey {
+    let password = into_match(
+        kdf::Password::from_slice(password.as_bytes()),
+        "Error: Could not use the password",
+    );
+    let salt = into_match(
+        kdf::Salt::from_slice(&salt_bytes),
+        "Error: Could not create the salt",
+    );
+    into_match(
+        kdf::derive_key(&password, &salt, 3, 8, 32),
+        "Error: Could not generate the secret key",
+    )
+}
+
+pub fn get_salt_bytes() -> [u8; 32] {
+    let mut salt_bytes = [0u8; 32];
+    match orion::util::secure_rand_bytes(&mut salt_bytes) {
+        Ok(_) => (),
+        Err(_) => {
+            eprintln!(
+                "{} Could not generate the random bytes for the salt",
+                "Error:".bright_red()
+            );
+            exit(1);
+        }
+    };
+    salt_bytes
+}

src/io.rs

@@ -0,0 +1,73 @@
+use colored::Colorize;
+use std::fs::{read_to_string, File};
+use std::io::{Read, Write};
+use std::process::exit;
+
+use crate::utils::into_match;
+
+pub fn write_plain(path: String, plaintext: String) {
+    let mut file = into_match(
+        File::create(&path),
+        &format!("Error: Could not create {}", path),
+    );
+    into_match(
+        file.write(&plaintext.into_bytes()),
+        &format!("Error: Could not write the data to {}", path),
+    );
+    into_match(
+        file.flush(),
+        &format!("Error: Could not flush data to {}", path),
+    );
+}
+
+pub fn write_cipher(path: String, salt_bytes: [u8; 32], ciphertext: Vec<u8>) {
+    let mut file = into_match(
+        File::create(&path),
+        &format!("Error: Could not create {}", path),
+    );
+    into_match(
+        file.write(&salt_bytes),
+        &format!("Error: Could write the salt to {}", path),
+    );
+    into_match(
+        file.write(&ciphertext),
+        &format!("Error: Could not write the ciphertext to {}", path),
+    );
+    into_match(
+        file.flush(),
+        &format!("Error: Could not flush data to {}", path),
+    );
+}
+
+pub fn read_plain(path: String) -> String {
+    into_match(
+        read_to_string(&path),
+        &format!("Error: Could not read {}", path),
+    )
+}
+
+pub fn read_cipher(path: String) -> (Vec<u8>, Vec<u8>) {
+    let mut file = into_match(
+        File::open(&path),
+        &format!("Error: Could not open {}", path),
+    );
+    let metadata = into_match(
+        File::metadata(&file),
+        &format!("Error: Could not read the metadata off of {}", path),
+    );
+    let mut data: Vec<u8> = vec![0u8; metadata.len() as usize];
+    match file.read(&mut data) {
+        Ok(_) => (),
+        Err(_) => {
+            eprintln!("{} Could not read {}", "Error:".bright_red(), path);
+            exit(1);
+        }
+    };
+
+    let mut salt_bytes: Vec<u8> = vec![0u8; 32];
+    salt_bytes.clone_from_slice(&data[..32]);
+    let mut cypher: Vec<u8> = vec![0u8; (metadata.len() - 32) as usize];
+    cypher.clone_from_slice(&data[32..]);
+
+    (salt_bytes, cypher)
+}

src/lib.rs

@@ -0,0 +1,3 @@
+pub mod crypto;
+pub mod io;
+pub mod utils;

src/main.rs

@@ -1,233 +1,44 @@
-// Program: rustcm-cli (0.1.3-alpha)
-// License: GNU GPL version 3
-// Author:  Arkaprabha Chakraborty
-//
-// Copyright (C) 2023 Arkaprabha Chakraborty
-
 use colored::Colorize;
-use orion::{aead, kdf};
-use passterm::prompt_password_tty;
 use std::env::args;
-use std::fs::{read_to_string, File};
-use std::io::{stdout, Read, Write};
 use std::process::exit;
 
+use rustcm_cli::{
+    crypto::{decrypt, encrypt, get_salt_bytes, get_secret_key},
+    io::{read_cipher, read_plain, write_cipher, write_plain},
+    utils::{get_password, into_match, to_array},
+};
+
 const PROGRAM_NAME: &str = "rustcm-cli";
 const PROGRAM_VERSION: &str = "0.1.3-alpha";
 
-fn get_password(prompt: &str) -> String {
-    into_match(
-        stdout().flush(),
-        format!("{} Could not flush date to stdout", "Error".bright_red()).as_str(),
-    );
-    let password = into_match(
-        prompt_password_tty(Some(prompt)),
-        format!("{} Could not read the password", "Error".bright_red()).as_str(),
-    );
-    //println!();
-
-    password
-}
-
-fn encrypt(plaintext: String, secret_key: orion::kdf::SecretKey) -> Vec<u8> {
-    let plaintext = plaintext.into_bytes();
-    match aead::seal(&secret_key, &plaintext) {
-        Ok(temp) => {
-            println!("{} Data was encrypted", "Success:".bright_green());
-            temp
-        }
-        Err(_) => {
-            eprintln!("{} Could not encrypt the data", "Error:".bright_red());
-            exit(1);
-        }
-    }
-}
-
-fn get_secret_key(salt_bytes: [u8; 32], password: String) -> orion::kdf::SecretKey {
-    let password = into_match(
-        kdf::Password::from_slice(password.as_bytes()),
-        format!("{} Could not use the password", "Error:".bright_red()).as_str(),
-    );
-    let salt = into_match(
-        kdf::Salt::from_slice(&salt_bytes),
-        format!("{} Could not create the salt", "Error:".bright_red()).as_str(),
-    );
-    into_match(
-        kdf::derive_key(&password, &salt, 3, 8, 32),
-        format!(
-            "{} Could not generate the secret key",
-            "Error:".bright_red()
-        )
-        .as_str(),
-    )
-}
-
-fn get_salt_bytes() -> [u8; 32] {
-    let mut salt_bytes = [0u8; 32];
-    match orion::util::secure_rand_bytes(&mut salt_bytes) {
-        Ok(_) => (),
-        Err(_) => {
-            eprintln!(
-                "{} Could not generate the random bytes for the salt",
-                "Error:".bright_red()
-            );
-            exit(1);
-        }
-    };
-
-    salt_bytes
-}
-
-fn into_match<T, E>(res: Result<T, E>, estr: &str) -> T {
-    let ok_val: T = match res {
-        Ok(temp) => temp,
-        Err(_) => {
-            eprintln!("{estr}");
-            exit(1)
-        }
-    };
-
-    ok_val
-}
-
-fn write_plain(path: String, plaintext: String) {
-    let mut file = into_match(
-        File::create(&path),
-        format!("{} Could not create {path}", "Error".bright_red()).as_str(),
-    );
-    into_match(
-        file.write(&plaintext.into_bytes()),
-        format!(
-            "{} Could not write the data to {path}",
-            "Error".bright_red()
-        )
-        .as_str(),
-    );
-    into_match(
-        file.flush(),
-        format!("{} Could not flush data to {path}", "Error".bright_red()).as_str(),
-    );
-}
-
-fn write_cipher(path: String, salt_bytes: [u8; 32], ciphertext: Vec<u8>) {
-    let mut file = into_match(
-        File::create(&path),
-        format!("{} Could not create {path}", "Error".bright_red()).as_str(),
-    );
-    into_match(
-        file.write(&salt_bytes),
-        format!("{} Could write the salt to {path}", "Error".bright_red()).as_str(),
-    );
-    into_match(
-        file.write(&ciphertext),
-        format!(
-            "{} Could not write the ciphertext to {path}",
-            "Error:".bright_red()
-        )
-        .as_str(),
-    );
-    into_match(
-        file.flush(),
-        format!("{} Could not flush data to {path}", "Error:".bright_red()).as_str(),
-    );
-}
-
-fn read_plain(path: String) -> String {
-    into_match(
-        read_to_string(&path),
-        format!("{} Could not read {path}", "Error:".bright_red()).as_str(),
-    )
-}
-
-fn read_cipher(path: String) -> (Vec<u8>, Vec<u8>) {
-    let mut file = into_match(
-        File::open(&path),
-        format!("{} Could not open {path}", "Error:".bright_red()).as_str(),
-    );
-    let metadata = into_match(
-        File::metadata(&file),
-        format!(
-            "{} Could not read the metadata off of {path}",
-            "Error:".bright_red()
-        )
-        .as_str(),
-    );
-    let mut data: Vec<u8> = vec![0u8; metadata.len() as usize];
-    match file.read(&mut data) {
-        Ok(_) => (),
-        Err(_) => {
-            eprintln!("{} Could not read {path}", "Error:".bright_red());
-            exit(1);
-        }
-    };
-
-    let mut salt_bytes: Vec<u8> = vec![0u8; 32];
-    salt_bytes.clone_from_slice(&data[..32]);
-    let mut cypher: Vec<u8> = vec![0u8; (metadata.len() - 32) as usize];
-    cypher.clone_from_slice(&data[32..]);
-
-    (salt_bytes, cypher)
-}
-
-fn decrypt(ciphertext: Vec<u8>, secret_key: orion::kdf::SecretKey) -> String {
-    let plaintext = match aead::open(&secret_key, &ciphertext) {
-        Ok(temp) => {
-            println!("{} Data was decrypted", "Success:".bright_green());
-            temp
-        }
-        Err(_) => {
-            eprintln!(
-                "{} Failed to decrypt the file, please check the password",
-                "Error:".bright_red()
-            );
-            exit(1);
-        }
-    };
-    match String::from_utf8(plaintext) {
-        Ok(temp) => temp,
-        Err(_) => {
-            eprintln!("{} Could not convert to String", "Error:".bright_red());
-            exit(1);
-        }
-    }
-}
-
-fn to_array(v: Vec<u8>) -> [u8; 32] {
-    let slice = v.as_slice();
-    let array: [u8; 32] = match slice.try_into() {
-        Ok(bytes) => bytes,
-        Err(_) => {
-            eprintln!(
-                "{} Expected a vector of length {} but it was {}",
-                "Error:".bright_red(),
-                32,
-                v.len()
-            );
-            exit(1);
-        }
-    };
-
-    array
-}
-
 fn main() {
     let mut args = args();
-    while args.next() != None {
+    args.next();
+
     let arg_str: String = match args.next() {
         Some(temp) => temp,
-            None => exit(0),
+        None => {
+            eprintln!(
+                "{} No arguments provided. Use --help for usage information.",
+                "Error:".bright_red()
+            );
+            exit(1);
+        }
     };
 
     match arg_str.as_str() {
-            "--help" | "-h" => {
+        "--help" | "-h" => handle_help(),
-                match args.next() {
+        "--version" | "-v" => handle_version(),
-                    Some(_) => {
+        "--encrypt" | "-e" => handle_encrypt(&mut args),
-                        eprintln!("{} Too many arguments", "Error:".bright_red());
+        "--decrypt" | "-d" => handle_decrypt(&mut args),
+        _ => {
+            eprintln!("{} Unrecognized argument", "Error:".bright_red());
             exit(1);
         }
-                    None => (),
     };
+}
 
+fn handle_help() {
     println!(
         "{} {}
 Rust Simple Text Cipher Machine.
@@ -264,17 +75,9 @@ COMMAND:
         "<path-to-output>".bright_magenta(),
     );
     exit(0);
-            }
+}
-
-            "--version" | "-v" => {
-                match args.next() {
-                    Some(_) => {
-                        eprintln!("{} Too many arguments", "Error:".bright_red());
-                        exit(1);
-                    }
-                    None => (),
-                };
 
+fn handle_version() {
     println!(
         "{} ({})
 Copyright (C) 2023 Arkaprabha Chakraborty
@@ -288,93 +91,50 @@ Written by Arkaprabha Chakraborty",
         "NO WARRANTY".bright_red()
     );
     exit(0);
-            }
+}
-            "--encrypt" | "-e" => {
-                let path: String = match args.next() {
-                    Some(temp) => temp,
-                    None => {
-                        eprintln!(
-                            "{} Was expecting two arguments but received none",
-                            "Error:".bright_red()
-                        );
-                        exit(1);
-                    }
-                };
-
-                let output: String = match args.next() {
-                    Some(temp) => temp,
-                    None => {
-                        eprintln!(
-                            "{} Was expecting two arguments but received one",
-                            "Error:".bright_red()
-                        );
-                        exit(1);
-                    }
-                };
-
-                match args.next() {
-                    Some(_) => {
-                        eprintln!("{} Too many arguments", "Error:".bright_red());
-                        exit(1);
-                    }
-                    None => (),
-                };
 
+fn handle_encrypt(args: &mut std::env::Args) {
+    let (path, output) = get_args(args);
     let plaintext: String = read_plain(path);
     let salt_bytes = get_salt_bytes();
     let password: String = get_password("Password: ");
     let secret_key = get_secret_key(salt_bytes, password);
     let ciphertext = encrypt(plaintext, secret_key);
-
     write_cipher(output, salt_bytes, ciphertext);
-
+    println!(
+        "{} Encryption completed successfully",
+        "Success:".bright_green()
+    );
     exit(0);
-            }
+}
-            "--decrypt" | "-d" => {
-                let path: String = match args.next() {
-                    Some(temp) => temp,
-                    None => {
-                        eprintln!(
-                            "{} Was expecting two arguments but received none",
-                            "Error:".bright_red()
-                        );
-                        exit(1);
-                    }
-                };
-
-                let output: String = match args.next() {
-                    Some(temp) => temp,
-                    None => {
-                        eprintln!(
-                            "{} Was expecting two arguments but received one",
-                            "Error:".bright_red()
-                        );
-                        exit(1);
-                    }
-                };
-
-                match args.next() {
-                    Some(_) => {
-                        eprintln!("{} Too many arguments", "Error:".bright_red());
-                        exit(1);
-                    }
-                    None => (),
-                };
 
+fn handle_decrypt(args: &mut std::env::Args) {
+    let (path, output) = get_args(args);
     let (salt_bytes, ciphertext) = read_cipher(path);
     let salt_bytes: [u8; 32] = to_array(salt_bytes);
     let password: String = get_password("Password: ");
     let secret_key = get_secret_key(salt_bytes, password);
     let plaintext = decrypt(ciphertext, secret_key);
-
     write_plain(output, plaintext);
-
+    println!(
+        "{} Decryption completed successfully",
+        "Success:".bright_green()
+    );
     exit(0);
-            }
+}
-            _ => {
+
-                eprintln!("{} Unrecognized argument", "Error:".bright_red());
+fn get_args(args: &mut std::env::Args) -> (String, String) {
+    let path: String = into_match(
+        Err(args.next()),
+        "Error: Was expecting two arguments but received none",
+    );
+    let output: String = into_match(
+        Err(args.next()),
+        "Error: Was expecting two arguments but received one",
+    );
+    if args.next().is_some() {
+        eprintln!("{} Too many arguments", "Error:".bright_red());
         exit(1);
     }
-        };
+    (path, output)
-    }
 }

src/utils.rs

@@ -0,0 +1,40 @@
+use colored::Colorize;
+use passterm::prompt_password_tty;
+use std::io::stdout;
+use std::io::Write;
+use std::process::exit;
+
+pub fn get_password(prompt: &str) -> String {
+    into_match(stdout().flush(), "Error: Could not flush date to stdout");
+    let password = into_match(
+        prompt_password_tty(Some(prompt)),
+        "Error: Could not read the password",
+    );
+    password
+}
+
+pub fn into_match<T, E>(res: Result<T, E>, estr: &str) -> T {
+    match res {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("{}", estr.bright_red());
+            exit(1)
+        }
+    }
+}
+
+pub fn to_array(v: Vec<u8>) -> [u8; 32] {
+    let slice = v.as_slice();
+    match slice.try_into() {
+        Ok(bytes) => bytes,
+        Err(_) => {
+            eprintln!(
+                "{} Expected a vector of length {} but it was {}",
+                "Error:".bright_red(),
+                32,
+                v.len()
+            );
+            exit(1);
+        }
+    }
+}