#!/bin/bash

# install or uninstall options to the script.
OPTION=$1

FIREWALLCMD="firewall-cmd"
FIREWALLCMD_DIRECT="${FIREWALLCMD} --direct"
FIREWALLCMD_QUIET="${FIREWALLCMD} --quiet"
FIREWALLCMD_QUIET_DIRECT="${FIREWALLCMD_QUIET} --direct"
IPTABLES="iptables -w"

aembit_agent_proxy_user_name="${AEMBIT_AGENT_PROXY_USER_NAME:-aembit_agent_proxy}"
aembit_group_id=$(id --group "${aembit_agent_proxy_user_name}")
agent_proxy_port=38080
container_cidr=$AEMBIT_DOCKER_CONTAINER_CIDR
# https://learn.microsoft.com/en-us/azure/virtual-network/what-is-ip-address-168-63-129-16
azure_platform_resources_ip="168.63.129.16"
# IP of the cloud-provider metadata services.
imds_ip="169.254.169.254"

# Rule Chains
aembit_raw_PREROUTING=aembit_raw_PREROUTING
aembit_raw_OUTPUT=aembit_raw_OUTPUT
aembit_nat_OUTPUT=aembit_nat_OUTPUT
aembit_nat_PREROUTING=aembit_nat_PREROUTING
aembit_filter_INPUT=aembit_filter_INPUT

agent_proxy_notification_endpoint="localhost:${AEMBIT_SERVICE_PORT}/iptables/loaded"

# Create all aembit rule chains.
create_all_aembit_chains() {
    add_ipv4_chain raw "${aembit_raw_PREROUTING}"
    add_ipv4_chain raw "${aembit_raw_OUTPUT}"
    add_ipv4_chain nat "${aembit_nat_PREROUTING}"
    add_ipv4_chain nat "${aembit_nat_OUTPUT}"
    add_ipv4_chain filter "${aembit_filter_INPUT}"
}

# Link all aembit rule chains to the corresponding tables
link_all_aembit_chains() {
    add_ipv4_rule raw PREROUTING -j "${aembit_raw_PREROUTING}"
    add_ipv4_rule raw OUTPUT -j "${aembit_raw_OUTPUT}"
    add_ipv4_rule nat PREROUTING -j "${aembit_nat_PREROUTING}"
    add_ipv4_rule nat OUTPUT -j "${aembit_nat_OUTPUT}"
    add_ipv4_rule filter INPUT -j "${aembit_filter_INPUT}"
}

# Remove the rule to link aembit rule chains and delete them.
unlink_and_delete_aembit_chains() {
    remove_ipv4_rule raw PREROUTING -j "${aembit_raw_PREROUTING}"
    remove_ipv4_chain raw "${aembit_raw_PREROUTING}"

    remove_ipv4_rule raw OUTPUT -j "${aembit_raw_OUTPUT}"
    remove_ipv4_chain raw "${aembit_raw_OUTPUT}"

    remove_ipv4_rule nat PREROUTING -j "${aembit_nat_PREROUTING}"
    remove_ipv4_chain nat "${aembit_nat_PREROUTING}"

    remove_ipv4_rule nat OUTPUT -j "${aembit_nat_OUTPUT}"
    remove_ipv4_chain nat "${aembit_nat_OUTPUT}"

    remove_ipv4_rule filter INPUT -j "${aembit_filter_INPUT}"
    remove_ipv4_chain filter "${aembit_filter_INPUT}"
}

# check if firewalld is active and running
firewalld_is_active_and_running() {
    if [ -n "${firewalld_status+x}" ]; then
	    return "${firewalld_status}"
    fi

    # Check if firewalld is installed
    if ! command -v firewalld &> /dev/null; then
	    echo "firewalld not installed."
	    firewalld_status=1
	    return "${firewalld_status}"
    fi

    # Check if firewalld is active and running
    if systemctl is-active --quiet firewalld && systemctl is-enabled --quiet firewalld; then
        echo "firewalld is active and running."
	    firewalld_status=0
	    return "${firewalld_status}"
    else
        echo "firewalld is either not active or not running."
	    firewalld_status=1
	    return "${firewalld_status}"
    fi
}

