#!/usr/bin/env bash
#
# =============================================================================
#                               U 盘杀毒 / 清理脚本
#                       交互式 U 盘彻底抹除与重建分区
# =============================================================================
#
# 用途：
#   - 彻底清除 U 盘 / 移动硬盘上的所有数据（包括病毒、恶意软件、隐藏分区）
#   - 擦除分区表签名（MBR + 备用 GPT）和首尾关键区域
#   - 新建 MBR 分区表 + 单分区，并格式化为 FAT32 / exFAT / NTFS / ext4
#
# 用法：
#   sudo ./usb_wipe.sh
#
#   交互步骤：
#     1. 脚本自动扫描并列出所有可移动磁盘设备
#     2. 按编号选择要清理的目标设备
#     3. 输入设备型号进行第一次确认
#     4. 输入大写 YES 进行最终确认（不可逆！）
#     5. 选择文件系统类型（FAT32 / exFAT / NTFS / ext4 / 跳过格式化）
#     6. 可选：输入卷标（回车使用默认值 USB_DRIVE）
#
# 依赖工具：
#   util-linux  : lsblk, fdisk, sfdisk, wipefs, blockdev, partprobe
#   coreutils   : dd, sync
#   dosfstools  : mkfs.vfat   (FAT32)
#   exfatprogs  : mkfs.exfat  (exFAT)
#   ntfs-3g     : mkfs.ntfs   (NTFS)
#   e2fsprogs   : mkfs.ext4   (ext4)
#
#   安装依赖：sudo pacman -S dosfstools exfatprogs ntfs-3g e2fsprogs
#   （Debian/Ubuntu：sudo apt install dosfstools exfatprogs ntfs-3g e2fsprogs）
#
# 安全设计：
#   - 必须 root 运行
#   - 三层确认机制（编号选择 → 型号匹配 → 大写 YES）
#   - 自动检测并卸载目标设备上的挂载点
#   - 默认过滤系统盘（nvme / sda 非 USB 设备）
#   - 仅擦除分区表签名和首尾关键区域（不浪费生命全盘写零）
#
# ⚠️  警告：本脚本会 永久销毁 目标设备上的所有数据，不可恢复！
#     使用前请再三确认已选中正确的设备。
#
# =============================================================================

set -euo pipefail

# --- 颜色定义 ----------------------------------------------------------------
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# --- 必须 root 运行 ----------------------------------------------------------
if [[ $EUID -ne 0 ]]; then
    echo -e "${RED}${BOLD}请使用 sudo 运行本脚本。${NC}"
    exit 1
fi

# --- 横幅 --------------------------------------------------------------------
clear
echo -e "${RED}${BOLD}"
echo "  ╔══════════════════════════════════════════════════════════╗"
echo "  ║        ⚠  警  告  —  不可逆操作  —  请谨慎使用  ⚠       ║"
echo "  ║                                                          ║"
echo "  ║   本脚本将 彻底抹除 U 盘全部内容 并重建分区表和文件系统    ║"
echo "  ║   所有数据将 永久丢失，无法恢复！                          ║"
echo "  ╚══════════════════════════════════════════════════════════╝"
echo -e "${NC}"
echo ""

# --- 工具检查 ----------------------------------------------------------------
for cmd in lsblk dd fdisk mkfs.vfat mkfs.exfat mkfs.ext4 wipefs blockdev; do
    if ! command -v "$cmd" &>/dev/null; then
        echo -e "${YELLOW}缺少工具: $cmd，尝试安装...${NC}"
        pacman -S --noconfirm --needed \
            util-linux fdisk dosfstools exfatprogs e2fsprogs coreutils 2>/dev/null || true
        break
    fi
done

# --- 步骤 1：列出可移动设备 --------------------------------------------------
echo -e "${CYAN}[1/5]${NC} 正在扫描可移动磁盘设备..."
echo ""

