#!/usr/bin/env bash
# network_isolation.sh - Configure network isolation for malware analysis VMs.
#
# Usage:
#   sudo bash network_isolation.sh --mode isolated [--interface eth0] [--inetsim]
#   sudo bash network_isolation.sh --mode capture  [--interface eth0] [--output /tmp/capture.pcap]
#   sudo bash network_isolation.sh --mode normal   [--interface eth0]
#   sudo bash network_isolation.sh --status
#
# Modes:
#   isolated  - Block all outbound traffic, enable DNS sinkhole, optionally start INetSim
#   capture   - Allow traffic but capture all packets
#   normal    - Restore default networking rules
#
# Options:
#   --interface IFACE    Network interface (default: auto-detect primary)
#   --inetsim            Start INetSim fake services in isolated mode
#   --dns-sinkhole IP    Redirect all DNS to this IP (default: 127.0.0.1)
#   --output FILE        PCAP output file for capture mode
#   --status             Show current isolation status and exit

set -euo pipefail

# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------
MODE=""
INTERFACE=""
ENABLE_INETSIM=false
DNS_SINKHOLE_IP="127.0.0.1"
CAPTURE_OUTPUT="/tmp/malware_capture_$(date +%Y%m%d_%H%M%S).pcap"
IPTABLES_BACKUP="/tmp/iptables_backup_analysis.rules"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
log_info()  { echo -e "${GREEN}[+]${NC} $*"; }
log_warn()  { echo -e "${YELLOW}[!]${NC} $*"; }
log_error() { echo -e "${RED}[-]${NC} $*"; }
die()       { log_error "$@"; exit 1; }

check_root() {
    [[ $EUID -eq 0 ]] || die "This script must be run as root."
}

detect_interface() {
    if [[ -z "$INTERFACE" ]]; then
        INTERFACE=$(ip route show default 2>/dev/null | awk '/default/ {print $5}' | head -1)
        if [[ -z "$INTERFACE" ]]; then
            INTERFACE="eth0"
        fi
    fi
    log_info "Using interface: $INTERFACE"
}

command_exists() { command -v "$1" &>/dev/null; }

# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --mode)
                MODE="$2"; shift 2 ;;
            --interface)
                INTERFACE="$2"; shift 2 ;;
            --inetsim)
                ENABLE_INETSIM=true; shift ;;
            --dns-sinkhole)
                DNS_SINKHOLE_IP="$2"; shift 2 ;;
            --output)
                CAPTURE_OUTPUT="$2"; shift 2 ;;
            --status)
                show_status; exit 0 ;;
            -h|--help)
                head -18 "$0" | tail -16; exit 0 ;;
            *)
                die "Unknown option: $1" ;;
        esac
    done

    if [[ -z "$MODE" ]]; then
        die "Mode is required. Use --mode isolated|capture|normal"
    fi
}

# ---------------------------------------------------------------------------
# Status
# ---------------------------------------------------------------------------
show_status() {
    echo "=== Network Isolation Status ==="
    echo ""

    echo "--- iptables rules (filter) ---"
    iptables -L -n --line-numbers 2>/dev/null || echo "  Unable to read iptables"
    echo ""

    echo "--- NAT rules ---"
    iptables -t nat -L -n 2>/dev/null || echo "  Unable to read NAT rules"
    echo ""

    echo "--- INetSim ---"
    if pgrep -x inetsim &>/dev/null; then
        echo "  INetSim is RUNNING (PID: $(pgrep -x inetsim))"
    else
        echo "  INetSim is NOT running"
    fi
    echo ""

    echo "--- DNS sinkhole ---"
    if iptables -t nat -L PREROUTING -n 2>/dev/null | grep -q "udp dpt:53"; then
        echo "  DNS sinkhole is ACTIVE"
    else
        echo "  DNS sinkhole is NOT active"
    fi
    echo ""

    echo "--- Active captures ---"
    if pgrep -x tcpdump &>/dev/null; then
        echo "  tcpdump is RUNNING (PID: $(pgrep -x tcpdump))"
    else
        echo "  No active packet capture"
    fi
}

