#!/bin/bash

# =============================================================================
# OVS Conntrack/NAT ICMP Test Script (CORRECTED VERSION)
# =============================================================================

# -- Configuration --
BRIDGE="br0"
CLIENT_NS="client"
SERVER_NS="server"
CLIENT_VETH="veth-client"
CLIENT_OVS_PORT="ovs-client"
CLIENT_IP="192.168.20.10/24"
CLIENT_MAC="00:00:00:00:20:10"
SERVER_VETH="veth-server"
SERVER_OVS_PORT="ovs-server"
SERVER_IP="192.168.10.20/24"
SERVER_MAC="00:00:00:00:10:20"

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

# =============================================================================
# Helper Functions
# =============================================================================

check_deps() {
    for cmd in ip ovs-vsctl ovs-ofctl tcpdump ovs-appctl; do
        if ! command -v $cmd &> /dev/null; then
            echo -e "${RED}ERROR: Command '$cmd' not found. Please ensure Open vSwitch and tcpdump are installed.${NC}"
            exit 1
        fi
    done
}

cleanup() {
    echo -e "${YELLOW}--- Cleaning up environment... ---${NC}"
    ip netns del $CLIENT_NS &> /dev/null
    ip netns del $SERVER_NS &> /dev/null
    ovs-vsctl --if-exists del-br $BRIDGE
    rm -f server.pcap client.pcap tcpdump_server.log tcpdump_client.log
    echo -e "${YELLOW}--- Cleanup complete ---${NC}"
}

setup() {
    echo -e "${YELLOW}--- 1. Setting up test environment ---${NC}"
    ovs-vsctl --if-exists del-br $BRIDGE
    ovs-vsctl add-br $BRIDGE
    ovs-vsctl set bridge $BRIDGE protocols=OpenFlow15

    ip netns add $CLIENT_NS
    ip netns add $SERVER_NS

    ip link add $CLIENT_VETH type veth peer name $CLIENT_OVS_PORT
    ip link add $SERVER_VETH type veth peer name $SERVER_OVS_PORT

    ip link set $CLIENT_VETH netns $CLIENT_NS
    ip link set $SERVER_VETH netns $SERVER_NS

    ovs-vsctl add-port $BRIDGE $CLIENT_OVS_PORT
    ovs-vsctl add-port $BRIDGE $SERVER_OVS_PORT

    ip netns exec $CLIENT_NS ip addr add $CLIENT_IP dev $CLIENT_VETH
    ip netns exec $CLIENT_NS ip link set dev $CLIENT_VETH address $CLIENT_MAC
    ip netns exec $CLIENT_NS ip link set dev $CLIENT_VETH up
    ip netns exec $CLIENT_NS ip link set dev lo up

    ip netns exec $SERVER_NS ip addr add $SERVER_IP dev $SERVER_VETH
    ip netns exec $SERVER_NS ip link set dev $SERVER_VETH address $SERVER_MAC
    ip netns exec $SERVER_NS ip link set dev $SERVER_VETH up
    ip netns exec $SERVER_NS ip link set dev lo up

    ip link set $CLIENT_OVS_PORT up
    ip link set $SERVER_OVS_PORT up

    echo "Environment setup complete."
}

load_flows() {
    echo -e "${YELLOW}--- 2. Loading OpenFlow rules ---${NC}"
    ovs-ofctl -O OpenFlow15 del-flows $BRIDGE
    ovs-ofctl -O OpenFlow15 add-flows $BRIDGE - <<'EOF'
table=0,ip,actions=ct(table=1,zone=42,nat)
table=1,in_port=ovs-client,ip,ct_state=+trk+new,actions=ct(commit,table=2,zone=42,nat(dst=192.168.10.20))
table=1,icmp,ct_state=+trk+rel-rpl,actions=ct(commit,table=2,zone=42,nat)
table=1,ip,actions=resubmit(,2)
table=2,in_port=ovs-client,ip,ct_state=+trk+new,actions=output:ovs-server
table=2,in_port=ovs-client,icmp,ct_state=+trk+rel,actions=output:ovs-server
table=2,in_port=ovs-server,icmp,ct_state=+trk+rel,actions=output:ovs-client
table=2,in_port=ovs-server,ip,ct_state=+trk+rpl,actions=output:ovs-client
EOF
    echo "Flow rules loaded."
    ovs-ofctl -O OpenFlow15 dump-flows $BRIDGE
}

# =============================================================================
# Main Logic
# =============================================================================

check_deps
trap cleanup EXIT
setup
load_flows

# --- Test Part 1: ICMP Error -> Server ---
echo -e "\n${YELLOW}--- 3. Starting Test Part 1: ICMP Error -> Server ---${NC}"

echo "Starting tcpdump in 'server' namespace..."
ip netns exec $SERVER_NS tcpdump -n -l -U -i $SERVER_VETH -w server.pcap 2>tcpdump_server.log &
TCPDUMP_SERVER_PID=$!
sleep 2

