diff --git a/src/cron.rs b/src/cron.rs index 6bcef68..ca34130 100644 --- a/src/cron.rs +++ b/src/cron.rs @@ -6,8 +6,18 @@ use std::os::unix::fs::MetadataExt; use std::path::Path; use std::path::PathBuf; use std::process::Command; +use std::time::{SystemTime, UNIX_EPOCH}; -/// Get the path to the user's crontab file +const MANAGED_BEGIN: &str = "# mihoto: begin auto-update"; +const MANAGED_END: &str = "# mihoto: end auto-update"; + +#[derive(Clone, Debug, Eq, PartialEq)] +struct CrontabSnapshot { + content: String, + existed: bool, +} + +/// Get the path to the staging file used by `crontab `. fn crontab_path() -> PathBuf { let run_dir = env::var("XDG_RUNTIME_DIR").unwrap_or_else(|_| { // Use current user's UID as fallback @@ -17,6 +27,20 @@ fn crontab_path() -> PathBuf { PathBuf::from(run_dir).join("mihoro-crontab") } +/// Get the path for a best-effort pre-write crontab backup. +fn crontab_backup_path() -> PathBuf { + let base = env::var("XDG_STATE_HOME") + .map(PathBuf::from) + .or_else(|_| env::var("XDG_CACHE_HOME").map(PathBuf::from)) + .or_else(|_| env::var("HOME").map(|home| PathBuf::from(home).join(".local/state"))) + .unwrap_or_else(|_| env::temp_dir()); + let secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs()) + .unwrap_or(0); + base.join("mihoro").join(format!("crontab-{}.bak", secs)) +} + /// Get the mihoro binary path from current executable fn mihoro_bin_path() -> Result { env::current_exe()? @@ -25,22 +49,257 @@ fn mihoro_bin_path() -> Result { .ok_or_else(|| anyhow!("Failed to get mihoro binary path")) } +fn quote_cron_command(command: &str) -> String { + format!("'{}'", command.replace('\'', "'\\''")) +} + +fn generate_cron_entry_for_bin(interval_hours: u16, bin_path: &str) -> String { + format!( + "0 */{} * * * {} update", + interval_hours, + quote_cron_command(bin_path) + ) +} + /// Generate cron entry for auto-update fn generate_cron_entry(interval_hours: u16) -> Result { let bin_path = mihoro_bin_path()?; Ok(format!( - "0 */{} * * * {} update\n", - interval_hours, bin_path + "{}\n", + generate_cron_entry_for_bin(interval_hours, &bin_path) )) } -/// Generate the crontab content with mihoro entry -fn generate_crontab(interval_hours: u16) -> Result { - let mihoro_entry = generate_cron_entry(interval_hours)?; - Ok(mihoro_entry) +fn build_managed_cron_block(interval_hours: u16, bin_path: &str) -> String { + format!( + "{}\n{}\n{}\n", + MANAGED_BEGIN, + generate_cron_entry_for_bin(interval_hours, bin_path), + MANAGED_END + ) +} + +fn remove_marked_cron_blocks(content: &str) -> (String, usize) { + let mut output = String::with_capacity(content.len()); + let mut pending_block: Option = None; + let mut removed = 0; + + for line in content.split_inclusive('\n') { + let trimmed = line.trim_end_matches(['\r', '\n']).trim(); + + if let Some(block) = pending_block.as_mut() { + block.push_str(line); + if trimmed == MANAGED_END { + pending_block = None; + removed += 1; + } + continue; + } + + if trimmed == MANAGED_BEGIN { + pending_block = Some(line.to_string()); + } else { + output.push_str(line); + } + } + + if let Some(block) = pending_block { + output.push_str(&block); + } + + (output, removed) +} + +fn count_managed_cron_blocks(content: &str) -> usize { + let (_, removed) = remove_marked_cron_blocks(content); + removed +} + +fn cron_command_after_schedule(line: &str) -> Option<&str> { + let mut fields = 0; + let mut in_field = false; + + for (idx, ch) in line.char_indices() { + if ch.is_whitespace() { + if in_field { + fields += 1; + in_field = false; + if fields == 5 { + return Some(line[idx..].trim_start()); + } + } + } else { + in_field = true; + } + } + + None +} + +fn is_legacy_current_binary_entry(line: &str, bin_path: &str) -> bool { + let trimmed = line.trim(); + if trimmed.is_empty() || trimmed.starts_with('#') { + return false; + } + + let Some(command) = cron_command_after_schedule(trimmed) else { + return false; + }; + command == format!("{} update", bin_path) + || command == format!("{} update", quote_cron_command(bin_path)) +} + +fn remove_legacy_current_binary_entries(content: &str, bin_path: &str) -> String { + let mut output = String::with_capacity(content.len()); + for line in content.split_inclusive('\n') { + let trimmed = line.trim_end_matches(['\r', '\n']); + if !is_legacy_current_binary_entry(trimmed, bin_path) { + output.push_str(line); + } + } + + output +} + +fn remove_managed_cron_content(content: &str, bin_path: &str) -> String { + let (without_blocks, _) = remove_marked_cron_blocks(content); + remove_legacy_current_binary_entries(&without_blocks, bin_path) +} + +fn merge_managed_cron_block(existing: &str, block: &str, bin_path: &str) -> String { + let mut merged = remove_managed_cron_content(existing, bin_path); + if !merged.is_empty() && !merged.ends_with('\n') { + merged.push('\n'); + } + merged.push_str(block); + merged +} + +fn is_no_crontab_message(stderr: &str) -> bool { + stderr.to_ascii_lowercase().contains("no crontab") +} + +fn snapshot_from_failed_crontab_list(stderr: &str) -> Result { + if is_no_crontab_message(stderr) { + return Ok(CrontabSnapshot { + content: String::new(), + existed: false, + }); + } + + anyhow::bail!("Failed to read crontab: {}", stderr.trim()); +} + +fn read_current_crontab() -> Result { + let output = Command::new("crontab").arg("-l").output()?; + if output.status.success() { + return Ok(CrontabSnapshot { + content: String::from_utf8_lossy(&output.stdout).into_owned(), + existed: true, + }); + } + + let stderr = String::from_utf8_lossy(&output.stderr); + snapshot_from_failed_crontab_list(&stderr) +} + +fn install_crontab(content: &str) -> Result<()> { + let crontab_file = crontab_path(); + fs::write(&crontab_file, content)?; + let status = Command::new("crontab").arg(&crontab_file).status()?; + if !status.success() { + anyhow::bail!("Failed to install crontab"); + } + Ok(()) +} + +fn remove_crontab() -> Result<()> { + let output = Command::new("crontab").arg("-r").output()?; + if output.status.success() { + return Ok(()); + } + + let stderr = String::from_utf8_lossy(&output.stderr); + if is_no_crontab_message(&stderr) { + return Ok(()); + } + + anyhow::bail!("Failed to remove crontab: {}", stderr.trim()); +} + +fn write_crontab_content(content: &str) -> Result<()> { + if content.is_empty() { + remove_crontab() + } else { + install_crontab(content) + } +} + +fn write_crontab_backup(snapshot: &CrontabSnapshot) -> Result { + let backup_path = crontab_backup_path(); + if let Some(parent) = backup_path.parent() { + fs::create_dir_all(parent)?; + } + fs::write(&backup_path, &snapshot.content)?; + Ok(backup_path) +} + +fn restore_crontab(snapshot: &CrontabSnapshot) -> Result<()> { + if snapshot.existed && !snapshot.content.is_empty() { + install_crontab(&snapshot.content) + } else { + remove_crontab() + } +} + +fn verify_auto_update_enabled(expected_block: &str, bin_path: &str) -> Result<()> { + let snapshot = read_current_crontab()?; + if count_managed_cron_blocks(&snapshot.content) != 1 { + anyhow::bail!("Expected exactly one mihoto-managed cron block after enable"); + } + if !snapshot.content.contains(expected_block) { + anyhow::bail!("Installed crontab does not contain the expected mihoto cron block"); + } + if remove_legacy_current_binary_entries(&snapshot.content, bin_path) != snapshot.content { + anyhow::bail!("Installed crontab still contains a legacy mihoto cron entry"); + } + Ok(()) +} + +fn verify_auto_update_disabled(bin_path: &str) -> Result<()> { + let snapshot = read_current_crontab()?; + if count_managed_cron_blocks(&snapshot.content) != 0 { + anyhow::bail!("Expected no mihoto-managed cron block after disable"); + } + if remove_legacy_current_binary_entries(&snapshot.content, bin_path) != snapshot.content { + anyhow::bail!("Crontab still contains a legacy mihoto cron entry after disable"); + } + Ok(()) } -/// Enable auto-update by installing cron job +fn find_mihoto_cron_entry(content: &str, bin_path: &str) -> Option { + let mut inside_block = false; + for line in content.lines() { + let trimmed = line.trim(); + if trimmed == MANAGED_BEGIN { + inside_block = true; + continue; + } + if trimmed == MANAGED_END { + inside_block = false; + continue; + } + if inside_block && !trimmed.is_empty() && !trimmed.starts_with('#') { + return Some(line.to_string()); + } + if is_legacy_current_binary_entry(line, bin_path) { + return Some(line.to_string()); + } + } + None +} + +/// Enable auto-update by installing or refreshing only the mihoto-managed cron block. pub fn enable_auto_update(interval_hours: u16, prefix: &str) -> Result<()> { if interval_hours == 0 { println!( @@ -54,19 +313,18 @@ pub fn enable_auto_update(interval_hours: u16, prefix: &str) -> Result<()> { anyhow::bail!("Auto-update interval must be between 1 and 24 hours"); } - let crontab_content = generate_crontab(interval_hours)?; - let crontab_file = crontab_path(); - - // Write crontab to runtime directory for reference - fs::write(&crontab_file, crontab_content)?; - - // Install crontab using crontab command - let status = std::process::Command::new("crontab") - .arg(&crontab_file) - .status()?; + let bin_path = mihoro_bin_path()?; + let snapshot = read_current_crontab()?; + let _backup_path = write_crontab_backup(&snapshot)?; + let block = build_managed_cron_block(interval_hours, &bin_path); + let merged = merge_managed_cron_block(&snapshot.content, &block, &bin_path); - if !status.success() { - anyhow::bail!("Failed to install crontab"); + write_crontab_content(&merged)?; + if let Err(err) = verify_auto_update_enabled(&block, &bin_path) { + if let Err(rollback_err) = restore_crontab(&snapshot) { + anyhow::bail!("{}; rollback failed: {}", err, rollback_err); + } + return Err(err); } println!( @@ -83,33 +341,40 @@ pub fn enable_auto_update(interval_hours: u16, prefix: &str) -> Result<()> { Ok(()) } -/// Disable auto-update by removing cron job +/// Disable auto-update by removing only the mihoto-managed cron block. pub fn disable_auto_update(prefix: &str) -> Result<()> { - let crontab_file = crontab_path(); - - // Remove our crontab reference file - if crontab_file.exists() { - fs::remove_file(&crontab_file)?; + let bin_path = mihoro_bin_path()?; + let snapshot = read_current_crontab()?; + if !snapshot.existed && snapshot.content.is_empty() { + println!( + "{} Auto-update disabled (no active cron job)", + prefix.yellow() + ); + return Ok(()); } - // Install empty crontab to remove all entries - let status = std::process::Command::new("crontab").arg("-r").status(); + let cleaned = remove_managed_cron_content(&snapshot.content, &bin_path); + if cleaned == snapshot.content { + println!("{} Auto-update disabled", prefix.green().bold()); + return Ok(()); + } - match status { - Ok(status) if status.success() => { - println!("{} Auto-update disabled", prefix.green().bold()); - Ok(()) - } - Ok(_) => { - // crontab -r returns non-zero if no crontab exists, which is fine - println!( - "{} Auto-update disabled (no active cron job)", - prefix.yellow() - ); - Ok(()) + let _backup_path = write_crontab_backup(&snapshot)?; + write_crontab_content(&cleaned)?; + if let Err(err) = verify_auto_update_disabled(&bin_path) { + if let Err(rollback_err) = restore_crontab(&snapshot) { + anyhow::bail!("{}; rollback failed: {}", err, rollback_err); } - Err(e) => Err(anyhow!("Failed to disable crontab: {}", e)), + return Err(err); } + + let crontab_file = crontab_path(); + if crontab_file.exists() { + let _ = fs::remove_file(&crontab_file); + } + + println!("{} Auto-update disabled", prefix.green().bold()); + Ok(()) } /// Format Unix timestamp to local datetime string using date command @@ -130,19 +395,20 @@ fn format_datetime(secs: u64) -> String { /// Get current cron status pub fn get_cron_status(_prefix: &str, mihomo_config_path: &str) -> Result<()> { - let crontab_file = crontab_path(); + let bin_path = mihoro_bin_path()?; + let snapshot = read_current_crontab()?; - if !crontab_file.exists() { - println!("{} Auto-update is disabled", "status:".yellow().bold()); - return Ok(()); + match find_mihoto_cron_entry(&snapshot.content, &bin_path) { + Some(cron_entry) => { + println!("{} Auto-update is enabled", "status:".green().bold()); + println!("{} {}", "->".dimmed(), cron_entry.dimmed()); + } + None => { + println!("{} Auto-update is disabled", "status:".yellow().bold()); + return Ok(()); + } } - let content = fs::read_to_string(&crontab_file)?; - let cron_entry = content.lines().next().unwrap_or(""); - - println!("{} Auto-update is enabled", "status:".green().bold()); - println!("{} {}", "->".dimmed(), cron_entry.dimmed()); - // Show last updated time from mihomo config file let config_path = Path::new(mihomo_config_path); if let Ok(metadata) = fs::metadata(config_path) { @@ -171,8 +437,134 @@ mod tests { } #[test] - fn test_generate_crontab() { - let crontab = generate_crontab(6).unwrap(); - assert!(crontab.contains("0 */6 * * *")); + fn test_quote_cron_command_wraps_paths_and_escapes_single_quotes() { + assert_eq!( + quote_cron_command("/opt/Mihoro Bin/mihoro"), + "'/opt/Mihoro Bin/mihoro'" + ); + assert_eq!( + quote_cron_command("/opt/mihoro's bin/mihoro"), + "'/opt/mihoro'\\''s bin/mihoro'" + ); + } + + #[test] + fn test_build_managed_cron_block_uses_stable_markers() { + let block = build_managed_cron_block(6, "/usr/local/bin/mihoro"); + + assert_eq!( + block, + "# mihoto: begin auto-update\n\ +0 */6 * * * '/usr/local/bin/mihoro' update\n\ +# mihoto: end auto-update\n" + ); + } + + #[test] + fn test_merge_managed_cron_block_preserves_unrelated_entries() { + let existing = "\ +27 16 * * * \"/root/.acme.sh\"/acme.sh --cron --home \"/root/.acme.sh\" > /dev/null +0 3 * * * /opt/backup/run.sh +"; + let merged = merge_managed_cron_block( + existing, + &build_managed_cron_block(12, "/root/.local/bin/mihoro"), + "/root/.local/bin/mihoro", + ); + + assert!(merged.contains("\"/root/.acme.sh\"/acme.sh --cron")); + assert!(merged.contains("/opt/backup/run.sh")); + assert_eq!(count_managed_cron_blocks(&merged), 1); + assert!(merged.contains("# mihoto: begin auto-update")); + assert!(merged.contains("0 */12 * * * '/root/.local/bin/mihoro' update")); + assert!(merged.contains("# mihoto: end auto-update")); + } + + #[test] + fn test_merge_managed_cron_block_replaces_existing_block() { + let existing = "\ +# existing user note +# mihoto: begin auto-update +0 */6 * * * '/root/.local/bin/mihoro' update +# mihoto: end auto-update +15 4 * * * /opt/backup/run.sh +"; + let merged = merge_managed_cron_block( + existing, + &build_managed_cron_block(12, "/root/.local/bin/mihoro"), + "/root/.local/bin/mihoro", + ); + + assert_eq!(count_managed_cron_blocks(&merged), 1); + assert!(merged.contains("# existing user note")); + assert!(merged.contains("15 4 * * * /opt/backup/run.sh")); + assert!(merged.contains("0 */12 * * * '/root/.local/bin/mihoro' update")); + assert!(!merged.contains("0 */6 * * * '/root/.local/bin/mihoro' update")); + } + + #[test] + fn test_remove_managed_cron_content_preserves_unrelated_entries() { + let existing = "\ +# keep this comment +# mihoto: begin auto-update +0 */12 * * * '/root/.local/bin/mihoro' update +# mihoto: end auto-update +27 16 * * * /root/.acme.sh/acme.sh --cron > /dev/null +"; + let cleaned = remove_managed_cron_content(existing, "/root/.local/bin/mihoro"); + + assert_eq!( + cleaned, + "# keep this comment\n27 16 * * * /root/.acme.sh/acme.sh --cron > /dev/null\n" + ); + } + + #[test] + fn test_marker_text_inside_user_command_is_preserved() { + let existing = "10 1 * * * /usr/bin/printf '# mihoto: begin auto-update'\n"; + let cleaned = remove_managed_cron_content(existing, "/root/.local/bin/mihoro"); + + assert_eq!(cleaned, existing); + } + + #[test] + fn test_no_crontab_message_is_successfully_classified() { + assert!(is_no_crontab_message("no crontab for pectics\n")); + assert!(is_no_crontab_message("crontab: no crontab for root\n")); + assert!(!is_no_crontab_message("crontab: permission denied\n")); + } + + #[test] + fn test_non_no_crontab_list_failure_is_error() { + let err = snapshot_from_failed_crontab_list("crontab: permission denied\n") + .expect_err("permission failures must not be treated as an empty crontab"); + + assert!(err.to_string().contains("Failed to read crontab")); + } + + #[test] + fn test_disable_on_no_crontab_is_success_path() { + let cleaned = remove_managed_cron_content("", "/root/.local/bin/mihoro"); + assert!(cleaned.is_empty()); + assert_eq!(count_managed_cron_blocks(&cleaned), 0); + } + + #[test] + fn test_legacy_current_binary_entry_is_migrated_but_similar_lines_remain() { + let existing = "\ +0 */6 * * * /root/.local/bin/mihoro update +3 2 * * * /root/.local/bin/mihoro update --dry-run +4 2 * * * /opt/other/mihoro update +"; + let merged = merge_managed_cron_block( + existing, + &build_managed_cron_block(12, "/root/.local/bin/mihoro"), + "/root/.local/bin/mihoro", + ); + + assert!(!merged.contains("0 */6 * * * /root/.local/bin/mihoro update\n")); + assert!(merged.contains("/root/.local/bin/mihoro update --dry-run")); + assert!(merged.contains("/opt/other/mihoro update")); + assert_eq!(count_managed_cron_blocks(&merged), 1); } }