#!/bin/bash
# Blowfish密钥提取脚本 - 在树莓派上以root运行
# 用法: sudo bash extract_blowfish_key.sh [PID]
#
# 安全保证:
# - 只读取内存,不修改任何寄存器或内存
# - GDB attach时间最短化 (<30秒)
# - 自动detach,即使中途异常

set -e

PID=${1:-$(pgrep -f Legged_sport | head -1)}
OUTDIR=/tmp/blowfish_dump_$$
mkdir -p $OUTDIR

if [ -z "$PID" ]; then
    echo "❌ 找不到 Legged_sport 进程"
    exit 1
fi

# 检查PID存在
if [ ! -d "/proc/$PID" ]; then
    echo "❌ PID $PID 不存在"
    exit 1
fi

echo "============================================"
echo "🐕 Blowfish密钥提取"
echo "============================================"
echo "PID: $PID"
echo "输出: $OUTDIR"
echo ""

# 检查当前狗的状态 (通过/tmp下的状态文件,如果有的话)
# 提示用户确认安全状态
echo "⚠️  确认狗的状态:"
echo "   - 趴下(damping)或在脚架上"
echo "   - 当前是高级模式 (未在执行运动指令)"
echo ""
read -p "确认安全状态? (yes/no): " confirm
if [ "$confirm" != "yes" ]; then
    echo "❌ 用户取消"
    rm -rf $OUTDIR
    exit 1
fi

# ===== 第1步: 保存进程内存映射 =====
echo ""
echo "[1/5] 保存进程内存映射..."
cat /proc/$PID/maps > $OUTDIR/maps.txt
echo "   ✅ /proc/$PID/maps -> $OUTDIR/maps.txt"

# 找出代码段范围 (anonymous r-x 段,是UPX解压后的代码)
ANON_RX=$(awk '/r-xp/ && !/\// {print $1}' $OUTDIR/maps.txt | head -1)
if [ -n "$ANON_RX" ]; then
    ANON_RX_START=$(echo $ANON_RX | cut -d- -f1)
    ANON_RX_END=$(echo $ANON_RX | cut -d- -f2)
    echo "   📌 匿名代码段(UPX解压): 0x$ANON_RX_START - 0x$ANON_RX_END"
fi

# 找堆段
HEAP=$(awk '/\[heap\]/ {print $1}' $OUTDIR/maps.txt | head -1)
if [ -n "$HEAP" ]; then
    HEAP_START=$(echo $HEAP | cut -d- -f1)
    HEAP_END=$(echo $HEAP | cut -d- -f2)
    echo "   📌 堆: 0x$HEAP_START - 0x$HEAP_END"
fi

# ===== 第2步: 准备GDB脚本 (只读,不修改) =====
echo ""
echo "[2/5] 生成GDB脚本..."

cat > $OUTDIR/gdb_script.gdb <<'GDBSCRIPT'
# GDB脚本: 只读模式提取Blowfish密钥
# 严格规定: 不修改任何寄存器、不写入任何内存、不调用任何函数

set pagination off
set print pretty on
set logging redirect on
set logging overwrite on

# 1. 进程基本信息
printf "=== 进程信息 ===\n"
info proc mappings

printf "\n=== 寄存器状态 ===\n"
info registers

printf "\n=== 调用栈 ===\n"
bt 5

# 2. 不设断点,只dump关键内存

# 完成
printf "\n=== 完成,准备detach ===\n"
GDBSCRIPT

echo "   ✅ $OUTDIR/gdb_script.gdb"

# ===== 第3步: GDB attach (最短时间) =====
echo ""
echo "[3/5] GDB attach (目标<5秒)..."

START_TS=$(date +%s)

# 用-batch模式,自动执行后立即退出
timeout 30 gdb -batch \
    -p $PID \
    -x $OUTDIR/gdb_script.gdb \
    > $OUTDIR/gdb_output.txt 2>&1

END_TS=$(date +%s)
ELAPSED=$((END_TS - START_TS))
echo "   ✅ GDB完成,耗时 ${ELAPSED}s"

# 检查狗是否还在线
if ! kill -0 $PID 2>/dev/null; then
    echo "   ⚠️  警告: Legged_sport进程消失!"
fi