echo "Injecting UDP (Client -> Server) to establish connection..."
# FIX: Use hardcoded hex packet data directly to ensure correct format and avoid variable expansion issues.
ovs-ofctl -O OpenFlow15 packet-out $BRIDGE "in_port=$CLIENT_OVS_PORT,packet=00000000102000000000201008004500001C000040000A11C762C0A8140AC0A814140001000200080000,actions=resubmit(,0)"

echo "Injecting UDP (Server -> Client) reply..."
# FIX: Same as above
ovs-ofctl -O OpenFlow15 packet-out $BRIDGE "in_port=$SERVER_OVS_PORT,packet=00000000201000000000102008004500001C000020000A11D162C0A80A14C0A8140A0002000100080000,actions=resubmit(,0)"

echo "Injecting forged ICMP error (Router -> Server)..."
# FIX: Same as above
ovs-ofctl -O OpenFlow15 packet-out $BRIDGE "in_port=$CLIENT_OVS_PORT,packet=000000001020000000002000080045000038011F0000FF011140C0A81401C0A814140304f778000005784500001C000020000A11e762C0A81414C0A8140A0002000100080000,actions=resubmit(,0)"

sleep 1
kill $TCPDUMP_SERVER_PID
wait $TCPDUMP_SERVER_PID 2>/dev/null

echo "Verifying results..."
CT_DUMP=$(ovs-appctl dpctl/dump-conntrack | grep "zone=42")
if echo "$CT_DUMP" | grep -q "udp,orig=(src=192.168.20.10,dst=192.168.20.20,sport=1,dport=2),reply=(src=192.168.10.20,dst=192.168.20.10,sport=2,dport=1)"; then
    echo -e "  [${GREEN}PASS${NC}] Conntrack entry created correctly."
else
    echo -e "  [${RED}FAIL${NC}] Conntrack entry incorrect or not found."
    echo "Current Conntrack Table:"
    echo "$CT_DUMP"
fi

PCAP_DUMP=$(ip netns exec $SERVER_NS tcpdump -n -r server.pcap icmp 2>/dev/null)
if echo "$PCAP_DUMP" | grep -q "192.168.20.1 > 192.168.10.20: ICMP"; then
    echo -e "  [${GREEN}PASS${NC}] Server received the correct ICMP error packet (DNAT successful)."
else
    echo -e "  [${RED}FAIL${NC}] Server did not receive the expected ICMP error packet."
    echo "ICMP traffic captured on server:"
    echo "$PCAP_DUMP"
fi

# --- Test Part 2: ICMP Error -> Client ---
echo -e "\n${YELLOW}--- 4. Starting Test Part 2: ICMP Error -> Client ---${NC}"

echo "Flushing Conntrack table..."
ovs-appctl dpctl/flush-conntrack zone=42

echo "Starting tcpdump in 'client' namespace..."
ip netns exec $CLIENT_NS tcpdump -n -l -U -i $CLIENT_VETH -w client.pcap 2>tcpdump_client.log &
TCPDUMP_CLIENT_PID=$!
sleep 2

echo "Injecting UDP (Client -> Server) to establish new connection..."
# FIX: Same as above
ovs-ofctl -O OpenFlow15 packet-out $BRIDGE "in_port=$CLIENT_OVS_PORT,packet=00000000102000000000201008004500001C000040000A11C762C0A8140AC0A814140001000200080000,actions=resubmit(,0)"

echo "Injecting forged ICMP error (Router -> Client)..."
# FIX: Same as above
ovs-ofctl -O OpenFlow15 packet-out $BRIDGE "in_port=$SERVER_OVS_PORT,packet=000000002010000000002000080045000038011F0000FF01114AC0A81401C0A8140A0304F778000005784500001C000040000A11D162C0A8140AC0A80A140001000200080000,actions=resubmit(,0)"

sleep 1
kill $TCPDUMP_CLIENT_PID
wait $TCPDUMP_CLIENT_PID 2>/dev/null

echo "Verifying results..."
CT_DUMP2=$(ovs-appctl dpctl/dump-conntrack | grep "zone=42")
if echo "$CT_DUMP2" | grep -q "udp,orig=(src=192.168.20.10,dst=192.168.20.20,sport=1,dport=2),reply=(src=192.168.10.20,dst=192.168.20.10,sport=2,dport=1)"; then
    echo -e "  [${GREEN}PASS${NC}] Conntrack entry created correctly."
else
    echo -e "  [${RED}FAIL${NC}] Conntrack entry incorrect or not found."
    echo "Current Conntrack Table:"
    echo "$CT_DUMP2"
fi

PCAP_DUMP2=$(ip netns exec $CLIENT_NS tcpdump -n -r client.pcap icmp 2>/dev/null)
if echo "$PCAP_DUMP2" | grep -q "192.168.20.20 > 192.168.20.10: ICMP"; then
    echo -e "  [${GREEN}PASS${NC}] Client received the correct ICMP error packet (Reverse NAT successful)."
else
    echo -e "  [${RED}FAIL${NC}] Client did not receive the expected ICMP error packet."
    echo "ICMP traffic captured on client:"
    echo "$PCAP_DUMP2"
fi

echo -e "\n${YELLOW}--- Test complete ---${NC}"