#!/bin/bash

set -euo pipefail

readonly IPT=iptables

# ------------------------- Утилиты -------------------------
require_root() {
    if [[ ${EUID:-$(id -u)} -ne 0 ]]; then
        echo "[!] Пожалуйста, запустите скрипт от имени root (sudo)." >&2
        exit 1
    fi
}

ensure_dependencies() {
    if ! command -v ${IPT} >/dev/null 2>&1; then
        echo "[!] iptables не найден. Установите: apt update && apt install -y iptables" >&2
        exit 1
    fi
    if ! command -v netfilter-persistent >/dev/null 2>&1; then
        echo "[i] Устанавливаем iptables-persistent (без интерактивного режима)..."
        DEBIAN_FRONTEND=noninteractive apt-get update >/dev/null 2>&1 || true
        DEBIAN_FRONTEND=noninteractive apt-get install -y iptables-persistent netfilter-persistent >/dev/null 2>&1 || true
    fi
}

validate_port() {
    local port="$1"
    [[ "$port" =~ ^[0-9]+$ ]] || return 1
    (( port >= 1 && port <= 65535 ))
}

validate_protocol() {
    local proto="$1"
    case "$proto" in
        tcp|udp|both) return 0 ;;
        *) return 1 ;;
    esac
}

validate_ip() {
    local ip="$1"
    # Принимаем IPv4 CIDR или одиночный IPv4
    if [[ "$ip" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}(/([0-9]|[1-2][0-9]|3[0-2]))?$ ]]; then
        IFS=/ read -r base _ <<<"$ip"
        IFS=. read -r a b c d <<<"$base"
        for o in "$a" "$b" "$c" "$d"; do
            (( o >= 0 && o <= 255 )) || return 1
        done
        return 0
    fi
    return 1
}

# === Сохранение правил ===
save_rules() {
    echo "[+] Сохраняем правила..."
    # Используем iptables-persistent для сохранения правил
    netfilter-persistent save >/dev/null 2>&1 || true
    echo "[+] Правила сохранены."
}

# === Показать правила ===
show_rules() {
    echo "========== IPTABLES RULES =========="
    iptables -L -n -v --line-numbers
    echo "========== DOCKER-USER RULES =========="
    if iptables -S DOCKER-USER >/dev/null 2>&1; then
        iptables -L DOCKER-USER -n -v --line-numbers
    else
        echo "(chain DOCKER-USER not present)"
    fi
}

parse_port_rules() {
    local chain="$1"
    local action_filter="${2:-}" # ACCEPT or DROP
    local src_filter="${3:-}" # -s IP
    local rules
    
    rules=$(iptables -S "$chain" 2>/dev/null)
    if [[ -z "$rules" ]]; then
        return
    fi

    echo "# Chain: $chain"
    echo "# Protocol | Port    | Action | Source IP"
    echo "-------------------------------------------"
    echo "$rules" | while read -r rule; do
        local proto="-" port="-" action="-" src="-"
        
        [[ "$rule" =~ -p\ (tcp|udp) ]] && proto="${BASH_REMATCH[1]}"
        [[ "$rule" =~ --dport\ ([0-9]+) ]] && port="${BASH_REMATCH[1]}"
        [[ "$rule" =~ -j\ (ACCEPT|DROP) ]] && action="${BASH_REMATCH[1]}"
        [[ "$rule" =~ -s\ ([0-9\./]+) ]] && src="${BASH_REMATCH[1]}"

        # Skip if no port rule or if action/src filter doesn't match
        if [[ "$port" == "-" ]]; then continue; fi
        if [[ -n "$action_filter" && "$action" != "$action_filter" ]]; then continue; fi
        if [[ -n "$src_filter" && "$src" != "$src_filter" ]]; then continue; fi

        printf "%-10s | %-7s | %-6s | %s\n" "$proto" "$port" "$action" "$src"
    done
}

# === Показать открытые порты ===
show_open_ports() {
    echo "\n========== ОТКРЫТЫЕ ПОРТЫ (ACCEPT) =========="
    parse_port_rules INPUT ACCEPT
    parse_port_rules DOCKER-USER ACCEPT
    echo "============================================="
    echo "[i] Если порт открыт, но не работает, убедитесь, что служба запущена и слушает этот порт (например, командой 'ss -tulpn | grep <PORT>')."
}

# === Показать закрытые порты (из DROP) ===
show_closed_ports() {
    echo "\n========== ЗАКРЫТЫЕ ПОРТЫ (DROP) =========="
    parse_port_rules INPUT DROP
    parse_port_rules DOCKER-USER DROP
    echo "==========================================="
}

show_all_ports() {
    echo "\n========== ВСЕ ПОРТОВЫЕ ПРАВИЛА =========="
    parse_port_rules INPUT
    parse_port_rules DOCKER-USER
    echo "=========================================="
}

