#!/bin/bash
# ********************************************************
# @author: Lin, Chao <linchaochao@kylinos.cn>
# @create time: 2026-01-08 10:56:59
# @last modified: 2026-01-08 10:56:59
# @description:
# ********************************************************

set -euo pipefail

show_help() {
    cat <<EOF
用法:
  $0 [-r] <target1> [target2] ...

说明:
  - 默认行为（无 -r）：
      * 所有参数均为目标（可为 ELF 可执行文件或目录）。
      * 脚本会处理每个目标：
          - 若是目录：递归查找 ELF 可执行文件（排除 .so/.a 等）；
          - 若是文件：直接处理该文件（需为合法 ELF 可执行）。
      * 每个匹配的文件将被重命名为 .real，并替换为指向 dispatcher 的符号链接。

  - 使用 -r（revert 模式）：
      * 恢复被替换的文件：删除符号链接，将 .real 文件恢复原名。
      * 参数可为目录或单个文件（符号链接或 .real 文件）。

Dispatcher 路径:
  - 默认: /usr/bin/dispatcher
  - 可通过环境变量 DISPATCHER 覆盖，例如：
        DISPATCHER=/opt/my-dispatcher $0 /usr/local/bin/myapp

选项:
  -r       启用 revert 模式（恢复原文件）
    -f FILE  将处理过的文件以 files=xxx,yyy,zzz 形式追加到 FILE
  -h       显示此帮助

示例:
  # 使用默认 dispatcher 替换单个文件
  $0 /usr/local/bin/myapp

  # 使用自定义 dispatcher（通过环境变量）
  DISPATCHER=./my-dispatcher $0 /bin/ls

  # Revert
  $0 -r /usr/local/bin/myapp

EOF
}

# —————— 配置 dispatcher 路径 ——————
if [ -n "${DISPATCHER:-}" ]; then
    DISPATCHER_PATH="$DISPATCHER"
else
    DISPATCHER_PATH="/usr/bin/dispatcher"
fi
readonly DISPATCHER_PATH

REVERT_MODE=false
OUTPUT_FILE=""
processed_files=()

# 解析选项
while getopts "rhf:" opt; do
    case "$opt" in
        r) REVERT_MODE=true ;;
        f) OUTPUT_FILE="$OPTARG" ;;
        h) show_help; exit 0 ;;
        *) echo "错误: 无效选项 -$OPTARG" >&2; show_help; exit 1 ;;
    esac
done
shift $((OPTIND - 1))

