From 1772f4f9ed2b9d2cd1b2f263520cc01eef780612 Mon Sep 17 00:00:00 2001 From: Arshdeep54 Date: Thu, 19 Dec 2024 22:01:27 +0530 Subject: [PATCH] fix cargo test --- src/auth.rs | 5 ++++- src/lib/logger.rs | 12 ++---------- src/lib/utils.rs | 5 +++++ src/ssh.rs | 3 ++- src/su.rs | 4 ++-- src/sudo.rs | 5 +++-- tests/logger_tests.rs | 8 +++++--- 7 files changed, 23 insertions(+), 19 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index b289ecf..a37e335 100644 --- 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 index 75106bd..59455b8 100644 --- 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 index 36c8590..1e0351a 100644 --- 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 index 73bffc2..ecf9c2e 100644 --- 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 index 779db47..65d7f09 100644 --- 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 index 442317d..51445a5 100644 --- 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 index c92a83d..523fff0 100644 --- 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"); } } \ No newline at end of file