# Helper for 'Press Enter to continue...'
press_enter() {
    read -rp $'Нажмите Enter, чтобы продолжить...'
}

# === Блокировка IP ===
block_ip() {
    read -p "Введите IP для блокировки: " ip
    if ! validate_ip "$ip"; then echo "[!] Неверный IP-адрес."; press_enter; return; fi
    iptables -A INPUT -s "$ip" -j DROP
    iptables -A DOCKER-USER -s "$ip" -j DROP
    echo "[+] IP $ip заблокирован."
    save_rules
}

# === Разблокировка IP ===
unblock_ip() {
    read -p "Введите IP для разблокировки: " ip
    if ! validate_ip "$ip"; then echo "[!] Неверный IP-адрес."; press_enter; return; fi
    iptables -D INPUT -s "$ip" -j DROP 2>/dev/null
    iptables -D DOCKER-USER -s "$ip" -j DROP 2>/dev/null
    echo "[+] IP $ip разблокирован."
    save_rules
}

# === Ensure DOCKER-USER chain exists and hooked ===
ensure_docker_user_chain() {
    iptables -S DOCKER-USER >/dev/null 2>&1 || iptables -N DOCKER-USER
    iptables -C FORWARD -j DOCKER-USER >/dev/null 2>&1 || iptables -I FORWARD -j DOCKER-USER
}

# === Меню управления портами ===
ports_menu() {
    while true; do
        cat <<MENU
================ PORT MANAGER ================
1) Показать открытые порты
2) Показать закрытые порты
3) Показать все порты
4) Вернуться в главное меню
MENU
        read -p "Выберите действие: " choice
        case $choice in
            1) show_open_ports; press_enter ;;
            2) show_closed_ports; press_enter ;;
            3) show_all_ports; press_enter ;;
            4) break ;;
            *) echo "Неверный выбор." ;;
        esac
    done
}

# === Главное меню ===
main_menu() {
    while true; do
        cat <<MENU
================ NETWORK MANAGER ================
1) Открыть порт
2) Закрыть порт
3) Добавить IP к порту
4) Убрать IP с порта
5) Показать правила
6) Порты (подменю)
7) Заблокировать IP
8) Разблокировать IP
9) Выход
MENU
        read -p "Выберите действие: " action
        case $action in
            1) new_open_port; press_enter ;;
            2) new_close_port; press_enter ;;
            3) new_add_ip_to_port; press_enter ;;
            4) new_remove_ip_from_port; press_enter ;;
            5) show_rules; press_enter ;;
            6) ports_menu ;;
            7) block_ip; press_enter ;;
            8) unblock_ip; press_enter ;;
            9) echo "Выход"; exit 0 ;;
            *) echo "Неверный выбор." ;;
        esac
    done
}

# --- Helper functions for managing iptables rules ---
add_port_rule() {
    local chain="$1" proto="$2" port="$3" ip="$4" action="$5"
    if [[ -z "$ip" ]]; then
        ${IPT} -A "$chain" -p "$proto" --dport "$port" -j "$action"
    else
        ${IPT} -A "$chain" -p "$proto" --dport "$port" -s "$ip" -j "$action"
    fi
}

delete_port_rule() {
    local chain="$1" proto="$2" port="$3" ip="$4" action="$5"
    if [[ -z "$ip" ]]; then
        ${IPT} -D "$chain" -p "$proto" --dport "$port" -j "$action" 2>/dev/null || true
    else
        ${IPT} -D "$chain" -p "$proto" --dport "$port" -s "$ip" -j "$action" 2>/dev/null || true
    fi
}

new_open_port() {
    read -p "Введите номер порта: " port
    if ! validate_port "$port"; then echo "[!] Неверный номер порта."; press_enter; return; fi
    read -p "Протокол (tcp/udp/both): " proto
    proto=$(echo "$proto" | tr '[:upper:]' '[:lower:]')
    if ! validate_protocol "$proto"; then echo "[!] Неверный протокол."; press_enter; return; fi
    read -p "Ограничить доступ по IP? (y/n): " limit

    limit=$(echo "$limit" | tr '[:upper:]' '[:lower:]')

    if [ "$proto" == "both" ]; then
        protos=("tcp" "udp")
    else
        protos=("$proto")
    fi

    for p in "${protos[@]}"; do
        # Удаляем ВСЕ существующие правила для этого порта из внешних цепочек
        # Удаляем правила ACCEPT
        while iptables -D INPUT -p "$p" --dport "$port" -j ACCEPT 2>/dev/null; do :; done
        while iptables -D DOCKER-USER -p "$p" --dport "$port" -j ACCEPT 2>/dev/null; do :; done
        # Удаляем правила DROP
        while iptables -D INPUT -p "$p" --dport "$port" -j DROP 2>/dev/null; do :; done
        while iptables -D DOCKER-USER -p "$p" --dport "$port" -j DROP 2>/dev/null; do :; done

        if [ "$limit" == "y" ]; then
            read -p "Введите IP-адреса (через пробел): " ips
            for ip in $ips; do
                if ! validate_ip "$ip"; then echo "[!] Пропускаем неверный IP: $ip"; continue; fi
                iptables -I INPUT 1 -p "$p" --dport "$port" -s "$ip" -j ACCEPT
                iptables -I DOCKER-USER 1 -p "$p" --dport "$port" -s "$ip" -j ACCEPT
            done
        else
            iptables -I INPUT 1 -p "$p" --dport "$port" -j ACCEPT
            iptables -I DOCKER-USER 1 -p "$p" --dport "$port" -j ACCEPT
        fi
    done
    echo "[+] Порт $port открыт."
    save_rules
}

