From e62ec8ea3054df9189f16869d7f01eb5e544aac3 Mon Sep 17 00:00:00 2001
From: trivernis <trivernis@gmail.com>
Date: Thu, 20 Feb 2020 21:36:49 +0100
Subject: [PATCH] Improve error handling

---
 src/main.rs | 42 +++++++++++++++++++++---------------------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 4275534..5597850 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,6 +9,7 @@ use rayon::prelude::*;
 use rpassword;
 use rpassword::read_password_from_tty;
 use std::fs::File;
+use std::io;
 use std::io::{Read, Write};
 use structopt::StructOpt;
 
@@ -94,17 +95,18 @@ fn main() {
 fn encrypt(_opts: &Opts, args: &Encrypt) {
     let input = (*args.input).parse().unwrap();
     let output = (*args.output).parse().unwrap();
-    let data: Vec<u8> = read_file_binary(input);
+    let data: Vec<u8> = read_file_binary(input).expect("Failed to read input file!");
 
     if let Some(output_checksum) = (args.clone()).output_checksum {
         let checksum = sha_checksum(&data);
         let checksum_b64 = base64::encode(checksum.as_slice());
-        write_file(output_checksum, checksum_b64.as_bytes());
+        write_file(output_checksum, checksum_b64.as_bytes())
+            .expect("Failed to write checksum file!");
     }
     let pass = read_password_from_tty(Some("Password: ")).unwrap();
     let key = create_key(pass);
     let enc_data = encrypt_data(data.as_slice(), key.as_slice());
-    write_file(output, enc_data.as_slice());
+    write_file(output, enc_data.as_slice()).expect("Failed to write output file!");
 }
 
 /// Decrypts a des encrypted file.
@@ -113,15 +115,15 @@ fn decrypt(_opts: &Opts, args: &Decrypt) {
     let input = (*args.input).parse().unwrap();
     let output = (*args.output).parse().unwrap();
     let dictionary = args.dictionary.clone();
-    let data = read_file_binary(input);
+    let data = read_file_binary(input).expect("Failed to read input file!");
 
     if let Some(input_checksum) = (args.clone()).input_checksum {
-        let bin_content = read_file_binary(input_checksum);
+        let bin_content = read_file_binary(input_checksum).expect("Failed to read checksum file!");
         let data_checksum = base64::decode(bin_content.as_slice()).unwrap();
 
         if let Some(dict) = dictionary {
             println!("Reading dictionary...");
-            let dictionary = read_file(dict);
+            let dictionary = read_file(dict).expect("Failed to read dictionary file!");
             let lines = dictionary.lines().collect::<Vec<&str>>();
 
             let pw_table: Vec<PassKey> = lines
@@ -138,7 +140,7 @@ fn decrypt(_opts: &Opts, args: &Decrypt) {
 
             println!("Starting multithreaded decryption...");
             if let Some(dec_data) = decrypt_with_dictionary(&data, pw_table, &data_checksum) {
-                write_file(output, &dec_data);
+                write_file(output, &dec_data).expect("Failed to write output file!");
                 println!("Finished!");
             } else {
                 println!("No password found!");
@@ -146,7 +148,7 @@ fn decrypt(_opts: &Opts, args: &Decrypt) {
         } else {
             println!("Starting brute force multithreaded decryption...");
             if let Some(dec_data) = decrypt_brute_brute_force(&data, &data_checksum) {
-                write_file(output, &dec_data);
+                write_file(output, &dec_data).expect("Failed to write output file!");
                 println!("Finished!");
             } else {
                 println!("No fitting key found. (This should have been impossible)")
@@ -156,14 +158,14 @@ fn decrypt(_opts: &Opts, args: &Decrypt) {
         let pass = read_password_from_tty(Some("Password: ")).unwrap();
         let key = create_key(pass);
         let result = decrypt_data(&data, key.as_slice());
-        write_file(output, &result);
+        write_file(output, &result).expect("Failed to write output file!");
     }
 }
 
 /// Creates a dictionary from an input file and writes it to the output file
 fn create_dictionary(_opts: &Opts, args: &CreateDictionary) {
     let input = (*args.input).parse().unwrap();
-    let contents = read_file(input);
+    let contents = read_file(input).expect("Failed to read input file!");
     let lines = contents.lines().collect::<Vec<&str>>();
     println!("Parsing {} passwords...", lines.len());
 
@@ -187,23 +189,21 @@ fn create_dictionary(_opts: &Opts, args: &CreateDictionary) {
 }
 
 /// Reads a file to the end and returns the content as byte array
-fn read_file_binary(filename: String) -> Vec<u8> {
-    let mut fin = File::open(filename).unwrap();
+fn read_file_binary(filename: String) -> Result<Vec<u8>, io::Error> {
     let mut data: Vec<u8> = vec![];
-    fin.read_to_end(&mut data).unwrap();
-    data
+    File::open(filename)?.read_to_end(&mut data)?;
+    Ok(data)
 }
 
 /// Reads a file to the end and returns the contents as a string
-fn read_file(filename: String) -> String {
-    let mut fin = File::open(filename).unwrap();
+fn read_file(filename: String) -> Result<String, io::Error> {
     let mut contents = String::new();
-    fin.read_to_string(&mut contents).unwrap();
-    contents
+    File::open(filename)?.read_to_string(&mut contents)?;
+    Ok(contents)
 }
 
 /// writes binary data to a file
-fn write_file(filename: String, data: &[u8]) {
-    let mut fout = File::create(filename).unwrap();
-    fout.write(data).unwrap();
+fn write_file(filename: String, data: &[u8]) -> Result<(), io::Error> {
+    File::create(filename)?.write(data)?;
+    Ok(())
 }