#!/bin/sh
set -eu

action=${1:-check}
trusted_proxy_cidrs=${TRUSTED_PROXY_CIDRS:-${CADDY_SOURCE_CIDR:-}}
image_port=${IMAGE_PORT:-8191}
chain=SAM_IMAGE_INGRESS

case "$image_port" in
    ''|*[!0-9]*) echo "IMAGE_PORT must be numeric" >&2; exit 2 ;;
esac

require_root() {
    if [ "$(id -u)" -ne 0 ]; then
        echo "Run this action as root." >&2
        exit 1
    fi
}

case "$action" in
    check)
        iptables -S DOCKER-USER 2>/dev/null | grep -F "$chain" || true
        iptables -S "$chain" 2>/dev/null || true
        ;;
    apply)
        require_root
        if [ -z "$trusted_proxy_cidrs" ]; then
            echo "TRUSTED_PROXY_CIDRS is required" >&2
            exit 2
        fi
        rule_count=0
        for trusted_proxy_cidr in $(printf '%s' "$trusted_proxy_cidrs" | tr ',' ' '); do
            case "$trusted_proxy_cidr" in
                *[!0-9A-Fa-f:./]*)
                    echo "Invalid trusted proxy CIDR: $trusted_proxy_cidr" >&2
                    exit 2
                    ;;
            esac
            rule_count=$((rule_count + 1))
        done
        if [ "$rule_count" -eq 0 ]; then
            echo "TRUSTED_PROXY_CIDRS must contain at least one CIDR" >&2
            exit 2
        fi
        iptables -n -L DOCKER-USER >/dev/null
        iptables -n -L "$chain" >/dev/null 2>&1 || iptables -N "$chain"
        iptables -F "$chain"
        iptables -A "$chain" -j DROP
        for trusted_proxy_cidr in $(printf '%s' "$trusted_proxy_cidrs" | tr ',' ' '); do
            iptables -I "$chain" 1 -s "$trusted_proxy_cidr" -j ACCEPT
        done
        iptables -C DOCKER-USER -p tcp -m conntrack --ctorigdstport "$image_port" -j "$chain" 2>/dev/null \
            || iptables -I DOCKER-USER 1 -p tcp -m conntrack --ctorigdstport "$image_port" -j "$chain"
        ;;
    remove)
        require_root
        while iptables -C DOCKER-USER -p tcp -m conntrack --ctorigdstport "$image_port" -j "$chain" 2>/dev/null; do
            iptables -D DOCKER-USER -p tcp -m conntrack --ctorigdstport "$image_port" -j "$chain"
        done
        if iptables -n -L "$chain" >/dev/null 2>&1; then
            iptables -F "$chain"
            iptables -X "$chain"
        fi
        ;;
    *)
        echo "Usage: $0 [check|apply|remove]" >&2
        exit 2
        ;;
esac