#!/bin/bash

# sudo権限で実行されているか確認し、そうでなければsudoで再実行
if [ "$(id -u)" -ne 0 ]; then
    exec sudo "$0" "$@"
fi

# グローバル変数
RULES_FILE="/tmp/iptables_forward_rules"
QUIET=false
PROTO="tcp"   # デフォルトプロトコル

# コンテナ環境検出関数
is_container() {
    # systemd-detect-virt を使用してコンテナ環境を検出
    if command -v systemd-detect-virt &> /dev/null; then
        if systemd-detect-virt --quiet | grep -qE "container|vm"; then
            return 0  # コンテナまたは仮想マシン環境
        fi
    fi

    # lxc-checkconfig を使用して LXC コンテナを検出 (systemd-detect-virt がない場合)
    if command -v lxc-checkconfig &> /dev/null; then
        if lxc-checkconfig 2>&1 | grep -q "Running in an LXC container"; then
            return 0 # LXC コンテナ環境
        fi
    fi

    # その他 (LXD など) のコンテナ環境検出方法を追加可能

    return 1 # コンテナ環境ではない
}


# ヘルプ表示
show_help() {
    cat << EOF
Usage: ipf [OPTIONS] [PORT:IP:PORT | -L | -d RULE_NUMBER]

Examples:
  ipf 11434:10.1.1.2:11434          # Forward local port 11434 to 10.1.1.2:11434
  ipf -L                            # List all rules with numbers
  ipf -d 1                          # Delete rule number 1
  ipf -d 1 -q                       # Delete rule number 1 (quiet mode)
  ipf -v                            # Show current iptables rules by iptables-save
  ipf -e                            # Edit all rules in editor and restore
  ipf -f                            # Enable IP forwarding
  ipf -t 1                          # Test rule number 1
  ipf -t 11434:10.10.1.2:11434      # Test rule by specification
  ipf -p udp 11434:10.1.1.2:11434   # Forward using UDP

Options:
  -h, --help     Show this help message
  -L, -l         List all rules
  -d NUM         Delete rule by number
  -q             Quiet mode (no output)
  -v             Show current iptables rules (via iptables-save)
  -e             Edit all rules in editor and restore
  -f             Enable IP forwarding and localnet routing
  -t ARG          Test connectivity for rule number or PORT:IP:PORT
  -p PROTO        Specify protocol (tcp|udp) for the following rule
  --version      Show version information

NOTE:
  All port forwarding is performed using DNAT only.  No MASQUERADE (SNAT) is applied,
  so the original source IP of the client is preserved.  This ensures that
  fail2ban logs the correct client IP on the destination server.
EOF
}

# バージョン情報を表示
show_version() {
    cat << EOF
ipf ver.1.0.4
Date: 2025-10-26
Created by: qwen3/gpt-oss/gemini-2.5-pro and krasherjoe

---
MIT License

Copyright (c) 2025 krasherjoe

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
---

** WARNING / 注意 **
This script operates with administrative privileges (sudo) to modify your system's
firewall (iptables) rules. Incorrect use, especially with the -e (edit) option,
can disrupt your network connectivity or overwrite existing security rules.
Please use with caution and understand the changes you are making.

このスクリプトは管理者権限(sudo)で動作し、システムのファイアウォール(iptables)ルールを
直接変更します。特に -e (編集) オプションの誤った使用は、ネットワーク接続を中断させたり、
既存のセキュリティルールを上書きする危険性があります。
コンテナ環境ではMASQUERADEを使用するため、クライアントの元のIPアドレスが隠蔽されます。
これにより、fail2banなどのツールが不正な行為者を正しく識別し、ブロックできなくなる可能性があります。
セキュリティへの影響を考慮してください。
内容をよく理解した上で、注意して使用してください。
EOF
}

# エラーメッセージを表示
error() {
    echo "Error: $*" >&2
    logger -p user.err "ipf: $*"
    exit 1
}

# ルールの保存
save_rules() {
    iptables-save > "$RULES_FILE"
}

# ルールの読み込み
load_rules() {
    if [[ -f "$RULES_FILE" ]]; then
        iptables-restore < "$RULES_FILE"
    fi
}

# ルール番号と内容をリスト表示
list_rules() {
    local i=1
    echo "Forwarding rules:"
    iptables -t nat -L PREROUTING -n -v --line-numbers | grep -E '^[0-9]+' | while read -r line; do
        echo "$i: $line"
        ((i++))
    done
}

