#!/bin/bash

if [[ ${DEBUG} -gt 0 ]]; then set -x; fi

set -e

# begin setup_prerouting_external
setup_prerouting_external() {

    echo "Setting load balancing rules (external traffic)"

    echo "Rule 1: ~33% -> POD1A_IP"
    sudo iptables -t nat -A PREROUTING -p tcp -d $NODE1_IP --dport $NODE_PORT \
        -m statistic --mode random --probability 0.33 \
        -j DNAT --to-destination $POD1A_IP:$POD_PORT \
        -m comment --comment "NODEPORT_PREROUTING_EXTERNAL_POD1A"

    echo "Rule 2: ~33% -> POD1B_IP"
    sudo iptables -t nat -A PREROUTING -p tcp -d $NODE1_IP --dport $NODE_PORT \
        -m statistic --mode random --probability 0.5 \
        -j DNAT --to-destination $POD1B_IP:$POD_PORT \
        -m comment --comment "NODEPORT_PREROUTING_EXTERNAL_POD1B"

    echo "Rule 3: Remaining -> POD2_IP"
    sudo iptables -t nat -A PREROUTING -p tcp -d $NODE1_IP --dport $NODE_PORT \
        -j DNAT --to-destination $POD2_IP:$POD_PORT \
        -m comment --comment "NODEPORT_PREROUTING_EXTERNAL_POD2"

    echo "Completed."
}
# end setup_prerouting_external

# begin setup_postrouting_external
setup_postrouting_external() {

    echo "Setting postrouting rules (external traffic)."

    BRIDGE_CIDR=$(ip a show br0 | grep 'inet ' | awk '{print $2}')
    
    echo "Rule 4. Masquerade for local node"
    sudo iptables -t nat -A POSTROUTING -s $BRIDGE_CIDR ! -d $BRIDGE_CIDR \
        -j MASQUERADE -m comment --comment "NODEPORT_POSTROUTING_MASQ_EXTERNAL_LOCAL"

    echo "Rule 5. Masquerade for remote nodes"
    sudo iptables -t nat -A POSTROUTING -p tcp -d $POD2_IP --dport $POD_PORT \
        -j MASQUERADE -m comment --comment "NODEPORT_POSTROUTING_MASQ_EXTERNAL_REMOTE"

    echo "Completed."
}
# end setup_postrouting_external

# begin setup_static_routing
setup_static_routing() {

    echo "Setting up static routes."
    
    echo "Adding static route for $POD2_IP via $NODE2_IP"
    sudo ip route add $POD2_IP via $NODE2_IP

    echo "Completed."
}
# end setup_static_routing

# begin setup_output_internal
setup_output_internal() {

    echo "Setting load balancing rules (internal traffic)."

    echo "Rule 6: ~33% -> POD1A_IP"
    sudo iptables -t nat -A OUTPUT -p tcp -d 127.0.0.1 --dport $NODE_PORT \
        -m statistic --mode random --probability 0.33 \
        -j DNAT --to-destination $POD1A_IP:$POD_PORT \
        -m comment --comment "NODEPORT_OUTPUT_INTERNAL_POD1A"

    echo "Rule 7: ~33% -> POD1B_IP"
    sudo iptables -t nat -A OUTPUT -p tcp -d 127.0.0.1 --dport $NODE_PORT \
        -m statistic --mode random --probability 0.5 \
        -j DNAT --to-destination $POD1B_IP:$POD_PORT \
        -m comment --comment "NODEPORT_OUTPUT_INTERNAL_POD1B"

    echo "Rule 8: Remainder -> POD2_IP"
    sudo iptables -t nat -A OUTPUT -p tcp -d 127.0.0.1 --dport $NODE_PORT \
        -j DNAT --to-destination $POD2_IP:$POD_PORT \
        -m comment --comment "NODEPORT_OUTPUT_INTERNAL_POD2"

    echo "Completed."
}
# end setup_output_internal

# begin setup_postrouting_internal
setup_postrouting_internal() {

    echo "Setting postrouting (internal)."

    echo "Setting route_localnet for bridge $BRIDGE_NAME"
    sudo sysctl -w net.ipv4.conf.$BRIDGE_NAME.route_localnet=1
    
    BRIDGE_CIDR=$(ip a show br0 | grep 'inet ' | awk '{print $2}')

    echo "Rule 9. Masquerade for local node"
    sudo iptables -t nat -A POSTROUTING -d $BRIDGE_CIDR -j MASQUERADE \
        -m comment --comment "NODEPORT_POSTROUTING_MASQ_INTERNAL_LOCAL"

    echo "Rule 10. Masquerade for remote nodes"
    sudo iptables -t nat -A POSTROUTING -p tcp -s 127.0.0.1 \
        -d $POD2_SUBNET --dport $POD_PORT -j SNAT --to-source $NODE1_IP \
        -m comment --comment "NODEPORT_POSTROUTING_SNAT_INTERNAL"

    echo "Completed."
}
# end setup_postrouting_internal