# Add ipv4 chain using firewalld direct
# args: table, chain name
function fwd_direct_add_ipv4_chain() {
    local table=$1
    local chain=$2
    ${FIREWALLCMD_QUIET_DIRECT} --add-chain ipv4 "${table}" "${chain}"
}

# Add ipv4 rule using firewalld direct
# args: table, chain name, iptable rule
function fwd_direct_add_ipv4_rule() {
    local table=$1
    local chain=$2
    shift
    shift
    ${FIREWALLCMD_QUIET_DIRECT} --add-rule ipv4 "${table}" "${chain}" 0 "$@"
}

# Remove ipv4 rule using firewalld direct
# args: table, chain name, iptable rule
function fwd_direct_remove_ipv4_rule() {
    local table=$1
    local chain=$2
    shift
    shift
    ${FIREWALLCMD_QUIET_DIRECT} --remove-rule ipv4 "${table}" "${chain}" 0 "$@"
}

# Remove ipv4 chain using firewalld direct
# args: table, chain name
function fwd_direct_remove_ipv4_chain() {
    local table=$1
    local chain=$2
    ${FIREWALLCMD_QUIET_DIRECT} --remove-rules ipv4 "${table}" "${chain}"
    ${FIREWALLCMD_QUIET_DIRECT} --remove-chain ipv4 "${table}" "${chain}"
}

# Add firewalld rich rule
function fwd_add_rich_rule() {
    rule="rule $*"
    ${FIREWALLCMD_QUIET} --add-rich-rule="${rule}"
}

# Remove firewalld rich rule
function fwd_remove_rich_rule() {
    rule="rule $*"
    ${FIREWALLCMD_QUIET} --remove-rich-rule="${rule}"
}

# Query active ipv4 rules using firewalld direct
# args: table, chain name
function fwd_direct_query_ipv4_rules() {
    local table=$1
    local chain=$2
    ${FIREWALLCMD_DIRECT} --get-rules ipv4 "${table}" "${chain}"
}

# Formats a firewalld rule for deletion.
#
# firewall-d output format doesn't include the table or chain at the beginning of the rule,
# but does include a leading '0', so we need to remove it and prepend the table and chain.
function format_fwd_delete_rule() {
    local rule=$1
    echo "nat ${aembit_nat_OUTPUT} ${rule#0 }"
}

# Iptables - Add ipv4 chain
# args: table, chain name
function iptables_add_ipv4_chain() {
    local table=$1
    local chain=$2
    ${IPTABLES} -t "${table}" -N "${chain}"
}

# Add ipv4 rule using iptables
# args: table, chain name, iptable rule
function iptables_add_ipv4_rule() {
    local table=$1
    local chain=$2
    shift
    shift
    ${IPTABLES} -t "${table}" -A "${chain}" "$@"
}

# Remove ipv4 rule using iptables
# args: table, chain name, iptable rule
function iptables_remove_ipv4_rule() {
    local table=$1
    local chain=$2
    shift
    shift
    ${IPTABLES} -t "${table}" -D "${chain}" "$@"
}

# Remove ipv4 chain using iptables
# args: table, chain name
function iptables_remove_ipv4_chain() {
    local table=$1
    local chain=$2
    # Delete all rules
    ${IPTABLES} -t "${table}" -F "${chain}"
    # Delete the chain
    ${IPTABLES} -t "${table}" -X "${chain}"
}

# Query iptables for active ipv4 rules
# args: table, chain name
function iptables_query_ipv4_rules() {
    local table=$1
    local chain=$2
    ${IPTABLES} -t "${table}" -S "${chain}"
}

# Formats an iptables rule for deletion.
#
# iptables removal expects "<table> <chain> [...]", but the format output by iptables -S is "-A <chain> [...]"
function format_iptables_delete_rule() {
    local rule=$1
    echo "${rule//-A /nat }"
}