new_close_port() {
    read -p "Введите номер порта: " port
    if ! validate_port "$port"; then echo "[!] Неверный номер порта."; press_enter; return; fi
    read -p "Протокол (tcp/udp/both): " proto
    proto=$(echo "$proto" | tr '[:upper:]' '[:lower:]')
    if ! validate_protocol "$proto"; then echo "[!] Неверный протокол."; press_enter; return; fi

    if [ "$proto" == "both" ]; then
        protos=("tcp" "udp")
    else
        protos=("$proto")
    fi

    for p in "${protos[@]}"; do
        # Удаляем ВСЕ существующие правила для этого порта из внешних цепочек
        # Удаляем правила ACCEPT
        while iptables -D INPUT -p "$p" --dport "$port" -j ACCEPT 2>/dev/null; do :; done
        while iptables -D DOCKER-USER -p "$p" --dport "$port" -j ACCEPT 2>/dev/null; do :; done
        # Удаляем правила DROP
        while iptables -D INPUT -p "$p" --dport "$port" -j DROP 2>/dev/null; do :; done
        while iptables -D DOCKER-USER -p "$p" --dport "$port" -j DROP 2>/dev/null; do :; done

        # Добавляем новое правило DROP в НАЧАЛО внешних цепочек
        iptables -I INPUT 1 -p "$p" --dport "$port" -j DROP
        iptables -I DOCKER-USER 1 -p "$p" --dport "$port" -j DROP
    done
    echo "[+] Порт $port закрыт."
    save_rules
}

new_add_ip_to_port() {
    read -p "Введите номер порта: " port
    if ! validate_port "$port"; then echo "[!] Неверный номер порта."; press_enter; return; fi
    read -p "Протокол (tcp/udp): " proto
    proto=$(echo "$proto" | tr '[:upper:]' '[:lower:]')
    if ! validate_protocol "$proto"; then echo "[!] Неверный протокол."; press_enter; return; fi
    read -p "Введите IP-адреса (через пробел): " ips

    for ip in $ips; do
        if ! validate_ip "$ip"; then echo "[!] Пропускаем неверный IP: $ip"; continue; fi
        # Удаляем DROP для этого IP/порта из внешних цепочек
        while iptables -D INPUT -p "$proto" --dport "$port" -s "$ip" -j DROP 2>/dev/null; do :; done
        while iptables -D DOCKER-USER -p "$proto" --dport "$port" -s "$ip" -j DROP 2>/dev/null; do :; done
        # Добавляем ACCEPT в начало внешних цепочек
        iptables -I INPUT 1 -p "$proto" --dport "$port" -s "$ip" -j ACCEPT
        iptables -I DOCKER-USER 1 -p "$proto" --dport "$port" -s "$ip" -j ACCEPT
        echo "[+] Добавлен доступ для $ip на порт $port/$proto"
    done
    save_rules
}

new_remove_ip_from_port() {
    read -p "Введите номер порта: " port
    if ! validate_port "$port"; then echo "[!] Неверный номер порта."; press_enter; return; fi
    read -p "Протокол (tcp/udp): " proto
    proto=$(echo "$proto" | tr '[:upper:]' '[:lower:]')
    if ! validate_protocol "$proto"; then echo "[!] Неверный протокол."; press_enter; return; fi
    read -p "Введите IP-адрес: " ip
    if ! validate_ip "$ip"; then echo "[!] Неверный IP-адрес."; press_enter; return; fi

    # Удаляем ACCEPT для этого IP/порта из внешних цепочек
    while iptables -D INPUT -p "$proto" --dport "$port" -s "$ip" -j ACCEPT 2>/dev/null; do :; done
    while iptables -D DOCKER-USER -p "$proto" --dport "$port" -s "$ip" -j ACCEPT 2>/dev/null; do :; done
    echo "[+] Убран доступ для $ip с порта $port/$proto"
    save_rules
}

require_root
ensure_dependencies
ensure_docker_user_chain

main_menu