# ルールを削除
delete_rule() {
    local num=$1
    if ! [[ "$num" =~ ^[0-9]+$ ]]; then
        error "Invalid rule number: $num"
    fi

    # ユーザーへの表示用に整形されたルール行を取得
    local rule_line_verbose=$(iptables -t nat -L PREROUTING --line-numbers | grep -E "^$num\s+" | head -n 1)
    if [[ -z "$rule_line_verbose" ]]; then
        error "No rule found with number: $num"
    fi

    # パース用に-nオプションを付けたルール行を取得
    local rule_line_numeric=$(iptables -t nat -L PREROUTING -n --line-numbers | grep -E "^\s*$num\s+" | head -n 1)
    if [[ -z "$rule_line_numeric" ]]; then
        error "Could not find numeric rule for number: $num" # Should not happen
    fi

    # ルールから詳細を抽出 (sedを使い、より堅牢に)
    local line_details=$(echo "$rule_line_numeric" | sed -n 's/.*\(tcp\|udp\).*dpt:\([0-9]*\).*to:\([0-9.]*\):\([0-9]*\).*/\1 \2 \3 \4/p')
    if [[ -z "$line_details" ]]; then
        error "Could not parse rule details from line: $rule_line_numeric"
    fi
    read -r proto local_port target_ip target_port <<< "$line_details"

    # 1. PREROUTING ルールを番号で削除
    if ! iptables -t nat -D PREROUTING "$num"; then
        error "Failed to delete PREROUTING rule $num"
    fi

    # 2. 対応する OUTPUT ルールをすべて削除
    while iptables -t nat -D OUTPUT -p "$proto" -m "$proto" --dport "$local_port" -j DNAT --to-destination "$target_ip:$target_port" >/dev/null 2>&1; do :; done

    # 3. 対応する FORWARD ルールをすべて削除
    while iptables -D FORWARD -p "$proto" -m "$proto" -d "$target_ip" --dport "$target_port" -j ACCEPT >/dev/null 2>&1; do :; done

    # 4. 確立済み通信を許可するルールを削除 (存在確認後に削除)
    if iptables -C FORWARD -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT 2>/dev/null; then
        iptables -D FORWARD -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT >/dev/null 2>&1
    fi

    if ! $QUIET; then
        echo "Deleted rule $num: $rule_line_verbose"
    fi
}

