#!/bin/bash
# Name: ipfn
# Version: 1.0.0
# Date: 2026-01-22
# Refactored for nftables compatibility (fixed inet dnat ambiguity)

# sudo権限チェック
if [ "$(id -u)" -ne 0 ]; then
    exec sudo "$0" "$@"
fi

TABLE_NAME="ipf_wrapper"
QUIET=false
PROTO="tcp"

# --- 初期化 ---
init_nft() {
    nft add table inet ${TABLE_NAME} 2>/dev/null
    nft add chain inet ${TABLE_NAME} prerouting { type nat hook prerouting priority dstnat \; } 2>/dev/null
    nft add chain inet ${TABLE_NAME} output { type nat hook output priority dstnat \; } 2>/dev/null
    nft add chain inet ${TABLE_NAME} postrouting { type nat hook postrouting priority srcnat \; } 2>/dev/null
    nft add chain inet ${TABLE_NAME} forward { type filter hook forward priority filter \; } 2>/dev/null
}

# --- ルール追加 ---
add_rule() {
    init_nft
    local port_ip_port=$1
    local local_port=$(echo "$port_ip_port" | cut -d':' -f1)
    local target_ip=$(echo "$port_ip_port" | cut -d':' -f2)
    local target_port=$(echo "$port_ip_port" | cut -d':' -f3)

    # UUID生成
    local uuid="ipf-id:$(cat /proc/sys/kernel/random/uuid)"
    local comment="comment \"$uuid\""

    # DNATファミリーの特定 (IPv4なら 'ip', IPv6なら 'ip6')
    local family="ip"
    if [[ "$target_ip" =~ : ]]; then family="ip6"; fi

    # 1. PREROUTING (dnat ip to ... または dnat ip6 to ... に修正)
    if ! nft add rule inet ${TABLE_NAME} prerouting "$PROTO" dport "$local_port" dnat $family to "$target_ip:$target_port" "$comment"; then
        echo "Error: Failed to add PREROUTING rule."
        exit 1
    fi

    # 2. OUTPUT (ローカル発パケット用)
    nft add rule inet ${TABLE_NAME} output "$PROTO" dport "$local_port" dnat $family to "$target_ip:$target_port" "$comment" >/dev/null

    # 3. FORWARD
    nft add rule inet ${TABLE_NAME} forward $family daddr "$target_ip" "$PROTO" dport "$target_port" ct state new,established,related accept "$comment" >/dev/null
    nft add rule inet ${TABLE_NAME} forward $family saddr "$target_ip" "$PROTO" sport "$target_port" ct state established,related accept "$comment" >/dev/null

    # 4. POSTROUTING (Masquerade)
    # コンテナ環境検出関数(is_container)が定義されている前提
    if command -v systemd-detect-virt &>/dev/null && systemd-detect-virt --quiet --container; then
        nft add rule inet ${TABLE_NAME} postrouting $family daddr "$target_ip" "$PROTO" dport "$target_port" masquerade "$comment" >/dev/null
    fi

    [[ $QUIET == false ]] && echo "ipfn: Forwarded :$local_port -> $target_ip:$target_port ($PROTO)"
}

# (他の関数 list_rules, delete_rule 等は前回の回答を継承)

# メイン実行部
case "$1" in
    -L|-l) # リスト表示処理へ ;;
    -d)    # 削除処理へ ;;
    -f)    # sysctl設定へ ;;
    *)     add_rule "$1" ;;
esac