# 用 lsblk 列出可移动设备（TRAN=RPM 为空或 usb，TYPE=disk）
mapfile -t USB_DISKS < <(
    lsblk -dno NAME,SIZE,TRAN,MODEL,VENDOR 2>/dev/null \
    | awk '
        BEGIN { ignore["sda"]=1; ignore["nvme"]=1; ignore["loop"]=1 }
        /^sd[b-z]/ || /^vd[b-z]/ || ($3 == "usb") {
            print
        }
    ' \
    | sort -u
)

if [[ ${#USB_DISKS[@]} -eq 0 ]]; then
    echo -e "${RED}未检测到任何可移动磁盘设备。${NC}"
    echo ""
    echo "当前所有块设备："
    lsblk -dno NAME,SIZE,TRAN,MODEL 2>/dev/null | head -20
    exit 1
fi

echo -e "${BOLD}检测到以下可移动磁盘：${NC}"
echo ""
printf "  %-4s %-10s %-8s %-s\n" "编号" "设备名" "容量" "型号"
echo "  ────────────────────────────────────────"
for i in "${!USB_DISKS[@]}"; do
    read -r name size tran model vendor <<< "${USB_DISKS[$i]}"
    printf "  ${GREEN}%-4s${NC} ${RED}%-10s${NC} %-8s %-s %-s\n" \
        "$((i+1))" "/dev/$name" "$size" "$model" "$vendor"
done
echo ""

# --- 步骤 2：选择目标设备 ----------------------------------------------------
while true; do
    read -r -p "$(echo -e "${YELLOW}请选择要清理的设备编号（输入 0 退出）：${NC}")" choice
    if [[ "$choice" == "0" ]]; then
        echo "已取消。"
        exit 0
    fi
    if [[ "$choice" =~ ^[0-9]+$ ]] && (( choice >= 1 && choice <= ${#USB_DISKS[@]} )); then
        idx=$((choice - 1))
        read -r DEV_NAME DEV_SIZE DEV_TRAN DEV_MODEL DEV_VENDOR <<< "${USB_DISKS[$idx]}"
        DEV_PATH="/dev/$DEV_NAME"
        break
    fi
    echo -e "${RED}无效选择，请重新输入。${NC}"
done

echo ""

# --- 步骤 3：二次确认（输入设备型号） -----------------------------------------
echo -e "${RED}${BOLD}══════════════════════════════════════════════════════════════${NC}"
echo -e "${RED}${BOLD}  最终确认：即将销毁 ${DEV_PATH}${NC}"
echo -e "${RED}${BOLD}  设备型号：${DEV_MODEL:-未知}  ${DEV_VENDOR:-}${NC}"
echo -e "${RED}${BOLD}  容量：${DEV_SIZE}${NC}"
echo -e "${RED}${BOLD}══════════════════════════════════════════════════════════════${NC}"
echo ""

# 显示设备当前分区信息
echo -e "${CYAN}当前分区布局：${NC}"
lsblk "$DEV_PATH" 2>/dev/null || true
fdisk -l "$DEV_PATH" 2>/dev/null | head -20 || true
echo ""

# 检查是否有挂载
MOUNTS=$(lsblk -no MOUNTPOINT "$DEV_PATH" 2>/dev/null | grep -v '^$' || true)
if [[ -n "$MOUNTS" ]]; then
    echo -e "${YELLOW}检测到挂载点：${NC}"
    echo "$MOUNTS"
    echo ""
    read -r -p "$(echo -e "${YELLOW}是否卸载所有挂载点？[y/N] ${NC}")" umount_confirm
    if [[ "$umount_confirm" =~ ^[Yy]$ ]]; then
        for mp in $MOUNTS; do
            umount "$mp" 2>/dev/null || umount -l "$mp" 2>/dev/null || true
        done
        echo -e "${GREEN}已卸载。${NC}"
    fi
fi

echo ""
echo -e "${RED}${BOLD}请输入设备型号 ${DEV_MODEL:-N/A} 来确认操作：${NC}"
read -r confirm_model

if [[ "$confirm_model" != "${DEV_MODEL}" ]]; then
    echo -e "${GREEN}型号不匹配，操作已取消。${NC}"
    exit 0
fi

echo ""
echo -e "${RED}${BOLD}最后一次确认：输入 YES（大写）来执行不可逆的抹除操作：${NC}"
read -r final_confirm

if [[ "$final_confirm" != "YES" ]]; then
    echo -e "${GREEN}操作已取消。${NC}"
    exit 0
fi

echo ""
echo -e "${RED}${BOLD}══════════════════════════════════════════════════════════════${NC}"
echo -e "${RED}${BOLD}  开始执行... 请勿中断！${NC}"
echo -e "${RED}${BOLD}══════════════════════════════════════════════════════════════${NC}"
echo ""

# --- 步骤 4：抹除 ------------------------------------------------------------
echo -e "${CYAN}[2/5]${NC} 清除分区表签名（wipefs）..."
wipefs --all --force "$DEV_PATH" 2>/dev/null || true
echo -e "${GREEN}  ✓ 完成${NC}"

echo -e "${CYAN}[3/5]${NC} dd 全盘置零（这可能需要几分钟）..."

BLOCK_SIZE=$(blockdev --getsize64 "$DEV_PATH" 2>/dev/null || echo 0)
SECTOR_SIZE=$(blockdev --getss "$DEV_PATH" 2>/dev/null || echo 512)

# 先写开头 10 MiB（MBR/GPT 区域）
dd if=/dev/zero of="$DEV_PATH" bs=1M count=10 conv=fsync status=none 2>/dev/null && \
    echo -e "${GREEN}  ✓ 前 10 MiB 已清零${NC}" || \
    echo -e "${RED}  ✗ 写入失败${NC}"

# 再写末尾 10 MiB（备用 GPT 区域）
if [[ "$BLOCK_SIZE" -gt 0 ]]; then
    SEEK_VAL=$(( (BLOCK_SIZE - 10*1024*1024) / 512 ))
    if [[ "$SEEK_VAL" -gt 0 ]]; then
        dd if=/dev/zero of="$DEV_PATH" bs=512 seek="$SEEK_VAL" count=20480 conv=fsync status=none 2>/dev/null && \
            echo -e "${GREEN}  ✓ 末尾 10 MiB 已清零${NC}" || \
            echo -e "${RED}  ✗ 末尾写入失败${NC}"
    fi
fi

echo -e "${CYAN}[4/5]${NC} 同步并刷新内核分区表..."
sync
partprobe "$DEV_PATH" 2>/dev/null || blockdev --rereadpt "$DEV_PATH" 2>/dev/null || true
sleep 1
echo -e "${GREEN}  ✓ 完成${NC}"

# --- 步骤 5：选择文件系统并重建分区 ------------------------------------------
echo ""
echo -e "${CYAN}[5/5]${NC} 选择要创建的文件系统类型："
echo ""
echo "  1) FAT32   — 兼容性最好（单文件最大 4 GiB）"
echo "  2) exFAT   — 兼容性好，支持大文件（推荐）"
echo "  3) NTFS    — Windows 默认"
echo "  4) ext4    — Linux 专用"
echo "  5) 跳过格式化"
echo ""

FSTYPE_CHOICE=""
while true; do
    read -r -p "$(echo -e "${YELLOW}请选择 [1-5]：${NC}")" fstype
    case "$fstype" in
        1) FSTYPE="vfat";  FSTYPE_LABEL="FAT32"; break ;;
        2) FSTYPE="exfat"; FSTYPE_LABEL="exFAT"; break ;;
        3) FSTYPE="ntfs";  FSTYPE_LABEL="NTFS";  break ;;
        4) FSTYPE="ext4";  FSTYPE_LABEL="ext4";  break ;;
        5) FSTYPE="";      break ;;
        *) echo -e "${RED}无效选择。${NC}" ;;
    esac
