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 rules = iptables.list("filter", &get_chain(docker));
    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() -> Vec<String> {
    let iptables = iptables::new(false).unwrap();
    let rules = iptables.list("filter", &get_chain(false));
    if rules.is_err() {
        return vec![];
    }

    let rgx = get_regex_for_ip();
    rules
        .unwrap()
        .iter()
        .filter(|r| r.contains("-A INPUT") && 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 rules = iptables.list("filter", &get_chain(docker));
    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("-A INPUT -s")
                    && 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);
    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();
    iptables.delete("filter", &get_chain(docker), &secure_port_rule(port))
}

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();
    iptables.delete(
        "filter",
        &get_chain(docker),
        &allow_ip_for_port_rule(port, ip),
    )
}

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_extrac_ip() {
        let regex = get_regex_for_ip();
        let input = "-A INPUT -j DROP";
        assert_eq!(extract_ip(&regex, input), None);
    }
}