if [ $# -eq 0 ]; then
    echo "错误: 至少需要一个目标（文件或目录）。" >&2
    show_help
    exit 1
fi

# 非 revert 模式下验证 dispatcher
if [ "$REVERT_MODE" = false ]; then
    if [ ! -x "$DISPATCHER_PATH" ]; then
        echo "错误: dispatcher 不存在或不可执行: $DISPATCHER_PATH" >&2
        echo "提示: 可通过环境变量 DISPATCHER 指定路径。" >&2
        exit 1
    fi
fi

# ========================
# Revert 模式
# ========================
if [ "$REVERT_MODE" = true ]; then
    echo "【REVERT 模式】恢复被替换的可执行文件..."
    for target in "$@"; do
        if [ -d "$target" ]; then
            while IFS= read -r -d '' link; do
                [ -L "$link" ] || continue
                real_file="$link.real"
                if [ -f "$real_file" ]; then
                    echo "恢复: $link"
                    rm -f "$link"
                    mv "$real_file" "$link"
                fi
            done < <(find "$target" -type l -print0 2>/dev/null || true)
        elif [ -e "$target" ]; then
            if [ -L "$target" ]; then
                real_file="$target.real"
                if [ -f "$real_file" ]; then
                    echo "恢复: $target"
                    rm -f "$target"
                    mv "$real_file" "$target"
                    continue
                fi
            elif [[ "$target" == *.real ]] && [ -f "$target" ]; then
                orig="${target%.real}"
                if [ -L "$orig" ]; then
                    echo "恢复: $orig"
                    rm -f "$orig"
                    mv "$target" "$orig"
                    continue
                fi
            fi
            echo "跳过（非被替换文件）: $target" >&2
        else
            echo "警告: 目标不存在，跳过: $target" >&2
        fi
    done

# ========================
# 正常模式： dispatcher
# ========================
else
    echo "使用 dispatcher: $DISPATCHER_PATH"
    echo "处理目标: $*"

    process_elf_file() {
        local file="$1"
        [[ "$file" == *.real ]] && return

        case "$file" in
            *.so|*.so.*|*.a|*.la|*.o) return ;;
        esac

        # 必须为 ELF
        if ! file "$file" 2>/dev/null | grep -q 'ELF'; then
            return
        fi

        # 忽略 ARM / aarch64 二进制
        if file "$file" 2>/dev/null | grep -qiE 'arm|aarch64'; then
            return
        fi

        # 要求可执行位（对符号链接会根据目标判断）
        if [ ! -x "$file" ]; then
            return
        fi

        # 仅处理 EXEC，或 DYN 且包含 INTERP（PIE）；否则视为共享库并跳过
        elf_type=$(LC_ALL=C readelf -h "$file" 2>/dev/null | awk '/Type:/ {print $2}' || echo "")
        if [ "$elf_type" = "EXEC" ]; then
            :
        elif [ "$elf_type" = "DYN" ]; then
            if ! LC_ALL=C readelf -l "$file" 2>/dev/null | grep -q 'INTERP'; then
                return
            fi
        else
            return
        fi

        echo "处理: $file"
        # 原子替换：先移动原文件到 .real，再创建临时符号链接，最后原子重命名为目标名
        if ! mv "$file" "$file.real"; then
            echo "错误: 无法移动 $file 到 $file.real" >&2
            continue
        fi

        tmp_link="${file}.tmp.$$"
        if ln -s "$DISPATCHER_PATH" "$tmp_link"; then
            if mv -Tf "$tmp_link" "$file"; then
                processed_files+=("$file")
            else
                echo "错误: 无法将临时链接移动为 $file，尝试恢复原文件" >&2
                rm -f "$tmp_link"
                if ! mv "$file.real" "$file"; then
                    echo "严重: 恢复失败: $file" >&2
                fi
            fi
        else
            echo "错误: 无法创建临时链接 $tmp_link，尝试恢复原文件" >&2
            if ! mv "$file.real" "$file"; then
                echo "严重: 恢复失败: $file" >&2
            fi
        fi
    }

    for target in "$@"; do
        if [ -d "$target" ]; then
            # 临时关闭 errexit，避免 read 在 EOF 时导致脚本因 set -euo pipefail 退出
            set +e
            while IFS= read -r -d '' f; do
                process_elf_file "$f"
            done < <(find "$target" \( -type f -o -type l \) -print0 2>/dev/null || true)
            set -e
        elif [ -f "$target" ] || [ -L "$target" ]; then
            process_elf_file "$target"
        else
            echo "警告: 跳过无效目标（非文件/目录）: $target" >&2
        fi
    done
    # 如果指定了输出文件，追加处理过的文件列表
    if [ -n "$OUTPUT_FILE" ] && [ "${#processed_files[@]}" -gt 0 ]; then
        if ! touch "$OUTPUT_FILE" 2>/dev/null; then
            echo "错误: 无法写入到输出文件: $OUTPUT_FILE" >&2
        else
            total=${#processed_files[@]}
            if [ "$total" -eq 1 ]; then
                printf 'files=%s\n' "${processed_files[0]}" >> "$OUTPUT_FILE"
            else
                printf 'files=%s,\n' "${processed_files[0]}" >> "$OUTPUT_FILE"
                i=1
                while [ $i -lt $total ]; do
                    if [ $i -eq $((total-1)) ]; then
                        printf ' %s\n' "${processed_files[$i]}" >> "$OUTPUT_FILE"
                    else
                        printf ' %s,\n' "${processed_files[$i]}" >> "$OUTPUT_FILE"
                    fi
                    i=$((i+1))
                done
            fi
        fi
    fi
fi

echo "完成。"
