Error handling, and fix database overwrite flaw

Author: replydev

src/argument_functions.rs

@@ -44,11 +44,9 @@ pub fn import(args: Vec<String>){
 
 pub fn add(args: Vec<String>){
     if args.len() == 5{
-        if database_loader::add_element(&args[2],&args[3],&args[4]){
-            println!("Success");
-        }
-        else{
-            println!("Invalid values");
+        match database_loader::add_element(&args[2],&args[3],&args[4]){
+            Ok(()) => println!("Success"),
+            Err(e) => println!("An error occurred: {}",e)
         }
     }
     else{
@@ -59,11 +57,10 @@ pub fn add(args: Vec<String>){
 pub fn remove(args: Vec<String>){
     if args.len() == 3{
         let id = args[2].parse::<usize>().unwrap();
-        if database_loader::remove_element_from_db(id) {
-            println!("ok");
-        }
-        else{
-            println!("{} is a wrong index", id);
+
+        match database_loader::remove_element_from_db(id){
+            Ok(()) => println!("Success"),
+            Err(e) => println!("An error has occurred: {}",e)
         }
     }
     else{
@@ -103,7 +100,7 @@ pub fn export(args: Vec<String>){
 
 pub fn json(args: Vec<String>){
     if args.len() == 2{
-        println!("{}",otp_helper::get_json_results())
+        println!("{}",otp_helper::get_json_results().expect("Failed to get json results"));
     }
     else{
         println!("Invalid argument, type cotp --json");

src/database_loader.rs

@@ -66,18 +66,18 @@ impl OTPElement {
     }
 }
 
-pub fn read_from_file() -> Vec<OTPElement>{
+pub fn read_from_file() -> Result<Vec<OTPElement>,String>{
     let mut encrypted_contents = read_to_string(&get_db_path()).unwrap();
     //rust close files at the end of the function
     let contents = cryptograpy::decrypt_string(&mut encrypted_contents, &cryptograpy::prompt_for_passwords("Password: "));
     match contents {
         Ok(contents) => {
             let vector: Vec<OTPElement> = serde_json::from_str(&contents).unwrap();
-            return vector;
+            return Ok(vector);
         },
         Err(e) => {
             println!("{}", e);
-            return Vec::new();
+            return Err(String::from("Cannot decrypt existing database"));
         }
     }
 }
@@ -88,28 +88,39 @@ pub fn check_secret(secret: &str) -> bool{
     return upper_secret.chars().all(char::is_alphanumeric);
 }
 
-pub fn add_element(secret: &String,issuer: &String,label: &String) -> bool{
+pub fn add_element(secret: &String,issuer: &String,label: &String) -> Result<(),String>{
     if !check_secret(&secret){
-        return false;
+        return Err(String::from("Bad secret"))
     }
     let otp_element = OTPElement::new(secret.to_string(), issuer.to_string(), label.to_string(),6, String::from("TOTP"), String::from("SHA1"),String::from("Default"),0,0,30,vec![]);
-    let mut elements = read_from_file();
+    let mut elements;
+    match read_from_file(){
+        Ok(result) => elements = result,
+        Err(e) => return Err(e)
+    }
     elements.push(otp_element);
     overwrite_database(elements);
-    true
+    Ok(())
 }
 
-pub fn remove_element_from_db(mut id: usize) -> bool{
+pub fn remove_element_from_db(mut id: usize) -> Result<(),String>{
     if id == 0{
-        return false;
+        return Err(String::from("0 is a bad index"));
     }
     //user inserts numbers starting from 1, so we will decrement the value becouse we use array indexes instead
     id -= 1;
 
-    let mut elements: Vec<OTPElement> = read_from_file();
+    let mut elements: Vec<OTPElement>;
+
+    match read_from_file(){
+        Ok(result) => elements = result,
+        Err(e) => {
+            return Err(e);
+        }
+    }
 
     if id >= elements.len(){
-        return false;
+        return Err(format!("{} is a bad index",id+1));
     }
 
     for i in 0..elements.len(){
@@ -119,15 +130,21 @@ pub fn remove_element_from_db(mut id: usize) -> bool{
         }
     }
     overwrite_database(elements);
-    true
+    Ok(())
 }
 
 pub fn edit_element(mut id: usize, secret: &str,issuer: &str,label: &str) -> Result<(), String> {
     if id == 0{
         return Err(String::from("Invalid element"));
     }
     id -= 1;
-    let mut elements: Vec<OTPElement> = read_from_file();
+
+    let mut elements: Vec<OTPElement>;
+    match read_from_file(){
+        Ok(result) => elements = result,
+        Err(e) => return Err(String::from("Cannot decrypt existing database"))
+    }
+    
 
     if id >= elements.len() {
         return Err(String::from("Invalid element"));

src/otp_helper.rs

@@ -23,7 +23,14 @@ impl JsonResult {
 }
 
 pub fn show_codes(){
-    let elements: Vec<database_loader::OTPElement> = database_loader::read_from_file();
+    let elements: Vec<database_loader::OTPElement>;
+    match database_loader::read_from_file(){
+        Ok(result) => elements = result,
+        Err(e) => {
+            println!("An error as occurred: {}",e);
+            return;
+        }
+    }
     for i in 0..elements.len() {
         print_totp(i,&elements[i]);
     }
@@ -49,9 +56,13 @@ fn get_good_otp_code(element: &database_loader::OTPElement) -> String {
     s_otp
 }
 
-pub fn get_json_results() -> String{
-    let elements: Vec<database_loader::OTPElement> = database_loader::read_from_file();
+pub fn get_json_results() -> Result<String,String>{
+    let elements: Vec<database_loader::OTPElement>;
 
+    match database_loader::read_from_file(){
+        Ok(result) => elements = result,
+        Err(e) => return Err(e)
+    }
     let mut results: Vec<JsonResult> = Vec::new();
 
     for i in 0..elements.len() {
@@ -61,5 +72,5 @@ pub fn get_json_results() -> String{
 
     let json_string: &str = &serde_json::to_string_pretty(&results).unwrap();
 
-    json_string.to_string()
+    Ok(json_string.to_string())
 }