1
0

WIP moving code into lib

This commit is contained in:
2024-12-31 03:57:40 +01:00
parent 41abf04a63
commit 20f1038ff8
9 changed files with 448 additions and 420 deletions

View File

@ -2,14 +2,14 @@ use structopt::StructOpt;
#[derive(Debug, StructOpt)]
#[structopt(
name = "martillo-maldito",
about = "A simple iptables wrapper, including a ban server"
name = "martillo_maldito",
about = "A IPTables wrapper, including a ban service"
)]
pub enum Cli {
#[structopt(about = "Initialize ban server")]
BanServer,
#[structopt(about = "Initialize ban service, monitoring SSH logs for login attempts")]
BanService,
#[structopt(about = "List all banned ips")]
ListBannedIps {
GetBannedIps {
#[structopt(name = "Docker", short = "d", long = "docker")]
docker: bool,
},
@ -30,51 +30,42 @@ pub enum Cli {
#[structopt(name = "Docker", short = "d", long = "docker")]
docker: bool,
},
#[structopt(about = "Ban port")]
#[structopt(about = "Secure a port")]
SecurePort {
#[structopt(name = "Port to ban", short = "p", long = "port")]
port: u16,
#[structopt(name = "Position", short = "P", long = "position")]
position: Option<usize>,
#[structopt(name = "Docker", short = "d", long = "docker")]
docker: bool,
#[structopt(name = "Position", short = "P", long = "position")]
position: Option<i32>,
},
#[structopt(about = "Unban port")]
#[structopt(about = "Unsecure a port")]
UnsecurePort {
#[structopt(name = "Port to unban", short = "p", long = "port")]
port: u16,
#[structopt(name = "Docker", short = "d", long = "docker")]
docker: bool,
},
#[structopt(about = "Allow ip and port")]
#[structopt(about = "Allow an IP for port")]
AllowIpForPort {
#[structopt(name = "Ip to allow", short = "i", long = "ip")]
#[structopt(name = "IP to allow", short = "i", long = "ip")]
ip: String,
#[structopt(name = "Port to allow", short = "p", long = "port")]
port: u16,
#[structopt(name = "Docker", short = "d", long = "docker")]
docker: bool,
#[structopt(name = "Position", short = "P", long = "position")]
position: Option<i32>,
},
#[structopt(about = "Allow port for only an ip")]
OnlyIpForPort {
#[structopt(name = "Ip to allow", short = "i", long = "ip")]
ip: String,
#[structopt(name = "Port to allow", short = "p", long = "port")]
port: u16,
position: Option<usize>,
#[structopt(name = "Docker", short = "d", long = "docker")]
docker: bool,
},
#[structopt(about = "Removes an allow port for an ip")]
#[structopt(about = "Removes an allowance IP for a port")]
RemoveAllowIpPort {
#[structopt(name = "Ip to remove", short = "i", long = "ip")]
#[structopt(name = "IP to remove", short = "i", long = "ip")]
ip: String,
#[structopt(name = "Port to remove", short = "p", long = "port")]
port: u16,
#[structopt(name = "Docker", short = "d", long = "docker")]
docker: bool,
},
#[structopt(about = "Saves the IPTables configuration")]
SaveIPTables,
#[structopt(about = "Saves the configuration")]
SaveRules,
}

View File

@ -1,7 +0,0 @@
use std::process::Command;
pub fn save_iptables() -> std::io::Result<std::process::Output> {
Command::new("iptables-save")
.args(["-f", "/etc/iptables/rules.v4"])
.output()
}

View File

@ -1,229 +0,0 @@
use iptables::IPTables;
use regex::Regex;
use std::collections::HashMap;
pub fn is_port_secured(port: u16, docker: bool) -> bool {
let iptables = iptables::new(false).unwrap();
let rules = iptables.list("filter", &get_chain(docker));
if rules.is_err() {
return false;
}
for rule in rules.unwrap() {
if rule.contains(&format!("-p tcp -m tcp --dport {} -j DROP", port)) {
return true;
}
}
false
}
pub fn list_secured_ports(docker: bool) -> Vec<u16> {
let iptables = iptables::new(false).unwrap();
let chain = get_chain(docker);
let rules = iptables.list("filter", &chain);
if rules.is_err() {
return vec![];
}
let rgx = get_regex_for_port();
rules
.unwrap()
.iter()
.filter(|r| r.contains("-p tcp -m tcp --dport") && r.contains("-j DROP"))
.map(|r| extract_port(&rgx, r).unwrap())
.collect()
}
pub fn list_banned_ips(docker: bool) -> Vec<String> {
let iptables = iptables::new(false).unwrap();
let chain = get_chain(docker);
let rules = iptables.list("filter", &chain);
if rules.is_err() {
return vec![];
}
let rgx = get_regex_for_ip();
rules
.unwrap()
.iter()
.filter(|r| {
r.contains(&format!("-A {}", chain)) && r.contains("-j DROP") && r.contains("-s")
})
.map(|r| extract_ip(&rgx, r).unwrap())
.collect()
}
pub fn map_secured_ports_allowed_ips(docker: bool) -> HashMap<u16, Vec<String>> {
let mut result: HashMap<u16, Vec<String>> = HashMap::new();
let secured_ports = list_secured_ports(docker);
if secured_ports.is_empty() {
return result;
}
let iptables = iptables::new(false).unwrap();
let chain = get_chain(docker);
let rules = iptables.list("filter", &chain);
if rules.is_err() {
return result;
}
let rules = rules.unwrap();
let rgx = get_regex_for_ip();
for port in secured_ports {
let ips = rules
.iter()
.filter(|r| {
r.contains(&format!("-A {} -s", chain))
&& r.contains(&format!("-p tcp -m tcp --dport {} -j ACCEPT", port))
})
.map(|r| extract_ip(&rgx, r).unwrap())
.collect();
result.insert(port, ips);
}
result
}
fn secure_port_rule(port: u16) -> String {
format!("-p tcp --dport {} -j DROP", port)
}
pub fn secure_port(
port: u16,
docker: bool,
position: Option<i32>,
) -> Result<(), Box<dyn std::error::Error>> {
let iptables = iptables::new(false).unwrap();
let table = "filter";
let chain = get_chain(docker);
let rule = secure_port_rule(port);
let position = if docker && position.is_none() {
let all_docker_rules = iptables.list("filter", &chain).unwrap();
Some(all_docker_rules.len() as i32 - 1)
} else {
position
};
if let Some(position) = position {
insert_unique(&iptables, table, &chain, &rule, position)
} else {
append_unique(&iptables, table, &chain, &rule)
}
}
pub fn unsecure_port(port: u16, docker: bool) -> Result<(), Box<dyn std::error::Error>> {
let iptables = iptables::new(false).unwrap();
let rule = secure_port_rule(port);
iptables.delete("filter", &get_chain(docker), &rule)
}
fn allow_ip_for_port_rule(port: u16, ip: &str) -> String {
format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip)
}
pub fn allow_ip_for_port(
ip: &str,
port: u16,
docker: bool,
position: Option<i32>,
) -> Result<(), Box<dyn std::error::Error>> {
let iptables = iptables::new(false).unwrap();
let table = "filter";
let chain = get_chain(docker);
let rule = allow_ip_for_port_rule(port, ip);
if let Some(position) = position {
insert_unique(&iptables, table, &chain, &rule, position)
} else {
append_unique(&iptables, table, &chain, &rule)
}
}
pub fn remove_allow_ip_for_port(
ip: &str,
port: u16,
docker: bool,
) -> Result<(), Box<dyn std::error::Error>> {
let iptables = iptables::new(false).unwrap();
let rule = allow_ip_for_port_rule(port, ip);
iptables.delete("filter", &get_chain(docker), &rule)
}
fn get_chain(docker: bool) -> String {
if docker {
"DOCKER-USER".to_string()
} else {
"INPUT".to_string()
}
}
fn append_unique(
iptables: &IPTables,
table: &str,
chain: &str,
rule: &str,
) -> Result<(), Box<dyn std::error::Error>> {
iptables.append_unique(table, chain, rule)
}
fn insert_unique(
iptables: &IPTables,
table: &str,
chain: &str,
rule: &str,
position: i32,
) -> Result<(), Box<dyn std::error::Error>> {
iptables.insert_unique(table, chain, rule, position)
}
fn extract_ip(regex: &Regex, input: &str) -> Option<String> {
regex
.captures(input)
.and_then(|caps| caps.get(0).map(|m| m.as_str().to_string()))
}
fn extract_port(regex: &Regex, input: &str) -> Option<u16> {
regex
.captures(input)
.and_then(|caps| caps.get(1).map(|m| m.as_str().parse::<u16>().unwrap()))
}
fn get_regex_for_ip() -> Regex {
Regex::new(r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b").unwrap()
}
fn get_regex_for_port() -> Regex {
Regex::new(r"--dport\s+(\d+)").unwrap()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn correct_extract_ip() {
let regex = get_regex_for_ip();
let input = "-A INPUT -s 81.69.255.132/32 -j DROP";
assert_eq!(extract_ip(&regex, input), Some("81.69.255.132".to_string()));
}
#[test]
fn no_match_extract_ip() {
let regex = get_regex_for_ip();
let input = "-A INPUT -j DROP";
assert_eq!(extract_ip(&regex, input), None);
}
#[test]
fn docker_extract_ip() {
let regex = get_regex_for_ip();
let input = "-A DOCKER -d 172.18.0.2/32 ! -i br-127d33df48a4 -o br-127d33df48a4 -p tcp -m tcp --dport 8078 -j ACCEPT";
assert_eq!(extract_ip(&regex, input), Some("172.18.0.2".to_string()));
}
}

3
src/lib.rs Normal file
View File

@ -0,0 +1,3 @@
pub mod cli;
pub mod login_attempt;
pub mod martillo_maldito;

View File

@ -1,15 +1,6 @@
pub mod cli;
pub mod iptables_save;
pub mod iptables_wrapper;
pub mod login_attempt;
use cli::Cli;
use iptables_wrapper::{
allow_ip_for_port, is_port_secured, list_banned_ips, list_secured_ports,
map_secured_ports_allowed_ips, remove_allow_ip_for_port, secure_port, unsecure_port,
};
use env_logger::Builder;
use log::{error, info};
use login_attempt::LoginAttempt;
use martillo_maldito::{cli::Cli, login_attempt::LoginAttempt, martillo_maldito::MartilloMaldito};
use std::{
collections::HashMap,
io::BufRead,
@ -21,71 +12,80 @@ use structopt::StructOpt;
#[tokio::main]
async fn main() {
start_logger();
match Cli::from_args() {
Cli::BanServer => {
if let Err(error) = start_ban_server().await {
error!("ban server: {error}");
Cli::BanService => {
start_logger();
if let Err(err) = start_ban_service().await {
error!(err = err.to_string().as_str();
"Ban service"
);
}
}
Cli::ListBannedIps { docker } => {
println!(
"{}",
serde_json::to_string(&list_banned_ips(docker)).unwrap()
);
Cli::GetBannedIps { docker } => {
let banned_ips = MartilloMaldito::ipv4(docker).get_banned_ips();
println!("{}", serde_json::to_string(&banned_ips).unwrap());
}
Cli::ListSecuredPorts { docker } => {
println!(
"{}",
serde_json::to_string(&list_secured_ports(docker)).unwrap()
);
let secured_ports = MartilloMaldito::ipv4(docker).get_secured_ports();
println!("{}", serde_json::to_string(&secured_ports).unwrap());
}
Cli::MapSecuredPortsAllowedIps { docker } => {
let secured_ports_with_allowed_ips =
MartilloMaldito::ipv4(docker).get_secured_ports_with_allowed_ips();
println!(
"{}",
serde_json::to_string(&map_secured_ports_allowed_ips(docker)).unwrap()
serde_json::to_string(&secured_ports_with_allowed_ips).unwrap()
);
}
Cli::IsPortSecured { port, docker } => {
println!("{}", is_port_secured(port, docker));
let is_port_secured = MartilloMaldito::ipv4(docker).is_port_secured(port);
println!("{}", is_port_secured);
}
Cli::SecurePort {
port,
docker,
position,
} => println!("{}", secure_port(port, docker, position).is_ok()),
Cli::UnsecurePort { port, docker } => println!("{}", unsecure_port(port, docker).is_ok()),
docker,
} => {
let port_secured = MartilloMaldito::ipv4(docker).secure_port(port, position);
println!("{}", port_secured.is_ok())
}
Cli::UnsecurePort { port, docker } => {
let port_unsecured = MartilloMaldito::ipv4(docker).unsecure_port(port);
println!("{}", port_unsecured.is_ok())
}
Cli::AllowIpForPort {
ip,
port,
docker,
position,
} => println!("{}", allow_ip_for_port(&ip, port, docker, position).is_ok()),
Cli::OnlyIpForPort { ip, port, docker } => {
let allowed = allow_ip_for_port(&ip, port, docker, Some(1));
let secured = secure_port(port, docker, Some(2));
println!("{}", allowed.is_ok() && secured.is_ok());
docker,
} => {
let allowed_ip = MartilloMaldito::ipv4(docker).allow_ip_for_port(&ip, port, position);
println!("{}", allowed_ip.is_ok())
}
Cli::RemoveAllowIpPort { ip, port, docker } => {
println!("{}", remove_allow_ip_for_port(&ip, port, docker).is_ok())
let removed_allow_ip =
MartilloMaldito::ipv4(docker).remove_allow_ip_for_port(&ip, port);
println!("{}", removed_allow_ip.is_ok())
}
Cli::SaveIPTables => {
println!("{}", iptables_save::save_iptables().is_ok())
Cli::SaveRules => {
println!("{}", MartilloMaldito::save_rules().is_ok())
}
}
}
async fn start_ban_server() -> std::io::Result<()> {
async fn start_ban_service() -> std::io::Result<()> {
let seconds_iptables = Duration::from_secs(60);
info!("Saving IPTables every {} secs", seconds_iptables.as_secs());
info!(every_seconds = seconds_iptables.as_secs();
"Saving IPTables"
);
spawn(move || loop {
sleep(seconds_iptables);
iptables_save::save_iptables().expect("Failed to save IPTables");
MartilloMaldito::save_rules().expect("Failed to save rules");
});
let mut child: Child = Command::new("journalctl")
let child: Child = Command::new("journalctl")
.arg("-D")
.arg("/var/log/journal")
.arg("-u")
@ -95,17 +95,18 @@ async fn start_ban_server() -> std::io::Result<()> {
.spawn()
.expect("Failed to start journalctl");
let stdout = child.stdout.as_mut().expect("Failed to capture stdout");
let stdout = child.stdout.expect("Failed to capture stdout");
let mut reader = std::io::BufReader::new(stdout);
let iptables = iptables::new(false).unwrap();
let martillo_maldito = MartilloMaldito::ipv4(false);
let mut login_attempts: HashMap<String, usize> = HashMap::new();
info!("Start reading logins from SSH");
loop {
let mut line = String::new();
if let Err(err) = reader.read_line(&mut line) {
error!("Reading line: {}", err);
error!(err = err.to_string().as_str();
"Reading line"
);
continue;
}
@ -114,9 +115,9 @@ async fn start_ban_server() -> std::io::Result<()> {
}
if let Some(login_attempt) = LoginAttempt::capture(&line) {
info!(
"Login attempt from {} using {} user",
login_attempt.ip, login_attempt.user
info!(ip = login_attempt.ip.as_str(),
user = login_attempt.user.as_str();
"Login attempt",
);
match login_attempts.get_mut(&login_attempt.ip) {
@ -124,15 +125,10 @@ async fn start_ban_server() -> std::io::Result<()> {
*count += 1;
if *count == 3 {
if iptables
.append_unique(
"filter",
"INPUT",
&format!("--source {} -j DROP", login_attempt.ip),
)
.is_ok()
{
info!("IP {} banned", login_attempt.ip);
if martillo_maldito.ban_ip(&login_attempt.ip).is_ok() {
info!(ip = login_attempt.ip.as_str();
"Banned IP"
);
}
login_attempts.remove(&login_attempt.ip);
@ -147,11 +143,5 @@ async fn start_ban_server() -> std::io::Result<()> {
}
fn start_logger() {
unsafe {
std::env::set_var(
"RUST_LOG",
std::env::var("RUST_LOG").unwrap_or("INFO".to_string()),
);
}
logfmt_logger::init();
Builder::from_default_env().init();
}

243
src/martillo_maldito.rs Normal file
View File

@ -0,0 +1,243 @@
use regex::Regex;
use std::{collections::HashMap, process::Command};
pub struct MartilloMaldito {
iptables: iptables::IPTables,
chain: String,
}
impl MartilloMaldito {
pub fn ipv4(docker: bool) -> MartilloMaldito {
MartilloMaldito {
iptables: iptables::new(false).unwrap(),
chain: Self::get_chain(docker).to_string(),
}
}
pub fn ipv6(docker: bool) -> MartilloMaldito {
MartilloMaldito {
iptables: iptables::new(true).unwrap(),
chain: Self::get_chain(docker).to_string(),
}
}
pub fn save_rules() -> std::io::Result<std::process::Output> {
Command::new("iptables-save")
.args(["-f", "/etc/iptables/rules.v4"])
.output()
}
pub fn is_port_secured(&self, port: u16) -> bool {
let rules = self.get_rules();
if rules.is_err() {
return false;
}
for rule in rules.unwrap() {
if rule.contains(&format!("-p tcp -m tcp --dport {} -j DROP", port)) {
return true;
}
}
false
}
pub fn get_secured_ports(&self) -> Vec<u16> {
let rules = self.get_rules();
if rules.is_err() {
return vec![];
}
let port_regex = iptables_regex_for_port();
rules
.unwrap()
.iter()
.filter(|r| r.contains("-p tcp -m tcp --dport") && r.contains("-j DROP"))
.map(|r| extract_port(&port_regex, r).unwrap())
.collect()
}
pub fn get_banned_ips(&self) -> Vec<String> {
let rules = self.get_rules();
if rules.is_err() {
return vec![];
}
let ip_regex = iptables_regex_for_ip();
rules
.unwrap()
.iter()
.filter(|r| {
r.contains(&format!("-A {}", self.chain))
&& r.contains("-j DROP")
&& r.contains("-s")
})
.map(|r| extract_ip(&ip_regex, r).unwrap())
.collect()
}
pub fn get_secured_ports_with_allowed_ips(&self) -> HashMap<u16, Vec<String>> {
let mut result: HashMap<u16, Vec<String>> = HashMap::new();
let secured_ports = self.get_secured_ports();
if secured_ports.is_empty() {
return result;
}
let rules = self.get_rules();
if rules.is_err() {
return result;
}
let rules = rules.unwrap();
let ip_regex = iptables_regex_for_ip();
for port in secured_ports {
let ips = rules
.iter()
.filter(|r| {
r.contains(&format!("-A {} -s", self.chain))
&& r.contains(&format!("-p tcp -m tcp --dport {} -j ACCEPT", port))
})
.map(|r| extract_ip(&ip_regex, r).unwrap())
.collect();
result.insert(port, ips);
}
result
}
pub fn ban_ip(&self, ip: &str) -> Result<(), Box<dyn std::error::Error>> {
self.append_unique("filter", &format!("-s {} -j DROP", ip))
}
pub fn secure_port(
&self,
port: u16,
position: Option<usize>,
) -> Result<(), Box<dyn std::error::Error>> {
let table = "filter";
let rule = secure_port_rule(port);
let position = if self.chain == "DOCKER-USER" && position.is_none() {
let all_docker_rules = self.get_rules()?;
Some(all_docker_rules.len() - 1)
} else {
position
};
if let Some(position) = position {
self.insert_unique(table, &rule, position)
} else {
self.append_unique(table, &rule)
}
}
pub fn unsecure_port(&self, port: u16) -> Result<(), Box<dyn std::error::Error>> {
let rule = secure_port_rule(port);
self.iptables.delete("filter", &self.chain, &rule)
}
pub fn allow_ip_for_port(
&self,
ip: &str,
port: u16,
position: Option<usize>,
) -> Result<(), Box<dyn std::error::Error>> {
let table = "filter";
let rule = allow_ip_for_port_rule(port, ip);
if let Some(position) = position {
self.insert_unique(table, &rule, position)
} else {
self.append_unique(table, &rule)
}
}
pub fn remove_allow_ip_for_port(
&self,
ip: &str,
port: u16,
) -> Result<(), Box<dyn std::error::Error>> {
let rule = allow_ip_for_port_rule(port, ip);
self.iptables.delete("filter", &self.chain, &rule)
}
fn get_rules(&self) -> Result<Vec<String>, Box<dyn std::error::Error>> {
self.iptables.list("filter", &self.chain)
}
fn append_unique(&self, table: &str, rule: &str) -> Result<(), Box<dyn std::error::Error>> {
self.iptables.append_unique(table, &self.chain, rule)
}
fn insert_unique(
&self,
table: &str,
rule: &str,
position: usize,
) -> Result<(), Box<dyn std::error::Error>> {
self.iptables
.insert_unique(table, &self.chain, rule, position as i32)
}
fn get_chain(docker: bool) -> &'static str {
if docker {
"DOCKER-USER"
} else {
"INPUT"
}
}
}
fn secure_port_rule(port: u16) -> String {
format!("-p tcp --dport {} -j DROP", port)
}
fn allow_ip_for_port_rule(port: u16, ip: &str) -> String {
format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip)
}
fn extract_ip(regex: &Regex, input: &str) -> Option<String> {
regex
.captures(input)
.and_then(|caps| caps.get(0).map(|m| m.as_str().to_string()))
}
fn extract_port(regex: &Regex, input: &str) -> Option<u16> {
regex
.captures(input)
.and_then(|caps| caps.get(1).map(|m| m.as_str().parse::<u16>().unwrap()))
}
fn iptables_regex_for_ip() -> Regex {
Regex::new(r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b").unwrap()
}
fn iptables_regex_for_port() -> Regex {
Regex::new(r"--dport\s+(\d+)").unwrap()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn correct_extract_ip() {
let regex = iptables_regex_for_ip();
let input = "-A INPUT -s 81.69.255.132/32 -j DROP";
assert_eq!(extract_ip(&regex, input), Some("81.69.255.132".to_string()));
}
#[test]
fn no_match_extract_ip() {
let regex = iptables_regex_for_ip();
let input = "-A INPUT -j DROP";
assert_eq!(extract_ip(&regex, input), None);
}
#[test]
fn docker_extract_ip() {
let regex = iptables_regex_for_ip();
let input = "-A DOCKER -d 172.18.0.2/32 ! -i br-127d33df48a4 -o br-127d33df48a4 -p tcp -m tcp --dport 8078 -j ACCEPT";
assert_eq!(extract_ip(&regex, input), Some("172.18.0.2".to_string()));
}
}