function query_ipv4_rules() {
    # squelch the output, we just need the status
    _fwd_output=$(firewalld_is_active_and_running)
    fwd_status=$?
    if [ "$fwd_status" -eq 0 ]; then
        fwd_direct_query_ipv4_rules "$@"
    else
        iptables_query_ipv4_rules "$@"
    fi
}

function add_ipv4_chain() {
    if firewalld_is_active_and_running; then
        fwd_direct_add_ipv4_chain "$@"
    else
        iptables_add_ipv4_chain "$@"
    fi
}

function remove_ipv4_chain() {
    if firewalld_is_active_and_running; then
        fwd_direct_remove_ipv4_chain "$@"
    else
        iptables_remove_ipv4_chain "$@"
    fi
}

function add_ipv4_rule() {
    if firewalld_is_active_and_running; then
        fwd_direct_add_ipv4_rule "$@"
    else
        iptables_add_ipv4_rule "$@"
    fi
}

function remove_ipv4_rule() {
    if firewalld_is_active_and_running; then
        fwd_direct_remove_ipv4_rule "$@"
    else
        iptables_remove_ipv4_rule "$@"
    fi
}

# Formats a firewall rule for deletion.
function format_delete_rule() {
    local rule=$1
    # squelch the output, we just need the status
    _fwd_output=$(firewalld_is_active_and_running)
    fwd_status=$?
    if [ "$fwd_status" -eq 0 ]; then
        format_fwd_delete_rule "${rule}"
    else
        format_iptables_delete_rule "${rule}"
    fi

}

function save_and_reload_daemon() {
    if firewalld_is_active_and_running; then
	    # save firewalld rules and reload the daemon
	    ${FIREWALLCMD_QUIET} --runtime-to-permanent
	    ${FIREWALLCMD_QUIET} --reload
    fi
}

# Filter rules to allow TCP traffic to the AP.
function allow_tcp_traffic_to_ap() {
    # Allow local traffic to be sent to the Agent Proxy traffic port over loopback.
    add_ipv4_rule filter "${aembit_filter_INPUT}" -p tcp --dport "${agent_proxy_port}" -i lo -j ACCEPT

    if [ -n "$container_cidr" ]; then
        # Allow traffic from local containers to be sent to the Agent proxy traffic port.
        add_ipv4_rule filter "${aembit_filter_INPUT}" -p tcp -s  "${container_cidr}" --dport "${agent_proxy_port}" -j ACCEPT
    fi
}

# Default filter rule to block external traffic to the AP.
function default_block_tcp_traffic_to_ap() {
    add_ipv4_rule filter "${aembit_filter_INPUT}" -p tcp --dport "${agent_proxy_port}" -j REJECT
}

# Redirect outbound TCP traffic to the AP.
function redirect_outbound_tcp_traffic_to_ap() {
    # Rules to ignore outbound traffic rules to avoid redirectig them to the AP.

    # Ignore traffic from AP.
    add_ipv4_rule nat "${aembit_nat_OUTPUT}" -m owner --gid-owner "${aembit_group_id}" -j RETURN

    # Ignore TCP loopback.
    add_ipv4_rule nat "${aembit_nat_OUTPUT}" -p tcp -o lo -j RETURN

    # Ignore outbound traffic to the Azure communication channel.
    # See https://learn.microsoft.com/en-us/azure/virtual-network/what-is-ip-address-168-63-129-16.
    add_ipv4_rule nat "${aembit_nat_OUTPUT}" -p tcp --destination "${azure_platform_resources_ip}" -j RETURN
    add_ipv4_rule nat "${aembit_nat_OUTPUT}" -p tcp --destination "${imds_ip}" -j RETURN

    if [ -z "${AEMBIT_STEERING_ALLOWED_HOSTS}" ]; then
        # Redirect all new TCP connections to the Agent Proxy if we didn't specify custom steering.
        #
        # We use the "--syn" flag to forward new connections only, since,
        # prior to adding NAT rules, conntrack is not loaded. Once we add NAT rules,
        # conntrack is loaded, but mistakenly identifies existing connections as new
        # connections, redirecting them erroneously to the Agent Proxy mid-connection.
        add_ipv4_rule nat "${aembit_nat_OUTPUT}" -p tcp --syn -m comment --comment global -j REDIRECT --to-port "${agent_proxy_port}"
    fi
}

