fix cargo test

sdslabs · Dec 19, 2024 · 1772f4f · 1772f4f
1 parent 6f9ca23
commit 1772f4f
Show file tree

Hide file tree

Showing 7 changed files with 23 additions and 19 deletions.
diff --git a/src/auth.rs b/src/auth.rs
@@ -8,6 +8,7 @@ use lib::init::init;
 use lib::keyhouse::{get_name, validate_user};
 use lib::notifier;
 use lib::logger;
+use lib::utils::AUTH_LOG_PATH;
 
 pub fn handle_auth(ssh_host_username: &str, ssh_key: &str) -> Result<()> {
     let config = read_config()?;
@@ -36,9 +37,10 @@ pub fn handle_auth(ssh_host_username: &str, ssh_key: &str) -> Result<()> {
                     match fork() {
                         Ok(ForkResult::Parent { .. }) => {}
                         Ok(ForkResult::Child) => {
-                            if let Err(e) = logger::log("auth", "SUCCESS", &format!("User: {}", name)) {
+                            if let Err(e) = logger::log(AUTH_LOG_PATH, "SUCCESS", &format!("User: {}", name)) {
                                 println!("Failed to log: {}", e);
                             }
+                            std::process::exit(0); 
                         }
                         Err(_) => println!("Fork failed"),
                     }
@@ -49,6 +51,7 @@ pub fn handle_auth(ssh_host_username: &str, ssh_key: &str) -> Result<()> {
                         name,
                         ssh_host_username.to_string(),
                     )?;
+                    std::process::exit(0); 
                 }
                 Err(_) => println!("Fork failed"),
             }

diff --git a/src/lib/logger.rs b/src/lib/logger.rs
@@ -3,15 +3,7 @@ use std::io::Write;
 use std::time::{SystemTime, UNIX_EPOCH};
 use std::io::Result;
 
-pub fn log(filetype: &str, status: &str, message: &str) -> Result<()> {
-    let filename = match filetype {
-        "ssh" => "/opt/watchdog/custom-logs/ssh.logs",
-        "sudo" => "/opt/watchdog/custom-logs/sudo.logs",
-        "su" => "/opt/watchdog/custom-logs/su.logs",
-        "auth" => "/opt/watchdog/custom-logs/auth.logs",
-        _ => return Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "Invalid filetype")),
-    };
-
+pub fn log(filepath: &str, status: &str, message: &str) -> Result<()> {
     let start = SystemTime::now();
     let since_the_epoch = start.duration_since(UNIX_EPOCH).expect("Time went backwards");
     let timestamp = since_the_epoch.as_secs();
@@ -21,7 +13,7 @@ pub fn log(filetype: &str, status: &str, message: &str) -> Result<()> {
     let mut file = OpenOptions::new()
         .append(true)
         .create(true)
-        .open(filename)?;
+        .open(filepath)?;
 
     file.write_all(log_message.as_bytes())?;
     Ok(())

diff --git a/src/lib/utils.rs b/src/lib/utils.rs
@@ -2,6 +2,11 @@ use std::fs;
 
 use crate::errors::*;
 
+pub const AUTH_LOG_PATH: &str = "/opt/watchdog/custom-logs/auth.logs";
+pub const SSH_LOG_PATH: &str = "/opt/watchdog/custom-logs/ssh.logs";
+pub const SUDO_LOG_PATH: &str = "/opt/watchdog/custom-logs/sudo.logs";
+pub const SU_LOG_PATH: &str = "/opt/watchdog/custom-logs/su.logs";
+
 pub fn clear_file(path: &str) -> Result<()> {
     fs::write(path, "")?;
     Ok(())

diff --git a/src/ssh.rs b/src/ssh.rs
@@ -11,6 +11,7 @@ use lib::keyhouse::get_name;
 use lib::notifier;
 use lib::utils::clear_file;
 use lib::logger;
+use lib::utils::SSH_LOG_PATH;
 
 pub fn handle_ssh() -> Result<()> {
     let pam_type = env::var("PAM_TYPE")
@@ -31,7 +32,7 @@ pub fn handle_ssh() -> Result<()> {
                 match fork() {
                     Ok(ForkResult::Parent { .. }) => {}
                     Ok(ForkResult::Child) => {
-                        if let Err(e) = logger::log("ssh", "SUCCESS", &format!("User: {}", name)) {
+                        if let Err(e) = logger::log(SSH_LOG_PATH, "SUCCESS", &format!("User: {}", name)) {
                             println!("Failed to log: {}", e);
                         }
                     }

diff --git a/src/su.rs b/src/su.rs
@@ -8,6 +8,7 @@ use lib::errors::*;
 use lib::init::init;
 use lib::notifier;
 use lib::logger;
+use lib::utils::SU_LOG_PATH;
 
 pub fn handle_su() -> Result<()> {
     let pam_type = env::var("PAM_TYPE")
@@ -29,8 +30,7 @@ pub fn handle_su() -> Result<()> {
                 match fork() {
                     Ok(ForkResult::Parent { .. }) => {}
                     Ok(ForkResult::Child) => {
-                        // Call the log function in this child process
-                        if let Err(e) = logger::log("su", "SUCCESS", &format!("User: {}", pam_user)) {
+                        if let Err(e) = logger::log(SU_LOG_PATH, "SUCCESS", &format!("User: {}", pam_user)) {
                             println!("Failed to log: {}", e);
                         }
                     }

diff --git a/src/sudo.rs b/src/sudo.rs
@@ -8,6 +8,8 @@ use lib::errors::*;
 use lib::init::init;
 use lib::notifier;
 use lib::logger;
+use lib::utils::SUDO_LOG_PATH;
+
 pub fn handle_sudo() -> Result<()> {
     let pam_type = env::var("PAM_TYPE")
                      .chain_err(|| "PAM_TYPE not set. If you are running this by `watchdog sudo`, please don't. It's an internal command, intended to be used by PAM.")?;
@@ -25,8 +27,7 @@ pub fn handle_sudo() -> Result<()> {
                 match fork() {
                     Ok(ForkResult::Parent { .. }) => {}
                     Ok(ForkResult::Child) => {
-                        // Call the log function in this child process
-                        if let Err(e) = logger::log("sudo", "SUCCESS", &format!("User: {}", pam_ruser)) {
+                        if let Err(e) = logger::log(SUDO_LOG_PATH, "SUCCESS", &format!("User: {}", pam_ruser)) {
                             println!("Failed to log: {}", e);
                         }
                     }

diff --git a/tests/logger_tests.rs b/tests/logger_tests.rs
@@ -6,17 +6,19 @@ mod tests {
 
     #[test]
     fn test_log() {
-        let filetype = "ssh";
+        let filepath = "test_ssh.logs";
         let status = "INFO";
         let message = "Test log message";
 
-        log(filetype, status, message).expect("Failed to write log");
+        log(filepath, status, message).expect("Failed to write log");
 
-        let mut file = fs::File::open("ssh.logs").expect("Failed to open log file");
+        let mut file = fs::File::open(filepath).expect("Failed to open log file");
         let mut contents = String::new();
         file.read_to_string(&mut contents).expect("Failed to read log file");
 
         assert!(contents.contains("Test log message"));
         assert!(contents.contains("INFO"));
+
+        fs::remove_file(filepath).expect("Failed to delete test log file");
     }
 }