diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..29876b3 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,2 @@ +pub const HOST: &str = "ortem-tun"; +pub const BASE_CONFIG: &str = concat!(env!("HOME"), "/.ssh/config"); diff --git a/src/main.rs b/src/main.rs index a7c349a..bca8e78 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,266 +1,14 @@ -use nix::{ - sys::signal::{SigHandler, Signal, signal}, - unistd::{ForkResult, close as fdclose, fork, getppid, setsid}, -}; -use reqwest::{Proxy, blocking::Client, header}; -use std::{ - env, - io::Result as IOResult, - net::ToSocketAddrs, - panic, - process::{Child, Command, Stdio, exit}, - sync::{Mutex, MutexGuard, OnceLock}, - thread, time, -}; - +mod config; +mod monitoring; mod ssh_config; -use ssh_config::TinySSHConfig; - -const HOST: &str = "ortem-tun"; -const BASE_CONFIG: &str = concat!(env!("HOME"), "/.ssh/config"); - -static mut DAEMONIZED: bool = false; - -struct ThresholdCaller { - // call Fn once after threshold is exceeded - cmd: Box, - arg: Option, - triggered: bool, - retries: usize, - max_retries: usize, -} - -impl ThresholdCaller { - fn new(cmd: impl Fn(&str) + Send + 'static, max_retries: usize) -> Self { - ThresholdCaller { - cmd: Box::new(cmd), - arg: None, - triggered: false, - retries: 0, - max_retries, - } - } - - pub fn get<'m>() -> MutexGuard<'m, Self> { - static FAIL_NOTIFIER: OnceLock> = OnceLock::new(); - FAIL_NOTIFIER - .get_or_init(|| { - Mutex::new(ThresholdCaller::new( - |info| notify(format!("Can't setup socks proxy: {}", info), false), - 5, - )) - }) - .lock() - .unwrap() - } - - pub fn call_once(&mut self, info: Option<&String>) { - //dbg!(self.retries, self.triggered); - if !self.triggered && self.retries >= self.max_retries { - let arg = info.or(self.arg.as_ref()); - (self.cmd)(arg.unwrap_or(&String::new())); - self.triggered = true; - } - } - - pub fn set_arg(&mut self, arg: Option) { - self.arg = arg; - } - - pub fn append_fail>(&mut self, msg: Option) { - self.inc(); - if let Some(m) = msg { - let msg = m.into(); - unsafe { - if !DAEMONIZED { - eprintln!("{}", &msg); - } - } - self.set_arg(Some(msg)); - } - } - - pub fn was_triggered(&self) -> bool { - self.triggered - } - - pub fn reset(&mut self) { - self.retries = 0; - self.triggered = false; - } - - pub fn inc(&mut self) { - self.retries += 1; - } -} - -pub fn daemonize() -> Result { - if getppid().as_raw() != 1 { - setsig(Signal::SIGTTOU, SigHandler::SigIgn); - setsig(Signal::SIGTTIN, SigHandler::SigIgn); - setsig(Signal::SIGTSTP, SigHandler::SigIgn); - } - for fd in 0..=2 { - fdclose(fd).ok(); - } - - unsafe { - match fork() { - Ok(ForkResult::Parent { .. }) => { - exit(0); - } - Ok(ForkResult::Child) => match setsid() { - Ok(pid) => Ok(pid.as_raw()), - Err(e) => Err(e.to_string()), - }, - Err(_) => exit(255), - } - } -} - -pub fn setsig(sig: Signal, hnd: SigHandler) { - unsafe { - signal(sig, hnd).unwrap(); - } -} - -fn spawn(cmd: &str, args: &[&str]) -> IOResult { - Command::new(cmd).args(args).spawn() -} - -fn _dbg() { - spawn("touch", &["/tmp/asdqwe"]).ok(); -} - -fn notify>(msg: S, critical: bool) { - let (urg, prefix) = if critical { - ("critical", "Отвалилась жопа: ") - } else { - ("normal", "") - }; - spawn( - "notify-send", - &[ - "-a", - env!("CARGO_PKG_NAME"), - "-t", - "5000", - "-u", - urg, - &format!("{prefix}{}", msg.into()), - ], - ) - .unwrap() - .wait() - .unwrap(); -} - -fn run_proxy(host: &str) -> IOResult { - Command::new("/usr/bin/ssh") - .args(&["-f", "-N", host, "-o", "ConnectTimeout=3"]) - .stderr(Stdio::piped()) - .spawn() -} - -fn pgrep(cmd: &str) -> Option { - let grep = Command::new("pgrep").args(&["-f", cmd]).output().unwrap(); - if grep.status.success() { - let output = String::from_utf8(grep.stdout).unwrap(); - output.lines().next().unwrap().parse().ok() - } else { - None - } -} - -fn process_present(cmd: &str) -> bool { - pgrep(cmd).is_some() -} +mod trigger; +mod utils; -fn kill_by_cmd(cmd: &str) { - if let Some(pid) = pgrep(cmd) { - Command::new("kill") - .args(&["-9", &pid.to_string()]) - .status() - .ok(); - } -} +use std::{env, panic, process::exit, thread}; -fn monitor_connection(ssh_host: &str) -> ! { - let config = TinySSHConfig::parse(BASE_CONFIG).expect("ssh config file has wrong format"); - let host = config.get_host(ssh_host).expect("no host found in config"); - assert!(host.can_be_proxy(), "expected DynamicForward entry"); - let proxy = Proxy::all(&format!( - "socks5://127.0.0.1:{}", - host.get("DynamicForward").unwrap() - )) - .unwrap(); - let mut headers = header::HeaderMap::new(); - headers.insert( - header::USER_AGENT, - header::HeaderValue::from_static("curl/7.74.0"), - ); - let client = Client::builder() - .proxy(proxy) - .timeout(time::Duration::from_secs(10)) - .default_headers(headers) - .build() - .unwrap(); - let proxy_ip = { - let host = host.get("HostName").unwrap(); - let host = format!("{host}:80"); - host.to_socket_addrs().unwrap().next().unwrap().ip() - }; - let killer = |msg| { - ThresholdCaller::get().append_fail(Some(msg)); - kill_by_cmd(ssh_host); - }; - loop { - if process_present(HOST) { - match client.get("https://2ip.ru/").send() { - Ok(resp) => match resp.text() { - Ok(received_ip) => { - if received_ip.trim() != proxy_ip.to_string() { - killer("wrong ip"); - } else { - if ThresholdCaller::get().was_triggered() { - notify("Соединение восстановлено", false); - } - ThresholdCaller::get().reset(); - } - } - Err(_) => { - killer("can't receive response"); - } - }, - Err(_) => { - killer("can't send request"); - } - } - } - thread::sleep(time::Duration::from_secs(10)); - } -} +use crate::{config::HOST, utils::notify}; -fn monitor_process() -> ! { - loop { - if !process_present(HOST) { - match run_proxy(HOST) { - Ok(result) => { - let exited = result.wait_with_output().unwrap(); - if !exited.status.success() { - ThresholdCaller::get() - .append_fail(Some(String::from_utf8(exited.stderr).unwrap())); - } - } - Err(e) => { - ThresholdCaller::get().append_fail(Some(e.to_string())); - } - } - ThresholdCaller::get().call_once(None); - } - thread::sleep(time::Duration::from_secs(1)); - } -} +static mut DAEMONIZED: bool = false; fn main() { if !cfg!(debug_assertions) { @@ -272,19 +20,20 @@ fn main() { let mut args = env::args(); if let Some(arg) = args.nth(1) { if arg == "-d" { - daemonize().unwrap(); + utils::daemonize().unwrap(); unsafe { DAEMONIZED = true; } } } - thread::spawn(|| monitor_connection(HOST)); - monitor_process(); + thread::spawn(|| monitoring::monitor_connection(HOST)); + monitoring::monitor_process(); } #[cfg(test)] mod tests { - use super::*; + use crate::{config::BASE_CONFIG, ssh_config::TinySSHConfig}; + use std::error::Error; type TestResult = Result<(), Box>; diff --git a/src/monitoring.rs b/src/monitoring.rs new file mode 100644 index 0000000..78f2ad6 --- /dev/null +++ b/src/monitoring.rs @@ -0,0 +1,117 @@ +use std::{net::ToSocketAddrs as _, thread, time::Duration}; + +use crate::{ + config::{BASE_CONFIG, HOST}, + ssh_config::TinySSHConfig, + trigger::Trigger, + utils::*, +}; +use reqwest::{Proxy, blocking::Client, header}; + +pub fn monitor_connection(ssh_host: &str) -> ! { + let config = TinySSHConfig::parse(BASE_CONFIG).expect("ssh config file has wrong format"); + let host = config + .get_host(ssh_host) + .expect("no host found in ssh config"); + assert!(host.can_be_proxy(), "expected DynamicForward entry"); + let proxy = Proxy::all(&format!( + "socks5://127.0.0.1:{}", + host.get("DynamicForward").unwrap() + )) + .unwrap(); + let mut headers = header::HeaderMap::new(); + headers.insert( + header::USER_AGENT, + header::HeaderValue::from_static("curl/7.74.0"), + ); + let client = Client::builder() + .proxy(proxy) + .timeout(Duration::from_secs(10)) + .default_headers(headers) + .build() + .unwrap(); + let kill_with = |msg| { + Trigger::get().append_fail(Some(msg)); + kill_by_cmd(ssh_host); + }; + let proxy_host = { + let host = host.get("HostName").unwrap(); + format!("{host}:80") + }; + let mut proxy_ip = None; + + 'mon: loop { + let Some(proxy_ip) = &proxy_ip else { + let ip = match resolve_host(&proxy_host) { + Ok(ip) => ip, + Err(err) => { + kill_with(err); + continue 'mon; + } + }; + proxy_ip = Some(ip); + continue 'mon; + }; + if process_present(HOST) { + if let Err(err) = verify_proxy_ip(&client, &proxy_ip) { + kill_with(err) + } else { + if Trigger::get().was_triggered() { + notify("Соединение восстановлено", false); + } + Trigger::get().reset(); + } + } + thread::sleep(Duration::from_secs(5)); + } +} + +fn verify_proxy_ip(client: &Client, expected_proxy_ip: &str) -> Result<(), String> { + let received_ip = client + .get("https://2ip.ru/") + .send() + .map_err(|e| format!("can't send request while checking ip: {}", e.to_string()))? + .text() + .map_err(|e| { + format!( + "can't receive response while checking ip: {}", + e.to_string() + ) + })?; + let received_ip = received_ip.trim(); + + if received_ip != expected_proxy_ip { + return Err(format!( + "wrong proxy ip: {received_ip}, expected: {expected_proxy_ip}" + )); + } + Ok(()) +} + +fn resolve_host(host: &str) -> Result { + host.to_socket_addrs() + .map_err(|e| e.to_string())? + .next() + .map(|ip| ip.ip().to_string()) + .ok_or_else(|| "No IPs returned".to_string()) +} + +pub fn monitor_process() -> ! { + loop { + if !process_present(HOST) { + match run_proxy(HOST) { + Ok(result) => { + let exited = result.wait_with_output().unwrap(); + if !exited.status.success() { + Trigger::get().append_fail(Some(String::from_utf8(exited.stderr).unwrap())); + } + } + Err(e) => { + Trigger::get().append_fail(Some(e.to_string())); + } + } + Trigger::get().call_once(None); + } + thread::sleep(Duration::from_secs(1)); + } +} diff --git a/src/trigger.rs b/src/trigger.rs new file mode 100644 index 0000000..3f3a865 --- /dev/null +++ b/src/trigger.rs @@ -0,0 +1,76 @@ +use std::sync::{Mutex, MutexGuard, OnceLock}; + +use crate::{DAEMONIZED, utils::notify}; + +pub struct Trigger { + // call Fn once after threshold is exceeded + cmd: Box, + arg: Option, + triggered: bool, + retries: usize, + max_retries: usize, +} + +impl Trigger { + fn new(cmd: impl Fn(&str) + Send + 'static, max_retries: usize) -> Self { + Trigger { + cmd: Box::new(cmd), + arg: None, + triggered: false, + retries: 0, + max_retries, + } + } + + pub fn get() -> MutexGuard<'static, Self> { + static FAIL_NOTIFIER: OnceLock> = OnceLock::new(); + FAIL_NOTIFIER + .get_or_init(|| { + Mutex::new(Trigger::new( + |info| notify(format!("Can't setup socks proxy: {}", info), false), + 5, + )) + }) + .lock() + .unwrap() + } + + pub fn call_once(&mut self, info: Option<&str>) { + //dbg!(self.retries, self.triggered); + if !self.triggered && self.retries >= self.max_retries { + let arg = info.or(self.arg.as_deref()); + (self.cmd)(arg.unwrap_or("")); + self.triggered = true; + } + } + + pub fn set_arg(&mut self, arg: Option) { + self.arg = arg; + } + + pub fn append_fail>(&mut self, msg: Option) { + self.inc(); + if let Some(m) = msg { + let msg = m.into(); + unsafe { + if !DAEMONIZED { + eprintln!("{}", &msg); + } + } + self.set_arg(Some(msg)); + } + } + + pub fn was_triggered(&self) -> bool { + self.triggered + } + + pub fn reset(&mut self) { + self.retries = 0; + self.triggered = false; + } + + pub fn inc(&mut self) { + self.retries += 1; + } +} diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..9323d75 --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,100 @@ +use std::{ + io, + process::{Child, Command, Stdio, exit}, +}; + +use nix::{ + sys::signal::{SigHandler, Signal, signal}, + unistd::{ForkResult, close as fdclose, fork, getppid, setsid}, +}; + +pub fn daemonize() -> Result { + if getppid().as_raw() != 1 { + setsig(Signal::SIGTTOU, SigHandler::SigIgn); + setsig(Signal::SIGTTIN, SigHandler::SigIgn); + setsig(Signal::SIGTSTP, SigHandler::SigIgn); + } + for fd in 0..=2 { + fdclose(fd).ok(); + } + + unsafe { + match fork() { + Ok(ForkResult::Parent { .. }) => { + exit(0); + } + Ok(ForkResult::Child) => match setsid() { + Ok(pid) => Ok(pid.as_raw()), + Err(e) => Err(e.to_string()), + }, + Err(_) => exit(255), + } + } +} + +pub fn setsig(sig: Signal, hnd: SigHandler) { + unsafe { + signal(sig, hnd).unwrap(); + } +} + +pub fn spawn(cmd: &str, args: &[&str]) -> io::Result { + Command::new(cmd).args(args).spawn() +} + +pub fn _dbg() { + spawn("touch", &["/tmp/asdqwe"]).ok(); +} + +pub fn notify>(msg: S, critical: bool) { + let (urg, prefix) = if critical { + ("critical", "Отвалилась жопа: ") + } else { + ("normal", "") + }; + spawn( + "notify-send", + &[ + "-a", + env!("CARGO_PKG_NAME"), + "-t", + "5000", + "-u", + urg, + &format!("{prefix}{}", msg.into()), + ], + ) + .unwrap() + .wait() + .unwrap(); +} + +pub fn run_proxy(host: &str) -> io::Result { + Command::new("/usr/bin/ssh") + .args(&["-f", "-N", host, "-o", "ConnectTimeout=3"]) + .stderr(Stdio::piped()) + .spawn() +} + +pub fn pgrep(cmd: &str) -> Option { + let grep = Command::new("pgrep").args(&["-f", cmd]).output().unwrap(); + if grep.status.success() { + let output = String::from_utf8(grep.stdout).unwrap(); + output.lines().next().unwrap().parse().ok() + } else { + None + } +} + +pub fn process_present(cmd: &str) -> bool { + pgrep(cmd).is_some() +} + +pub fn kill_by_cmd(cmd: &str) { + if let Some(pid) = pgrep(cmd) { + Command::new("kill") + .args(&["-9", &pid.to_string()]) + .status() + .ok(); + } +}