#!/usr/bin/env bash
set -euo pipefail

setup() {
    local xray_uid
    xray_uid="$(id -u xray)"

    iptables -t nat -N XRAY_OUTPUT 2>/dev/null || true
    iptables -t nat -F XRAY_OUTPUT

    iptables -t nat -A XRAY_OUTPUT -m owner --uid-owner "${xray_uid}" -j RETURN
    iptables -t nat -A XRAY_OUTPUT -d 0.0.0.0/8 -j RETURN
    iptables -t nat -A XRAY_OUTPUT -d 10.0.0.0/8 -j RETURN
    iptables -t nat -A XRAY_OUTPUT -d 127.0.0.0/8 -j RETURN
    iptables -t nat -A XRAY_OUTPUT -d 169.254.0.0/16 -j RETURN
    iptables -t nat -A XRAY_OUTPUT -d 172.16.0.0/12 -j RETURN
    iptables -t nat -A XRAY_OUTPUT -d 192.168.0.0/16 -j RETURN
    iptables -t nat -A XRAY_OUTPUT -d 224.0.0.0/4 -j RETURN
    iptables -t nat -A XRAY_OUTPUT -d 240.0.0.0/4 -j RETURN
    iptables -t nat -A XRAY_OUTPUT -p tcp -j REDIRECT --to-ports "${XRAY_REDIRECT_PORT}"

    iptables -t nat -C OUTPUT -p tcp -j XRAY_OUTPUT 2>/dev/null || \
        iptables -t nat -A OUTPUT -p tcp -j XRAY_OUTPUT
}

cleanup() {
    iptables -t nat -D OUTPUT -p tcp -j XRAY_OUTPUT 2>/dev/null || true
    iptables -t nat -F XRAY_OUTPUT 2>/dev/null || true
    iptables -t nat -X XRAY_OUTPUT 2>/dev/null || true
}

case "${1:-setup}" in
    setup)
        setup
        ;;
    cleanup)
        cleanup
        ;;
    *)
        echo "Usage: $0 [setup|cleanup]" >&2
        exit 2
        ;;
esac