all

Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "sockser"
+version = "0.1.0"
+authors = ["plazmoid <kronos44@mail.ru>"]
+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 = "*"

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<dyn Fn(&str) + Send + 'static>,
+    arg: Option<String>,
+    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<Mutex<ThresholdCaller>> = 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<String>) {
+        self.arg = arg;
+    }
+
+    pub fn append_fail<S: Into<String>>(&mut self, msg: Option<S>) {
+        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<i32, String> {
+    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<Child> {
+    Command::new(cmd).args(args).spawn()
+}
+
+fn _dbg() {
+    spawn("touch", &["/tmp/asdqwe"]).ok();
+}
+
+fn notify<S: Into<String>>(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<Child> {
+    Command::new("/usr/bin/ssh")
+        .args(&["-f", "-N", host, "-o", "ConnectTimeout=3"])
+        .stderr(Stdio::piped())
+        .spawn()
+}
+
+fn pgrep(cmd: &str) -> Option<i32> {
+    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<dyn Error>>;
+
+    #[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(())
+    }
+}

src/ssh_config.rs

@@ -0,0 +1,120 @@
+use std::collections::HashMap;
+use std::fs::read_to_string;
+
+struct LineLexer {
+    linenum: usize,
+    content: Vec<String>,
+    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<String> {
+        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<String> {
+        let l = self.next();
+        if l.is_none() || l.as_ref().unwrap().len() > 0 {
+            l
+        } else {
+            None
+        }
+    }
+
+    pub fn next_nonempty(&mut self) -> Option<String> {
+        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<String, String>,
+}
+
+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<TinySSHHost>,
+}
+
+impl TinySSHConfig {
+    pub fn parse<S: Into<String>>(filename: S) -> Result<Self, String> {
+        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<String> = 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)
+    }
+}