moving code from command server into cli
This commit is contained in:
44
src/cli.rs
Normal file
44
src/cli.rs
Normal file
@ -0,0 +1,44 @@
|
||||
use std::path::PathBuf;
|
||||
use structopt::StructOpt;
|
||||
|
||||
#[derive(Debug, StructOpt)]
|
||||
#[structopt(name = "martillo-maldito", about = "A simple iptables ban server")]
|
||||
pub enum Arguments {
|
||||
#[structopt(about = "Initialize ban server")]
|
||||
BanServer {
|
||||
#[structopt(name = "Ssh auth log file", short = "f", long = "ssh-file")]
|
||||
ssh_auth_log: PathBuf,
|
||||
#[structopt(name = "Iptables save file", short = "s", long = "iptables-save")]
|
||||
iptables_save: Option<PathBuf>,
|
||||
},
|
||||
|
||||
#[structopt(about = "Ban port")]
|
||||
BanPort {
|
||||
#[structopt(name = "Port to ban", short = "p", long = "port")]
|
||||
port: u16,
|
||||
},
|
||||
|
||||
#[structopt(about = "Unban port")]
|
||||
UnbanPort {
|
||||
#[structopt(name = "Port to unban", short = "p", long = "port")]
|
||||
port: u16,
|
||||
},
|
||||
|
||||
#[structopt(about = "Allow ip and port")]
|
||||
AllowIpPort {
|
||||
#[structopt(name = "Ip to allow", short = "i", long = "ip")]
|
||||
ip: String,
|
||||
|
||||
#[structopt(name = "Port to allow", short = "p", long = "port")]
|
||||
port: u16,
|
||||
},
|
||||
|
||||
#[structopt(about = "Remove ip and port")]
|
||||
RemoveIpPort {
|
||||
#[structopt(name = "Ip to remove", short = "i", long = "ip")]
|
||||
ip: String,
|
||||
|
||||
#[structopt(name = "Port to remove", short = "p", long = "port")]
|
||||
port: u16,
|
||||
},
|
||||
}
|
||||
@ -1,7 +1,7 @@
|
||||
use std::process::Command;
|
||||
use std::{path::Path, process::Command};
|
||||
|
||||
pub fn save_iptables() {
|
||||
pub fn save_iptables(path: &Path) {
|
||||
let _ = Command::new("iptables-save")
|
||||
.args(&["-f", "/host_iptables/rules.v4"])
|
||||
.args(["-f", path.to_str().unwrap()])
|
||||
.output();
|
||||
}
|
||||
|
||||
41
src/iptables_wrapper.rs
Normal file
41
src/iptables_wrapper.rs
Normal file
@ -0,0 +1,41 @@
|
||||
pub fn ban_port(port: u16) {
|
||||
let iptables = iptables::new(false).unwrap();
|
||||
let _ = iptables.append_unique(
|
||||
"filter",
|
||||
"INPUT",
|
||||
&format!("-p tcp --dport {} -j DROP", port),
|
||||
);
|
||||
|
||||
println!("banned port {}", port);
|
||||
}
|
||||
|
||||
pub fn unban_port(port: u16) {
|
||||
let iptables = iptables::new(false).unwrap();
|
||||
let _ = iptables.delete(
|
||||
"filter",
|
||||
"INPUT",
|
||||
&format!("-p tcp --dport {} -j DROP", port),
|
||||
);
|
||||
|
||||
println!("unbanned port {}", port);
|
||||
}
|
||||
|
||||
pub fn allow_ip_port(ip: &str, port: u16) {
|
||||
let iptables = iptables::new(false).unwrap();
|
||||
let _ = iptables.append_unique(
|
||||
"filter",
|
||||
"INPUT",
|
||||
&format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip),
|
||||
);
|
||||
println!("allowed {} to access {}", ip, port);
|
||||
}
|
||||
|
||||
pub fn remove_ip_port(ip: &str, port: u16) {
|
||||
let iptables = iptables::new(false).unwrap();
|
||||
let _ = iptables.delete(
|
||||
"filter",
|
||||
"INPUT",
|
||||
&format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip),
|
||||
);
|
||||
println!("removed access {} to {}", ip, port);
|
||||
}
|
||||
60
src/main.rs
60
src/main.rs
@ -1,37 +1,57 @@
|
||||
pub mod cli;
|
||||
pub mod iptables_save;
|
||||
pub mod iptables_wrapper;
|
||||
pub mod login_attempt;
|
||||
pub mod tpc_command_server;
|
||||
|
||||
use crate::tpc_command_server::start_tcp_command_server;
|
||||
use cli::Arguments;
|
||||
use iptables_wrapper::{allow_ip_port, ban_port, remove_ip_port, unban_port};
|
||||
use linemux::MuxedLines;
|
||||
use login_attempt::LoginAttempt;
|
||||
use std::path::PathBuf;
|
||||
use std::thread;
|
||||
use std::{collections::HashMap, thread::sleep, time::Duration};
|
||||
use structopt::StructOpt;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> std::io::Result<()> {
|
||||
async fn main() {
|
||||
let opts = Arguments::from_args();
|
||||
match opts {
|
||||
Arguments::BanServer {
|
||||
ssh_auth_log,
|
||||
iptables_save,
|
||||
} => {
|
||||
let _ = start_ban_server(ssh_auth_log, iptables_save).await;
|
||||
}
|
||||
Arguments::BanPort { port } => ban_port(port),
|
||||
Arguments::UnbanPort { port } => unban_port(port),
|
||||
Arguments::AllowIpPort { ip, port } => allow_ip_port(&ip, port),
|
||||
Arguments::RemoveIpPort { ip, port } => remove_ip_port(&ip, port),
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_ban_server(
|
||||
ssh_auth_log: PathBuf,
|
||||
iptables_save: Option<PathBuf>,
|
||||
) -> std::io::Result<()> {
|
||||
let iptables = iptables::new(false).unwrap();
|
||||
let mut lines = MuxedLines::new()?;
|
||||
lines.add_file("/host_ssh/auth.log").await?;
|
||||
lines.add_file(&ssh_auth_log).await?;
|
||||
let mut login_attempts: HashMap<String, usize> = HashMap::new();
|
||||
|
||||
let seconds_iptables = Duration::from_secs(60);
|
||||
println!(
|
||||
"starting iptables-save, run every {} seconds",
|
||||
seconds_iptables.as_secs()
|
||||
);
|
||||
if let Some(iptables_save) = iptables_save {
|
||||
let seconds_iptables = Duration::from_secs(60);
|
||||
println!(
|
||||
"starting iptables-save, run every {} seconds",
|
||||
seconds_iptables.as_secs()
|
||||
);
|
||||
thread::spawn(move || loop {
|
||||
sleep(seconds_iptables);
|
||||
iptables_save::save_iptables(&iptables_save);
|
||||
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");
|
||||
println!("listening to changes over {}", ssh_auth_log.display());
|
||||
while let Ok(Some(line)) = lines.next_line().await {
|
||||
if let Some(login_attempt) = LoginAttempt::capture(line.line()) {
|
||||
println!(
|
||||
|
||||
@ -1,103 +0,0 @@
|
||||
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:9999").unwrap();
|
||||
println!("listening on port 9999 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) {
|
||||
for line in BufReader::new(stream).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 {
|
||||
"banport" => {
|
||||
if let Some(port) = arguments.get(0) {
|
||||
let iptables = iptables::new(false).unwrap();
|
||||
let _ = iptables.append_unique(
|
||||
"filter",
|
||||
"INPUT",
|
||||
&format!("-p tcp --dport {} -j DROP", port),
|
||||
);
|
||||
|
||||
format!("banned port {} for all ips", port)
|
||||
} else {
|
||||
"missing args for banport: port".to_string()
|
||||
}
|
||||
}
|
||||
"unbanport" => {
|
||||
if let Some(port) = arguments.get(0) {
|
||||
let iptables = iptables::new(false).unwrap();
|
||||
let _ = iptables.delete(
|
||||
"filter",
|
||||
"INPUT",
|
||||
&format!("-p tcp --dport {} -j DROP", port),
|
||||
);
|
||||
|
||||
format!("unbanned port {}", port)
|
||||
} else {
|
||||
"missing args for unbanport: port".to_string()
|
||||
}
|
||||
}
|
||||
"allowipport" => {
|
||||
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),
|
||||
);
|
||||
format!("allowed {} to access {}", ip, port)
|
||||
} else {
|
||||
"missing args for allowipport: ip and port".to_string()
|
||||
}
|
||||
}
|
||||
"removeipport" => {
|
||||
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),
|
||||
);
|
||||
format!("rm {} access to {}", ip, port)
|
||||
} else {
|
||||
"missing args for rmipport: ip and port".to_string()
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
format!("unknown command: {}", command)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user