#!/bin/bash
# Author: Florent Tainturier
# Contact: flotain41@gmail.com
# Version 2.0
# Script complet pour gérer la blacklist et la whitelist avec résolution de noms, pour IPv4 et IPv6

set -euo pipefail

# Variables
workplace=/root/log
blacklistipv4=$workplace/blacklistipv4.txt
blacklistipv6=$workplace/blacklistipv6.txt
whitelistfile=$workplace/whitelist.txt
logfile=$workplace/firewall.log

# Fonction de logging
function log() {
    local message="$1"
    local level="${2:-INFO}"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [$level] $message" >> "$logfile"
    echo "[$level] $message"
}

# Vérification des privilèges root
function check_root() {
    if [ "$(id -u)" != "0" ]; then
        echo "Ce script doit être exécuté en tant que root" 1>&2
        exit 1
    fi
}

# Fonctions de validation
function is_valid_ipv4() {
    local ip=$1
    local stat=1
    if [[ $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        OIFS=$IFS
        IFS='.'
        ip=($ip)
        IFS=$OIFS
        [[ ${ip[0]} -le 255 && ${ip[1]} -le 255 && ${ip[2]} -le 255 && ${ip[3]} -le 255 ]]
        stat=$?
    fi
    return $stat
}

function is_valid_ipv6() {
    local ip=$1
    local stat=1
    if [[ $ip =~ ^([0-9a-fA-F]{0,4}:){1,7}[0-9a-fA-F]{0,4}$ ]]; then
        stat=0
    fi
    return $stat
}

function is_valid_domain() {
    local domain=$1
    [[ $domain =~ ^([a-zA-Z0-9](([a-zA-Z0-9-]){0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}$ ]]
}

# Fonctions de résolution
function resolve_ip_to_name() {
    local ip="$1"
    local name=$(dig +short -x "$ip" 2>/dev/null | sed 's/\.$//')
    if [[ -z "$name" ]]; then
        echo "$ip"
    else
        echo "$ip ($name)"
    fi
}

function resolve_name_to_ip() {
    local name="$1"
    local ip=$(dig +short "$name" 2>/dev/null | head -n1)
    if [[ -z "$ip" ]]; then
        echo "$name"
    else
        echo "$ip"
    fi
}

# Fonction pour ajouter à la whitelist
function add_to_whitelist() {
    local target="$1"
    local resolved_entry

    if is_valid_ipv4 "$target" || is_valid_ipv6 "$target"; then
        resolved_entry=$(resolve_ip_to_name "$target")
    elif is_valid_domain "$target"; then
        local ip=$(resolve_name_to_ip "$target")
        if [[ -n "$ip" ]]; then
            resolved_entry="$target ($ip)"
        else
            resolved_entry="$target"
        fi
    else
        log "Adresse ou domaine invalide : $target" "ERROR"
        return
    fi

    log "Ajout de $resolved_entry à la whitelist."
    echo "$resolved_entry" >> "$whitelistfile"
}

# Fonction pour supprimer de la whitelist
function remove_from_whitelist() {
    local target="$1"
    if grep -q "^$target" "$whitelistfile"; then
        sed -i "/^$target/d" "$whitelistfile"
        log "Suppression de $target de la whitelist."
    else
        log "L'adresse $target n'est pas dans la whitelist." "WARNING"
    fi
}

# Fonction pour ajouter à la blacklist IPv4
function add_to_blacklist_ipv4() {
    local target="$1"
    local port="$2"
    local proto="$3"

    if grep -q "^$target" "$whitelistfile"; then
        log "L'adresse $target est dans la whitelist. Elle ne peut pas être blacklistée." "WARNING"
        return
    fi

    local resolved_entry
    if is_valid_ipv4 "$target"; then
        resolved_entry="$target"
    elif is_valid_domain "$target"; then
        resolved_entry=$(resolve_name_to_ip "$target")
        if ! is_valid_ipv4 "$resolved_entry"; then
            log "Le domaine $target ne résout pas en une adresse IPv4 valide." "ERROR"
            return
        fi
    else
        log "Adresse ou domaine invalide : $target" "ERROR"
        return
    fi

    log "Ajout de $target ($resolved_entry) à la blacklist IPv4."
    echo "$target ($resolved_entry)" >> "$blacklistipv4"

    if [[ -n "$port" && -n "$proto" ]]; then
        if iptables -A INPUT -s "$resolved_entry" -p "$proto" --dport "$port" -j DROP; then
            log "$target ($resolved_entry) a été ajouté à iptables pour le port $port et le protocole $proto."
        else
            log "Erreur lors de l'ajout de $target ($resolved_entry) à iptables." "ERROR"
        fi
    else
        if iptables -A INPUT -s "$resolved_entry" -j DROP; then
            log "$target ($resolved_entry) a été ajouté aux règles iptables."
        else
            log "Erreur lors de l'ajout de $target ($resolved_entry) à iptables." "ERROR"
        fi
    fi
}

# Fonction pour ajouter à la blacklist IPv6
function add_to_blacklist_ipv6() {
    local target="$1"
    local port="$2"
    local proto="$3"

    if grep -q "^$target" "$whitelistfile"; then
        log "L'adresse $target est dans la whitelist. Elle ne peut pas être blacklistée." "WARNING"
        return
    fi

    local resolved_entry
    if is_valid_ipv6 "$target"; then
        resolved_entry="$target"
    elif is_valid_domain "$target"; then
        resolved_entry=$(resolve_name_to_ip "$target")
        if ! is_valid_ipv6 "$resolved_entry"; then
            log "Le domaine $target ne résout pas en une adresse IPv6 valide." "ERROR"
            return
        fi
    else
        log "Adresse ou domaine invalide : $target" "ERROR"
        return
    fi

    log "Ajout de $target ($resolved_entry) à la blacklist IPv6."
    echo "$target ($resolved_entry)" >> "$blacklistipv6"

    if [[ -n "$port" && -n "$proto" ]]; then
        if ip6tables -A INPUT -s "$resolved_entry" -p "$proto" --dport "$port" -j DROP; then
            log "$target ($resolved_entry) a été ajouté à ip6tables pour le port $port et le protocole $proto."
        else
            log "Erreur lors de l'ajout de $target ($resolved_entry) à ip6tables." "ERROR"
        fi
    else
        if ip6tables -A INPUT -s "$resolved_entry" -j DROP; then
            log "$target ($resolved_entry) a été ajouté aux règles ip6tables."
        else
            log "Erreur lors de l'ajout de $target ($resolved_entry) à ip6tables." "ERROR"
        fi
    fi
}

# Fonction pour débloquer une adresse IPv4
function unblacklist_ipv4() {
    local target="$1"
    local ip_to_unblock

    if is_valid_ipv4 "$target"; then
        ip_to_unblock="$target"
    elif is_valid_domain "$target"; then
        ip_to_unblock=$(resolve_name_to_ip "$target")
        if ! is_valid_ipv4 "$ip_to_unblock"; then
            log "Le domaine $target ne résout pas en une adresse IPv4 valide." "ERROR"
            return
        fi
    else
        log "Adresse ou domaine invalide : $target" "ERROR"
        return
    fi

    if grep -q "$ip_to_unblock" "$blacklistipv4"; then
        sed -i "/$ip_to_unblock/d" "$blacklistipv4"
        
        # Supprimer toutes les règles iptables correspondant à cette adresse IP
        iptables-save | grep -v " $ip_to_unblock " | iptables-restore
        
        # Vérifier si des règles subsistent pour cette adresse IP
        if iptables -L INPUT -n | grep -q "$ip_to_unblock"; then
            # Si des règles subsistent, les supprimer manuellement
            iptables -L INPUT --line-numbers -n | grep "$ip_to_unblock" | awk '{print $1}' | sort -r | xargs -I {} iptables -D INPUT {}
        fi
        
        log "Suppression de $target ($ip_to_unblock) de la blacklist IPv4 et de toutes les règles iptables associées."
    else
        log "L'adresse $target ($ip_to_unblock) n'est pas dans la blacklist IPv4." "WARNING"
    fi
}

# Fonction pour débloquer une adresse IPv6
function unblacklist_ipv6() {
    local target="$1"
    local ip_to_unblock

    if is_valid_ipv6 "$target"; then
        ip_to_unblock="$target"
    elif is_valid_domain "$target"; then
        ip_to_unblock=$(resolve_name_to_ip "$target")
        if ! is_valid_ipv6 "$ip_to_unblock"; then
            log "Le domaine $target ne résout pas en une adresse IPv6 valide." "ERROR"
            return
        fi
    else
        log "Adresse ou domaine invalide : $target" "ERROR"
        return
    fi

    if grep -q "$ip_to_unblock" "$blacklistipv6"; then
        sed -i "/$ip_to_unblock/d" "$blacklistipv6"
        
        # Supprimer toutes les règles ip6tables correspondant à cette adresse IP
        ip6tables-save | grep -v " $ip_to_unblock " | ip6tables-restore
        
        # Vérifier si des règles subsistent pour cette adresse IP
        if ip6tables -L INPUT -n | grep -q "$ip_to_unblock"; then
            # Si des règles subsistent, les supprimer manuellement
            ip6tables -L INPUT --line-numbers -n | grep "$ip_to_unblock" | awk '{print $1}' | sort -r | xargs -I {} ip6tables -D INPUT {}
        fi
        
        log "Suppression de $target ($ip_to_unblock) de la blacklist IPv6 et de toutes les règles ip6tables associées."
    else
        log "L'adresse $target ($ip_to_unblock) n'est pas dans la blacklist IPv6." "WARNING"
    fi
}

# Fonctions d'affichage
function show_whitelist() {
    echo "Contenu de la whitelist :"
    cat "$whitelistfile" || echo "La whitelist est vide."
}

function show_blacklist_ipv4() {
    echo "Contenu de la blacklist IPv4 :"
    cat "$blacklistipv4" || echo "La blacklist IPv4 est vide."
}

function show_blacklist_ipv6() {
    echo "Contenu de la blacklist IPv6 :"
    cat "$blacklistipv6" || echo "La blacklist IPv6 est vide."
}

function show_iptables() {
    echo "Règles iptables :"
    iptables -L INPUT -n || echo "Erreur lors de la récupération des règles iptables."
}

function show_ip6tables() {
    echo "Règles ip6tables :"
    ip6tables -L INPUT -n || echo "Erreur lors de la récupération des règles ip6tables."
}

function show_all_lists() {
    show_whitelist
    echo ""
    show_blacklist_ipv4
    echo ""
    show_blacklist_ipv6
    echo ""
    show_iptables
    echo ""
    show_ip6tables
}

# Fonction pour sauvegarder les règles
function save_iptables_rules() {
    # Créer le répertoire s'il n'existe pas
    mkdir -p /etc/iptables

    # Sauvegarder les règles IPv4
    if iptables-save > /etc/iptables/rules.v4; then
        log "Règles iptables sauvegardées avec succès dans /etc/iptables/rules.v4."
    else
        log "Erreur lors de la sauvegarde des règles iptables." "ERROR"
    fi

    # Sauvegarder les règles IPv6
    if ip6tables-save > /etc/iptables/rules.v6; then
        log "Règles ip6tables sauvegardées avec succès dans /etc/iptables/rules.v6."
    else
        log "Erreur lors de la sauvegarde des règles ip6tables." "ERROR"
    fi

    # Définir les permissions appropriées
    chmod 600 /etc/iptables/rules.v4 /etc/iptables/rules.v6

    log "Les règles ont été sauvegardées. Pour les restaurer au démarrage, vous devrez configurer votre système pour exécuter :"
    log "iptables-restore < /etc/iptables/rules.v4"
    log "ip6tables-restore < /etc/iptables/rules.v6"
}

# Fonction principale
function main() {
    check_root

    case "$1" in
        --add-whitelist)
            target="$2"
            add_to_whitelist "$target"
            ;;
        --remove-whitelist)
            target="$2"
            remove_from_whitelist "$target"
            ;;
        --blacklist)
            target="$2"
            port="${3:-}"
            proto="${4:-}"
            if is_valid_ipv4 "$target" || (is_valid_domain "$target" && is_valid_ipv4 "$(resolve_name_to_ip "$target")"); then
                add_to_blacklist_ipv4 "$target" "$port" "$proto"
            elif is_valid_ipv6 "$target" || (is_valid_domain "$target" && is_valid_ipv6 "$(resolve_name_to_ip "$target")"); then
                add_to_blacklist_ipv6 "$target" "$port" "$proto"
            else
                log "Adresse ou domaine invalide : $target" "ERROR"
            fi
            ;;
        --unblacklist)
            target="$2"
            if is_valid_ipv4 "$target" || (is_valid_domain "$target" && is_valid_ipv4 "$(resolve_name_to_ip "$target")"); then
                unblacklist_ipv4 "$target"
            elif is_valid_ipv6 "$target" || (is_valid_domain "$target" && is_valid_ipv6 "$(resolve_name_to_ip "$target")"); then
                unblacklist_ipv6 "$target"
            else
                log "Adresse ou domaine invalide : $target" "ERROR"
            fi
            ;;
        --show)
            show_all_lists
            ;;
        --save-rules)
            save_iptables_rules
            ;;
        *)
            echo "Usage: $0 --add-whitelist <IP/Domaine> | --remove-whitelist <IP/Domaine> | --blacklist <IP/Domaine> [port] [proto] | --unblacklist <IP/Domaine> | --show | --save-rules"
            exit 1
            ;;
    esac
}

# Exécution du script
main "$@"