commit cfcbf638d69e5c25c29c02d62fa2eac6f8256fa7 Author: plazmoid Date: Mon May 31 17:24:17 2021 +0500 all diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..415b4b0 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "sockser" +version = "0.1.0" +authors = ["plazmoid "] +edition = "2018" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +nix = "*" +reqwest = { version = "*", features = ["blocking", "socks"]} +once_cell = "*" \ No newline at end of file diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..0703311 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,290 @@ +use nix::{ + sys::signal::{signal, SigHandler, Signal}, + unistd::{close as fdclose, fork, getppid, setsid, ForkResult}, +}; +use reqwest::{blocking::Client, header, Proxy}; +use std::{ + env, + io::Result as IOResult, + panic, + process::{exit, Child, Command, Stdio}, + sync::{Mutex, MutexGuard}, + thread, time, +}; + +use once_cell::sync::OnceCell; + +mod ssh_config; +use ssh_config::TinySSHConfig; + +const HOST: &str = "aws-tun"; +const BASE_CONFIG: &str = "/home/plazmoid/.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: OnceCell> = OnceCell::new(); + FAIL_NOTIFIER + .get_or_init(|| { + Mutex::new(ThresholdCaller::new( + |info| notify(format!("Can't setup socks proxy: {}", info), true), + 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 = if critical { "critical" } else { "normal" }; + spawn( + "notify-send", + &["-t", "5000", "-u", urg, &format!("sockser: {}", 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() +} + +fn kill_by_cmd(cmd: &str) { + if let Some(pid) = pgrep(cmd) { + Command::new("kill") + .args(&["-9", &pid.to_string()]) + .status() + .ok(); + } +} + +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()); + 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 = host.get("HostName").unwrap(); + 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 { + killer("wrong ip"); + } else { + if ThresholdCaller::get().was_triggered() { + notify("reconnected", false); + } + ThresholdCaller::get().reset(); + } + } + Err(_) => { + killer("can't receive response"); + } + }, + Err(_) => { + killer("can't send request"); + } + } + } + thread::sleep(time::Duration::from_secs(10)); + } +} + +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)); + } +} + +fn main() { + panic::set_hook(Box::new(|info| { + notify(info.to_string(), true); + exit(1); + })); + let mut args = env::args(); + if let Some(arg) = args.nth(1) { + if arg == "-d" { + daemonize().unwrap(); + unsafe { + DAEMONIZED = true; + } + } + } + thread::spawn(|| monitor_connection(HOST)); + monitor_process(); +} + +#[cfg(test)] +mod tests { + use super::*; + use std::error::Error; + type TestResult = Result<(), Box>; + + #[test] + fn test_ssh_config_parser() -> TestResult { + let config = TinySSHConfig::parse(BASE_CONFIG)?; + assert_eq!( + config.get_host("aws2").unwrap().get("HostName").unwrap(), + "35.177.229.131" + ); + Ok(()) + } + + #[test] + fn test_ssh_host_can_be_proxy() -> TestResult { + let config = TinySSHConfig::parse(BASE_CONFIG)?; + assert!(config.get_host("aws-tun").unwrap().can_be_proxy()); + Ok(()) + } +} diff --git a/src/ssh_config.rs b/src/ssh_config.rs new file mode 100644 index 0000000..9a15e13 --- /dev/null +++ b/src/ssh_config.rs @@ -0,0 +1,120 @@ +use std::collections::HashMap; +use std::fs::read_to_string; + +struct LineLexer { + linenum: usize, + content: Vec, + current: String, + previous: String, +} + +impl LineLexer { + pub fn new(data: String) -> Self { + Self { + linenum: 0, + content: data.lines().map(String::from).collect(), + current: String::new(), + previous: String::new(), + } + } + + pub fn next(&mut self) -> Option { + let line = self.content.get(self.linenum); + if line.is_some() { + self.previous = self.current.clone(); + self.current = line.unwrap().trim().to_string(); + self.linenum += 1; + } + line.map(|s| s.clone()) + } + + pub fn next_if_nonempty(&mut self) -> Option { + let l = self.next(); + if l.is_none() || l.as_ref().unwrap().len() > 0 { + l + } else { + None + } + } + + pub fn next_nonempty(&mut self) -> Option { + loop { + let l = self.next(); + if l.is_none() { + return l; + } + if l.as_ref().unwrap().len() > 0 { + return l; + } + } + } +} + +#[derive(Debug)] +pub struct TinySSHHost { + hostname: String, + entries: HashMap, +} + +impl TinySSHHost { + pub fn new(hostname: String) -> Self { + TinySSHHost { + hostname, + entries: HashMap::new(), + } + } + + pub fn get(&self, k: &str) -> Option<&String> { + self.entries.get(&k.to_string()) + } + + pub fn can_be_proxy(&self) -> bool { + self.get("DynamicForward").is_some() + } +} + +pub struct TinySSHConfig { + hosts: Vec, +} + +impl TinySSHConfig { + pub fn parse>(filename: S) -> Result { + let data = read_to_string(filename.into()).map_err(|e| e.to_string())?; + let mut tsc = Self { hosts: vec![] }; + let mut lexer = LineLexer::new(data); + loop { + let line = lexer.next_nonempty(); + if line.is_none() { + break; + } + let mut host = if lexer.current.starts_with("Host") { + match lexer.current.split_whitespace().nth(1) { + Some(h) => TinySSHHost::new(h.trim().to_string()), + None => return Err("No hostname found".to_string()), + } + } else { + return Err("Wrong section order".to_string()); + }; + loop { + let l = lexer.next_if_nonempty(); + if l.is_none() { + break; + } + let l = l.unwrap(); + let mut param: Vec = l.trim().split(' ').map(String::from).collect(); + if param.len() != 2 { + return Err(format!("Wrong param format: {}", l)); + } + let v = param.pop().unwrap(); + let k = param.pop().unwrap(); + host.entries.insert(k, v); + } + tsc.hosts.push(host); + } + Ok(tsc) + } + + pub fn get_host(&self, host: &str) -> Option<&TinySSHHost> { + self.hosts.iter().find(|h| h.hostname == host) + } +}