From a46ef2b053105670e20eb485e02c644a539480e2 Mon Sep 17 00:00:00 2001 From: Arkaprabha Chakraborty Date: Fri, 30 Aug 2024 12:53:24 +0530 Subject: [PATCH] Refactor: Modularize main logic and simplify argument handling --- src/crypto.rs | 72 ++++++++ src/io.rs | 73 ++++++++ src/lib.rs | 3 + src/main.rs | 448 ++++++++++++-------------------------------------- src/utils.rs | 40 +++++ 5 files changed, 292 insertions(+), 344 deletions(-) create mode 100644 src/crypto.rs create mode 100644 src/io.rs create mode 100644 src/lib.rs create mode 100644 src/utils.rs diff --git a/src/crypto.rs b/src/crypto.rs new file mode 100644 index 0000000..98631fc --- /dev/null +++ b/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 { + 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, 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 +} diff --git a/src/io.rs b/src/io.rs new file mode 100644 index 0000000..eb9ac2a --- /dev/null +++ b/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) { + 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, Vec) { + 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 = 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 = vec![0u8; 32]; + salt_bytes.clone_from_slice(&data[..32]); + let mut cypher: Vec = vec![0u8; (metadata.len() - 32) as usize]; + cypher.clone_from_slice(&data[32..]); + + (salt_bytes, cypher) +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..04e1643 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,3 @@ +pub mod crypto; +pub mod io; +pub mod utils; diff --git a/src/main.rs b/src/main.rs index 0f0e8cc..46b3f8a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,235 +1,46 @@ -// 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 { - 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(res: Result, 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) { - 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, Vec) { - 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 = 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 = vec![0u8; 32]; - salt_bytes.clone_from_slice(&data[..32]); - let mut cypher: Vec = vec![0u8; (metadata.len() - 32) as usize]; - cypher.clone_from_slice(&data[32..]); - - (salt_bytes, cypher) -} - -fn decrypt(ciphertext: Vec, 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; 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 { - let arg_str: String = match args.next() { - Some(temp) => temp, - None => exit(0), - }; + args.next(); - match arg_str.as_str() { - "--help" | "-h" => { - match args.next() { - Some(_) => { - eprintln!("{} Too many arguments", "Error:".bright_red()); - exit(1); - } - None => (), - }; + let arg_str: String = match args.next() { + Some(temp) => temp, + None => { + eprintln!( + "{} No arguments provided. Use --help for usage information.", + "Error:".bright_red() + ); + exit(1); + } + }; - println!( - "{} {} + match arg_str.as_str() { + "--help" | "-h" => handle_help(), + "--version" | "-v" => handle_version(), + "--encrypt" | "-e" => handle_encrypt(&mut args), + "--decrypt" | "-d" => handle_decrypt(&mut args), + _ => { + eprintln!("{} Unrecognized argument", "Error:".bright_red()); + exit(1); + } + }; +} + +fn handle_help() { + println!( + "{} {} Rust Simple Text Cipher Machine. USAGE: @@ -247,134 +58,83 @@ COMMAND: {}, {} {} {} Runs the program in decryption mode", - PROGRAM_NAME.bright_yellow(), - PROGRAM_VERSION.bright_blue(), - PROGRAM_NAME.bright_yellow(), - "-h".bright_cyan(), - "--help".bright_cyan(), - "-v".bright_cyan(), - "--version".bright_cyan(), - "-e".bright_cyan(), - "--encrypt".bright_cyan(), - "".bright_magenta(), - "".bright_magenta(), - "-d".bright_cyan(), - "--decrypt".bright_cyan(), - "".bright_magenta(), - "".bright_magenta(), - ); - exit(0); - } + PROGRAM_NAME.bright_yellow(), + PROGRAM_VERSION.bright_blue(), + PROGRAM_NAME.bright_yellow(), + "-h".bright_cyan(), + "--help".bright_cyan(), + "-v".bright_cyan(), + "--version".bright_cyan(), + "-e".bright_cyan(), + "--encrypt".bright_cyan(), + "".bright_magenta(), + "".bright_magenta(), + "-d".bright_cyan(), + "--decrypt".bright_cyan(), + "".bright_magenta(), + "".bright_magenta(), + ); + exit(0); +} - "--version" | "-v" => { - match args.next() { - Some(_) => { - eprintln!("{} Too many arguments", "Error:".bright_red()); - exit(1); - } - None => (), - }; - - println!( - "{} ({}) +fn handle_version() { + println!( + "{} ({}) 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 {}, to the extent permitted by law. Written by Arkaprabha Chakraborty", - PROGRAM_NAME.bright_yellow(), - PROGRAM_VERSION.bright_blue(), - "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 => (), - }; - - 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); - - 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 => (), - }; - - 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); - - exit(0); - } - _ => { - eprintln!("{} Unrecognized argument", "Error:".bright_red()); - exit(1); - } - }; - } + PROGRAM_NAME.bright_yellow(), + PROGRAM_VERSION.bright_blue(), + "NO WARRANTY".bright_red() + ); + exit(0); +} + +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); +} + +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); +} + +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) } diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..974d073 --- /dev/null +++ b/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(res: Result, estr: &str) -> T { + match res { + Ok(temp) => temp, + Err(_) => { + eprintln!("{}", estr.bright_red()); + exit(1) + } + } +} + +pub fn to_array(v: Vec) -> [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); + } + } +}