#!/bin/sh
set -eu

action=${1:-check}
source_cidr=${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 "$source_cidr" ]; then
            echo "CADDY_SOURCE_CIDR is required" >&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" -s "$source_cidr" -j ACCEPT
        iptables -A "$chain" -j DROP
        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