#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
usbvault —— 银河麒麟 Kylin V10 USB 外接存储设备加密格式化管理工具
====================================================================
用途 :
  以"单位统一编码"对 U 盘 / 移动硬盘做 LUKS2 整盘加密格式化。
  盘内所有数据密文落盘；插入同编码（已登记）的单位计算机自动解锁挂载；
  插入其他机器无法识别，重新格式化即密码学销毁全部数据。

卷结构 :
  /dev/sdX（裸盘，无分区表）
  └─ LUKS2 header (AES-XTS-plain64 / 512bit / Argon2id)
     └─ /dev/mapper/usbvault-<uuid8>
        └─ NTFS (LABEL=UNITVAULT)

编码派生 :
  K = PBKDF2-HMAC-SHA256(编码, 固定域盐, 600000 次, 64 字节)
  K 存于 /etc/usbvault/unit.key（root:0600）；编码明文不落盘、不写入 U 盘。

运行环境 : 银河麒麟桌面操作系统 V10（ID=kylin），Python 3.8+，仅标准库。
权限要求 : 除 help/version 外均需 root。
"""

import argparse
import fcntl
import getpass
import grp
import hashlib
import hmac
import json
import os
import pwd
import re
import secrets
import shutil
import stat as stat_mod
import subprocess
import sys
import time
from datetime import datetime

VERSION = "1.2.0"

# ---------- GUI 模式检测与 Hook ----------
GUI_MODE = os.environ.get("USBVAULT_GUI") == "1"
"""GUI 模式下 sys.exit() 会抛 RuntimeError 而非退出进程。"""

# 可被 GUI 层覆盖的钩子：设置为 callable 后替换对应函数的终端行为
HOOKS = {
    "ask": None,       # callable(prompt, default="") -> str
    "confirm": None,   # callable(prompt, assume_yes=False) -> bool
    "err": None,       # callable(msg)
    "warn": None,      # callable(msg)
    "step": None,      # callable(msg)
    "info": None,      # callable(msg)
    "audit": None,     # callable(action, device, result, detail, serial, vidpid)
}

# ---------- 机器可读输出（供 GUI 前端调用） ----------
# --json 模式下：所有人类可读输出改走 stderr，stdout 只保留一行 JSON 结果，
# 这样前端既能读进度，又能拿到结构化结果，且不影响现有纯文本用法。
JSON_MODE = False
_JSON_STDOUT = sys.stdout   # 在 main() 重定向 sys.stdout 之前保存真实 stdout
JSON_RESULT = {}            # 各命令通过 json_set() 填充的结果字段

def json_set(**kw):
    """在 JSON 模式下登记命令结果字段（非 JSON 模式为 no-op）。"""
    if JSON_MODE:
        JSON_RESULT.update(kw)

def emit_json(payload):
    """向真实 stdout 输出一行 JSON（非 JSON 模式静默）。"""
    if not JSON_MODE:
        return
    try:
        json.dump(payload, _JSON_STDOUT, ensure_ascii=False)
        _JSON_STDOUT.write("\n")
        _JSON_STDOUT.flush()
    except (OSError, ValueError):
        pass

def read_audit_tail(n=30):
    """读审计日志尾部 n 条（解析失败的行原样保留为字符串）。"""
    if not os.path.isfile(AUDIT_FILE):
        return []
    try:
        with open(AUDIT_FILE, "r", encoding="utf-8", errors="replace") as f:
            lines = [l for l in f.read().splitlines() if l.strip()]
    except OSError:
        return []
    out = []
    for line in lines[-n:]:
        try:
            out.append(json.loads(line))
        except ValueError:
            out.append({"raw": line})
    return out


def _exit(code):
    """统一退出：GUI 模式抛异常，CLI 模式 sys.exit()。"""
    if GUI_MODE:
        raise RuntimeError("exit:%d" % code)
    sys.exit(code)

# ---------- 路径（支持环境变量隔离，用于测试） ----------
def _p(env, default):
    return os.environ.get(env, default)

ETC_DIR     = _p("USBVAULT_ETC", "/etc/usbvault")
VAR_DIR     = _p("USBVAULT_VAR", "/var/lib/usbvault")
LOG_DIR     = _p("USBVAULT_LOG", "/var/log/usbvault")
MEDIA_DIR   = _p("USBVAULT_MEDIA", "/media/usbvault")
RUN_DIR     = _p("USBVAULT_RUN", "/run/usbvault")
KEY_FILE    = os.path.join(ETC_DIR, "unit.key")
CONFIG_FILE = os.path.join(ETC_DIR, "config.json")
ADMIN_HASH_FILE = os.path.join(ETC_DIR, "admin.hash")
AUDIT_FILE  = os.path.join(LOG_DIR, "audit.log")

# ---------- 密码参数 ----------
KDF_SALT    = b"usbvault-v1-unit-key"
KDF_ITERS   = 600000
KEY_LEN     = 64
CODE_RE     = re.compile(r"^[A-Za-z0-9]{12}$")
LUKS_CIPHER = "aes-xts-plain64"
LUKS_KEYSIZE = 512
FS_LABEL    = "UNITVAULT"
MAPPER_PREFIX = "usbvault-"

# ---------- 终端颜色（借鉴 usb-cleanup 风格） ----------
if sys.stdout.isatty():
    C_RST = "\033[0m"; C_RED = "\033[1;31m"; C_GRN = "\033[1;32m"
    C_YEL = "\033[1;33m"; C_BLU = "\033[1;34m"; C_CYA = "\033[1;36m"
else:
    C_RST = C_RED = C_GRN = C_YEL = C_BLU = C_CYA = ""

def ok(msg):  print("%s✔%s %s" % (C_GRN, C_RST, msg))
def warn(msg):
    if HOOKS["warn"]: return HOOKS["warn"](msg)
    print("%s⚠%s %s" % (C_YEL, C_RST, msg))
def err(msg):
    if HOOKS["err"]: return HOOKS["err"](msg)
    print("%s✘%s %s" % (C_RED, C_RST, msg), file=sys.stderr)
def step(msg):
    if HOOKS["step"]: return HOOKS["step"](msg)
    print("%s▶%s %s" % (C_CYA, C_RST, msg))
def info(msg):
    if HOOKS["info"]: return HOOKS["info"](msg)
    print("  %s" % msg)

os.environ["PATH"] = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"


# ================= 通用工具 =================

def require_root():
    if os.geteuid() != 0:
        err("需要 root 权限，请使用 sudo 执行。")
        _exit(1)

def run(cmd, input_bytes=None, timeout=900, check=False):
    """执行外部命令，返回 (returncode, stdout, stderr)。"""
    try:
        r = subprocess.run(
            cmd,
            input=input_bytes,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout,
        )
        rc, out, e = r.returncode, r.stdout.decode("utf-8", "replace"), r.stderr.decode("utf-8", "replace")
    except subprocess.TimeoutExpired:
        return 124, "", "timeout"
    except FileNotFoundError:
        return 127, "", "command not found: %s" % cmd[0]
    if check and rc != 0:
        raise RuntimeError("命令失败 [%s]: %s" % (" ".join(cmd), e.strip()))
    return rc, out, e

def run_stream(cmd, input_bytes=None, timeout=900):
    """流式执行外部命令，yield (stdout_line_or_None, stderr_line_or_None, event_dict)。
    event_dict = {"type": "start"|"stdout"|"stderr"|"exit", "rc": int_or_None}
    GUI 进度条可消费此生成器，实时获取 dd/cryptsetup/mkntfs 输出。
    """
    yield {"type": "start"}
    try:
        p = subprocess.Popen(
            cmd,
            stdin=subprocess.PIPE if input_bytes is not None else None,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            bufsize=1,
            universal_newlines=True,
        )
        if input_bytes is not None:
            p.stdin.write(input_bytes.decode("utf-8", "replace"))
            p.stdin.close()
        import select
        while True:
            ready, _, _ = select.select([p.stdout, p.stderr], [], [], 0.5)
            for stream in ready:
                line = stream.readline()
                if line:
                    key = "stdout" if stream is p.stdout else "stderr"
                    yield {"type": key, "line": line.rstrip("\n")}
            if p.poll() is not None:
                # 读取剩余
                for line in p.stdout:
                    yield {"type": "stdout", "line": line.rstrip("\n")}
                for line in p.stderr:
                    yield {"type": "stderr", "line": line.rstrip("\n")}
                break
    except FileNotFoundError:
        yield {"type": "exit", "rc": 127, "error": "command not found: %s" % cmd[0]}
        return
    yield {"type": "exit", "rc": p.returncode}
    if p.returncode != 0:
        raise RuntimeError("命令失败 [%s], rc=%d" % (" ".join(cmd), p.returncode))

def ask(prompt, default=""):
    """从 /dev/tty 读取输入（管道环境安全，借鉴 usb-cleanup）。"""
    if HOOKS["ask"]:
        return HOOKS["ask"](prompt, default)
    try:
        tty = open("/dev/tty", "r+")
    except OSError:
        tty = None
    try:
        if tty:
            full = prompt + (("[%s] " % default) if default else "")
            tty.write(full); tty.flush()
            line = tty.readline()
            return line.rstrip("\n") if line else default
        return input(prompt) or default
    finally:
        if tty:
            tty.close()

def confirm(prompt, assume_yes=False):
    if assume_yes:
        warn("[自动确认] %s" % prompt)
        return True
    if HOOKS["confirm"]:
        return HOOKS["confirm"](prompt, assume_yes=False)
    resp = ask("%s [y/N]: " % prompt)
    return resp.strip().lower() in ("y", "yes", "是")

def ask_code(prompt="请输入单位统一编码(12位字母数字)"):
    """读编码：测试/部署可用环境变量 USBVAULT_CODE 注入。
    编码不需要隐藏回显（非敏感密码），改用 ask() 从 /dev/tty 读取，
    避免 getpass.getpass() 在 sudo 环境下因 TTY 控制失败导致 EOFError。
    """
    env = os.environ.get("USBVAULT_CODE")
    if env is not None:
        code = env.strip()
        if not CODE_RE.match(code):
            err("环境变量 USBVAULT_CODE 不是合法的 12 位字母数字编码。")
            _exit(1)
        return code
    for _ in range(3):
        code = ask("%s: " % prompt).strip()
        if CODE_RE.match(code):
            return code
        err("编码必须为 12 位字母数字（区分大小写），请重新输入。")
    _exit(1)

def human_size(n):
    try:
        n = int(n)
    except (TypeError, ValueError):
        return "?"
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if n < 1024:
            return "%.1f%s" % (n, unit)
        n /= 1024.0
    return "%.1fPiB" % n


# ================= 密钥与配置 =================

def derive_key(code):
    return hashlib.pbkdf2_hmac("sha256", code.encode("utf-8"), KDF_SALT, KDF_ITERS, dklen=KEY_LEN)

def load_unit_key(required=True):
    if not os.path.isfile(KEY_FILE):
        if required:
            err("本机尚未登记单位编码（缺少 %s）。请先执行: sudo usbvault init" % KEY_FILE)
            _exit(1)
        return None
    with open(KEY_FILE, "rb") as f:
        key = f.read()
    if len(key) != KEY_LEN:
        err("单位密钥文件损坏（%s，长度 %d ≠ %d）。" % (KEY_FILE, len(key), KEY_LEN))
        _exit(1)
    return key

def key_fingerprint():
    key = load_unit_key(required=False)
    if key is None:
        return "未登记"
    return hashlib.sha256(key).hexdigest()[:8]

def write_unit_key(key):
    os.makedirs(ETC_DIR, mode=0o700, exist_ok=True)
    tmp = KEY_FILE + ".tmp"
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    try:
        os.write(fd, key)
    finally:
        os.close(fd)
    os.chmod(tmp, 0o600)
    os.replace(tmp, KEY_FILE)

def load_config():
    cfg = {
        "mount_uid": 0,
        "mount_gid": 100,        # users 组
        "fmask": "0111",
        "dmask": "0002",
        "noexec": False,
        "atime": "noatime",
        "guard_enabled": True,   # 设备管控默认开启（未登记编码时不实际生效）
        "deploy_role": "admin",  # 部署角色：admin 超管机 / client 受控客户端
    }
    if os.path.isfile(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                cfg.update(json.load(f))
        except (ValueError, OSError) as e:
            warn("配置文件读取失败，使用默认值: %s" % e)
    # gid 不存在时回退 root
    try:
        grp.getgrgid(int(cfg["mount_gid"]))
    except KeyError:
        cfg["mount_gid"] = 0
    return cfg

def save_config(cfg):
    """原子写回配置文件（root:0600）。"""
    os.makedirs(ETC_DIR, mode=0o700, exist_ok=True)
    tmp = CONFIG_FILE + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)
    os.chmod(tmp, 0o600)
    os.replace(tmp, CONFIG_FILE)


# ================= 部署角色与超管密码 =================

def is_client():
    """本机是否为受控客户端（client）。默认 admin，向后兼容旧部署。"""
    return load_config().get("deploy_role", "admin") == "client"

def admin_password_set():
    return os.path.isfile(ADMIN_HASH_FILE)

def hash_admin_password(pw, salt=None):
    """PBKDF2 派生密码哈希，返回 'salt_hex$hash_hex'。"""
    if salt is None:
        salt = secrets.token_bytes(16)
    dk = hashlib.pbkdf2_hmac("sha256", pw.encode("utf-8"), salt, 600000, dklen=32)
    return "%s$%s" % (salt.hex(), dk.hex())

def verify_admin_password(pw):
    """校验超管密码。未设置时返回 False（调用方需先查 admin_password_set）。"""
    try:
        with open(ADMIN_HASH_FILE, "r", encoding="utf-8") as f:
            stored = f.read().strip()
        salt_hex, hash_hex = stored.split("$", 1)
        expect = hash_admin_password(pw, bytes.fromhex(salt_hex)).split("$", 1)[1]
        return hmac.compare_digest(expect, hash_hex)
    except (OSError, ValueError):
        return False

_TTY_ERROR = None   # 最近一次密码读取失败的诊断信息，用于错误提示

def _read_hidden(stream, prompt, out=None):
    """在给定流上关闭回显读取一行（不关闭流）。
    提示信息默认写回同一流；当流只可读（如 sys.stdin）时由 out 指定输出。"""
    import termios
    if out is None:
        out = stream
    fd = stream.fileno()
    old = termios.tcgetattr(fd)
    new = termios.tcgetattr(fd)
    new[3] = new[3] & ~termios.ECHO
    try:
        out.write(prompt)
        out.flush()
        termios.tcsetattr(fd, termios.TCSADRAIN, new)
        line = stream.readline()
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old)
        try:
            out.write("\n")
            out.flush()
        except (OSError, ValueError):
            pass
    return line.rstrip("\n")

def _read_password_tty(prompt):
    """隐藏回显读取一行密码。
    顺序：USBVAULT_ADMIN_PASS → GUI 钩子 → /dev/tty → 标准输入（终端）。
    全部不可用时返回 None，原因记录在 _TTY_ERROR。"""
    global _TTY_ERROR
    env = os.environ.get("USBVAULT_ADMIN_PASS")
    if env is not None:
        return env
    if HOOKS["ask"]:
        return HOOKS["ask"](prompt, "")
    _TTY_ERROR = None
    try:
        tty = open("/dev/tty", "r+")
    except OSError as e:
        _TTY_ERROR = "无法打开 /dev/tty（%s）" % e
        tty = None
    if tty is not None:
        try:
            return _read_hidden(tty, prompt)
        except (OSError, ValueError, ImportError) as e:
            _TTY_ERROR = "/dev/tty 读取失败（%s）" % e
        finally:
            tty.close()
    if sys.stdin is not None and sys.stdin.isatty():
        try:
            return _read_hidden(sys.stdin, prompt, sys.stdout)
        except (OSError, ValueError, ImportError) as e:
            _TTY_ERROR = "%s；标准输入读取失败（%s）" % (_TTY_ERROR or "终端不可用", e)
    else:
        _TTY_ERROR = "%s；标准输入也不是终端" % (_TTY_ERROR or "终端不可用")
    return None

def require_admin(operation="此操作"):
    """敏感操作门卫：
    - client 机：直接拒绝（提示到超管机执行）
    - admin 机已设超管密码：/dev/tty 校验（最多 3 次），审计 admin-auth
    - admin 机未设密码：放行并提示建议设置
    校验失败/被拒绝时退出进程（GUI 模式抛异常由钩子拦截）。"""
    if is_client():
        err("%s被拒绝：本机为受控客户端，请到超管机（admin）执行。" % operation)
        audit("admin-auth", result="fail", detail="%s 拒绝（client 角色）" % operation)
        _exit(1)
    if not admin_password_set():
        warn("本机尚未设置超管密码，%s仅校验 root 权限。建议执行: sudo usbvault passwd" % operation)
        return
    for attempt in range(3):
        pw = _read_password_tty("请输入超管密码以继续（%s）: " % operation)
        if pw is None:
            err("无法读取超管密码：%s。" % (_TTY_ERROR or "无可用终端"))
            err("请在真实终端中执行，或用 USBVAULT_ADMIN_PASS=密码 注入。")
            audit("admin-auth", result="fail", detail="%s 无终端" % operation)
            _exit(1)
        if verify_admin_password(pw):
            audit("admin-auth", result="ok", detail=operation)
            return
        err("密码错误（剩余 %d 次机会）。" % (2 - attempt))
    audit("admin-auth", result="fail", detail="%s 密码连续错误" % operation)
    _exit(1)

def cmd_passwd(args):
    """设置/修改超管密码。已设置时需先验原密码（client 机同样允许改密）。
    环境变量：
      USBVAULT_ADMIN_PASS — 注入密码（旧密码验证，或无 admin.hash / --reset 时直接作为新密码）
      USBVAULT_NEW_PASS   — 显式注入新密码（避免与旧密码验证时 USBVAULT_ADMIN_PASS 的歧义）
    """
    require_root()
    reset = bool(getattr(args, "reset", False))
    had = admin_password_set()
    # 显式新密码优先；其余场景复用 USBVAULT_ADMIN_PASS
    new_pass_injected = os.environ.get("USBVAULT_NEW_PASS")

    if had and not reset:
        # 需要先验证旧密码——此时 USBVAULT_ADMIN_PASS 若设置，被视为旧密码
        old = _read_password_tty("请输入当前超管密码: ")
        if old is None:
            err("无法读取当前超管密码：%s。" % (_TTY_ERROR or "无可用终端"))
            err("忘记密码可执行: sudo usbvault passwd --reset（需 root，无需旧密码）")
            audit("passwd", result="fail", detail="原密码无法读取")
            return 1
        if not verify_admin_password(old):
            err("原密码校验失败，未做任何修改。")
            err("忘记密码可执行: sudo usbvault passwd --reset（需 root，无需旧密码）")
            audit("passwd", result="fail", detail="原密码校验失败")
            return 1

    if had and reset:
        warn("重置模式：跳过旧密码校验（需 root 权限），仅用于忘记密码时恢复。")

    # 新密码读取：优先用 USBVAULT_NEW_PASS；无 admin.hash 或 --reset 时 USBVAULT_ADMIN_PASS 也当新密码
    if new_pass_injected is not None:
        pw1 = new_pass_injected
    elif (not had or reset) and os.environ.get("USBVAULT_ADMIN_PASS") is not None:
        # 无 admin.hash 或 reset，USBVAULT_ADMIN_PASS 没有被旧密码验证消耗
        pw1 = os.environ["USBVAULT_ADMIN_PASS"]
    else:
        pw1 = _read_password_tty("请输入新超管密码（至少 8 位）: ")
        if pw1 is None:
            err("无法读取新超管密码：%s。" % (_TTY_ERROR or "无可用终端"))
            err("请在真实终端中执行，或用 USBVAULT_NEW_PASS=新密码 注入。")
            return 1
    if len(pw1) < 8:
        err("密码至少 8 位。")
        return 1
    # 二次确认：无显式注入时需终端验证
    if new_pass_injected is None and not ((not had or reset) and os.environ.get("USBVAULT_ADMIN_PASS") is not None):
        pw2 = _read_password_tty("请再次输入新超管密码: ")
        if pw1 != pw2:
            err("两次输入不一致。")
            return 1
    os.makedirs(ETC_DIR, mode=0o700, exist_ok=True)
    tmp = ADMIN_HASH_FILE + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        f.write(hash_admin_password(pw1))
    os.chmod(tmp, 0o600)
    os.replace(tmp, ADMIN_HASH_FILE)
    ok("超管密码已%s。" % ("重置" if reset else ("更新" if had else "设置")))
    audit("passwd", result="ok", detail="超管密码已重置（--reset）" if reset else "超管密码已设置/更新")
    return 0

def cmd_set_role(args):
    """切换部署角色。client 机拒绝执行；已设密码时需验证。"""
    require_root()
    role = getattr(args, "role", "")
    if role not in ("admin", "client"):
        err("用法: sudo usbvault set-role admin|client")
        return 1
    if is_client():
        err("本机为受控客户端，不允许自行切换角色（请在超管机重装或删除配置后操作）。")
        audit("set-role", result="fail", detail="client 拒绝切角色")
        return 1
    if admin_password_set():
        require_admin("切换部署角色")
    cfg = load_config()
    old = cfg.get("deploy_role", "admin")
    cfg["deploy_role"] = role
    save_config(cfg)
    ok("部署角色: %s → %s" % (old, role))
    if role == "client":
        warn("本机已切换为受控客户端：格式化/换码/关管控/临时放行将被拒绝。")
    audit("set-role", result="ok", detail="%s -> %s" % (old, role))
    return 0

def cmd_verify_admin(args):
    """隐藏命令：供 prerm / uninstall.sh 在卸载前校验超管密码。
    未设密码返回 0；已设密码校验通过返回 0，失败返回 1。"""
    require_root()
    if not admin_password_set():
        return 0
    for attempt in range(3):
        pw = _read_password_tty("卸载 USBVault 需要超管密码: " if attempt == 0
                                else "密码错误，请重试: ")
        if pw is None:
            warn("无终端交互且未注入 USBVAULT_ADMIN_PASS，无法验证超管密码。")
            audit("uninstall-attempt", result="fail", detail="无终端")
            return 1
        if verify_admin_password(pw):
            audit("uninstall-attempt", result="ok", detail="密码校验通过")
            return 0
    warn("超管密码连续错误，卸载已中止。")
    audit("uninstall-attempt", result="fail", detail="密码连续错误")
    return 1


# ================= 审计日志 =================

def audit(action, device="", result="ok", detail="", serial="", vidpid=""):
    if HOOKS["audit"]:
        HOOKS["audit"](action, device=device, result=result, detail=detail, serial=serial, vidpid=vidpid)
    rec = {
        "ts": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "user": pwd.getpwuid(os.getuid()).pw_name if os.geteuid() != 0 else "root",
        "action": action,
        "device": device,
        "serial": serial,
        "vidpid": vidpid,
        "result": result,
        "detail": detail,
    }
    try:
        os.makedirs(LOG_DIR, mode=0o750, exist_ok=True)
        with open(AUDIT_FILE, "a", encoding="utf-8") as f:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
        os.chmod(AUDIT_FILE, 0o640)
    except OSError as e:
        warn("审计日志写入失败: %s" % e)
    try:
        subprocess.run(["logger", "-t", "usbvault", "-p", "authpriv.info",
                        "%s %s %s %s" % (action, device, result, detail)],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except OSError:
        pass


# ================= 设备探测 =================

def lsblk_json():
    rc, out, e = run(["lsblk", "-J", "-b", "-o",
                      "NAME,PATH,SIZE,TYPE,FSTYPE,LABEL,UUID,MOUNTPOINT,TRAN,RM,SERIAL,MODEL"])
    if rc != 0:
        return []
    try:
        return json.loads(out).get("blockdevices", [])
    except ValueError:
        return []

def walk_devices(nodes, out=None):
    if out is None:
        out = []
    for d in nodes:
        out.append(d)
        if d.get("children"):
            walk_devices(d["children"], out)
    return out

def mounts_list():
    res = []
    try:
        with open("/proc/mounts", "r") as f:
            for line in f:
                parts = line.split()
                if len(parts) >= 2:
                    res.append((parts[0], parts[1]))
    except OSError:
        pass
    return res

def dev_mountpoints(devpath):
    """返回指定块设备（含其子设备/dm 子设备）的挂载点列表。"""
    names = {os.path.basename(devpath), devpath}
    out = []
    for src, dst in mounts_list():
        base = os.path.basename(src)
        if src == devpath or base in names or src.startswith(devpath):
            out.append(dst)
    return out

def udev_info(devpath):
    rc, out, _ = run(["udevadm", "info", "--query=property", "--name=" + devpath])
    props = {}
    if rc == 0:
        for line in out.splitlines():
            if "=" in line:
                k, v = line.split("=", 1)
                props[k] = v
    return props

def usb_disks(allow_loop=False):
    """返回候选 USB 整盘列表（测试模式可含 loop）。lsblk -J 键为小写。"""
    disks = []
    for d in lsblk_json():
        t = d.get("type")
        if t == "disk" and d.get("tran") == "usb":
            d["_loop"] = False
        elif allow_loop and t == "loop":
            d["_loop"] = True
        else:
            continue
        d["_children"] = d.get("children") or []
        disks.append(d)
    return disks

def disk_guard_errors(d, allow_loop=False):
    """格式化前的安全护栏检查，返回错误列表（空 = 通过）。"""
    errs = []
    path = d.get("path", "")
    if not allow_loop and not d.get("_loop"):
        if d.get("tran") != "usb":
            errs.append("设备总线不是 USB（tran=%s）" % d.get("tran"))
    if d.get("_loop") and not allow_loop:
        errs.append("loop 设备需显式 --allow-loop")
    # 系统盘/已挂载检查：任何子设备挂载即拒绝
    for c in walk_devices(d.get("_children", [])):
        mp = c.get("mountpoint")
        if mp:
            errs.append("包含已挂载分区 %s → %s（系统盘或使用中，禁止格式化）" % (c.get("name"), mp))
    if dev_mountpoints(path):
        errs.append("设备自身处于挂载状态")
    return errs

def disk_summary(d):
    props = udev_info(d.get("path", ""))
    return {
        "path": d.get("path"),
        "name": d.get("name"),
        "size": human_size(d.get("size")),
        "model": d.get("model") or props.get("ID_MODEL", "").replace("_", " ") or "?",
        "serial": props.get("ID_SERIAL_SHORT") or d.get("serial") or "",
        "vidpid": "%s:%s" % (props.get("ID_VENDOR_ID", ""), props.get("ID_MODEL_ID", "")),
        "fstype": d.get("fstype") or props.get("ID_FS_TYPE", ""),
        "loop": d.get("_loop", False),
    }


# ================= LUKS 操作 =================

def is_luks(devpath):
    rc, _, _ = run(["cryptsetup", "isLuks", "--type", "luks2", devpath])
    if rc != 0:
        rc, _, _ = run(["cryptsetup", "isLuks", devpath])
    return rc == 0

def luks_uuid(devpath):
    rc, out, _ = run(["cryptsetup", "luksUUID", devpath])
    return out.strip() if rc == 0 else ""

def mapper_name_for(uuid_str):
    return MAPPER_PREFIX + uuid_str.replace("-", "")[:8]

def luks_open(devpath, key_bytes, name):
    rc, _, e = run(["cryptsetup", "luksOpen", "--key-file", "-", devpath, name],
                   input_bytes=key_bytes, timeout=300)
    return rc, e

def luks_close(name):
    rc, _, e = run(["cryptsetup", "luksClose", name], timeout=120)
    return rc, e

def mapper_path(name):
    return "/dev/mapper/" + name

def active_mappers():
    d = "/dev/mapper"
    if not os.path.isdir(d):
        return []
    return sorted(f for f in os.listdir(d) if f.startswith(MAPPER_PREFIX))

def active_session_user():
    """返回当前活动图形会话的用户名与 uid；无则 (None, None)。"""
    try:
        rc, out, _ = run(["loginctl", "list-sessions", "--no-legend"], timeout=15)
        if rc != 0:
            return None, None
        best = None  # (uid, user, state)
        for line in out.splitlines():
            parts = line.split()
            if len(parts) < 3:
                continue
            sid, uid_s, user = parts[0], parts[1], parts[2]
            try:
                uid = int(uid_s)
            except ValueError:
                continue
            if uid < 1000 or uid >= 60000:
                continue
            rc2, st, _ = run(["loginctl", "show-session", sid, "-p", "State", "--value"],
                             timeout=10)
            state = st.strip() if rc2 == 0 else ""
            score = {"active": 2, "online": 1}.get(state, 0)
            if best is None or score > best[0]:
                best = (score, uid, user)
        if best:
            return best[2], best[1]
    except OSError:
        pass
    return None, None

def _run_as_user(user, cmd, timeout=120, env_extra=None):
    """以指定用户身份执行命令；图形会话需要 DBUS_SESSION_BUS_ADDRESS。"""
    env = os.environ.copy()
    try:
        pw = pwd.getpwnam(user)
        uid = pw.pw_uid
        runtime = "/run/user/%d" % uid
        bus = os.path.join(runtime, "bus")
        if os.path.exists(bus):
            env["DBUS_SESSION_BUS_ADDRESS"] = "unix:path=" + bus
            env["XDG_RUNTIME_DIR"] = runtime
        # runuser 不重置 HOME/USER，若沿用 root 的 HOME，udisksctl/gio 读取用户
        # 配置会失败（表现为 udisks2 路径挂载失败、回退到传统挂载）
        env["HOME"] = pw.pw_dir
        env["USER"] = user
        env["LOGNAME"] = user
        env["SHELL"] = pw.pw_shell or "/bin/bash"
    except KeyError:
        pass
    if env_extra:
        env.update(env_extra)
    full = ["runuser", "-u", user, "--"] + list(cmd)
    try:
        r = subprocess.run(full, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                           timeout=timeout, env=env)
        return (r.returncode,
                r.stdout.decode("utf-8", "replace"),
                r.stderr.decode("utf-8", "replace"))
    except (subprocess.TimeoutExpired, FileNotFoundError) as e:
        return 124, "", str(e)

def mount_opts(cfg):
    opts = [
        "uid=%d" % int(cfg["mount_uid"]),
        "gid=%d" % int(cfg["mount_gid"]),
        "fmask=%s" % cfg["fmask"],
        "dmask=%s" % cfg["dmask"],
        "allow_other",
        "utf8",
    ]
    if cfg.get("atime"):
        opts.append(str(cfg["atime"]))
    if cfg.get("noexec"):
        opts.append("noexec")
    return ",".join(opts)

def pick_mountpoint(uuid_str):
    """选择挂载点目录：优先 FS 标签，冲突时用 uuid8 后缀。"""
    base = FS_LABEL
    mp = os.path.join(MEDIA_DIR, base)
    if not os.path.ismount(mp):
        return mp
    suffix = uuid_str.replace("-", "")[:8]
    return os.path.join(MEDIA_DIR, "%s-%s" % (base, suffix))

def _udisks_mount(mapper, retries=12):
    """通过 udisksctl 以活动会话用户身份挂载，让桌面文件管理器识别为本地移动设备。
    成功返回挂载点；失败返回 None。
    注意：mapper 刚创建时 udisks2 尚未把该 dm 设备纳入其对象树（轮询周期约 2s），
    因此先轮询 udisksctl info 确认设备已被识别，再执行 mount。"""
    user, _ = active_session_user()
    if not user:
        warn("未检测到活动图形会话用户，改用传统挂载方式")
        return None
    dev = mapper_path(mapper)

    # 阶段一：等待 udisks2 注册该 dm 设备（最多约 24 秒）
    known = False
    for _ in range(retries):
        rc, _, _ = _run_as_user(user, ["udisksctl", "info", "-b", dev],
                                timeout=30)
        if rc == 0:
            known = True
            break
        time.sleep(2)
    if not known:
        warn("udisks2 未识别 %s，改用传统挂载方式" % dev)
        return None

    # 阶段二：挂载
    last_err = ""
    for attempt in range(retries):
        rc, out, e = _run_as_user(user,
            ["udisksctl", "mount", "-b", dev], timeout=90)
        if rc == 0:
            m = re.search(r"(?:Mounted|已挂载)\s+\S+\s+(?:at|于)\s+(.+?)[.\s]*$", out.strip())
            mp = (m.group(1).strip() if m else None) or find_mount_of_mapper(mapper)
            if mp:
                return mp
        last_err = (e + out).strip()
        # 已挂载直接成功
        if "already mounted" in last_err or "已挂载" in last_err:
            return find_mount_of_mapper(mapper)
        time.sleep(1)
    warn("udisksctl 挂载失败（%s），改用传统挂载方式" % (last_err[:120] or "超时"))
    return None

def _udisks_unmount(mapper, user=None):
    """通过 udisksctl 卸载（保持桌面侧状态一致）；失败返回 False。"""
    if user is None:
        user, _ = active_session_user()
    if not user:
        return False
    rc, _, _ = _run_as_user(user,
        ["udisksctl", "unmount", "-b", mapper_path(mapper)], timeout=90)
    return rc == 0

def do_mount(mapper, uuid_str):
    # 优先走 udisks2：挂载点落到 /media/<user>/LABEL，文件管理器正常识别
    # 仅在默认 MEDIA_DIR（未通过 USBVAULT_MEDIA 隔离，即真实使用场景）下启用；
    # 测试环境通过 USBVAULT_MEDIA 指到 /tmp 时使用传统 mount，便于断言挂载点。
    if os.environ.get("USBVAULT_MEDIA") is None:
        mp = _udisks_mount(mapper)
        if mp:
            return mp
    # 回退：直接 mount（无活动图形会话/udisks2 不可用时）
    cfg = load_config()
    mp = pick_mountpoint(uuid_str)
    os.makedirs(MEDIA_DIR, mode=0o755, exist_ok=True)
    os.makedirs(mp, mode=0o755, exist_ok=True)
    rc, _, e = run(["mount", "-t", "ntfs-3g", "-o", mount_opts(cfg),
                    mapper_path(mapper), mp], timeout=300)
    if rc != 0:
        err("挂载失败: %s" % e.strip())
        return None
    return mp

def do_unmount(mp, mapper=None):
    run(["sync"], timeout=300)
    if mapper:
        try:
            _udisks_unmount(mapper)
            if not os.path.ismount(mp):
                return True, ""
        except OSError:
            pass
    for i in range(3):
        rc, _, e = run(["umount", mp], timeout=120)
        if rc == 0:
            return True, ""
        time.sleep(1)
    rc, _, e = run(["umount", "-l", mp], timeout=120)
    return rc == 0, e.strip()

def find_mount_of_mapper(mapper):
    src = mapper_path(mapper)
    for s, dst in mounts_list():
        if s == src:
            return dst
    return None

def gen_recovery_key():
    """生成 160bit 恢复密钥，返回 (raw40, 分组显示)。"""
    raw = secrets.token_hex(20)  # 40 个十六进制字符
    shown = "-".join(raw[i:i+5] for i in range(0, 40, 5))
    return raw, shown


# ================= 设备管控（guard） =================
# 借鉴 USBGuard 的内核 authorized 机制：命中外接 USB 存储整盘时由 udev 85 规则
# 同步调用 guard-check 判定；非"本机同编码注册盘"写 authorized=0 整设备去授权
# （设备从系统消失，不生成块设备、不触发桌面提示）。鼠标/键盘等 HID 外设不产生
# sd* 块设备事件，从入口隔离、不受影响。

GUARD_ALLOW_TTL = 24 * 3600  # 临时放行标记有效期（秒）

def guard_enabled(cfg):
    return bool(cfg.get("guard_enabled", True))

def _usb_device_node_for(kname):
    """沿 /sys/class/block/<kname>/device 向上查找 usb_device 节点（含 idVendor）。
    找不到（如 loop 盘、内置读卡器）返回 None。"""
    node = os.path.realpath("/sys/class/block/%s/device" % kname)
    while node and node != "/":
        if os.path.isfile(os.path.join(node, "idVendor")):
            return node
        node = os.path.dirname(node)
    return None

def _usb_node_by_serial(serial):
    """按序列号在 /sys/bus/usb/devices/ 反查 usb_device 节点。
    用于设备被 authorized=0 阻断后块节点已消失时的恢复授权。"""
    if not serial:
        return None
    base = "/sys/bus/usb/devices"
    try:
        names = os.listdir(base)
    except OSError:
        return None
    for name in names:
        node = os.path.join(base, name)
        try:
            with open(os.path.join(node, "serial")) as f:
                if f.read().strip() == serial and \
                        os.path.isfile(os.path.join(node, "authorized")):
                    return node
        except OSError:
            continue
    return None

def _usb_set_authorized(node, val):
    """写 USB 设备 authorized（val: "0"/"1"），成功返回 True。"""
    try:
        with open(os.path.join(node, "authorized"), "w") as f:
            f.write(val)
        return True
    except OSError:
        return False

def guard_blocked_list():
    """读取当前被拦设备记录列表（/run/usbvault/guard-blocked-*.json）。"""
    out = []
    if not os.path.isdir(RUN_DIR):
        return out
    for f in sorted(os.listdir(RUN_DIR)):
        if not re.match(r"guard-blocked-.+\.json$", f):
            continue
        try:
            with open(os.path.join(RUN_DIR, f), "r", encoding="utf-8") as fh:
                rec = json.load(fh)
            if isinstance(rec, dict):
                out.append(rec)
        except (OSError, ValueError):
            continue
    return out

def guard_block(kname, serial="", vidpid="", reason=""):
    """阻断：USB 设备写 authorized=0 整设备去授权（从系统消失），并记录留痕。
    非 USB 设备（无 sysfs usb_device 节点，如 loop）无法去授权，仅记录。"""
    node = _usb_device_node_for(kname)
    deauth = _usb_set_authorized(node, "0") if node else False
    rec = {
        "ts": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "kname": kname,
        "serial": serial or kname,
        "vidpid": vidpid,
        "reason": reason,
        "deauthorized": deauth,
    }
    try:
        os.makedirs(RUN_DIR, exist_ok=True)
        with open(os.path.join(RUN_DIR, "guard-blocked-%s.json" % kname),
                  "w", encoding="utf-8") as f:
            json.dump(rec, f, ensure_ascii=False)
    except OSError:
        pass
    audit("guard-block", device="/dev/" + kname,
          result="ok" if deauth else "fail",
          detail=reason + ("" if deauth else "（无 USB 节点，未去授权）"),
          serial=serial, vidpid=vidpid)
    return deauth

def _disk_kname_of(kname):
    """分区名（sdb1）→ 整盘名（sdb）；整盘名原样返回；其他（loop0、dm-0）原样返回。"""
    m = re.match(r"^((?:sd[a-z]+|mmcblk\d+|nvme\d+n\d+))(p?\d+)$", kname)
    return m.group(1) if m else kname

def _disk_in_use(kname):
    """整盘或其任一分区当前已挂载时返回 True。
    用于保护系统盘/在用盘：从 U 盘启动或根分区在该盘上时绝不能被去授权。
    udev add 事件同步执行、早于 udisks2 挂载，正常情况下不会命中；
    开机 coldplug 阶段根分区已挂载，此判断可兜底。"""
    disk = _disk_kname_of(kname)
    try:
        with open("/proc/mounts", "r", encoding="utf-8", errors="replace") as f:
            for line in f:
                parts = line.split()
                if len(parts) < 2 or not parts[0].startswith("/dev/"):
                    continue
                name = os.path.basename(os.path.realpath(parts[0]))
                if _disk_kname_of(name) == disk:
                    return True
    except OSError:
        return False
    return False

def guard_is_allowed(kname, serial=""):
    """存在未过期（24h）的临时放行标记则消费之并返回 True。
    主键用序列号（重新枚举后内核设备名可能变化），无序列号时兜底 kname。"""
    p = os.path.join(RUN_DIR, "guard-allow-%s" % (serial or kname))
    try:
        st = os.stat(p)
    except OSError:
        return False
    fresh = (time.time() - st.st_mtime) < GUARD_ALLOW_TTL
    try:
        os.unlink(p)
    except OSError:
        pass
    return fresh

def guard_session_allowed(serial):
    """临时放行后的"本次会话"标记（24h）：一次性放行标记被整盘事件消费后，
    该盘分区事件/重新枚举仍应放行，不再重复阻断。"""
    if not serial:
        return False
    p = os.path.join(RUN_DIR, "guard-allowed-%s" % serial)
    try:
        st = os.stat(p)
    except OSError:
        return False
    return (time.time() - st.st_mtime) < GUARD_ALLOW_TTL

def guard_write_session_allow(serial):
    if not serial:
        return
    try:
        os.makedirs(RUN_DIR, exist_ok=True)
        with open(os.path.join(RUN_DIR, "guard-allowed-%s" % serial), "w") as f:
            f.write(str(time.time()))
    except OSError:
        pass

def cmd_guard_check(args):
    """由 udev 85 规则同步调用：判定外接 USB 存储是否放行（退出码恒为 0）。
    RUN 同步执行，判定在 udisks2 感知设备之前完成，无挂载窗口。"""
    require_root()
    kname = os.path.basename(args.kname)
    dev = "/dev/" + kname
    os.makedirs(RUN_DIR, exist_ok=True)
    with open(os.path.join(RUN_DIR, ".guard.lock"), "w") as lf:
        fcntl.flock(lf, fcntl.LOCK_EX)
        _guard_check_inner(kname, dev)
    return 0

def _guard_check_inner(kname, dev):
    # 分区事件（sdb1）按整盘（sdb）判定与阻断：阻断单位始终是整台 USB 设备
    disk = _disk_kname_of(kname)
    disk_dev = "/dev/" + disk
    dprops = udev_info(disk_dev) or udev_info(dev) or {}
    serial = dprops.get("ID_SERIAL_SHORT") or disk
    vidpid = "%s:%s" % (dprops.get("ID_VENDOR_ID", ""), dprops.get("ID_MODEL_ID", ""))
    # ① 未登记编码：管控不生效（避免未部署先锁死）
    if not os.path.isfile(KEY_FILE):
        return
    # ② 管控开关关闭
    if not guard_enabled(load_config()):
        return
    # ③ 系统盘 / 在用盘保护：整盘或任一分区已挂载（如从 U 盘启动）时绝不阻断
    if _disk_in_use(kname):
        audit("guard-skip", device=dev, result="skip",
              detail="整盘或分区正在挂载（系统盘保护）", serial=serial, vidpid=vidpid)
        return
    # ④ 临时放行：一次性标记（消费后写会话标记，供该盘后续分区/重插事件使用）
    if guard_session_allowed(serial):
        return
    if guard_is_allowed(kname, serial):
        guard_write_session_allow(serial)
        audit("guard-permit", device=dev, result="ok", detail="临时放行标记",
              serial=serial, vidpid=vidpid)
        return
    # ⑤ 密钥读取异常：fail-closed 阻断
    try:
        with open(KEY_FILE, "rb") as f:
            key = f.read()
    except OSError as e:
        guard_block(disk, serial, vidpid, "unit.key 读取失败（fail-closed）: %s" % e)
        return
    if len(key) != KEY_LEN:
        guard_block(disk, serial, vidpid, "unit.key 损坏（fail-closed）")
        return
    # ⑥ 同编码注册盘判定：只认整盘 LUKS2。分区盘（普通 U 盘/Ventoy/分区级加密）
    #    一律不可能是本系统的注册盘，直接阻断，不做 Argon2id 校验白耗 1 秒。
    if disk == kname:
        rc, _, _ = run(["cryptsetup", "open", "--test-passphrase", "--type", "luks2",
                        "--key-file", KEY_FILE, disk_dev], timeout=120)
        if rc == 0:
            return  # 放行，90 规则 automount 接管自动解锁挂载
        reason = "非本机同编码加密注册盘"
    else:
        reason = "非本机同编码加密注册盘（分区盘）"
    # ⑦ 阻断：写 authorized=0，整台 USB 设备从系统消失
    guard_block(disk, serial, vidpid, reason)

def cmd_guard(args):
    """设备管控：on/off/status/blocked/allow。"""
    require_root()
    action = getattr(args, "guard_action", None) or "status"
    kname = getattr(args, "kname", None)
    if action in ("off", "allow"):
        op_name = {"off": "关闭管控", "allow": "临时放行"}[action]
        require_admin(op_name)
    if action in ("on", "off"):
        cfg = load_config()
        cfg["guard_enabled"] = (action == "on")
        save_config(cfg)
        ok("设备管控已%s" % ("开启" if action == "on" else "关闭"))
        if action == "on" and not os.path.isfile(KEY_FILE):
            warn("本机尚未登记编码，管控暂不会实际生效（登记编码后自动生效）。")
        json_set(guard_enabled=(action == "on"))
        audit("guard", result="ok", detail="guard %s" % action)
        return 0
    if action == "status":
        en = guard_enabled(load_config())
        registered = os.path.isfile(KEY_FILE)
        print()
        print("%s设备管控（guard）%s" % (C_CYA, C_RST))
        print("─" * 62)
        print("管控开关 : %s" % ((C_GRN + "开启" + C_RST) if en else (C_RED + "关闭" + C_RST)))
        print("编码登记 : %s" % (("已登记（指纹 %s）" % key_fingerprint()) if registered
                                   else (C_RED + "未登记" + C_RST)))
        print("实际生效 : %s" % ("是 —— 仅识别同编码注册的外接 USB 存储" if (en and registered)
                                   else "否 —— 不阻断任何设备"))
        blocked = guard_blocked_list()
        print("被拦设备 : %d 台次（guard blocked 查看详情）" % len(blocked))
        json_set(guard_enabled=en, unit_registered=registered,
                 guard_blocked=len(blocked),
                 unit_fingerprint=key_fingerprint() if registered else "")
        return 0
    if action == "blocked":
        blocked = guard_blocked_list()
        json_set(blocked=blocked)
        if not blocked:
            info("当前没有被拦截的设备记录。")
            return 0
        for rec in blocked:
            print("%-19s %-7s %-22s %-11s %s" % (
                rec.get("ts", "?"), rec.get("kname", "?"),
                rec.get("serial") or "-", rec.get("vidpid") or "-",
                rec.get("reason", "")))
        return 0
    if action == "allow":
        if not kname:
            err("用法: sudo usbvault guard allow <设备名>（如 sdb；用 guard blocked 查看）")
            return 1
        kname = os.path.basename(kname)
        # 被拦记录里的序列号作为主键（重新枚举后内核设备名可能变化）
        serial = ""
        for rec in guard_blocked_list():
            if rec.get("kname") == kname:
                serial = rec.get("serial") or ""
                break
        os.makedirs(RUN_DIR, exist_ok=True)
        with open(os.path.join(RUN_DIR, "guard-allow-%s" % (serial or kname)), "w") as f:
            f.write(str(time.time()))
        ok("已临时放行 /dev/%s（序列号 %s；24 小时内、一次性插入有效）"
           % (kname, serial or "未知"))
        # 设备仍在线（authorized=0 状态）时立即恢复授权，触发重新枚举。
        # 被阻断后块节点已消失，_usb_device_node_for 查不到，需按序列号反查 USB 节点
        node = _usb_device_node_for(kname) or _usb_node_by_serial(serial)
        if node and _usb_set_authorized(node, "1"):
            info("设备已重新授权，系统正在重新识别…")
        else:
            info("请重新插拔设备使其生效。")
        try:
            os.unlink(os.path.join(RUN_DIR, "guard-blocked-%s.json" % kname))
        except OSError:
            pass
        json_set(allow_kname=kname, allow_serial=serial)
        audit("guard-allow", device="/dev/" + kname, result="ok",
              detail="临时放行 24h", serial=serial)
        return 0
    err("未知 guard 操作: %s（支持 on/off/status/blocked/allow）" % action)
    return 1


# ================= 命令实现 =================

def cmd_init(args):
    require_root()
    require_admin("登记/更新单位编码")
    step("单位编码登记")
    old_exists = os.path.isfile(KEY_FILE)
    if old_exists:
        warn("本机已登记编码（指纹 %s）。" % key_fingerprint())
        warn("重新登记将使所有以旧编码格式化的 U 盘无法自动解锁！")
        warn("若需换码并保留旧盘数据，请使用: sudo usbvault rekey")
        if not confirm("确定要重新登记吗", args.yes):
            err("已取消。")
            return 1
    code = ask_code()
    code2 = ask_code("请再次输入编码确认")
    if not hmac.compare_digest(code, code2):
        err("两次输入不一致，已取消。")
        audit("init", result="fail", detail="两次输入不一致")
        return 1
    key = derive_key(code)
    write_unit_key(key)
    os.chmod(ETC_DIR, 0o700)
    ok("已登记单位编码，密钥写入 %s (root:0600)" % KEY_FILE)
    info("编码明文未落盘。请在单位保密册上备份此编码！")
    json_set(unit_registered=True, unit_fingerprint=key_fingerprint())
    audit("init", result="ok", detail="编码登记/更新")
    return 0

def cmd_format(args):
    require_root()
    require_admin("加密格式化")
    dev = os.path.abspath(args.device)
    if not os.path.exists(dev):
        err("设备不存在: %s" % dev)
        return 1
    disks = usb_disks(allow_loop=args.allow_loop)
    d = next((x for x in disks if x.get("path") == dev), None)
    if d is None:
        # 给出具体原因
        cand = lsblk_json()
        raw = next((x for x in walk_devices(cand) if x.get("path") == dev), None)
        if raw is None:
            err("未找到块设备: %s" % dev)
        elif raw.get("type") == "loop" and not args.allow_loop:
            err("loop 设备需显式 --allow-loop（测试用）。")
        elif raw.get("type") != "disk":
            err("仅支持整盘格式化（%s 是 %s），请使用整个磁盘设备。" % (dev, raw.get("type")))
        else:
            err("设备 %s 不是 USB 存储（tran=%s）。" % (dev, raw.get("tran")))
        return 1
    errs = disk_guard_errors(d, allow_loop=args.allow_loop)
    if errs:
        for e in errs:
            err("安全检查未通过: %s" % e)
        audit("format", device=dev, result="reject", detail="; ".join(errs))
        return 1
    s = disk_summary(d)
    key = load_unit_key()

    if args.dry_run:
        step("［预览模式 dry-run］将执行以下操作，不改动任何数据：")
        for line in (
            "1. wipefs 清除 %s 及其子分区的所有签名" % dev,
            "2. LUKS2 整盘格式化（%s / %dbit / Argon2id）" % (LUKS_CIPHER, LUKS_KEYSIZE),
            "3. 生成恢复密钥与头部备份（%s）" % VAR_DIR,
            "4. 盘内创建 NTFS（LABEL=%s）" % FS_LABEL,
            "5. 解锁并挂载到 %s" % MEDIA_DIR,
        ):
            info(line)
        info("目标设备: %s  %s  %s  序列号 %s  当前类型 %s" %
             (dev, s["model"], s["size"], s["serial"] or "?", s["fstype"] or "空"))
        return 0

    print()
    warn("=" * 62)
    warn("即将加密格式化以下设备，盘上所有数据将被销毁：")
    warn("  设备   : %s" % dev)
    warn("  型号   : %s" % s["model"])
    warn("  容量   : %s" % s["size"])
    warn("  序列号 : %s" % (s["serial"] or "未知"))
    warn("  当前   : %s" % (s["fstype"] or "无文件系统"))
    if s["fstype"] == "crypto_LUKS":
        warn("  注意   : 该盘已是 LUKS 加密卷！")
    warn("=" * 62)
    if not s["loop"] and s["serial"] and not args.yes:
        expected_tail = s["serial"][-4:].lower()
        tail = ask("请输入序列号后 4 位以确认（序列号后 4 位是 %s）: " % expected_tail)
        # 容错：用户可能输入多于 4 位（如提示显示后 8 位时用户取后 5 位），
        # 统一取用户输入的最后 4 位与序列号后 4 位比较
        user_input = tail.strip().lower()
        user_tail = user_input[-4:] if len(user_input) >= 4 else user_input
        if user_tail != expected_tail:
            err("序列号确认不符（期望后 4 位 %s，实际 %s），已取消。" % (expected_tail, user_tail))
            audit("format", device=dev, result="reject", detail="序列号确认不符")
            return 1
    elif not confirm("确认继续", args.yes):
        err("已取消。")
        return 1

    # 1) 擦除旧签名（先子分区后整盘）
    children = [c.get("path") for c in d.get("_children", []) if c.get("path")]
    old_sigs = []
    step("擦除旧签名")
    for p in children + [dev]:
        rc, out, e = run(["wipefs", "-a", p])
        if out.strip():
            old_sigs.append(out.strip())
        if rc != 0:
            warn("wipefs %s: %s" % (p, e.strip()[:120]))
    if old_sigs:
        info("检测到旧数据签名（盘内可能有历史明文数据）")

    # 2) LUKS2 格式化
    step("创建 LUKS2 加密卷（Argon2id，约需十几秒）")
    rc, _, e = run(["cryptsetup", "luksFormat",
                    "--type", "luks2",
                    "--cipher", LUKS_CIPHER,
                    "--key-size", str(LUKS_KEYSIZE),
                    "--pbkdf", "argon2id",
                    "--batch-mode",
                    "--key-file", "-",
                    dev], input_bytes=key, timeout=600)
    if rc != 0:
        err("LUKS2 格式化失败: %s" % e.strip())
        audit("format", device=dev, result="fail", detail="luksFormat: " + e.strip()[:200])
        return 1
    uuid_str = luks_uuid(dev)
    mapper = mapper_name_for(uuid_str)
    ok("LUKS2 卷已创建 UUID=%s" % uuid_str)

    # 3) 恢复密钥 + 头部备份
    step("生成恢复密钥与头部备份")
    os.makedirs(VAR_DIR, mode=0o700, exist_ok=True)
    rec_raw, rec_shown = gen_recovery_key()
    json_set(device=dev, mapper=mapper, uuid=uuid_str,
             recovery_key=rec_shown,
             recovery_file=os.path.join(VAR_DIR, "recovery-%s.key" % mapper),
             header_backup=os.path.join(VAR_DIR, "header-%s.img" % mapper))
    rec_file = os.path.join(VAR_DIR, "recovery-%s.key" % mapper)
    with open(rec_file, "w") as f:
        f.write(rec_raw)
    os.chmod(rec_file, 0o600)
    rc, _, e = run(["cryptsetup", "luksAddKey", dev, rec_file, "--key-file", "-"],
                   input_bytes=key, timeout=600)
    if rc != 0:
        warn("恢复密钥槽添加失败（可继续，但编码丢失将无法恢复）: %s" % e.strip())
    else:
        ok("恢复密钥已保存: %s" % rec_file)
        info("恢复密钥（请离线抄存）: %s" % rec_shown)
    hdr_file = os.path.join(VAR_DIR, "header-%s.img" % mapper)
    rc, _, e = run(["cryptsetup", "luksHeaderBackup", dev,
                    "--header-backup-file", hdr_file], timeout=120)
    if rc == 0:
        os.chmod(hdr_file, 0o600)
        ok("LUKS 头部已备份: %s" % hdr_file)
    else:
        warn("头部备份失败: %s" % e.strip())

    # 4) 打开 + 覆写（历史明文盘） + NTFS
    step("解锁加密卷")
    rc, e = luks_open(dev, key, mapper)
    if rc != 0:
        err("解锁失败: %s" % e.strip())
        audit("format", device=dev, result="fail", detail="luksOpen: " + e.strip()[:200])
        return 1
    try:
        mpath = mapper_path(mapper)
        if old_sigs and not args.quick:
            warn("检测到历史数据，正在全盘密文覆写以清除历史明文（USB3 约 1-5 分钟）…")
            # LUKS2 头部占 16MB，预留 32MB 防止 NTFS 空间不足
            size_bytes = int(d.get("size", 0))
            safe_size = max(0, size_bytes - 32 * 1024 * 1024)
            cmd_dd = ["dd", "if=/dev/zero", "of=%s" % mpath,
                      "bs=4M", "count=%d" % (safe_size // (4 * 1024 * 1024)),
                      "conv=fsync"]
            if sys.stdout.isatty():
                cmd_dd.append("status=progress")
            rc, _, e = run(cmd_dd, timeout=3600)
            if rc != 0:
                warn("覆写中断（不影响加密卷可用性）: %s" % e.strip()[:120])
            else:
                ok("全盘覆写完成，历史明文已清除")
        else:
            info("全新介质，跳过覆写（--quick 语义）")
        step("创建 NTFS 文件系统 (LABEL=%s)" % FS_LABEL)
        rc, _, e = run(["mkntfs", "-f", "-L", FS_LABEL, mpath], timeout=1800)
        if rc != 0:
            err("NTFS 创建失败: %s" % e.strip())
            audit("format", device=dev, result="fail", detail="mkntfs: " + e.strip()[:200])
            return 1
        ok("NTFS 文件系统已创建")
    finally:
        rc, e = luks_close(mapper)
        if rc != 0:
            warn("临时关闭失败: %s" % e.strip())

    # 5) 解锁挂载
    if not args.no_mount:
        step("解锁并挂载")
        rc, e = luks_open(dev, key, mapper)
        if rc != 0:
            warn("自动挂载失败，可手动: sudo usbvault open")
        else:
            state_save(d.get("NAME", ""), mapper)
            mp = do_mount(mapper, uuid_str)
            if mp:
                ok("已挂载: %s → %s（本机所有用户可读写）" % (mapper, mp))
                json_set(mountpoint=mp)
    json_set(device=dev, mapper=mapper, uuid=uuid_str)
    ok("加密格式化完成")
    audit("format", device=dev, result="ok",
          serial=s["serial"], vidpid=s["vidpid"],
          detail="LUKS2+NTFS uuid=%s mapper=%s" % (uuid_str, mapper))
    return 0

def cmd_open(args):
    require_root()
    # 候选：参数指定 / 未锁定盘上的 LUKS
    if args.device:
        dev = os.path.abspath(args.device)
    else:
        cands = []
        for d in usb_disks(allow_loop=True):
            if d.get("fstype") == "crypto_LUKS":
                cands.append(d)
        if not cands:
            err("未发现 LUKS 加密 USB 设备。")
            return 1
        if len(cands) == 1:
            dev = cands[0]["path"]
        else:
            info("发现多个加密设备：")
            for i, c in enumerate(cands, 1):
                info("%d) %s  %s  %s" % (i, c["path"], human_size(c.get("size")), c.get("model", "")))
            idx = ask("选择 [1-%d]: " % len(cands)).strip()
            if not idx.isdigit() or not (1 <= int(idx) <= len(cands)):
                err("无效选择。")
                return 1
            dev = cands[int(idx) - 1]["path"]
    if not os.path.exists(dev):
        err("设备不存在: %s" % dev)
        return 1
    uuid_str = luks_uuid(dev)
    if not uuid_str:
        err("%s 不是 LUKS 加密卷。" % dev)
        return 1
    mapper = mapper_name_for(uuid_str)
    if os.path.exists(mapper_path(mapper)):
        warn("该卷已解锁（%s）。" % mapper)
        mp = find_mount_of_mapper(mapper)
        if mp:
            info("挂载点: %s" % mp)
        return 0
    # 先用本机密钥
    key = load_unit_key(required=False)
    rc, e = (luks_open(dev, key, mapper) if key else (1, "no key"))
    if rc != 0:
        # 回退：输入编码
        info("本机密钥无法解锁，请输入该盘的编码。")
        code = ask_code("请输入编码")
        rc, e = luks_open(dev, derive_key(code), mapper)
        if rc != 0:
            err("解锁失败（编码错误或非本单位盘）: %s" % e.strip()[:120])
            audit("open", device=dev, result="fail", detail="密钥不匹配")
            return 1
    state_save(os.path.basename(dev), mapper)
    mp = do_mount(mapper, uuid_str)
    if mp:
        ok("已解锁并挂载: %s → %s" % (mapper, mp))
        audit("open", device=dev, result="ok", detail="mapper=%s" % mapper)
    return 0

def cmd_close(args):
    require_root()
    targets = []
    if args.target and args.target != "all":
        targets = [args.target]
    else:
        targets = active_mappers()
        if not targets:
            info("当前没有已解锁的 usbvault 卷。")
            return 0
    rc_all = 0
    for t in targets:
        mapper = t if t.startswith(MAPPER_PREFIX) else None
        if mapper is None:
            # t 可能是内核设备名（如 sda），查状态文件
            st = state_load(t)
            mapper = st.get("mapper") if st else None
        if not mapper:
            err("无法确定卷名: %s" % t)
            rc_all = 1
            continue
        close_one(mapper, kname=None, manual=True)
    return rc_all

def close_one(mapper, kname=None, manual=False):
    mp = find_mount_of_mapper(mapper)
    step("锁定 %s" % mapper)
    if not os.path.exists(mapper_path(mapper)):
        # 幂等：卷已不在（例如重复停止服务、设备已拔出并清理完毕）
        info("卷 %s 已处于锁定状态" % mapper)
        if kname:
            state_clear(kname)
        return 0
    if mp:
        okc, e = do_unmount(mp, mapper=mapper)
        if okc:
            ok("已卸载 %s" % mp)
        else:
            warn("卸载失败: %s" % e)
    rc, e = luks_close(mapper)
    if rc != 0 and backing_gone(mapper):
        # 仅当底层块设备确已消失（真实拔盘）时才强制移除：
        # 此时数据已不可能再写入，移除僵尸映射不会造成数据损坏。
        rc2, _, e2 = run(["dmsetup", "remove", "-f", mapper_path(mapper)], timeout=60)
        if rc2 == 0:
            warn("已强制移除残留映射 %s（设备已拔出）" % mapper)
            rc, e = 0, ""
        else:
            e = e2 or e
    if rc == 0:
        ok("加密卷已锁定")
        audit("close", device=mapper, result="ok", detail="manual" if manual else "auto")
    else:
        err("锁定失败: %s" % e.strip()[:120])
        audit("close", device=mapper, result="fail", detail=e.strip()[:120])
        return 1
    if kname:
        flag = os.path.join(RUN_DIR, "manual-%s" % kname)
        try:
            os.makedirs(RUN_DIR, exist_ok=True)
            with open(flag, "w") as f:
                f.write("1")
        except OSError:
            pass
    if kname:
        state_clear(kname)
    return 0

def cmd_automount(args):
    """由 systemd 服务调用：自动解锁并挂载（幂等）。"""
    require_root()
    kname = args.kname
    dev = "/dev/" + kname
    lockf = os.path.join(RUN_DIR, "auto-%s.lock" % kname)
    os.makedirs(RUN_DIR, exist_ok=True)
    with open(lockf, "w") as lf:
        fcntl.flock(lf, fcntl.LOCK_EX)
        # 等待 udev 提供 ID_FS_TYPE
        props = {}
        for _ in range(20):
            props = udev_info(dev)
            if props.get("ID_FS_TYPE") == "crypto_LUKS":
                break
            time.sleep(0.5)
        if props.get("ID_FS_TYPE") != "crypto_LUKS":
            # udev 数据库可能尚未就绪（开机早期/测试环境）：直接探测 LUKS 签名
            rc, _, _ = run(["cryptsetup", "isLuks", dev], timeout=20)
            if rc != 0:
                return 0  # 非本单位加密盘，静默退出
        if os.path.exists(os.path.join(RUN_DIR, "manual-%s" % kname)):
            audit("automount", device=dev, result="skip", detail="手动锁定标志")
            return 0
        # 幂等：已解锁挂载且后端确实是当前设备时退出（防陈旧映射误判）
        cur_dev_num = dev_num(dev)
        for m in active_mappers():
            if find_mount_of_mapper(m):
                st = state_load(kname)
                if (st and st.get("mapper") == m
                        and cur_dev_num is not None
                        and mapper_backing_dev(m) == cur_dev_num):
                    return 0
        if not os.path.isfile(KEY_FILE):
            audit("automount", device=dev, result="fail", detail="本机未登记编码")
            return 1
        with open(KEY_FILE, "rb") as f:
            key = f.read()
        if len(key) != KEY_LEN:
            audit("automount", device=dev, result="fail", detail="密钥文件损坏")
            return 1
        uuid_str = luks_uuid(dev)
        if not uuid_str:
            audit("automount", device=dev, result="fail", detail="读取UUID失败")
            return 1
        mapper = mapper_name_for(uuid_str)
        need_open = True
        if os.path.exists(mapper_path(mapper)):
            cur = dev_num(dev)
            back = mapper_backing_dev(mapper)
            if cur is not None and back == cur:
                need_open = False   # 同一物理设备已解锁，直接复用
            else:
                # 拔盘后残留的僵尸映射（后端指向已消失的块设备），强制清理后重新解锁。
                # 否则会把已失效的 mapper 挂上去，表现为“U盘无法识别/读写 I/O 错误”。
                warn("发现陈旧映射 %s（后端 %s，当前设备 %s），强制清理后重新解锁"
                     % (mapper, back, cur))
                force_close_mapper(mapper)
        if need_open:
            rc, e = luks_open(dev, key, mapper)
            if rc != 0:
                audit("automount", device=dev, result="fail", detail="密钥不匹配或解锁失败")
                return 1
        mp = find_mount_of_mapper(mapper)
        if not mp:
            mp = do_mount(mapper, uuid_str)
        if mp:
            state_save(kname, mapper)
            audit("automount", device=dev, result="ok",
                  serial=props.get("ID_SERIAL_SHORT", ""),
                  vidpid="%s:%s" % (props.get("ID_VENDOR_ID", ""), props.get("ID_MODEL_ID", "")),
                  detail="mapper=%s mp=%s" % (mapper, mp))
        else:
            audit("automount", device=dev, result="fail", detail="挂载失败")
            return 1
    return 0

def cmd_lock(args):
    """设备拔出 / 服务停止时卸载并锁定（udev remove 事件调用）。"""
    require_root()
    kname = args.kname
    st = state_load(kname)
    mapper = st.get("mapper") if st else None
    if mapper and not os.path.exists(mapper_path(mapper)):
        # 映射已不存在（重复 remove / 已被清理）：清掉状态文件后继续兜底
        state_clear(kname)
        mapper = None
    if mapper and getattr(args, "udev", False) and not backing_gone(mapper):
        # udev remove 是异步事件：若该内核名已被新插入的设备复用，
        # 此时映射的后端仍在线，绝不能误锁正在使用的卷。
        audit("close", device=mapper, result="skip", detail="后端设备仍在线（设备名复用）")
        return 0
    if not mapper:
        # 兜底：找 backing 设备已消失的孤儿 mapper
        for m in active_mappers():
            if backing_gone(m):
                mapper = m
                break
    if mapper:
        return close_one(mapper, kname=None, manual=False)
    state_clear(kname)
    return 0

def dev_num(devpath):
    """返回块设备的 (major, minor)，设备不存在返回 None。"""
    try:
        st = os.stat(devpath)
    except OSError:
        return None
    if not stat_mod.S_ISBLK(st.st_mode):
        return None
    return (os.major(st.st_rdev), os.minor(st.st_rdev))

def mapper_backing_dev(mapper):
    """返回 mapper 依赖的底层块设备 (major, minor)，读不到返回 None。"""
    rc, out, _ = run(["dmsetup", "deps", mapper_path(mapper)], timeout=30)
    if rc != 0:
        return None
    m = re.search(r"\((\d+),\s*(\d+)\)", out)
    return (int(m.group(1)), int(m.group(2))) if m else None

def force_close_mapper(mapper):
    """强制清理 mapper：先卸载其所有挂载点，再 luksClose，失败则 dmsetup remove -f。
    用于处理拔盘后残留的 dm-error 僵尸映射（后端块设备已消失）。"""
    mp = find_mount_of_mapper(mapper)
    if mp:
        do_unmount(mp, mapper=mapper)
    rc, _ = luks_close(mapper)
    if rc == 0:
        return True
    back = mapper_backing_dev(mapper)
    if back is not None and not backing_gone(mapper):
        # 后端设备仍在线（例如同 UUID 的另一块盘），不能强删，避免破坏在用数据
        warn("映射 %s 仍被占用，无法清理" % mapper)
        return False
    # 后端已消失（或映射本身无依赖关系，属失效映射），用 dmsetup 强制移除
    rc, _, e = run(["dmsetup", "remove", "-f", mapper_path(mapper)], timeout=60)
    if rc == 0:
        warn("已强制移除陈旧映射 %s" % mapper)
        return True
    warn("清理陈旧映射 %s 失败: %s" % (mapper, e.strip()))
    return False

def backing_gone(mapper):
    """检查 mapper 的底层设备是否已消失（拔盘 / loop 释放）。"""
    rc, out, _ = run(["dmsetup", "deps", mapper_path(mapper)], timeout=30)
    if rc != 0:
        return False
    m = re.search(r"\((\d+),\s*(\d+)\)", out)
    if not m:
        return False
    syspath = "/sys/dev/block/%s:%s" % (m.group(1), m.group(2))
    if not os.path.exists(syspath):
        return True  # 底层块设备已整个消失（真实拔盘场景）
    try:
        # loop 设备：backing_file 为空说明镜像已 detach（测试场景）
        real = os.path.basename(os.path.realpath(syspath))
        bf = "/sys/block/%s/loop/backing_file" % real
        if os.path.isfile(bf) and not open(bf).read().strip():
            return True
    except OSError:
        pass
    return False

def cmd_cleanup_stale(args):
    require_root()
    n = 0
    for m in active_mappers():
        if find_mount_of_mapper(m):
            continue
        if backing_gone(m):
            rc, e = luks_close(m)
            if rc == 0:
                n += 1
    # 清理已消失设备的状态/标志文件
    if os.path.isdir(RUN_DIR):
        for f in os.listdir(RUN_DIR):
            mm = re.match(r"state-(.+)\.json$", f)
            if mm and not os.path.exists("/dev/" + mm.group(1)):
                try:
                    os.unlink(os.path.join(RUN_DIR, f))
                except OSError:
                    pass
    if n:
        ok("已清理 %d 个残留加密卷" % n)
    return 0

def state_path(kname):
    return os.path.join(RUN_DIR, "state-%s.json" % kname)

def state_save(kname, mapper):
    try:
        os.makedirs(RUN_DIR, exist_ok=True)
        with open(state_path(kname), "w") as f:
            json.dump({"mapper": mapper, "ts": time.time()}, f)
    except OSError:
        pass

def state_load(kname):
    try:
        with open(state_path(kname), "r") as f:
            return json.load(f)
    except (OSError, ValueError):
        return None

def state_clear(kname):
    try:
        os.unlink(state_path(kname))
    except OSError:
        pass

def cmd_ui_state(args):
    """GUI 数据源：一次性输出角色、编码、管控、设备与最近审计（供前端 --json 调用）。"""
    devices = []
    for d in usb_disks(allow_loop=False):
        try:
            s = disk_summary(d)
        except Exception:
            continue
        path = s.get("path", "") or ""
        fstype = s.get("fstype", "") or ""
        mapper = ""
        unlocked = False
        if fstype == "crypto_LUKS":
            uuid_str = luks_uuid(path)
            mapper = mapper_name_for(uuid_str) if uuid_str else ""
            if mapper and os.path.exists(mapper_path(mapper)):
                unlocked = True
                mp = find_mount_of_mapper(mapper) or ""
            else:
                mp = "已锁定"
        else:
            mp = d.get("mountpoint") or "未格式化"
        raw_size = d.get("size") or 0
        devices.append({
            "path": path,
            "kname": s.get("name") or os.path.basename(path),
            "model": s.get("model", "") or "",
            "serial": s.get("serial", "") or "",
            "vidpid": s.get("vidpid", "") or "",
            "size": raw_size,
            "size_human": human_size(raw_size),
            "fstype": fstype,
            "mountpoint": mp,
            "mapper": mapper,
            "unlocked": unlocked,
        })
    key_ok = os.path.isfile(KEY_FILE)
    cfg = load_config()
    json_set(
        version=VERSION,
        is_root=(os.geteuid() == 0),
        role="client" if is_client() else "admin",
        unit_registered=key_ok,
        unit_fingerprint=key_fingerprint() if key_ok else "",
        admin_password_set=admin_password_set(),
        guard_enabled=guard_enabled(cfg),
        guard_blocked=len(guard_blocked_list()),
        devices=devices,
        active_mappers=active_mappers(),
        audit=read_audit_tail(getattr(args, "lines", 20) or 20),
    )
    return 0


def cmd_status(args):
    key = load_unit_key(required=False)
    print()
    print("%sUSBVault 状态%s  v%s" % (C_CYA, C_RST, VERSION))
    print("─" * 62)
    if key:
        print("单位编码 : %s已登记%s（密钥指纹 %s）" % (C_GRN, C_RST, key_fingerprint()))
    else:
        print("单位编码 : %s未登记%s（sudo usbvault init）" % (C_RED, C_RST))
    inst = installed_files()
    print("udev规则 : %s" % ("已安装" if inst["udev"] else "未安装"))
    print("自动服务 : %s" % ("已安装" if inst["systemd"] else "未安装"))
    print("─" * 62)
    disks = usb_disks(allow_loop=True)
    if not disks:
        print("当前没有 USB 存储设备。")
    for d in disks:
        s = disk_summary(d)
        tag = "加密盘" if s["fstype"] == "crypto_LUKS" else (s["fstype"] or "未格式化")
        print("%-9s %-8s %-22s 序列号 %s" % (s["path"], s["size"], s["model"], s["serial"] or "?"))
        info("类型: %s" % tag)
        if s["fstype"] == "crypto_LUKS":
            uuid_str = luks_uuid(s["path"])
            mapper = mapper_name_for(uuid_str) if uuid_str else "?"
            if os.path.exists(mapper_path(mapper)):
                mp = find_mount_of_mapper(mapper)
                info("状态: %s已解锁%s 挂载: %s" % (C_GRN, C_RST, mp or "未挂载"))
            else:
                info("状态: 已锁定")
        print()
    ms = active_mappers()
    if ms:
        print("活动加密卷: %s" % ", ".join(ms))
    return 0

def cmd_repair(args):
    require_root()
    target = args.target
    mapper = None
    dev = None
    if target and os.path.exists(target):
        if target.startswith("/dev/mapper/"):
            mapper = os.path.basename(target)
        else:
            dev = target
    else:
        cands = [m for m in active_mappers() if find_mount_of_mapper(m)]
        locked = [d for d in usb_disks(allow_loop=True) if d.get("fstype") == "crypto_LUKS"]
        if cands:
            mapper = cands[0]
        elif locked:
            dev = locked[0]["path"]
        else:
            err("未发现加密卷。")
            return 1
    if mapper is None:
        uuid_str = luks_uuid(dev)
        if not uuid_str:
            err("%s 不是 LUKS 卷。" % dev)
            return 1
        mapper = mapper_name_for(uuid_str)
        key = load_unit_key(required=False)
        rc, e = (luks_open(dev, key, mapper) if key else (1, ""))
        if rc != 0:
            code = ask_code("请输入该盘编码")
            rc, e = luks_open(dev, derive_key(code), mapper)
            if rc != 0:
                err("解锁失败: %s" % e.strip()[:120])
                return 1
        was_locked = True
    else:
        was_locked = False
    mp = find_mount_of_mapper(mapper)
    if mp:
        warn("卷已挂载于 %s，修复脏标记需要先卸载。" % mp)
        if not confirm("卸载并修复", args.yes):
            return 1
        do_unmount(mp, mapper=mapper)
    step("运行 ntfsfix（清除脏标记并检查基本一致性）")
    rc, out, e = run(["ntfsfix", "-d", mapper_path(mapper)], timeout=600)
    print(out.strip() or e.strip())
    if rc == 0:
        ok("修复完成，重新挂载…")
        uuid_str = luks_uuid_of_mapper(mapper)
        mp = do_mount(mapper, uuid_str) if uuid_str else None
        if mp:
            ok("已重新挂载: %s" % mp)
        audit("repair", device=mapper, result="ok")
        return 0
    err("修复失败: %s" % e.strip()[:200])
    audit("repair", device=mapper, result="fail", detail=e.strip()[:200])
    if was_locked:
        luks_close(mapper)
    return 1

def luks_uuid_of_mapper(mapper):
    rc, out, _ = run(["cryptsetup", "status", mapper], timeout=30)
    if rc != 0:
        return ""
    m = re.search(r"UUID:\s+(\S+)", out)
    return m.group(1) if m else ""

def cmd_rekey(args):
    require_root()
    require_admin("更换单位编码（换码）")
    step("更换单位编码")
    warn("换码影响：")
    warn("  1. 本机密钥将更新，之后格式化的盘使用新编码。")
    warn("  2. 当前插入的、用旧编码加密的 U 盘将被同步换码（数据不变）。")
    warn("  3. 未插入的旧编码 U 盘将无法自动解锁！需插入后手动换码或重新格式化。")
    if not confirm("确定继续", args.yes):
        return 1
    old_key = load_unit_key()
    new_code = ask_code("请输入新编码")
    if ask_code("请再次输入新编码") != new_code:
        err("两次输入不一致。")
        return 1
    new_key = derive_key(new_code)
    if hmac.compare_digest(old_key, new_key):
        warn("新编码与旧编码派生密钥相同，无需换码。")
        return 0
    # 更新本机密钥
    write_unit_key(new_key)
    ok("本机密钥已更新（指纹 %s → %s）" % (key_fingerprint(), hashlib.sha256(new_key).hexdigest()[:8]))
    # 换码当前插入的盘
    changed, failed = 0, 0
    for d in usb_disks(allow_loop=True):
        if d.get("fstype") != "crypto_LUKS":
            continue
        dev = d["path"]
        uuid_str = luks_uuid(dev)
        mapper = mapper_name_for(uuid_str)
        # 已挂载的先卸载（luksChangeKey 不需要卸载，但稳妥起见保持挂载也可）
        tmp = "/dev/shm/usbvault-newkey-%d" % os.getpid()
        with open(tmp, "wb") as f:
            f.write(new_key)
        os.chmod(tmp, 0o600)
        try:
            rc, _, e = run(["cryptsetup", "luksChangeKey", dev, tmp, "--key-file", "-"],
                           input_bytes=old_key, timeout=600)
        finally:
            try:
                os.unlink(tmp)
            except OSError:
                pass
        if rc == 0:
            ok("已换码: %s (%s)" % (dev, mapper))
            changed += 1
            audit("rekey", device=dev, result="ok")
        else:
            # 可能是其他编码的盘
            warn("换码失败 %s（可能是其他编码的盘，跳过）: %s" % (dev, e.strip()[:100]))
            failed += 1
            audit("rekey", device=dev, result="fail", detail=e.strip()[:100])
    info("本次换码 %d 块盘，跳过 %d 块。" % (changed, failed))
    json_set(rekey_changed=changed, rekey_failed=failed,
             unit_fingerprint=key_fingerprint())
    if failed:
        info("跳过的盘如需保留数据，请插入后执行: sudo usbvault rekey --device <dev>（输入其编码）")
    return 0

def cmd_backup_header(args):
    require_root()
    dev = args.device or _single_luks_device()
    if not dev:
        err("未找到加密 U 盘，或设备不唯一，请用参数指定。")
        return 1
    uuid_str = luks_uuid(dev)
    if not uuid_str:
        err("%s 不是 LUKS 卷。" % dev)
        return 1
    os.makedirs(VAR_DIR, mode=0o700, exist_ok=True)
    f = os.path.join(VAR_DIR, "header-%s.img" % mapper_name_for(uuid_str))
    rc, _, e = run(["cryptsetup", "luksHeaderBackup", dev, "--header-backup-file", f])
    if rc != 0:
        err("备份失败: %s" % e.strip())
        return 1
    os.chmod(f, 0o600)
    ok("头部已备份: %s" % f)
    json_set(header_file=f, device=dev)
    audit("backup-header", device=dev, result="ok", detail=f)
    return 0

def cmd_restore_header(args):
    require_root()
    require_admin("恢复 LUKS 头部")
    dev = args.device or _single_luks_device()
    hdr = args.file
    if not dev or not os.path.exists(dev):
        err("设备无效: %s" % dev)
        return 1
    if not hdr or not os.path.isfile(hdr):
        err("头部备份文件无效: %s" % hdr)
        return 1
    warn("恢复头部将覆盖 %s 当前的 LUKS 头部！" % dev)
    warn("恢复后只能用备份头部对应的密钥（编码）解锁。")
    code = ask_code("请输入当前单位编码以确认")
    key = load_unit_key()
    if not hmac.compare_digest(derive_key(code), key):
        err("编码校验失败。")
        audit("restore-header", device=dev, result="reject", detail="编码校验失败")
        return 1
    if not confirm("确认恢复头部", args.yes):
        return 1
    rc, _, e = run(["cryptsetup", "luksHeaderRestore", dev,
                    "--header-backup-file", hdr, "--batch-mode"], timeout=300)
    if rc != 0:
        err("恢复失败: %s" % e.strip())
        audit("restore-header", device=dev, result="fail", detail=e.strip()[:200])
        return 1
    ok("头部已恢复自 %s" % hdr)
    audit("restore-header", device=dev, result="ok", detail=hdr)
    return 0

def _single_luks_device():
    cands = [d for d in usb_disks(allow_loop=True) if d.get("fstype") == "crypto_LUKS"]
    return cands[0]["path"] if len(cands) == 1 else None

def cmd_audit(args):
    require_root()
    if not os.path.isfile(AUDIT_FILE):
        info("暂无审计日志。")
        return 0
    lines = open(AUDIT_FILE, "r", encoding="utf-8", errors="replace").read().splitlines()
    if args.follow:
        import shlex
        os.execvp("tail", ["tail", "-n", "50", "-f", AUDIT_FILE])
    n = args.lines or 30
    json_set(records=read_audit_tail(n))
    for line in lines[-n:]:
        try:
            r = json.loads(line)
            print("%s  %-9s %-28s %-10s %s" %
                  (r.get("ts"), r.get("result"), r.get("action") + " " + r.get("device", ""),
                   r.get("serial", "") or "-", r.get("detail", "")))
        except ValueError:
            print(line)
    return 0

def cmd_clean_trace(args):
    require_root()
    step("清理本机使用痕迹（仅针对 %s 挂载点）" % MEDIA_DIR)
    n = 0
    homes = set()
    for pw in pwd.getpwall():
        home = pw.pw_dir
        if os.path.isdir(home) and (pw.pw_uid >= 1000 or pw.pw_name == "root"):
            homes.add(home)
    for home in sorted(homes):
        bk = os.path.join(home, ".config/gtk-3.0/bookmarks")
        if not os.path.isfile(bk):
            continue
        try:
            with open(bk, "r", encoding="utf-8", errors="replace") as f:
                lines = f.read().splitlines()
            keep = [l for l in lines if MEDIA_DIR not in l]
            if len(keep) != len(lines):
                if args.dry_run:
                    info("[预览] 将从 %s 移除 %d 条记录" % (bk, len(lines) - len(keep)))
                else:
                    with open(bk, "w", encoding="utf-8") as f:
                        f.write("\n".join(keep) + ("\n" if keep else ""))
                    ok("已清理 %s（%d 条）" % (bk, len(lines) - len(keep)))
                n += len(lines) - len(keep)
        except OSError as e:
            warn("处理 %s 失败: %s" % (bk, e))
    info("共处理 %d 条记录。" % n)
    audit("clean-trace", result="ok", detail="%d 条" % n)
    return 0

def _udev_has_guard():
    """管控检查自 1.2.0-5 起并入 90 号规则，按内容判断是否生效。"""
    try:
        with open("/etc/udev/rules.d/90-usbvault.rules", "r", errors="ignore") as f:
            return "guard-check" in f.read()
    except OSError:
        return False

def installed_files():
    return {
        "udev": os.path.isfile("/etc/udev/rules.d/90-usbvault.rules"),
        "guard_udev": _udev_has_guard(),
        "systemd": os.path.isfile("/etc/systemd/system/usbvault-auto@.service"),
        "polkit": os.path.isfile(
            "/etc/polkit-1/localauthority/50-local.d/50-usbvault.pkla"),
    }

def cmd_doctor(args):
    require_root()
    checks = []
    def chk(name, cond, note=""):
        checks.append((name, bool(cond), note))

    rc, out, _ = run(["cryptsetup", "--version"])
    chk("cryptsetup", rc == 0 and "2." in out, out.splitlines()[0] if out else "缺失")
    rc, out, _ = run(["cryptsetup", "luksDump", "--help"])
    chk("LUKS2 支持", rc == 0 and "luks2" in out.lower())
    rc, out, _ = run(["mkntfs", "--version"])
    chk("mkntfs (ntfsprogs)", rc == 0 or "v" in out.lower(), out.splitlines()[0] if out else "缺失")
    chk("ntfs-3g 挂载", shutil.which("mount.ntfs") is not None or os.path.exists("/sbin/mount.ntfs"))
    rc, _, _ = run(["modprobe", "dm_crypt"], timeout=30)
    chk("dm_crypt 内核模块", rc == 0 or os.path.exists("/sys/module/dm_crypt"))
    key = load_unit_key(required=False)
    chk("单位编码登记", key is not None, "指纹 %s" % key_fingerprint())
    inst = installed_files()
    chk("udev 规则", inst["udev"], "/etc/udev/rules.d/90-usbvault.rules")
    chk("udev 管控规则", inst["guard_udev"],
        "/etc/udev/rules.d/90-usbvault.rules 内置 guard-check")
    chk("设备管控（guard）", True,
        "已开启（仅识别同编码注册盘）" if guard_enabled(load_config())
        else "已关闭（sudo usbvault guard on 开启）")
    chk("部署角色", True,
        "受控客户端（client，敏感操作已禁用）" if is_client() else "超管机（admin）")
    chk("超管密码", True,
        "已设置（敏感操作需二次认证）" if admin_password_set()
        else "未设置（建议: sudo usbvault passwd）")
    chk("systemd 自动服务", inst["systemd"], "usbvault-auto@.service")
    rc, out, _ = run(["systemctl", "is-enabled", "usbvault-cleanup.service"])
    chk("开机清理服务", out.strip() == "enabled")
    chk("udisks2（UKUI 挂载抑制）", shutil.which("udisksctl") is not None)
    chk("polkit 挂载授权规则", inst["polkit"],
        "/etc/polkit-1/localauthority/50-local.d/50-usbvault.pkla")

    print()
    print("%sUSBVault 环境自检%s v%s" % (C_CYA, C_RST, VERSION))
    print("─" * 62)
    bad = 0
    for name, cond, note in checks:
        mark = (C_GRN + "✔" + C_RST) if cond else (C_RED + "✘" + C_RST)
        if not cond:
            bad += 1
        print(" %s %-24s %s" % (mark, name, note or ""))
    print("─" * 62)
    json_set(checks=[{"name": n, "ok": c, "note": note} for n, c, note in checks],
             checks_failed=bad)
    if bad:
        warn("%d 项未通过。若为 udev/服务项，请执行 sudo usbvault 安装目录下 install.sh" % bad)
        return 1
    ok("全部通过")
    return 0


# ================= 交互菜单 =================

MENU = [
    ("1", "登记 / 重新登记单位编码", "cmd_init"),
    ("2", "加密格式化 USB 存储设备", "cmd_format"),
    ("3", "手动解锁并挂载", "cmd_open"),
    ("4", "锁定并弹出（卸载+关卷）", "cmd_close"),
    ("5", "查看设备与挂载状态", "cmd_status"),
    ("6", "更换单位编码（换码）", "cmd_rekey"),
    ("7", "备份 LUKS 头部", "cmd_backup_header"),
    ("8", "恢复 LUKS 头部", "cmd_restore_header"),
    ("9", "修复 NTFS 卷（脏标记）", "cmd_repair"),
    ("10", "查看审计日志", "cmd_audit"),
    ("11", "环境自检", "cmd_doctor"),
    ("12", "清理本机使用痕迹", "cmd_clean_trace"),
    ("13", "设备管控（guard on/off/status）", "cmd_guard"),
    ("14", "设置/修改超管密码", "cmd_passwd"),
    ("15", "切换部署角色（admin/client）", "cmd_set_role"),
]

def cmd_menu(args):
    require_root()
    while True:
        print()
        print("%s╔══════════════════════════════════════════════╗%s" % (C_CYA, C_RST))
        print("%s║   USBVault —— USB 存储加密管控工具 v%s     ║%s" % (C_CYA, VERSION, C_RST))
        print("%s╚══════════════════════════════════════════════╝%s" % (C_CYA, C_RST))
        for k, label, _ in MENU:
            print("  %2s) %s" % (k, label))
        print("   0) 退出")
        choice = ask("请选择: ").strip()
        if choice == "0":
            return 0
        entry = next((m for m in MENU if m[0] == choice), None)
        if not entry:
            continue
        try:
            ns = argparse.Namespace(yes=False, dry_run=False, allow_loop=False,
                                    quick=False, no_mount=False, lines=30, follow=False,
                                    device=None, target=None, file=None, kname=None,
                                    guard_action=None, role=None)
            if entry[2] == "cmd_set_role":
                ns.role = ask("部署角色 [admin/client]: ").strip()
            if entry[2] == "cmd_format":
                dev = ask("请输入设备路径（如 /dev/sdb，留空自动选择）: ").strip()
                if not dev:
                    cands = [d for d in usb_disks() if d.get("fstype") != "crypto_LUKS"]
                    if not cands:
                        err("未发现可格式化的 USB 设备。")
                        continue
                    for i, c in enumerate(cands, 1):
                        info("%d) %s %s %s" % (i, c["path"], human_size(c.get("size")), c.get("model", "")))
                    idx = ask("选择 [1-%d]: " % len(cands)).strip()
                    if not idx.isdigit() or not (1 <= int(idx) <= len(cands)):
                        continue
                    dev = cands[int(idx) - 1]["path"]
                ns.device = dev
            if entry[2] == "cmd_guard":
                act = ask("guard 操作 [status/on/off/blocked/allow]（默认 status）: ").strip() or "status"
                ns.guard_action = act
                if act == "allow":
                    ns.kname = ask("请输入要临时放行的设备名（如 sdb）: ").strip()
            rc = globals()[entry[2]](ns)
            if rc:
                warn("命令返回码 %d" % rc)
        except KeyboardInterrupt:
            print()
            warn("已中断。")
        except SystemExit:
            pass


# ================= 入口 =================

def build_parser():
    ap = argparse.ArgumentParser(
        prog="usbvault",
        description="银河麒麟 V10 USB 外接存储加密格式化管理工具",
        epilog="示例: sudo usbvault init && sudo usbvault format /dev/sdb")
    ap.add_argument("--version", action="version", version="usbvault " + VERSION)
    ap.add_argument("--json", action="store_true",
                    help="机器可读输出：stdout 仅输出一行 JSON，人类可读信息走 stderr（GUI 调用）")
    sub = ap.add_subparsers(dest="cmd")

    p = sub.add_parser("init", help="登记单位编码（12位字母数字）")
    p.add_argument("--yes", action="store_true", help="自动确认（自动化/测试用）")
    p.set_defaults(func=cmd_init)

    p = sub.add_parser("format", help="加密格式化 USB 设备（LUKS2+NTFS）")
    p.add_argument("device", help="整盘设备，如 /dev/sdb")
    p.add_argument("--allow-loop", action="store_true", help="允许 loop 设备（测试用）")
    p.add_argument("--quick", action="store_true", help="跳过历史数据全盘覆写")
    p.add_argument("--no-mount", action="store_true", help="格式化后不自动挂载")
    p.add_argument("--dry-run", action="store_true", help="仅预览将执行的操作")
    p.add_argument("--yes", action="store_true", help="自动确认（自动化/测试用）")
    p.set_defaults(func=cmd_format)

    p = sub.add_parser("open", help="手动解锁并挂载")
    p.add_argument("device", nargs="?", help="LUKS 设备（默认自动发现）")
    p.set_defaults(func=cmd_open)

    p = sub.add_parser("close", help="锁定（卸载+关闭加密卷）")
    p.add_argument("target", nargs="?", default="all", help="mapper名/设备名/all")
    p.set_defaults(func=cmd_close)

    p = sub.add_parser("status", help="查看设备与挂载状态")
    p.set_defaults(func=cmd_status)

    p = sub.add_parser("ui-state", help=argparse.SUPPRESS)
    p.add_argument("--lines", type=int, default=20, help="审计日志条数")
    p.set_defaults(func=cmd_ui_state)

    p = sub.add_parser("rekey", help="更换单位编码")
    p.add_argument("--device", help="仅处理指定设备")
    p.add_argument("--yes", action="store_true", help="自动确认（自动化/测试用）")
    p.set_defaults(func=cmd_rekey)

    p = sub.add_parser("repair", help="修复 NTFS 卷（ntfsfix）")
    p.add_argument("target", nargs="?", help="设备或 mapper 路径")
    p.add_argument("--yes", action="store_true")
    p.set_defaults(func=cmd_repair)

    p = sub.add_parser("backup-header", help="备份 LUKS 头部")
    p.add_argument("device", nargs="?")
    p.set_defaults(func=cmd_backup_header)

    p = sub.add_parser("restore-header", help="恢复 LUKS 头部")
    p.add_argument("file", help="头部备份文件")
    p.add_argument("device", nargs="?")
    p.add_argument("--yes", action="store_true")
    p.set_defaults(func=cmd_restore_header)

    p = sub.add_parser("audit", help="查看审计日志")
    p.add_argument("--lines", type=int, default=30)
    p.add_argument("--follow", action="store_true")
    p.set_defaults(func=cmd_audit)

    p = sub.add_parser("doctor", help="环境自检")
    p.set_defaults(func=cmd_doctor)

    p = sub.add_parser("clean-trace", help="清理本机挂载使用痕迹")
    p.add_argument("--dry-run", action="store_true")
    p.set_defaults(func=cmd_clean_trace)

    p = sub.add_parser("guard", help="USB 存储设备管控（on/off/status/blocked/allow）")
    p.add_argument("guard_action", nargs="?", default="status",
                   choices=["on", "off", "status", "blocked", "allow"],
                   help="on/off 管控开关；status 状态；blocked 被拦列表；allow 临时放行")
    p.add_argument("kname", nargs="?", help="allow 操作的设备名（如 sdb）")
    p.set_defaults(func=cmd_guard)

    p = sub.add_parser("passwd", help="设置/修改超管密码（敏感操作二次认证）")
    p.add_argument("--reset", action="store_true",
                   help="忘记密码时重置（root，无需旧密码）")
    p.set_defaults(func=cmd_passwd)

    p = sub.add_parser("set-role", help="切换部署角色（admin 超管机 / client 受控客户端）")
    p.add_argument("role", choices=["admin", "client"])
    p.set_defaults(func=cmd_set_role)

    p = sub.add_parser("verify-admin", help=argparse.SUPPRESS)
    p.set_defaults(func=cmd_verify_admin)

    # 内部命令（systemd 调用）
    p = sub.add_parser("automount", help=argparse.SUPPRESS)
    p.add_argument("kname")
    p.set_defaults(func=cmd_automount)

    # 内部命令（udev 85 规则调用）
    p = sub.add_parser("guard-check", help=argparse.SUPPRESS)
    p.add_argument("kname")
    p.set_defaults(func=cmd_guard_check)

    p = sub.add_parser("lock", help=argparse.SUPPRESS)
    p.add_argument("kname")
    p.add_argument("--udev", action="store_true",
                   help="由 udev remove 事件调用：后端设备仍在线时跳过（防设备名复用误锁）")
    p.set_defaults(func=cmd_lock)

    p = sub.add_parser("cleanup-stale", help=argparse.SUPPRESS)
    p.set_defaults(func=cmd_cleanup_stale)

    p = sub.add_parser("menu", help="交互式菜单")
    p.set_defaults(func=cmd_menu)

    return ap

def main(argv):
    global JSON_MODE
    ap = build_parser()
    args = ap.parse_args(argv)
    if getattr(args, "json", False):
        JSON_MODE = True
        # 人类可读输出全部转到 stderr，stdout 只留 emit_json 写的那一行 JSON
        sys.stdout = sys.stderr
    if not getattr(args, "func", None):
        if os.geteuid() == 0:
            rc = cmd_menu(args)
        else:
            ap.print_help()
            print("\n提示：交互菜单需 root，请执行 sudo usbvault")
            rc = 0
    else:
        rc = args.func(args)
    if JSON_MODE:
        payload = {"ok": rc == 0, "rc": rc}
        payload.update(JSON_RESULT)
        emit_json(payload)
    return rc

if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv[1:]))
    except KeyboardInterrupt:
        print()
        sys.exit(130)