# Removes a list of rules from iptables/firewalld.
#
# This function expects arguments to be a list of rules (as strings), like:
# remove_ipv4_rules "rule1" "rule2" "rule3" ...
function remove_ipv4_rules() {
    active_rules_for_host=("$@")
    for rule in "${active_rules_for_host[@]}"; do
        rule=$(format_delete_rule "${rule}")
        echo "removing: ${rule}"
        IFS=' ' read -ra rule_parts <<< "${rule}"
        remove_ipv4_rule "${rule_parts[@]}"
    done
}

# Adds a list of rules to iptables/firewalld.
#
# Args:
# $1: The table to add the rule to,
# $2: The hostname to add as the rule's comment.
# $@: The list of rules to add.
function add_ipv4_rules() {
    for ip in "${incoming_ips_for_host[@]}"; do

        if [[ "${table}" == "${aembit_nat_PREROUTING}" ]] && [[ -n "${container_cidr}" ]]; then
            # add a source IP to the rule if we're dealing with the PREROUTING table and we have a container CIDR
            rule_parts=("nat" "${table}" "-d" "${ip}" "-p" "tcp" "--syn" "-s" "${container_cidr}" "-m" "comment" "--comment" "${hostname}" "-j" "REDIRECT" "--to-port" "${agent_proxy_port}")
        else
            rule_parts=("nat" "${table}" "-d" "${ip}" "-p" "tcp" "--syn" "-m" "comment" "--comment" "${hostname}" "-j" "REDIRECT" "--to-port" "${agent_proxy_port}")
        fi

        echo "adding: ${rule_parts[*]}"
        add_ipv4_rule "${rule_parts[@]}"
    done
}


