3 files changed, 28 insertions, 19 deletions

diff --git a/src/main.rs b/src/main.rs
index c6e8201..850d43b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -14,7 +14,6 @@ const ERROR_CONFIG:         u8 = 2;
 const ERROR_NO_USER:        u8 = 3;
 const ERROR_NOT_AUTHORIZED: u8 = 4;
 const ERROR_AUTH_FAILED:    u8 = 5;
-const ERROR_RUN_ROOT:       u8 = 6;
 
 
 fn main() -> ExitCode {
@@ -74,19 +73,12 @@ fn main() -> ExitCode {
         }
     };
 
-
     // authenticate the user
     if !validate(&user.name, persist) {
         eprintln!("Authentication failed.");
         return ExitCode::from(ERROR_AUTH_FAILED);
     }
 
-    // set the uid and gid of the process to root to run the command as root
-    if !unistd::setuid(unistd::geteuid()).is_ok() || !unistd::setgid(unistd::getegid()).is_ok() {
-        eprintln!("Failed to set root permissions");
-        return ExitCode::from(ERROR_RUN_ROOT);
-    };
-
     // execute the passed command
     let start = 1 + flags.arg_count;
     let err = exec::execvp(&args[start], &args[start..]);
@@ -128,6 +120,7 @@ struct Config {
 /// * `false` - If the user failed to authenticate
 fn validate(user: &str, persist: bool) -> bool {
     if persist && persist::get_persist(user) {
+        secure::elevate_privilages(0, 0);
         return true;
     }
     let input = match rpassword::prompt_password(format!("crab ({}) password: ", user)) {
@@ -142,6 +135,7 @@ fn validate(user: &str, persist: bool) -> bool {
     if !auth.authenticate().is_ok() || !auth.open_session().is_ok() {
         return false;
     }
+    secure::elevate_privilages(0, 0);
     if persist {
         persist::set_persist(user);
     }
@@ -222,7 +216,7 @@ fn config(path: &str) -> Option<Config> {
     let mut identitys = vec![];
     for (line_num, line) in file.split("\n").enumerate() {
         let args: Vec<&str> = line.split(" ").collect();
-        if line.trim() == "" {
+        if line.starts_with("#") || line.trim() == "" {
             continue;
         }
         if args.len() < 2 {
diff --git a/src/persist.rs b/src/persist.rs
index fe4ce69..0e0d1fd 100644
--- a/src/persist.rs
+++ b/src/persist.rs
@@ -34,7 +34,7 @@ pub fn get_persist(user: &str) -> bool {
 pub fn set_persist(user: &str) {
   let mut json = match get_persist_config() {
       Some(data) => data,
-      None => serde_json::from_str("{}").unwrap()
+      None => return
   };
   json[user] = Value::from(now());
   let session = match get_current_session() {
@@ -43,8 +43,8 @@ pub fn set_persist(user: &str) {
   };
   match secure::write_file(PERSIST_PATH, &format!("{}", session), &json.to_string())  {
     Ok(_) => {},
-    Err(_) => {
-        eprintln!("crab: An Internal Has Error")
+    Err(e) => {
+        eprintln!("crab: An Internal Has Error: {}", e);
     }
   };
 }
@@ -78,7 +78,9 @@ fn get_persist_config() -> Option<Value> {
   };
   let data = match secure::read_file(PERSIST_PATH, &format!("{}", session)) {
       Some(data) => data,
-      None => "{}".to_string()
+      None => {
+        "{}".to_string()
+      }
   };
   let json: Value = match serde_json::from_str(&data) {
       Ok(data) => data,
diff --git a/src/secure.rs b/src/secure.rs
index 39339a0..f463048 100644
--- a/src/secure.rs
+++ b/src/secure.rs
@@ -1,5 +1,5 @@
-use std::{os::{unix::prelude::PermissionsExt, linux::fs::MetadataExt}, fs, io};
-use nix::unistd;
+use std::{os::{unix::prelude::PermissionsExt, linux::fs::MetadataExt}, fs, io::{self, ErrorKind}};
+use nix::unistd::{self, Uid, Gid};
 
 
 /// Writes a file securly to a specified path with given data
@@ -11,10 +11,10 @@ use nix::unistd;
 /// A ``io::Result<()>`` if the write succeded or failed
 pub fn write_file(dir: &str, file: &str, data: &str) -> Result<(), io::Error> {
   fs::create_dir_all(dir)?;
-  set_file_permissions(0, 0, 0o100600, dir)?;
+  set_file_permissions(0, 0, 0o600, dir)?;
   let path = path(dir, file);
   fs::write(&path, "")?;
-  set_file_permissions(0, 0, 0o100600, &path)?;
+  set_file_permissions(0, 0, 0o600, &path)?;
   fs::write(&path, data)?;
   Ok(())
 }
@@ -41,6 +41,17 @@ pub fn read_file(dir: &str, file: &str) -> Option<String> {
 }
 
 
+pub fn elevate_privilages(uid: u32, gid: u32) -> bool {
+  if unistd::setuid(Uid::from(uid)).is_err() {
+    return false;
+  } 
+  if unistd::setgid(Gid::from(gid)).is_err() {
+    return false;
+  }
+  true
+}
+
+
 /// Sets the permission for a secure file
 /// #### Arguments
 /// * `uid` - The user to own the file
@@ -70,7 +81,9 @@ fn set_file_permissions(uid: u32, gid: u32, mode: u32, path: &str) -> Result<(),
 fn check_file_permissions(uid: u32, gid: u32, mode: u32, path: &str) -> bool {
   let metadata = match fs::metadata(path) {
     Ok(data) => data,
-    Err(_) => return false
+    Err(e) => {
+      return e.kind() == ErrorKind::NotFound;
+    }
   };
   let perms = metadata.permissions();
   return perms.mode() == mode && metadata.st_uid() == uid && metadata.st_gid() == gid;
@@ -79,5 +92,5 @@ fn check_file_permissions(uid: u32, gid: u32, mode: u32, path: &str) -> bool {
 
 /// Get the path of a file given a directory and file name
 fn path(dir: &str, file: &str) -> String {
-  return format!("{}/{}", dir, file);
+  return format!("{}/{}.persist", dir, file);
 }
\ No newline at end of file