1nf-multi-chaos

#!/bin/bash

set -euo pipefail

# Check for root privileges
if [[ $EUID -ne 0 ]]; then
	echo "Error: This script must be run as root" >&2
	exit 1
fi

show_usage() {
	echo "Usage:"
	echo "  $0 <table_name> <duration_seconds>"
	echo ""
	echo "Example:"
	echo "  $0 my-forward 30"
	exit 1
}

if [[ $# -ne 2 ]]; then
	show_usage
fi

TABLE_NAME="$1"
DURATION="$2"

# Validate duration is a positive integer
if ! [[ "$DURATION" =~ ^[0-9]+$ ]] || [[ "$DURATION" -lt 1 ]]; then
	echo "Error: Duration must be a positive integer (seconds)" >&2
	exit 1
fi

# Verify table exists
if ! nft list table ip "$TABLE_NAME" >/dev/null 2>&1; then
	echo "Error: Table '$TABLE_NAME' does not exist" >&2
	exit 1
fi

# Parse source IPs and port from existing table
NFT_OUTPUT=$(nft list table ip "$TABLE_NAME")
DST_IP=$(echo "$NFT_OUTPUT" | grep "masquerade" | sed 's/.*ip daddr \([0-9.]*\).*/\1/')

# Check if this is port-specific or all-ports forwarding
PORT=""
if echo "$NFT_OUTPUT" | grep -q "tcp dport"; then
	PORT=$(echo "$NFT_OUTPUT" | grep "tcp dport" | head -1 | sed 's/.*tcp dport \([0-9]*\).*/\1/')
	# Extract source IPs from tcp dport rules
	mapfile -t SOURCE_IPS < <(echo "$NFT_OUTPUT" | grep "tcp dport" | sed 's/.*ip daddr \([0-9.]*\).*/\1/' | sort -u)
else
	# All-ports forwarding: extract source IPs from ip daddr rules without port
	mapfile -t SOURCE_IPS < <(echo "$NFT_OUTPUT" | grep "dnat to" | sed 's/.*ip daddr \([0-9.]*\).*/\1/' | sort -u)
fi

if [[ -z "$DST_IP" ]]; then
	echo "Error: Could not parse destination IP from table structure" >&2
	exit 1
fi

if [[ ${#SOURCE_IPS[@]} -eq 0 ]]; then
	echo "Error: No source IPs found in table" >&2
	exit 1
fi

# Timestamp helper
log() {
	echo "$(date '+%Y-%m-%d %H:%M:%S') | $1"
}

# Cleanup function to restore original state
cleanup() {
	log "Interrupted - removing chaos filter chains..."
	# Delete the chaos filter chains if they exist
	nft delete chain ip "$TABLE_NAME" prerouting_chaos 2>/dev/null || true
	nft delete chain ip "$TABLE_NAME" output_chaos 2>/dev/null || true
	log "Chaos chains removed, NAT configuration preserved"
	exit 0
}

# Set trap for interrupt
trap cleanup INT TERM

log "Starting chaos injection for table: $TABLE_NAME"
if [[ -n "$PORT" ]]; then
	log "Configuration: port=$PORT, sources=${#SOURCE_IPS[@]} IPs, dst=$DST_IP, duration=${DURATION}s per IP"
else
	log "Configuration: all ports, sources=${#SOURCE_IPS[@]} IPs, dst=$DST_IP, duration=${DURATION}s per IP"
fi

# Apply chaos for a source IP
apply_chaos() {
	local src_ip=$1
	local current=$2
	local total=$3

	# Phase 1: DROP - create filter chains and add drop rules
	log "[$current/$total] $src_ip: Phase 1 - DROP (${DURATION}s)"

	# Delete any existing chaos chains first
	nft delete chain ip "$TABLE_NAME" prerouting_chaos 2>/dev/null || true
	nft delete chain ip "$TABLE_NAME" output_chaos 2>/dev/null || true

	# Create new filter chains with drop rules for the target IP
	if [[ -n "$PORT" ]]; then
		# Port-specific forwarding
		nft "
table ip $TABLE_NAME {
    chain prerouting_chaos {
        type filter hook prerouting priority -100; policy accept;
        tcp dport $PORT ip daddr $src_ip drop
    }
    chain output_chaos {
        type filter hook output priority -100; policy accept;
        tcp dport $PORT ip daddr $src_ip drop
    }
}
"
	else
		# All-ports forwarding
		nft "
table ip $TABLE_NAME {
    chain prerouting_chaos {
        type filter hook prerouting priority -100; policy accept;
        ip daddr $src_ip drop
    }
    chain output_chaos {
        type filter hook output priority -100; policy accept;
        ip daddr $src_ip drop
    }
}
"
	fi
	sleep "$DURATION"

	# Phase 2: UNREACHABLE - recreate chains with reject rules
	log "[$current/$total] $src_ip: Phase 2 - UNREACHABLE (${DURATION}s)"

	# Delete existing chaos chains
	nft delete chain ip "$TABLE_NAME" prerouting_chaos 2>/dev/null || true
	nft delete chain ip "$TABLE_NAME" output_chaos 2>/dev/null || true

	# Create new filter chains with reject rules
	if [[ -n "$PORT" ]]; then
		# Port-specific forwarding
		nft "
table ip $TABLE_NAME {
    chain prerouting_chaos {
        type filter hook prerouting priority -100; policy accept;
        tcp dport $PORT ip daddr $src_ip reject with icmp host-unreachable
    }
    chain output_chaos {
        type filter hook output priority -100; policy accept;
        tcp dport $PORT ip daddr $src_ip reject with icmp host-unreachable
    }
}
"
	else
		# All-ports forwarding
		nft "
table ip $TABLE_NAME {
    chain prerouting_chaos {
        type filter hook prerouting priority -100; policy accept;
        ip daddr $src_ip reject with icmp host-unreachable
    }
    chain output_chaos {
        type filter hook output priority -100; policy accept;
        ip daddr $src_ip reject with icmp host-unreachable
    }
}
"
	fi
	sleep "$DURATION"

	# Phase 3: RESTORE - remove chaos chains
	log "[$current/$total] $src_ip: RESTORED"
	nft delete chain ip "$TABLE_NAME" prerouting_chaos 2>/dev/null || true
	nft delete chain ip "$TABLE_NAME" output_chaos 2>/dev/null || true
}

# Apply chaos to all sources sequentially
TOTAL_IPS=${#SOURCE_IPS[@]}
for i in "${!SOURCE_IPS[@]}"; do
	apply_chaos "${SOURCE_IPS[$i]}" "$((i + 1))" "$TOTAL_IPS"
done

log "Chaos injection complete - processed $TOTAL_IPS source IP(s)"