cleanup() {

    echo "Teardown of external traffic"

    BRIDGE_NAME="br0"
    BRIDGE_CIDR=$(ip a show $BRIDGE_NAME 2>/dev/null | grep 'inet ' | awk '{print $2}')

    echo "Rule 1: ~33% -> POD1A_IP"
    sudo iptables -t nat -D PREROUTING -p tcp -d $NODE1_IP --dport $NODE_PORT \
        -m statistic --mode random --probability 0.33 \
        -j DNAT --to-destination $POD1A_IP:$POD_PORT \
        -m comment --comment "NODEPORT_PREROUTING_EXTERNAL_POD1A"

    echo "Rule 2: ~33% -> POD1B_IP"
    sudo iptables -t nat -D PREROUTING -p tcp -d $NODE1_IP --dport $NODE_PORT \
        -m statistic --mode random --probability 0.5 \
        -j DNAT --to-destination $POD1B_IP:$POD_PORT \
        -m comment --comment "NODEPORT_PREROUTING_EXTERNAL_POD1B"

    echo "Rule 3: Remaining -> POD2_IP"
    sudo iptables -t nat -D PREROUTING -p tcp -d $NODE1_IP --dport $NODE_PORT \
        -j DNAT --to-destination $POD2_IP:$POD_PORT \
        -m comment --comment "NODEPORT_PREROUTING_EXTERNAL_POD2"

    echo "Rule 4. Masquerade for local node"
    sudo iptables -t nat -D POSTROUTING -s $BRIDGE_CIDR ! -d $BRIDGE_CIDR -j MASQUERADE \
        -m comment --comment "NODEPORT_POSTROUTING_MASQ_EXTERNAL_LOCAL"

    echo "Rule 5. Masquerade for remote nodes"
    sudo iptables -t nat -D POSTROUTING -p tcp -d $POD2_IP --dport $POD_PORT -j MASQUERADE \
        -m comment --comment "NODEPORT_POSTROUTING_MASQ_EXTERNAL_REMOTE"

    echo "Teardown of internal traffic"

    echo "Rule 6: ~33% -> POD1A_IP"
    sudo iptables -t nat -D OUTPUT -p tcp -d 127.0.0.1 --dport $NODE_PORT \
        -m statistic --mode random --probability 0.33 \
        -j DNAT --to-destination $POD1A_IP:$POD_PORT \
        -m comment --comment "NODEPORT_OUTPUT_INTERNAL_POD1A"

    echo "Rule 7: ~33% -> POD1B_IP"
    sudo iptables -t nat -D OUTPUT -p tcp -d 127.0.0.1 --dport $NODE_PORT \
        -m statistic --mode random --probability 0.5 \
        -j DNAT --to-destination $POD1B_IP:$POD_PORT \
        -m comment --comment "NODEPORT_OUTPUT_INTERNAL_POD1B"

    echo "Rule 8: Remainder -> POD2_IP"
    sudo iptables -t nat -D OUTPUT -p tcp -d 127.0.0.1 --dport $NODE_PORT \
        -j DNAT --to-destination $POD2_IP:$POD_PORT \
        -m comment --comment "NODEPORT_OUTPUT_INTERNAL_POD2"

    echo "Rule 9. Masquerade for local node"
    sudo iptables -t nat -D POSTROUTING -d $BRIDGE_CIDR -j MASQUERADE \
        -m comment --comment "NODEPORT_POSTROUTING_MASQ_INTERNAL_LOCAL"

    echo "Rule 10. Masquerade for remote nodes"
    sudo iptables -t nat -D POSTROUTING -p tcp -s 127.0.0.1 \
        -d $POD2_SUBNET --dport $POD_PORT -j SNAT --to-source $NODE1_IP \
        -m comment --comment "NODEPORT_POSTROUTING_SNAT_INTERNAL"

    echo "Deleting static route for $POD2_IP via $NODE2_IP"
    sudo ip route delete $POD2_IP via $NODE2_IP

    echo "Complete."
}

status() {

    echo "iptables NAT rules:"
    sudo iptables -t nat -vL
}

usage() {
    echo "Usage: $0 <command>"
    echo "Commands:"
    echo "  static_routing       Sets up static routing"
    echo "  prerouting_external  Sets up PREROUTING NAT"
    echo "  postrouting_external Sets up POSTROUTING for external traffic"
    echo "  output_internal      Sets up OUTPUT NAT rules for internal traffic"
    echo "  postrouting_internal Sets up POSTROUTING for internal traffic"
    echo "  setup                Sets up all NAT/routing rules"
    echo "  status               Shows the status of the iptables NAT rules"
    echo "  cleanup              Removes the NodePort configuration"
}

if [[ $# -eq 0 ]]; then
    usage
fi

case "$1" in
    static_routing)
        setup_static_routing
        ;;
    prerouting_external)
        setup_prerouting_external
        ;;
    postrouting_external)
        setup_postrouting_external
        ;;               
    output_internal)
        setup_output_internal
        ;;                 
    postrouting_internal)
        setup_postrouting_internal
        ;;   
    setup)
        setup_static_routing
        setup_prerouting_external
        setup_postrouting_external
        setup_output_internal
        setup_postrouting_internal
        ;;
    status)
        status
        ;;        
    cleanup)
        set +e
        cleanup
        ;;
    *)
        echo "Invalid command: $1"
        usage
        exit 1
        ;;
esac