# Manipulates outbound traffic rules to redirect traffic to the Agent Proxy.
#
# Args:
# $1: the table to manipulate,
# $2: the hostname to redirect traffic for,
# $@: list of IPs which the hostname resolved to.
function redirect_outbound_tcp_traffic() {
    table="$1"
    hostname="$2"
    shift
    shift
    incoming_ips_for_host=("$@")

    active_rules=$(query_ipv4_rules nat "${table}" | grep -e "-m comment --comment \"\?${hostname}\"\?")
    # Remove all quotes from the active rules (otherwise we end up double-quoting comments in the rules, which breaks the comparison below)
    active_rules=${active_rules//\"}

    echo -e "active rules:\n${active_rules}"
    echo -e "incoming IPs: ${incoming_ips_for_host[*]}"

    IFS=$'\n' read -rd '' -a active_rules_for_host <<< "${active_rules}"

    # We need to do three things:
    # 1. Remove any existing rules that are stale (i.e. DNS resolution for this hostname doesn't include the ip anymore),
    # 2. Add any new rules for IPs that are not already in the active rules,
    # 3. Don't remove rules that are still valid (i.e. there's a rule for the incoming IP in the active rules).
    #
    # We can do this by taking the intersection of the incoming IPs and the active rules, and then removing the IPs/rules that are in both.
    # At the end, whatever's left in active_rules_for_host gets removed, and whatever's in incoming_ips_for_hosts gets added.

    # For each index of the incoming_ips...
    for ip_index in "${!incoming_ips_for_host[@]}"; do
        ip="${incoming_ips_for_host[$ip_index]}"

        # ...and for each index of the active_rules...
        for rule_index in "${!active_rules_for_host[@]}"; do
            rule="${active_rules_for_host[$rule_index]}"

            # ...if the rule contains the fragment `-d ${ip}`, unset the rule and the IP to remove them from their respective arrays
            if [[ "$rule" == *"-d ${ip}"* ]]; then
                unset 'active_rules_for_host[rule_index]'
                unset 'incoming_ips_for_host[ip_index]'
                break
            fi
        done
    done

    remove_ipv4_rules "${active_rules_for_host[@]}"

    add_ipv4_rules "${table}" "${hostname}" "${incoming_ips_for_host[@]}"
}

# Manipulates outbound traffic rules to redirect traffic to the Agent Proxy.
#
# Args:
# $1: hostname,
# $@: list of IPs which the hostname resolved to.
function redirect_outbound_tcp_traffic_for_ips_to_ap() {
    hostname="$1"
    shift

    # redirect host-based traffic
    redirect_outbound_tcp_traffic "${aembit_nat_OUTPUT}" "${hostname}" "$@"

    # also redirect container-based traffic if Agent Proxy was installed with a container CIDR
    if [ -n "$container_cidr" ]; then
        redirect_outbound_tcp_traffic "${aembit_nat_PREROUTING}" "${hostname}" "$@"
    fi
}

# Redirect TCP traffic from containers to the AP.
function redirect_container_traffic_to_ap() {
    if [ -z "${AEMBIT_STEERING_ALLOWED_HOSTS}" ]; then
        # REDIRECT tcp traffic from container CIDR to the AP if we didn't specify custom steering
        add_ipv4_rule nat "${aembit_nat_PREROUTING}" -p tcp --syn -s "${container_cidr}" -j REDIRECT --to-port "${agent_proxy_port}"
    fi
}

# TCP related rules.
function tcp_proxy_rules() {
    #
    # filter rules
    #
    allow_tcp_traffic_to_ap

    # Disallow the rest of the traffic (coming from external interfaces) to be sent to the Agent Proxy traffic port.
    default_block_tcp_traffic_to_ap

    # Rules to redirect traffic to AP.

    # Redirect outbound TCP traffic to AP.
    redirect_outbound_tcp_traffic_to_ap

    # Redirect pre-routing TCP from local containers.
    if [ -n "$container_cidr" ]; then
        redirect_container_traffic_to_ap
    fi
}

# Notify Agent Proxy that iptables/firewalld rules were installed.
function notify_ap() {
    max_attempts=10

    for ((i=1; i<=max_attempts; i++)); do
        status_code=$(curl -s -o /dev/null -w "%{http_code}" -X POST "${agent_proxy_notification_endpoint}")
        if [ "$status_code" -eq 200 ]; then
            echo "Notified Agent Proxy."
            break
        else
            echo "Couldn't notify Agent Proxy, received HTTP $status_code ($i/$max_attempts)."
            sleep 1
        fi
    done

    if ((i > max_attempts)); then
        echo "Couldn't notify Agent Proxy after $max_attempts attempts."
    fi
}

# Install chain rules for steering traffic to the Agent Proxy.
#
# These function wrappers add the iptables or firewalld specific rules
# based upon whether firewalld is active and running.
function install_traffic_steering_rules_to_ap() {
    create_all_aembit_chains

    tcp_proxy_rules

    link_all_aembit_chains

    save_and_reload_daemon

    notify_ap
}


# Uninstall chain rules for steering traffic to the Agent Proxy.
#
# These function wrappers remove the iptables or firewalld specific rules
# based upon whether firewalld is active and running.
function uninstall_traffic_steering_rules_to_ap() {
    unlink_and_delete_aembit_chains
    save_and_reload_daemon
}

case "${OPTION}" in
    install)
        echo "Installing traffic steering rules to the Agent Proxy."
        install_traffic_steering_rules_to_ap
        ;;
    uninstall)
        echo "Uninstalling traffic steering rules."
        uninstall_traffic_steering_rules_to_ap
        ;;
    update)
        shift
        # Remaining args are expected to be <hostname> <ip>+, e.g.
        # ./rules.sh update graph.microsoft.com 10.10.10.10 11.11.11.11 [...]
        echo "Updating traffic steering rules"
        redirect_outbound_tcp_traffic_for_ips_to_ap "$@"
        ;;
    *)
        echo "Invalid option - ${OPTION}"
        ;;
esac
