Update src/main.rs: Added some error handling

src/main.rs

@@ -5,85 +5,206 @@
 // Copyright (C) 2023 Arkaprabha Chakraborty
 
 use orion::{aead, kdf};
-use passterm;
+use passterm::read_password;
 use std::env::args;
 use std::fs::read_to_string;
 use std::fs::File;
 use std::io::{stdout, Read, Write};
+use std::process::exit;
 
 const PROGRAM_NAME: &str = "rustcm-cli";
 const PROGRAM_VERSION: &str = "0.1.0-alpha";
 
 pub fn get_password(prompt: &str) -> String {
     print!("{}", prompt);
-    stdout().flush().unwrap();
+    match stdout().flush() {
-    let password: String = passterm::read_password().unwrap();
+        Ok(_) => (),
+        Err(_) => {
+            eprintln!("Error: Could not flush stdout");
+            exit(0);
+        }
+    };
+
+    let password: String = match read_password() {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not read the password");
+            exit(0)
+        }
+    };
     println!();
+
     password
 }
 
 pub fn encrypt(plaintext: String, secret_key: orion::kdf::SecretKey) -> Vec<u8> {
     let plaintext = plaintext.into_bytes();
-    let ciphertext: Vec<u8> =
+    let ciphertext: Vec<u8> = match aead::seal(&secret_key, &plaintext) {
-        aead::seal(&secret_key, &plaintext).expect("Error: Could not encrypt the data");
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not encrypt the data");
+            exit(0);
+        }
+    };
 
     ciphertext
 }
 
-pub fn get_secret_key(presalt: [u8; 32], password: String) -> orion::kdf::SecretKey {
+pub fn get_secret_key(salt_bytes: [u8; 32], password: String) -> orion::kdf::SecretKey {
-    let password = kdf::Password::from_slice(password.as_bytes()).unwrap();
+    let password = match kdf::Password::from_slice(password.as_bytes()) {
-    let salt = kdf::Salt::from_slice(&presalt).unwrap();
+        Ok(temp) => temp,
-    let secret_key =
+        Err(_) => {
-        kdf::derive_key(&password, &salt, 3, 8, 32).expect("Error: Could not create secret key");
+            eprintln!("Error: Could not generate the password");
+            exit(0)
+        }
+    };
+    let salt = match kdf::Salt::from_slice(&salt_bytes) {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not generate the salt");
+            exit(0)
+        }
+    };
+    let secret_key = match kdf::derive_key(&password, &salt, 3, 8, 32) {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not generate the secret key");
+            exit(0)
+        }
+    };
 
     secret_key
 }
 
-pub fn get_presalt() -> [u8; 32] {
+pub fn get_salt_bytes() -> [u8; 32] {
-    let mut presalt = [0u8; 32];
+    let mut salt_bytes = [0u8; 32];
-    orion::util::secure_rand_bytes(&mut presalt).expect("Error: Could not get presalt");
+    match orion::util::secure_rand_bytes(&mut salt_bytes) {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not generate the salt");
+            exit(0)
+        }
+    };
 
-    presalt
+    salt_bytes
 }
 
-pub fn get_salt(presalt: [u8; 32]) -> orion::kdf::Salt {
+pub fn get_salt(salt_bytes: [u8; 32]) -> orion::kdf::Salt {
-    let salt = kdf::Salt::from_slice(&presalt).expect("Error: Could not create salt");
+    let salt = match kdf::Salt::from_slice(&salt_bytes) {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not generate the salt");
+            exit(0)
+        }
+    };
 
     salt
 }
 
 pub fn write_plain(path: String, plaintext: String) {
-    let mut file = File::create(path).expect("Error: Could not create the file");
+    let mut file = match File::create(path) {
-    file.write(&plaintext.into_bytes())
+        Ok(temp) => temp,
-        .expect("Error: Could not write_cipher plaintext to the file");
+        Err(_) => {
-    file.flush().unwrap();
+            eprintln!("Error: Could not create the file");
+            exit(0);
+        }
+    };
+
+    match file.write(&plaintext.into_bytes()) {
+        Ok(_) => (),
+        Err(_) => {
+            eprintln!("Error: Could not write the data to the file");
+            exit(0);
+        }
+    };
+
+    match file.flush() {
+        Ok(_) => (),
+        Err(_) => {
+            eprintln!("Error: Could not flush the file");
+            exit(0);
+        }
+    };
 }
 
-pub fn write_cipher(path: String, presalt: [u8; 32], ciphertext: Vec<u8>) {
+pub fn write_cipher(path: String, salt_bytes: [u8; 32], ciphertext: Vec<u8>) {
-    let mut file = File::create(path).expect("Error: Could not create the file");
+    let mut file = match File::create(path) {
-    file.write(&presalt).unwrap();
+        Ok(temp) => temp,
-    file.write(&ciphertext)
+        Err(_) => {
-        .expect("Error: Could not write_cipher presalt to the file");
+            eprintln!("Error: Could not create the file");
-    file.flush().unwrap();
+            exit(0);
+        }
+    };
+
+    match file.write(&salt_bytes) {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could write the salt_bytes to the file");
+            exit(0);
+        }
+    };
+
+    match file.write(&ciphertext) {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not write the ciphertext to the file");
+            exit(0);
+        }
+    };
+
+    match file.flush() {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not flush the file");
+            exit(0);
+        }
+    };
 }
 
 pub fn read_plain(path: String) -> String {
-    read_to_string(path).unwrap()
+    let file_str: String = match read_to_string(path) {
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not read the file");
+            exit(0);
+        }
+    };
+
+    file_str
 }
 
 pub fn read_cipher(path: String) -> (Vec<u8>, Vec<u8>) {
-    let mut file = File::open(path).expect("Error: Could not open the file");
+    let mut file = match File::open(path) {
-    let metadata = File::metadata(&file).expect("Error: Could not read the metadata off the file");
+        Ok(temp) => temp,
-    let mut data: Vec<u8> = vec![0u8; metadata.len() as usize];
+        Err(_) => {
-    file.read(&mut data)
+            eprintln!("Error: Could not open the file");
-        .expect("Error: Could not read from the file");
+            exit(0);
+        }
+    };
 
-    let mut presalt: Vec<u8> = vec![0u8; 32];
+    let metadata = match File::metadata(&file) {
-    presalt.clone_from_slice(&data[..32]);
+        Ok(temp) => temp,
+        Err(_) => {
+            eprintln!("Error: Could not read the metadata off the file");
+            exit(0);
+        }
+    };
+
+    let mut data: Vec<u8> = vec![0u8; metadata.len() as usize];
+    match file.read(&mut data) {
+        Ok(_) => (),
+        Err(_) => {
+            eprintln!("Error: Could not read the file");
+            exit(0);
+        }
+    };
+
+    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..]);
 