# ===== 第4步: dump堆内存 (用/proc/PID/mem直接读取,比GDB快) =====
echo ""
echo "[4/5] 直接读取堆内存 (不用GDB)..."

if [ -n "$HEAP_START" ]; then
    HEAP_SIZE=$((0x$HEAP_END - 0x$HEAP_START))
    HEAP_OFFSET=$((0x$HEAP_START))
    echo "   读取堆 ${HEAP_SIZE} 字节..."

    # 用dd直接从/proc/$PID/mem读 (bs=65536大块读取,比bs=1快1000x)
    dd if=/proc/$PID/mem of=$OUTDIR/heap.bin \
       bs=65536 count=$HEAP_SIZE skip=$HEAP_OFFSET \
       iflag=skip_bytes,count_bytes 2>/dev/null || \
       echo "   ⚠️ /proc/mem读取失败,可能需要ptrace"

    if [ -f $OUTDIR/heap.bin ]; then
        HEAPACTUAL=$(stat -c%s $OUTDIR/heap.bin)
        echo "   ✅ heap.bin ($HEAPACTUAL 字节)"
    fi
fi

# 也dump匿名代码段 (UPX解压后的代码,含运行时S-box)
if [ -n "$ANON_RX_START" ]; then
    RX_SIZE=$((0x$ANON_RX_END - 0x$ANON_RX_START))
    RX_OFFSET=$((0x$ANON_RX_START))
    echo "   读取匿名代码段 ${RX_SIZE} 字节..."

    dd if=/proc/$PID/mem of=$OUTDIR/anon_code.bin \
       bs=65536 count=$RX_SIZE skip=$RX_OFFSET \
       iflag=skip_bytes,count_bytes 2>/dev/null

    if [ -f $OUTDIR/anon_code.bin ]; then
        CODEACTUAL=$(stat -c%s $OUTDIR/anon_code.bin)
        echo "   ✅ anon_code.bin ($CODEACTUAL 字节)"
    fi
fi

# 也dump所有匿名rw段 (含运行时S-box和P-array)
echo "   读取所有匿名rw段..."
RW_COUNT=0
awk '/rw-p/ && !/\// && !/\[stack\]/ && !/\[vvar\]/ {print $1}' $OUTDIR/maps.txt | while read range; do
    START=$(echo $range | cut -d- -f1)
    END=$(echo $range | cut -d- -f2)
    SIZE=$((0x$END - 0x$START))
    OFFSET=$((0x$START))

    # 只dump合理大小的段 (1KB - 64MB)
    if [ $SIZE -gt 1024 ] && [ $SIZE -lt 67108864 ]; then
        OUTFILE=$OUTDIR/rw_${START}.bin
        dd if=/proc/$PID/mem of=$OUTFILE \
           bs=65536 count=$SIZE skip=$OFFSET \
           iflag=skip_bytes,count_bytes 2>/dev/null

        if [ -f $OUTFILE ] && [ -s $OUTFILE ]; then
            echo "      rw_${START}.bin ($SIZE 字节)"
            RW_COUNT=$((RW_COUNT + 1))
        else
            rm -f $OUTFILE
        fi
    fi
done

# ===== 第5步: 打包 =====
echo ""
echo "[5/5] 打包输出..."
cd /tmp
TARBALL=blowfish_dump_$(date +%Y%m%d_%H%M%S).tar.gz
tar czf $TARBALL -C $OUTDIR .
mv $TARBALL /tmp/blowfish_dump.tar.gz
echo "   ✅ /tmp/blowfish_dump.tar.gz ($(stat -c%s /tmp/blowfish_dump.tar.gz) 字节)"

# 检查狗是否还在线
echo ""
if kill -0 $PID 2>/dev/null; then
    echo "✅ Legged_sport仍在运行 (PID $PID)"
else
    echo "❌ Legged_sport已退出!"
fi

echo ""
echo "============================================"
echo "完成!"
echo "============================================"
echo "数据位置: /tmp/blowfish_dump.tar.gz"
echo "中间文件: $OUTDIR/"
echo ""
echo "从Mac拉取:"
echo "  scp pi@192.168.123.161:/tmp/blowfish_dump.tar.gz ./data/lowlevel_captures/"