use std::fs; use std::process::Command; use std::thread; use std::time::Duration; use anyhow::{Context, Result}; use clap::Parser; use detmir_core::{exit_codes, now_utc_rfc3339}; use serde::Serialize; const DEFAULT_HEARTBEAT_FILE: &str = "/opt/infra-admin/.state/tsj_guardian_heartbeat"; const DEFAULT_SERVICE_NAME: &str = "tsj-guardian-bot.service"; const DEFAULT_GOST_SERVICE_NAME: &str = "gost-tg.service"; const DEFAULT_GOST_PATTERN: &str = "/usr/local/bin/gost -L http+socks5://127.0.0.1:11090 -F socks5+wss://gw.example.local:4443"; #[derive(Debug, Parser)] #[command(about = "TSJ Guardian bot heartbeat watchdog and gost duplicate guard.")] struct Cli { #[arg(long, default_value = DEFAULT_HEARTBEAT_FILE)] heartbeat_file: String, #[arg(long, default_value_t = 180)] max_age_seconds: i64, #[arg(long, default_value = DEFAULT_SERVICE_NAME)] service_name: String, #[arg(long, default_value = DEFAULT_GOST_SERVICE_NAME)] gost_service_name: String, #[arg(long, default_value = DEFAULT_GOST_PATTERN)] gost_dup_pattern: String, #[arg(long)] apply: bool, #[arg(long)] json: bool, #[arg(long)] loop_forever: bool, #[arg(long, default_value_t = 60)] interval_seconds: u64, } impl Cli { fn apply_env(mut self) -> Self { self.heartbeat_file = env_string("HEARTBEAT_FILE").unwrap_or(self.heartbeat_file); self.max_age_seconds = env_string("MAX_AGE_SEC") .and_then(|value| value.parse().ok()) .unwrap_or(self.max_age_seconds); self.service_name = env_string("SERVICE_NAME").unwrap_or(self.service_name); self.gost_service_name = env_string("GOST_SERVICE_NAME").unwrap_or(self.gost_service_name); self.gost_dup_pattern = env_string("GOST_DUP_PATTERN").unwrap_or(self.gost_dup_pattern); self } } #[derive(Debug, Clone, Serialize, PartialEq, Eq)] struct WatchdogAction { action: String, target: String, status: String, } #[derive(Debug, Serialize)] struct WatchdogReport { ok: bool, dry_run: bool, generated_at_utc: String, heartbeat_file: String, heartbeat_age_seconds: Option, heartbeat_status: String, gost_main_pid: Option, gost_pids: Vec, actions: Vec, } fn env_string(name: &str) -> Option { std::env::var(name).ok().filter(|value| !value.is_empty()) } fn main_pid(service_name: &str) -> Option { let output = Command::new("systemctl") .args(["show", "-p", "MainPID", "--value", service_name]) .output() .ok()?; if !output.status.success() { return None; } String::from_utf8_lossy(&output.stdout) .trim() .parse::() .ok() .filter(|pid| *pid > 1) } fn matching_pids(pattern: &str) -> Vec { let output = match Command::new("ps").args(["-eo", "pid=,args="]).output() { Ok(output) => output, Err(_) => return Vec::new(), }; if !output.status.success() { return Vec::new(); } let self_pid = std::process::id(); let mut pids = String::from_utf8_lossy(&output.stdout) .lines() .filter_map(|line| { let line = line.trim_start(); let (pid, args) = line.split_once(char::is_whitespace)?; let pid = pid.parse::().ok()?; if pid != self_pid && args.contains(pattern) { Some(pid) } else { None } }) .collect::>(); pids.sort_unstable(); pids.dedup(); pids } fn choose_keep_pid(main_pid: Option, pids: &[u32]) -> Option { if pids.is_empty() { return None; } match main_pid { Some(pid) if pids.contains(&pid) => Some(pid), _ => pids.first().copied(), } } fn signal_pid(pid: u32, signal: &str) -> Result<()> { let status = Command::new("kill") .args([format!("-{signal}"), pid.to_string()]) .status() .with_context(|| format!("failed to signal pid {pid}"))?; if status.success() { Ok(()) } else { Err(anyhow::anyhow!("kill -{signal} {pid} exited with {status}")) } } fn pid_alive(pid: u32) -> bool { Command::new("kill") .args(["-0", &pid.to_string()]) .status() .map(|status| status.success()) .unwrap_or(false) } fn restart_service(service_name: &str) -> Result<()> { let status = Command::new("systemctl") .args(["restart", service_name]) .status() .with_context(|| format!("failed to restart {service_name}"))?; if status.success() { Ok(()) } else { Err(anyhow::anyhow!( "systemctl restart {service_name} exited with {status}" )) } } fn heartbeat_status(path: &str, max_age_seconds: i64, now: i64) -> (String, Option) { let raw = match fs::read_to_string(path) { Ok(raw) => raw, Err(_) => return ("missing".to_string(), None), }; let heartbeat = match raw.trim().parse::() { Ok(value) => value, Err(_) => return ("invalid".to_string(), None), }; let age = now.saturating_sub(heartbeat); if age > max_age_seconds { ("stale".to_string(), Some(age)) } else { ("fresh".to_string(), Some(age)) } } fn run_once(cli: &Cli) -> Result { let mut actions = Vec::new(); let gost_main_pid = main_pid(&cli.gost_service_name); let gost_pids = matching_pids(&cli.gost_dup_pattern); if let Some(keep_pid) = choose_keep_pid(gost_main_pid, &gost_pids) { for pid in gost_pids.iter().copied().filter(|pid| *pid != keep_pid) { if cli.apply { let status = signal_pid(pid, "TERM") .map(|_| "terminated") .unwrap_or("term-failed"); actions.push(WatchdogAction { action: "kill-term".to_string(), target: pid.to_string(), status: status.to_string(), }); } else { actions.push(WatchdogAction { action: "kill-term".to_string(), target: pid.to_string(), status: "planned".to_string(), }); } } if cli.apply && gost_pids.len() > 1 { thread::sleep(Duration::from_secs(2)); for pid in gost_pids.iter().copied().filter(|pid| *pid != keep_pid) { if pid_alive(pid) { let status = signal_pid(pid, "KILL") .map(|_| "killed") .unwrap_or("kill-failed"); actions.push(WatchdogAction { action: "kill-kill".to_string(), target: pid.to_string(), status: status.to_string(), }); } } } } let now = chrono_like_now_epoch(); let (hb_status, hb_age) = heartbeat_status(&cli.heartbeat_file, cli.max_age_seconds, now); if hb_status != "fresh" { if cli.apply { let status = restart_service(&cli.service_name) .map(|_| "restarted") .unwrap_or("restart-failed"); actions.push(WatchdogAction { action: "restart-service".to_string(), target: cli.service_name.clone(), status: status.to_string(), }); } else { actions.push(WatchdogAction { action: "restart-service".to_string(), target: cli.service_name.clone(), status: "planned".to_string(), }); } } Ok(WatchdogReport { ok: true, dry_run: !cli.apply, generated_at_utc: now_utc_rfc3339(), heartbeat_file: cli.heartbeat_file.clone(), heartbeat_age_seconds: hb_age, heartbeat_status: hb_status, gost_main_pid, gost_pids, actions, }) } fn chrono_like_now_epoch() -> i64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|duration| duration.as_secs() as i64) .unwrap_or(0) } fn print_report(report: &WatchdogReport, json: bool) -> Result<()> { if json { println!("{}", serde_json::to_string_pretty(report)?); } else { println!( "tsj-guardian-watchdog: heartbeat={} age={:?} dry_run={}", report.heartbeat_status, report.heartbeat_age_seconds, report.dry_run ); for action in &report.actions { println!("{} {} {}", action.action, action.target, action.status); } } Ok(()) } fn run(cli: Cli) -> Result<()> { if cli.loop_forever { loop { let report = run_once(&cli)?; print_report(&report, cli.json)?; thread::sleep(Duration::from_secs(cli.interval_seconds)); } } let report = run_once(&cli)?; print_report(&report, cli.json)?; Ok(()) } fn main() -> Result<()> { let cli = Cli::parse().apply_env(); match run(cli) { Ok(()) => Ok(()), Err(err) => { eprintln!("{err:#}"); std::process::exit(exit_codes::ERROR); } } } #[cfg(test)] mod tests { use super::*; #[test] fn keeps_main_pid_when_present() { assert_eq!(choose_keep_pid(Some(20), &[10, 20, 30]), Some(20)); } #[test] fn keeps_first_pid_when_main_missing() { assert_eq!(choose_keep_pid(Some(99), &[10, 20, 30]), Some(10)); } #[test] fn detects_heartbeat_states() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("heartbeat"); assert_eq!( heartbeat_status(path.to_str().unwrap(), 180, 1_000), ("missing".to_string(), None) ); fs::write(&path, "bad\n").unwrap(); assert_eq!( heartbeat_status(path.to_str().unwrap(), 180, 1_000), ("invalid".to_string(), None) ); fs::write(&path, "900\n").unwrap(); assert_eq!( heartbeat_status(path.to_str().unwrap(), 180, 1_000), ("fresh".to_string(), Some(100)) ); fs::write(&path, "700\n").unwrap(); assert_eq!( heartbeat_status(path.to_str().unwrap(), 180, 1_000), ("stale".to_string(), Some(300)) ); } }