diff --git a/src/main.rs b/src/main.rs index 7ec37bf..a56ce20 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,11 @@ pub mod iptables_save; pub mod login_attempt; +pub mod tpc_command_server; +use crate::tpc_command_server::start_tcp_command_server; use linemux::MuxedLines; use login_attempt::LoginAttempt; +use std::thread; use std::{collections::HashMap, thread::sleep, time::Duration}; #[tokio::main] @@ -17,12 +20,15 @@ async fn main() -> std::io::Result<()> { "starting iptables-save, run every {} seconds", seconds_iptables.as_secs() ); - tokio::spawn(async move { - loop { - sleep(seconds_iptables); - iptables_save::save_iptables(); - println!("saved iptables rules"); - } + + thread::spawn(move || loop { + sleep(seconds_iptables); + iptables_save::save_iptables(); + println!("saved iptables rules"); + }); + + thread::spawn(|| { + start_tcp_command_server(); }); println!("listening to changes over /host_ssh/auth.log"); diff --git a/src/tpc_command_server.rs b/src/tpc_command_server.rs new file mode 100644 index 0000000..8f87560 --- /dev/null +++ b/src/tpc_command_server.rs @@ -0,0 +1,91 @@ +use std::io::{BufRead, BufReader, Write}; +use std::net::{TcpListener, TcpStream}; +use std::thread; + +pub fn start_tcp_command_server() { + let listener = TcpListener::bind("127.0.0.1:9632").unwrap(); + println!("listening on port 9632 for tcp commands"); + + for stream in listener.incoming() { + match stream { + Ok(stream) => { + thread::spawn(move || handle_client(&stream)); + } + Err(e) => { + eprintln!("err with tcp conn: {}", e); + } + } + } +} + +fn handle_client(mut stream: &TcpStream) { + let reader = BufReader::new(stream); + + for line in reader.lines() { + let buffer = match line { + Ok(data) => data, + Err(_) => return, + }; + + if buffer.is_empty() { + break; + } + + let mut parts = buffer.trim().split_whitespace(); + if let Some(command) = parts.next() { + let arguments: Vec<&str> = parts.collect(); + let response = handle_command(command, arguments); + stream.write_all(response.as_bytes()).unwrap(); + } else { + stream.write_all("invalid command".as_bytes()).unwrap(); + } + } +} + +fn handle_command(command: &str, arguments: Vec<&str>) -> String { + match command { + "ban" => { + if let (Some(ip), Some(port)) = (arguments.get(0), arguments.get(1)) { + let iptables = iptables::new(false).unwrap(); + let _ = iptables.append_unique( + "filter", + "INPUT", + &format!("-s {} -p tcp --dport {} -j ACCEPT", ip, port), + ); + + let _ = iptables.append_unique( + "filter", + "INPUT", + &format!("-p tcp --dport {} -j DROP", port), + ); + + format!("banned port {}, only {} allowed", port, ip) + } else { + "missing args for ban: ip and port".to_string() + } + } + "unban" => { + if let (Some(ip), Some(port)) = (arguments.get(0), arguments.get(1)) { + let iptables = iptables::new(false).unwrap(); + let _ = iptables.delete( + "filter", + "INPUT", + &format!("-s {} -p tcp --dport {} -j ACCEPT", ip, port), + ); + + let _ = iptables.delete( + "filter", + "INPUT", + &format!("-p tcp --dport {} -j DROP", port), + ); + + format!("unbanned port {}, used for {}", port, ip) + } else { + "missing args for unban: ip and port".to_string() + } + } + _ => { + format!("unknown command: {}", command) + } + } +}