# ---------------------------------------------------------------------------
# Isolation mode
# ---------------------------------------------------------------------------
apply_isolation() {
    log_info "Applying network isolation..."

    # Backup current rules
    iptables-save > "$IPTABLES_BACKUP"
    log_info "iptables rules backed up to $IPTABLES_BACKUP"

    # Flush existing rules
    iptables -F
    iptables -X
    iptables -t nat -F
    iptables -t nat -X

    # Default policy: drop all
    iptables -P INPUT DROP
    iptables -P FORWARD DROP
    iptables -P OUTPUT DROP

    # Allow loopback
    iptables -A INPUT -i lo -j ACCEPT
    iptables -A OUTPUT -o lo -j ACCEPT

    # Allow established connections (for management SSH if needed)
    iptables -A INPUT -m state --state ESTABLISHED,RELATED -j ACCEPT
    iptables -A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT

    # Allow DHCP (to keep VM network interface up)
    iptables -A OUTPUT -p udp --dport 67:68 -j ACCEPT
    iptables -A INPUT -p udp --sport 67:68 -j ACCEPT

    # DNS sinkhole: redirect all DNS queries to sinkhole IP
    iptables -t nat -A PREROUTING -p udp --dport 53 -j DNAT --to-destination "${DNS_SINKHOLE_IP}:53"
    iptables -t nat -A PREROUTING -p tcp --dport 53 -j DNAT --to-destination "${DNS_SINKHOLE_IP}:53"
    # Allow local DNS so sinkhole works
    iptables -A OUTPUT -p udp --dport 53 -d "$DNS_SINKHOLE_IP" -j ACCEPT
    iptables -A OUTPUT -p tcp --dport 53 -d "$DNS_SINKHOLE_IP" -j ACCEPT

    # Log dropped packets (limited rate)
    iptables -A INPUT -m limit --limit 5/min -j LOG --log-prefix "ISOLATION_IN_DROP: "
    iptables -A OUTPUT -m limit --limit 5/min -j LOG --log-prefix "ISOLATION_OUT_DROP: "

    log_info "Network isolation applied - all outbound traffic blocked"

    # Start INetSim if requested
    if $ENABLE_INETSIM; then
        start_inetsim
    fi

    # Set up DNS sinkhole with dnsmasq
    setup_dns_sinkhole
}

setup_dns_sinkhole() {
    if ! command_exists dnsmasq; then
        log_warn "dnsmasq not installed. Installing..."
        apt-get install -y -qq dnsmasq 2>/dev/null || {
            log_warn "Could not install dnsmasq. DNS sinkhole will use iptables NAT only."
            return
        }
    fi

    # Configure dnsmasq as a sinkhole
    local conf="/etc/dnsmasq.d/sinkhole.conf"
    cat > "$conf" <<DNSEOF
# Malware analysis DNS sinkhole
address=/#/${DNS_SINKHOLE_IP}
listen-address=127.0.0.1
bind-interfaces
no-resolv
log-queries
log-facility=/var/log/dns_sinkhole.log
DNSEOF

    systemctl restart dnsmasq 2>/dev/null || service dnsmasq restart 2>/dev/null || {
        log_warn "Could not start dnsmasq"
        return
    }
    log_info "DNS sinkhole active - all queries resolve to $DNS_SINKHOLE_IP"
    log_info "DNS query log: /var/log/dns_sinkhole.log"
}

start_inetsim() {
    if ! command_exists inetsim; then
        log_warn "INetSim not installed. Install with: apt-get install inetsim"
        return
    fi

    if pgrep -x inetsim &>/dev/null; then
        log_warn "INetSim is already running"
        return
    fi

    log_info "Starting INetSim..."
    inetsim --data-dir /var/lib/inetsim --log-dir /var/log/inetsim &
    sleep 2
    if pgrep -x inetsim &>/dev/null; then
        log_info "INetSim started (PID: $(pgrep -x inetsim))"
    else
        log_warn "INetSim failed to start"
    fi
}

# ---------------------------------------------------------------------------
# Capture mode
# ---------------------------------------------------------------------------
apply_capture() {
    log_info "Applying capture mode on $INTERFACE..."

    if ! command_exists tcpdump; then
        die "tcpdump is required for capture mode. Install with: apt-get install tcpdump"
    fi

    # Kill any existing capture
    pkill -x tcpdump 2>/dev/null || true

    # Start packet capture in background
    tcpdump -i "$INTERFACE" -w "$CAPTURE_OUTPUT" -U &
    local pid=$!
    log_info "Packet capture started (PID: $pid)"
    log_info "Output: $CAPTURE_OUTPUT"
    log_info "Stop capture with: kill $pid"
}

# ---------------------------------------------------------------------------
# Normal mode (restore)
# ---------------------------------------------------------------------------
apply_normal() {
    log_info "Restoring normal networking..."

    # Stop INetSim
    if pgrep -x inetsim &>/dev/null; then
        log_info "Stopping INetSim..."
        pkill -x inetsim 2>/dev/null || true
    fi

    # Stop tcpdump
    if pgrep -x tcpdump &>/dev/null; then
        log_info "Stopping tcpdump..."
        pkill -x tcpdump 2>/dev/null || true
    fi

    # Remove DNS sinkhole config
    rm -f /etc/dnsmasq.d/sinkhole.conf 2>/dev/null
    systemctl restart dnsmasq 2>/dev/null || service dnsmasq restart 2>/dev/null || true

    # Restore iptables
    if [[ -f "$IPTABLES_BACKUP" ]]; then
        iptables-restore < "$IPTABLES_BACKUP"
        log_info "iptables rules restored from backup"
    else
        # Reset to permissive defaults
        iptables -F
        iptables -X
        iptables -t nat -F
        iptables -t nat -X
        iptables -P INPUT ACCEPT
        iptables -P FORWARD ACCEPT
        iptables -P OUTPUT ACCEPT
        log_info "iptables reset to default ACCEPT policy"
    fi

    log_info "Normal networking restored"
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
main() {
    parse_args "$@"
    check_root
    detect_interface

    case "$MODE" in
        isolated) apply_isolation ;;
        capture)  apply_capture ;;
        normal)   apply_normal ;;
        *)        die "Invalid mode: $MODE. Use isolated|capture|normal" ;;
    esac
}

main "$@"