PySpOS/reset_slot.py at master · ViudiraTech/PySpOS · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
'''
 *
 *      reset_slot.py
 *      Standalone slot wipe tool.
 *
 *      2026/9/25 By GoutouStdio
 *      Copyright (C) 2022-2026 GoutouStdio, based on the MIT license.
 *
 */
'''

import argparse
import os
import shutil
import sys

import secure_boot


# Hardened terminal prompts live in src/ttyutil.py; without src/ the raw
# input fallback applies.
try:
    sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "src"))
    import ttyutil as _ttyutil
except ImportError:
    _ttyutil = None


# Ask to continue, healing a raw terminal left by a crashed session first.
def _ask_continue(prompt):
    if _ttyutil is None:
        return input(prompt).strip().lower()
    _ttyutil.ensure_sane_tty()
    return _ttyutil.read_choice(prompt, valid=("y", "n"), default="n",
                                max_retries=5)


# Resolve a slot name to its directory, rejecting anything else.
def _safe_slot_path(root_dir, slot):
    if slot not in secure_boot.SLOTS:
        raise ValueError("槽位必须是 slot_a 或 slot_b")
    root = os.path.realpath(root_dir)
    path = os.path.realpath(os.path.join(root, slot))
    if os.path.commonpath((root, path)) != root or os.path.basename(path) != slot:
        raise RuntimeError("槽位路径越界")
    return path


# Erase the contents of the named slots.
def wipe_slots(root_dir, slots, assume_yes=False):
    root_dir = os.path.abspath(root_dir)
    if secure_boot.read_locked(root_dir):
        raise RuntimeError("当前是 LOCKED 信任域，拒绝删除槽位")
    if not assume_yes:
        if _ask_continue(f"将删除槽位 {', '.join(slots)}，继续？(y/n): ") not in ("y", "yes"):
            print("已取消")
            return False
    for slot in slots:
        path = _safe_slot_path(root_dir, slot)
        if os.path.islink(path):
            os.unlink(path)
        elif os.path.isdir(path):
            shutil.rmtree(path)
        elif os.path.exists(path):
            raise RuntimeError(f"槽位路径不是目录: {path}")
        print(f"已删除槽位: {slot}")
    print("policy、rollback state 和 etc 未删除；下次启动将从 src 重建。")
    return True


# Parse arguments and wipe the requested slots.
def main(argv=None):
    parser = argparse.ArgumentParser(description="清空 PySpOS 开发槽位")
    parser.add_argument("--root-dir", default=os.path.dirname(os.path.abspath(__file__)))
    parser.add_argument("--slot", action="append", choices=secure_boot.SLOTS)
    parser.add_argument("--all", action="store_true")
    parser.add_argument("--yes", action="store_true")
    args = parser.parse_args(argv)
    slots = list(secure_boot.SLOTS) if args.all else (args.slot or [])
    if not slots:
        parser.error("请指定 --slot 或 --all")
    try:
        return 0 if wipe_slots(args.root_dir, slots, args.yes) else 1
    except Exception as exc:
        print(f"清空失败: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())