-    (presalt, cypher)
+    (salt_bytes, cypher)
 }
 
 pub fn decrypt(ciphertext: Vec<u8>, secret_key: orion::kdf::SecretKey) -> String {
@@ -104,11 +225,14 @@ where
     let slice = v.as_slice();
     let array: [T; 32] = match slice.try_into() {
         Ok(bytes) => bytes,
-        Err(_) => panic!(
+        Err(_) => {
-            "Error: Expected a Vec of length {} but it was {}",
+            eprintln!(
+                "Error: Expected a vector of length {} but it was {}",
                 32,
                 v.len()
-        ),
+            );
+            exit(0);
+        }
     };
 
     array
@@ -117,11 +241,18 @@ where
 fn main() {
     let mut args = args();
     while args.next() != None {
-        let arg_str = args.next().unwrap();
+        let arg_str: String = match args.next() {
-        if arg_str.eq("--help") || arg_str.eq("-h") {
+            Some(temp) => temp,
+            None => {
+                exit(0);
+            }
+        };
+
+        match arg_str.as_str() {
+            "--help" | "-h" => {
                 println!(
                     "{PROGRAM_NAME} {PROGRAM_VERSION}
-Rust Simple Text Cipher Machine. Encrypts or decrypts files using the ChaCha20-Poly1305 algorithm.
+Rust Simple Text Cipher Machine.
 
 USAGE:
     rustcm-cli [COMMAND]
@@ -138,36 +269,77 @@ COMMAND:
 
     -d, --decrypt <input-path> <output-path>
         Runs the program in decryption mode."
-            );
+                )
-        } else if arg_str.eq("--version") || arg_str.eq("-v") {
+            }
+
+            "--version" | "-v" => {
                 println!(
-                "rustcm-cli (0.1.0)
+                    "rustcm-cli (0.1.0-alpha)
 Copyright (C) 2023 Arkaprabha Chakraborty
 License GPLv3: GNU GPL version 3
 This is free software: you are free to change and redistribute it.
 There is NO WARRANTY, to the extent permitted by law.
 
 Written by Arkaprabha Chakraborty"
-            )
+                );
             }
-        if arg_str.eq("--encrypt") || arg_str.eq("-e") {
+            "encrypt" | "-e" => {
+                let path: String = match args.next() {
+                    Some(temp) => temp,
+                    None => {
+                        eprintln!("Error: Was expecting two arguments but received none");
+                        exit(0);
+                    }
+                };
+
+                let output: String = match args.next() {
+                    Some(temp) => temp,
+                    None => {
+                        eprintln!("Error: Was expecting two arguments but received one");
+                        exit(0);
+                    }
+                };
+
                 let password: String = get_password("Password: ");
-            let path: String = args.next().unwrap();
+                let plaintext: String = read_plain(path);
-            let output: String = args.next().unwrap();
+                let salt_bytes = get_salt_bytes();
-            let plaintext: String = read_plain(path.clone());
+                let secret_key = get_secret_key(salt_bytes, password);
-            let presalt = get_presalt();
-            let secret_key = get_secret_key(presalt, password.clone());
                 let ciphertext = encrypt(plaintext, secret_key);
-            write_cipher(output, presalt, ciphertext);
+
-        } else if arg_str.eq("--decrypt") || arg_str.eq("-d") {
+                write_cipher(output, salt_bytes, ciphertext);
+
+                println!("File encrypted successfully")
+            }
+            "--decrypt" | "-d" => {
+                let path: String = match args.next() {
+                    Some(temp) => temp,
+                    None => {
+                        eprintln!("Error: Was expecting two arguments but received none");
+                        exit(0);
+                    }
+                };
+
+                let output: String = match args.next() {
+                    Some(temp) => temp,
+                    None => {
+                        eprintln!("Error: Was expecting two arguments but received one");
+                        exit(0);
+                    }
+                };
+
                 let password: String = get_password("Password: ");
-            let path: String = args.next().unwrap();
+                let (salt_bytes, ciphertext) = read_cipher(path);
-            let output: String = args.next().unwrap();
+                let salt_bytes: [u8; 32] = convert_to_array(salt_bytes);
-            let (presalt, ciphertext) = read_cipher(path.clone());
+                let secret_key = get_secret_key(salt_bytes, password);
-            let presalt: [u8; 32] = convert_to_array(presalt);
-            let secret_key = get_secret_key(presalt, password);
                 let plaintext = decrypt(ciphertext, secret_key);
+
                 write_plain(output, plaintext);
-        }
+
+                println!("File decryption was successful")
+            }
+            _ => {
+                eprintln!("Error: Unrecognized argument");
+            }
+        };
     }
 }