disk_cleaner.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Windows 磁盘清理工具（侧重 C 盘）。
默认仅扫描并报告可释放空间；使用 --apply 才会删除。
部分目录需要管理员权限。
"""

from __future__ import annotations

import argparse
import ctypes
import os
import subprocess
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Iterable


# ---------------------------------------------------------------------------
# 基础工具
# ---------------------------------------------------------------------------

def is_admin() -> bool:
    try:
        return bool(ctypes.windll.shell32.IsUserAnAdmin())
    except Exception:
        return False


def human_bytes(n: int) -> str:
    if n < 0:
        n = 0
    units = ("B", "KB", "MB", "GB", "TB")
    v = float(n)
    for u in units:
        if v < 1024.0 or u == units[-1]:
            if u == "B":
                return f"{int(v)} {u}"
            return f"{v:.2f} {u}"
        v /= 1024.0
    return f"{n} B"


def iter_files(root: Path) -> Iterable[Path]:
    try:
        for dirpath, _dirnames, filenames in os.walk(root, topdown=True, onerror=None):
            base = Path(dirpath)
            for name in filenames:
                yield base / name
    except (OSError, PermissionError):
        return


def dir_size(root: Path) -> tuple[int, int]:
    """返回 (可统计的字节数, 无法访问的文件数)。"""
    total = 0
    errors = 0
    for p in iter_files(root):
        try:
            if p.is_file():
                total += p.stat().st_size
        except (OSError, PermissionError):
            errors += 1
    return total, errors


def safe_rmtree_contents(root: Path, on_progress: Callable[[], None] | None = None) -> tuple[int, int]:
    """删除 root 下所有子项，保留 root 目录本身。"""
    removed_bytes = 0
    failures = 0
    if not root.exists() or not root.is_dir():
        return 0, 0

    for child in list(root.iterdir()):
        try:
            if child.is_file() or child.is_symlink():
                try:
                    removed_bytes += child.stat().st_size
                except OSError:
                    pass
                child.unlink(missing_ok=True)
            elif child.is_dir():
                b, f = safe_rmtree_contents(child, on_progress)
                removed_bytes += b
                failures += f
                try:
                    child.rmdir()
                except OSError:
                    failures += 1
        except (OSError, PermissionError):
            failures += 1
        if on_progress:
            on_progress()
    return removed_bytes, failures


def empty_recycle_bin() -> tuple[bool, str]:
    try:
        r = subprocess.run(
            [
                "powershell.exe",
                "-NoProfile",
                "-NonInteractive",
                "-Command",
                "Clear-RecycleBin -Force -ErrorAction Stop",
            ],
            capture_output=True,
            text=True,
            timeout=600,
        )
        if r.returncode == 0:
            return True, ""
        msg = (r.stderr or r.stdout or "").strip() or f"exit {r.returncode}"
        return False, msg
    except Exception as e:
        return False, str(e)


def disable_hibernation() -> tuple[bool, str]:
    """关闭休眠并删除 hiberfil.sys（需管理员）。"""
    try:
        r = subprocess.run(
            ["powercfg.exe", "/h", "off"],
            capture_output=True,
            text=True,
            timeout=60,
        )
        if r.returncode == 0:
            return True, ""
        msg = (r.stderr or r.stdout or "").strip() or f"exit {r.returncode}"
        return False, msg
    except Exception as e:
        return False, str(e)


def file_size_safe(path: Path) -> int:
    try:
        return path.stat().st_size
    except OSError:
        return 0


# ---------------------------------------------------------------------------
# 清理目标
# ---------------------------------------------------------------------------

@dataclass
class CleanTarget:
    key: str
    title: str
    path: Path                               # 主显示路径
    extra_paths: list[Path] = field(default_factory=list)
    needs_admin_hint: bool = False
    note: str = ""
    info_only: bool = False                  # 仅展示，不参与勾选与清理
    label_only_path: bool = False            # 若 True，path 仅作显示，不参与清理/统计
    _size_cache: tuple[int, int] | None = field(default=None, repr=False)

    # ---- 路径处理 ----
    def resolve(self) -> Path:
        try:
            return self.path.expanduser().resolve()
        except Exception:
            return self.path

    def all_paths(self) -> list[Path]:
        """返回参与统计与清理的所有路径（不含 label_only_path 时的 self.path）。"""
        if self.label_only_path:
            sources: list[Path] = list(self.extra_paths)
        else:
            sources = [self.path, *self.extra_paths]
        seen: set[str] = set()
        out: list[Path] = []
        for p in sources:
            try:
                rp = p.expanduser().resolve()
            except Exception:
                rp = p
            k = str(rp).lower()
            if k in seen:
                continue
            seen.add(k)
            out.append(rp)
        return out

    def display_paths(self) -> list[Path]:
        """用于 UI 显示的路径：始终包含主路径在前。"""
        seen: set[str] = set()
        out: list[Path] = []
        for p in [self.path, *self.extra_paths]:
            try:
                rp = p.expanduser().resolve()
            except Exception:
                rp = p
            k = str(rp).lower()
            if k in seen:
                continue
            seen.add(k)
            out.append(rp)
        return out

    # ---- 体积 ----
    def size(self) -> tuple[int, int]:
        if self._size_cache is not None:
            return self._size_cache

        if self.key == "thumbcache":
            total, errors = 0, 0
            for p in self.all_paths():
                if p.exists():
                    for f in p.glob("thumbcache_*.db"):
                        try:
                            total += f.stat().st_size
                        except OSError:
                            errors += 1
            self._size_cache = (total, errors)
            return self._size_cache

        total = 0
        errors = 0
        for p in self.all_paths():
            if not p.exists():
                continue
            try:
                if p.is_file():
                    total += p.stat().st_size
                else:
                    b, e = dir_size(p)
                    total += b
                    errors += e
            except (OSError, PermissionError):
                errors += 1
        self._size_cache = (total, errors)
        return self._size_cache

    # ---- 清理 ----
    def clean(self, dry_run: bool) -> tuple[int, int]:
        if self.info_only:
            return 0, 0
        if self.key == "thumbcache":
            total_b, total_f = 0, 0
            for p in self.all_paths():
                b, f = clean_thumbcache_only(p, dry_run=dry_run)
                total_b += b
                total_f += f
            return total_b, total_f
        if self.key == "recycle_bin":
            if dry_run:
                return self.size()
            ok, _ = empty_recycle_bin()
            return (0, 0) if ok else (0, 1)
        if self.key == "hibernation":
            if dry_run:
                return self.size()
            ok, _ = disable_hibernation()
            return (0, 0) if ok else (0, 1)

        total_b, total_f = 0, 0
        for p in self.all_paths():
            if not p.exists():
                continue
            if dry_run:
                if p.is_file():
                    total_b += file_size_safe(p)
                else:
                    b, _ = dir_size(p)
                    total_b += b
                continue
            if p.is_file():
                try:
                    b = p.stat().st_size
                    p.unlink()
                    total_b += b
                except OSError:
                    total_f += 1
            else:
                b, f = safe_rmtree_contents(p)
                total_b += b
                total_f += f
        return total_b, total_f


# ---------------------------------------------------------------------------
# 浏览器缓存路径搜索
# ---------------------------------------------------------------------------

CHROMIUM_CACHE_SUBDIRS = (
    "Cache",
    "Code Cache",
    "GPUCache",
    "ShaderCache",
    Path("Service Worker") / "CacheStorage",
    Path("Service Worker") / "ScriptCache",
)


def _filter_existing(paths: Iterable[Path]) -> list[Path]:
    return [p for p in paths if p.exists()]


def find_chromium_cache_dirs(user_data_root: Path) -> list[Path]:
    if not user_data_root.exists():
        return []
    results: list[Path] = []
    for profile in user_data_root.iterdir():
        try:
            if not profile.is_dir():
                continue
        except OSError:
            continue
        name = profile.name
        if not (
            name == "Default"
            or name.startswith("Profile ")
            or name in ("Guest Profile", "System Profile")
        ):
            continue
        for sub in CHROMIUM_CACHE_SUBDIRS:
            results.append(profile / sub)
    return _filter_existing(results)


def find_firefox_cache_dirs(local: Path) -> list[Path]:
    base = local / "Mozilla" / "Firefox" / "Profiles"
    if not base.exists():
        return []
    results: list[Path] = []
    for profile in base.iterdir():
        try:
            if not profile.is_dir():
                continue
        except OSError:
            continue
        for sub in ("cache2", "startupCache", "OfflineCache", "thumbnails"):
            results.append(profile / sub)
    # Firefox 部分缓存也可能在 Roaming 下，但常用的全在 Local 里。
    return _filter_existing(results)


def _on_drive(p: Path, drive: str) -> bool:
    try:
        return str(p).upper().startswith(drive.upper())
    except Exception:
        return False


# ---------------------------------------------------------------------------
# 构造目标列表
# ---------------------------------------------------------------------------

def build_targets(drive: str) -> list[CleanTarget]:
    drive = drive.rstrip("\\/").upper()
    if len(drive) != 2 or drive[1] != ":":
        raise ValueError("驱动器格式应为 C: 或 D:")

    windir = Path(os.environ.get("WINDIR", f"{drive}\\Windows"))
    local = Path(os.environ.get("LOCALAPPDATA", ""))
    targets: list[CleanTarget] = []
    system_drive = (os.environ.get("SystemDrive", "C:") or "C:").upper()
    is_system_drive = drive == system_drive

    # ---- 临时目录 ----
    for env in ("TEMP", "TMP"):
        raw = os.environ.get(env)
        if raw:
            p = Path(raw)
            if _on_drive(p, drive):
                targets.append(
                    CleanTarget(
                        key=f"user_{env.lower()}",
                        title=f"用户临时目录 ({env})",
                        path=p,
                        note="当前登录用户的临时文件，通常可安全清理。",
                    )
                )

    if is_system_drive:
        targets.append(
            CleanTarget(
                key="windows_temp",
                title="Windows\\Temp",
                path=windir / "Temp",
                needs_admin_hint=True,
                note="系统级临时目录；部分文件可能被服务占用。",
            )
        )

        targets.append(
            CleanTarget(
                key="prefetch",
                title="Prefetch（可选）",
                path=windir / "Prefetch",
                needs_admin_hint=True,
                note="预读取数据；删除后可能短暂影响启动预测，一般可清理。",
            )
        )

        targets.append(
            CleanTarget(
                key="softdist_download",
                title="Windows Update 缓存 (SoftwareDistribution\\Download)",
                path=windir / "SoftwareDistribution" / "Download",
                needs_admin_hint=True,
                note="Windows 更新已下载的安装包；建议先停止 wuauserv 服务。",
            )
        )

        program_data = Path(os.environ.get("ProgramData", f"{drive}\\ProgramData"))
        targets.append(
            CleanTarget(
                key="wer",
                title="Windows 错误报告 (WER)",
                path=program_data / "Microsoft" / "Windows" / "WER",
                extra_paths=[
                    local / "Microsoft" / "Windows" / "WER" if local else program_data,
                ],
                needs_admin_hint=True,
                note="错误报告与崩溃转储，可清理。",
            )
        )

    # ---- 用户级缓存 ----
    if local.exists() and _on_drive(local, drive):
        targets.append(
            CleanTarget(
                key="inet_cache",
                title="IE/旧版 Edge 缓存 (INetCache)",
                path=local / "Microsoft" / "Windows" / "INetCache",
                note="旧式 Web 缓存。",
            )
        )
        targets.append(
            CleanTarget(
                key="thumbcache",
                title="缩略图缓存",
                path=local / "Microsoft" / "Windows" / "Explorer",
                note="仅清理 thumbcache_*.db（缩略图会自动重建）。",
            )
        )
        targets.append(
            CleanTarget(
                key="crash_dumps",
                title="本机崩溃转储 (CrashDumps)",
                path=local / "CrashDumps",
                note="应用崩溃 dump 文件。",
            )
        )

        # Chrome
        chrome_root = local / "Google" / "Chrome" / "User Data"
        chrome_dirs = find_chromium_cache_dirs(chrome_root)
        if chrome_dirs:
            targets.append(
                CleanTarget(
                    key="chrome_cache",
                    title=f"Chrome 缓存（{len(chrome_dirs)} 项）",
                    path=chrome_root,
                    extra_paths=chrome_dirs,
                    label_only_path=True,
                    note="仅清理各 Profile 下的 Cache/Code Cache/GPUCache 等；建议先关闭 Chrome。",
                )
            )

        # Edge
        edge_root = local / "Microsoft" / "Edge" / "User Data"
        edge_dirs = find_chromium_cache_dirs(edge_root)
        if edge_dirs:
            targets.append(
                CleanTarget(
                    key="edge_cache",
                    title=f"Edge 缓存（{len(edge_dirs)} 项）",
                    path=edge_root,
                    extra_paths=edge_dirs,
                    label_only_path=True,
                    note="仅清理各 Profile 下的 Cache/Code Cache/GPUCache 等；建议先关闭 Edge。",
                )
            )

        # Firefox
        firefox_dirs = find_firefox_cache_dirs(local)
        if firefox_dirs:
            targets.append(
                CleanTarget(
                    key="firefox_cache",
                    title=f"Firefox 缓存（{len(firefox_dirs)} 项）",
                    path=local / "Mozilla" / "Firefox" / "Profiles",
                    extra_paths=firefox_dirs,
                    label_only_path=True,
                    note="仅清理各 Profile 下的 cache2 等；建议先关闭 Firefox。",
                )
            )

    # ---- Delivery Optimization ----
    if is_system_drive:
        targets.append(
            CleanTarget(
                key="delivery_opt",
                title="传递优化缓存 (Delivery Optimization)",
                path=windir
                / "ServiceProfiles"
                / "NetworkService"
                / "AppData"
                / "Local"
                / "Microsoft"
                / "Windows"
                / "DeliveryOptimization"
                / "Cache",
                needs_admin_hint=True,
                note="Windows 更新 P2P 缓存；常需管理员权限。",
            )
        )

    # ---- 回收站 ----
    targets.append(
        CleanTarget(
            key="recycle_bin",
            title="回收站",
            path=Path(f"{drive}\\$Recycle.Bin"),
            needs_admin_hint=True,
            note="清空所有分区回收站条目（调用 Clear-RecycleBin）。",
        )
    )

    # ---- 系统大文件：休眠 / 分页 ----
    if is_system_drive:
        targets.append(
            CleanTarget(
                key="hibernation",
                title="休眠文件 hiberfil.sys（关闭休眠以删除）",
                path=Path(f"{drive}\\hiberfil.sys"),
                needs_admin_hint=True,
                note="执行 powercfg /h off；之后将无法使用休眠/快速启动。",
            )
        )

        targets.append(
            CleanTarget(
                key="pagefile_info",
                title="分页文件 pagefile.sys（仅信息）",
                path=Path(f"{drive}\\pagefile.sys"),
                info_only=True,
                note="由系统管理；如需缩小请在“系统属性 → 性能 → 高级 → 虚拟内存”中调整。",
            )
        )
        targets.append(
            CleanTarget(
                key="swapfile_info",
                title="交换文件 swapfile.sys（仅信息）",
                path=Path(f"{drive}\\swapfile.sys"),
                info_only=True,
                note="UWP 应用使用；通常不建议手动清理。",
            )
        )

    # 去重
    seen: set[str] = set()
    unique: list[CleanTarget] = []
    for t in targets:
        try:
            k = (t.key, str(t.resolve()).lower())
        except Exception:
            k = (t.key, str(t.path).lower())
        if k in seen:
            continue
        seen.add(k)
        unique.append(t)
    return unique


def clean_thumbcache_only(root: Path, dry_run: bool) -> tuple[int, int]:
    removed = 0
    fails = 0
    if not root.exists():
        return 0, 0
    for p in root.glob("thumbcache_*.db"):
        try:
            if dry_run:
                removed += p.stat().st_size
            else:
                b = p.stat().st_size
                p.unlink()
                removed += b
        except OSError:
            fails += 1
    return removed, fails


def run_dism_reset_base(dry_run: bool) -> None:
    cmd = [
        "DISM.exe",
        "/Online",
        "/Cleanup-Image",
        "/StartComponentCleanup",
        "/ResetBase",
    ]
    print("\n[DISM] 命令:", " ".join(cmd))
    print("说明：会压缩/清理组件存储，耗时较长，需要管理员；/ResetBase 不可轻易撤销。")
    if dry_run:
        print("[DISM] 模拟运行：未执行。")
        return
    if not is_admin():
        print("[DISM] 需要管理员权限，已跳过。")
        return
    subprocess.run(cmd, check=False)


def configure_stdio_utf8() -> None:
    """让中文在 Windows 控制台正常显示。"""
    for stream in (sys.stdout, sys.stderr):
        reconf = getattr(stream, "reconfigure", None)
        if callable(reconf):
            try:
                reconf(encoding="utf-8", errors="replace")
            except Exception:
                pass


# ---------------------------------------------------------------------------
# CLI 入口
# ---------------------------------------------------------------------------

def main() -> int:
    configure_stdio_utf8()
    parser = argparse.ArgumentParser(
        description="Windows 磁盘清理（默认只扫描，--apply 才删除）",
    )
    parser.add_argument(
        "--drive",
        default=os.environ.get("SystemDrive", "C:"),
        help="目标盘符，默认系统盘",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="实际执行删除；不加此参数仅估算空间",
    )
    parser.add_argument(
        "--targets",
        default="all",
        help="逗号分隔的目标 key，或 all。key 见扫描输出第一列。",
    )
    parser.add_argument(
        "--dism",
        action="store_true",
        help="（危险/慢）在 --apply 且管理员下尝试 DISM 组件存储清理",
    )
    args = parser.parse_args()

    try:
        all_targets = build_targets(args.drive)
    except ValueError as e:
        print(e, file=sys.stderr)
        return 2

    selected = all_targets
    if args.targets.strip().lower() != "all":
        keys = {k.strip().lower() for k in args.targets.split(",") if k.strip()}
        selected = [t for t in all_targets if t.key.lower() in keys]
        missing = keys - {t.key.lower() for t in selected}
        if missing:
            print("未知 target:", ", ".join(sorted(missing)), file=sys.stderr)

    dry_run = not args.apply
    print("模式:", "【模拟】仅统计" if dry_run else "【执行】将尝试删除")
    print("目标盘:", args.drive.upper())
    print("管理员:", "是" if is_admin() else "否（部分目录可能无法清理）")
    print()

    total_bytes = 0
    for t in selected:
        sz, errn = t.size()
        flag = " [建议管理员]" if t.needs_admin_hint and not is_admin() else ""
        kind = "(信息)" if t.info_only else ""
        print(f"{t.key:18} {human_bytes(sz):>12}  无法读≈{errn}  {t.title}{flag} {kind}")
        for p in t.display_paths():
            print(f"    路径: {p}")
        if t.note:
            print(f"    说明: {t.note}")
        if not t.info_only:
            total_bytes += sz

    print()
    print("估算可清理（扫描到的文件体积，实际释放可能略少）:", human_bytes(total_bytes))

    if dry_run:
        print("\n若要真正清理，请使用: python disk_cleaner.py --apply --drive", args.drive)
        if args.dism:
            run_dism_reset_base(dry_run=True)
        return 0

    print("\n开始清理…")
    freed = 0
    fail_total = 0
    for t in selected:
        if t.info_only:
            continue
        b, f = t.clean(dry_run=False)
        freed += b
        fail_total += f
        print(f"  {t.key}: 处理约 {human_bytes(b)}，失败项 {f}")

    print()
    print("完成。粗略释放/处理:", human_bytes(freed), "失败项:", fail_total)

    if args.dism:
        run_dism_reset_base(dry_run=False)

    return 0


if __name__ == "__main__":
    if os.name != "nt":
        print("本脚本仅适用于 Windows。", file=sys.stderr)
        sys.exit(1)
    raise SystemExit(main())