# ポートフォワーディングルールを追加
add_rule() {
    # IPフォワーディングが有効かチェック
    if [[ $(sysctl -n net.ipv4.ip_forward) -ne 1 ]]; then
        error "IP forwarding is disabled. Please enable it by running: ipf -f"
    fi

    local container_env=$(is_container)
    if [[ $container_env -eq 0 ]]; then
        echo -e "\e[33mDetected container environment.  MASQUERADE is enabled, which hides the original client IP address.  This will prevent tools like fail2ban from correctly identifying and blocking malicious actors. Consider the security implications.\e[0m" # コンテナ環境であることを示すメッセージ
    fi

    # 確立済みの通信は許可する (戻りのパケットのため)
    iptables -I FORWARD 1 -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

    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)

    # IPアドレス検証
    if ! [[ "$target_ip" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; then
        error "Invalid IP address: $target_ip"
    fi

    # ポート番号検証
    if ! [[ "$local_port" =~ ^[0-9]+$ ]] || ! [[ "$target_port" =~ ^[0-9]+$ ]]; then
        error "Invalid port number"
    fi
    if (( local_port < 1 || local_port > 65535 || target_port < 1 || target_port > 65535 )); then
        error "Port number out of range (1-65535)"
    fi

    # iptablesでルールを追加
    # 外部からのパケットを対象
    iptables -t nat -A PREROUTING -p "$PROTO" -m "$PROTO" --dport "$local_port" -j DNAT --to-destination "$target_ip:$target_port"
    # ローカルで生成されたパケットを対象
    iptables -t nat -A OUTPUT -p "$PROTO" -m "$PROTO" --dport "$local_port" -j DNAT --to-destination "$target_ip:$target_port"

    if [[ $container_env -eq 0 ]]; then
        # コンテナ環境ではMASQUERADEを使用
        iptables -t nat -A POSTROUTING -p "$PROTO" -m "$PROTO" -d "$target_ip" --dport "$target_port" -j MASQUERADE
    fi

    # 転送されるパケットを許可する
    iptables -A FORWARD -p "$PROTO" -m "$PROTO" -d "$target_ip" --dport "$target_port" -j ACCEPT

    if ! $QUIET; then
        echo "Rule added: port $local_port -> $target_ip:$target_port (proto=$PROTO)"
    fi
}

# ルールを編集して復元
edit_rules() {
    local temp_file
    temp_file=$(mktemp -t iptables.rules.XXXXXX) || error "Failed to create temporary file"
    chmod 0600 "$temp_file"
    trap 'rm -f "$temp_file"' EXIT

    # 現在のルールを一時ファイルに保存
    iptables-save > "$temp_file"

    # エディタを決定
    local editor=${EDITOR:-nano}
    if ! command -v "$editor" > /dev/null; then
        editor=vi
    fi

    # nanoで編集
    if ! "$editor" "$temp_file"; then
        error "Editor closed without saving or an error occurred."
    fi

    # 編集後の内容でリストア
    if iptables-restore < "$temp_file"; then
        echo "iptables rules restored successfully from your edits."
        echo "--- Displaying new rules ---"
        iptables-save
    else
        error "Failed to restore iptables rules. Please check for syntax errors in your edits."
    fi
}

# 疎通確認用関数
test_rule() {
    local arg=$1
    local local_port target_ip target_port

    if [[ "$arg" =~ ^[0-9]+$ ]]; then
        # rule number → ルール行を取得
        local line=$(iptables -t nat -L PREROUTING -n --line-numbers | grep "^$arg ")
        if [[ -z "$line" ]]; then
            error "No rule found with number: $arg"
        fi
        local_port=$(echo "$line" | awk '{print $8}'|cut -d ':' -f2)
        target_ip=$(echo "$line" | awk -F: '{print $1}')
        target_port=$(echo "$line" | awk -F: '{print $2}')
    else
        # 文字列形式
        local_port=$(echo "$arg" | cut -d':' -f1)
        target_ip=$(echo "$arg" | cut -d':' -f2)
        target_port=$(echo "$arg" | cut -d':' -f3)
    fi

    echo "Testing connectivity to local port \"$(tput setaf 3)nc -z -w 5 127.0.0.1 $local_port$(tput sgr0)\""
    if nc -z -w 5 127.0.0.1 "$local_port" 2>/dev/null; then
        echo "$(tput setaf 6)Connection successful.$(tput sgr0)"
    else
        echo "$(tput setaf 1)Connection failed.$(tput sgr0)"
    fi
}

# メイン処理
main() {
    # オプション解析
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -h|--help)
                show_help
                exit 0
                ;;
            -L|-l)
                list_rules
                exit 0
                ;;
            -v)
                iptables-save | grep -v '^#' | less -R
                exit 0
                ;;
            -e)
                edit_rules
                exit 0
                ;;
            -f)
                if sysctl -w net.ipv4.ip_forward=1 > /dev/null && \
                   sysctl -w net.ipv4.conf.all.route_localnet=1 > /dev/null && \
                   sysctl -w net.ipv4.conf.default.route_localnet=1 > /dev/null; then
                    echo "IP forwarding and localnet routing enabled."
                else
                    error "Failed to enable kernel parameters."
                fi
                exit 0
                ;;
            --version)
                show_version
                exit 0
                ;;
            -d)
                if [[ $# -lt 2 ]]; then
                    error "Missing rule number for -d option"
                fi
                delete_rule "$2"
                exit 0
                ;;
            -t)
                if [[ $# -lt 2 ]]; then
                    error "Missing argument for -t option"
                fi
                test_rule "$2"
                exit 0
                ;;
            -q)
                QUIET=true
                shift
                ;;
            -p)
                if [[ $# -lt 2 ]]; then
                    error "Missing protocol after -p"
                fi
                if [[ "$2" != "tcp" && "$2" != "udp" ]]; then
                    error "Unsupported protocol: $2"
                fi
                PROTO="$2"
                shift 2
                ;;
            -*)
                error "Unknown option: $1"
                ;;
            *)
                break
                ;;
        esac
    done

    # 引数が残っていれば、ルール追加処理
    if [[ $# -eq 0 ]]; then
        show_help
        exit 1
    fi

    local arg=$1

    # ルール追加処理
    if [[ "$arg" =~ ^[0-9]+:[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+:[0-9]+$ ]]; then
        add_rule "$arg"
    else
        error "Invalid format. Expected: PORT:IP:PORT or -L or -d RULE_NUMBER"
    fi
}

# 実行
main "$@"