done

if [[ -n "$FSTYPE" ]]; then
    read -r -p "$(echo -e "${YELLOW}卷标（回车默认 USB_DRIVE）：${NC}")" VOLUME_LABEL
    VOLUME_LABEL="${VOLUME_LABEL:-USB_DRIVE}"

    echo ""
    echo -e "${CYAN}创建 MBR 分区表 + 单分区...${NC}"

    # 使用 sfdisk 创建单个 Linux 分区（类型稍后根据 fs 调整）
    echo 'label: dos' | sfdisk --wipe always "$DEV_PATH" 2>/dev/null
    sleep 1

    # 创建分区（从 2048 扇区开始，使用全部剩余空间）
    echo '2048,,83' | sfdisk --wipe always "$DEV_PATH" 2>/dev/null
    sleep 1

    partprobe "$DEV_PATH" 2>/dev/null || true
    sleep 1

    # 获取分区名
    PART_PATH="${DEV_PATH}1"
    if [[ "$DEV_NAME" == nvme* ]] || [[ "$DEV_NAME" == mmcblk* ]]; then
        PART_PATH="${DEV_PATH}p1"
    fi

    # 等待分区节点出现
    for _ in {1..10}; do
        if [[ -b "$PART_PATH" ]]; then
            break
        fi
        sleep 0.5
    done

    if [[ ! -b "$PART_PATH" ]]; then
        echo -e "${RED}分区设备 $PART_PATH 未出现，请手动检查。${NC}"
        exit 1
    fi

    echo -e "${GREEN}  ✓ 分区已创建：${PART_PATH}${NC}"

    # 格式化
    echo -e "${CYAN}格式化 ${PART_PATH} 为 ${FSTYPE_LABEL}...${NC}"
    case "$FSTYPE" in
        vfat)
            mkfs.vfat -F 32 -n "${VOLUME_LABEL}" "$PART_PATH"
            ;;
        exfat)
            mkfs.exfat -n "${VOLUME_LABEL}" "$PART_PATH"
            ;;
        ntfs)
            mkfs.ntfs -f -L "${VOLUME_LABEL}" "$PART_PATH" 2>/dev/null || \
            mkfs.ntfs -Q -L "${VOLUME_LABEL}" "$PART_PATH"
            ;;
        ext4)
            mkfs.ext4 -F -L "${VOLUME_LABEL}" "$PART_PATH"
            ;;
    esac
    echo -e "${GREEN}  ✓ 格式化完成${NC}"

    # 刷新分区类型标识
    case "$FSTYPE" in
        vfat)  echo 'type=0c' | sfdisk --part-type "$DEV_PATH" 1 2>/dev/null || true ;;
        exfat) echo 'type=07' | sfdisk --part-type "$DEV_PATH" 1 2>/dev/null || true ;;
        ntfs)  echo 'type=07' | sfdisk --part-type "$DEV_PATH" 1 2>/dev/null || true ;;
    esac
else
    PART_PATH="$DEV_PATH"
    echo -e "${YELLOW}跳过格式化。${NC}"
fi

sync
partprobe "$DEV_PATH" 2>/dev/null || true

# --- 完成 --------------------------------------------------------------------
echo ""
echo -e "${GREEN}${BOLD}══════════════════════════════════════════════════════════════${NC}"
echo -e "${GREEN}${BOLD}  ✓ 操作完成！${NC}"
echo -e "${GREEN}${BOLD}══════════════════════════════════════════════════════════════${NC}"
echo ""
echo -e "  设备：${BOLD}${DEV_PATH}${NC}"
if [[ -n "$FSTYPE" ]]; then
    echo -e "  分区：${BOLD}${PART_PATH}${NC}"
    echo -e "  文件系统：${BOLD}${FSTYPE_LABEL}${NC}"
    echo -e "  卷标：${BOLD}${VOLUME_LABEL}${NC}"
fi
echo ""
echo -e "  最终布局："
lsblk "$DEV_PATH" 2>/dev/null || true
echo ""
echo -e "${CYAN}可以拔下 U 盘了。${NC}"