SysDevCleaner/sysdevcleaner/core/util.py

"""Низкоуровневые утилиты: размеры, пути, защита, корзина."""
from __future__ import annotations

import ctypes
import os
import shutil
import stat
import string
import subprocess
import sys
import threading
from typing import Callable, Iterable

IS_WIN = sys.platform == "win32"

# ---------------------------------------------------------------- размеры ----

_UNITS = ("Б", "КБ", "МБ", "ГБ", "ТБ", "ПБ")


def human(size: float) -> str:
    """1536 -> '1.5 КБ'."""
    size = float(size or 0)
    neg = size < 0
    size = abs(size)
    i = 0
    while size >= 1024 and i < len(_UNITS) - 1:
        size /= 1024.0
        i += 1
    txt = f"{size:.0f} {_UNITS[i]}" if (i == 0 or size >= 100) else f"{size:.1f} {_UNITS[i]}"
    return ("-" if neg else "") + txt


class Cancelled(Exception):
    pass


class Cancel:
    """Кооперативная отмена длительных обходов."""

    def __init__(self) -> None:
        self._e = threading.Event()

    def cancel(self) -> None:
        self._e.set()

    def reset(self) -> None:
        self._e.clear()

    @property
    def cancelled(self) -> bool:
        return self._e.is_set()

    def check(self) -> None:
        if self._e.is_set():
            raise Cancelled()


def dir_stats(
    path: str,
    cancel: "Cancel | None" = None,
    max_entries: int = 400_000,
    follow_links: bool = False,
) -> tuple[int, int, int, float]:
    """Размер каталога: (байты, файлов, папок, max_mtime). Устойчиво к ошибкам."""
    total = files = dirs = 0
    newest = 0.0
    seen = 0
    stack = [path]
    while stack:
        if cancel is not None and cancel.cancelled:
            break
        cur = stack.pop()
        try:
            it = os.scandir(cur)
        except (OSError, ValueError):
            continue
        with it:
            while True:
                try:
                    entry = next(it)
                except StopIteration:
                    break
                except OSError:
                    break
                seen += 1
                if seen > max_entries:
                    stack.clear()
                    break
                try:
                    if entry.is_dir(follow_symlinks=follow_links):
                        dirs += 1
                        stack.append(entry.path)
                    elif entry.is_file(follow_symlinks=follow_links):
                        st = entry.stat(follow_symlinks=False)
                        files += 1
                        total += st.st_size
                        if st.st_mtime > newest:
                            newest = st.st_mtime
                except OSError:
                    continue
    if not newest:
        try:
            newest = os.stat(path).st_mtime
        except OSError:
            newest = 0.0
    return total, files, dirs, newest


def path_size(path: str, cancel: "Cancel | None" = None) -> tuple[int, int, int, float]:
    """Размер файла или каталога в том же формате, что и dir_stats."""
    try:
        st = os.stat(path)
    except OSError:
        return 0, 0, 0, 0.0
    if stat.S_ISDIR(st.st_mode):
        return dir_stats(path, cancel)
    return st.st_size, 1, 0, st.st_mtime


# ------------------------------------------------------------------ пути ----


def env(name: str) -> str:
    return os.environ.get(name, "") or ""


def expand(path: str) -> str:
    return os.path.expandvars(os.path.expanduser(path))


def known_paths() -> dict[str, str]:
    home = os.path.expanduser("~")
    local = env("LOCALAPPDATA") or os.path.join(home, "AppData", "Local")
    roaming = env("APPDATA") or os.path.join(home, "AppData", "Roaming")
    return {
        "HOME": home,
        "LOCAL": local,
        "LOCALLOW": os.path.join(os.path.dirname(local), "LocalLow") if local else "",
        "ROAMING": roaming,
        "TEMP": env("TEMP") or os.path.join(local, "Temp"),
        "WINDIR": env("WINDIR") or "C:\\Windows",
        "SYSTEMDRIVE": (env("SYSTEMDRIVE") or "C:") + "\\",
        "PROGRAMDATA": env("PROGRAMDATA") or "C:\\ProgramData",
        "PROGRAMFILES": env("PROGRAMFILES") or "C:\\Program Files",
        "PROGRAMFILES86": env("ProgramFiles(x86)") or "C:\\Program Files (x86)",
        "PUBLIC": env("PUBLIC") or "C:\\Users\\Public",
    }


def drives() -> list[str]:
    """Список смонтированных дисков ('C:\\', 'D:\\')."""
    if not IS_WIN:
        return ["/"]
    out = []
    try:
        mask = ctypes.windll.kernel32.GetLogicalDrives()
    except Exception:
        mask = 0
    for i, letter in enumerate(string.ascii_uppercase):
        if mask >> i & 1:
            root = f"{letter}:\\"
            try:
                dtype = ctypes.windll.kernel32.GetDriveTypeW(ctypes.c_wchar_p(root))
            except Exception:
                dtype = 3
            if dtype in (2, 3, 4):  # removable / fixed / network
                out.append(root)
    return out or ["C:\\"]


def disk_usage(root: str) -> dict:
    try:
        u = shutil.disk_usage(root)
        return {
            "root": root,
            "total": u.total,
            "used": u.used,
            "free": u.free,
            "percent": (u.used / u.total * 100) if u.total else 0.0,
        }
    except OSError:
        return {"root": root, "total": 0, "used": 0, "free": 0, "percent": 0.0}


# --------------------------------------------------------------- защита -----


def _norm(p: str) -> str:
    return os.path.normcase(os.path.abspath(p)).rstrip("\\/")


def _protected_roots() -> set[str]:
    k = known_paths()
    raw = [
        k["WINDIR"],
        os.path.join(k["WINDIR"], "System32"),
        os.path.join(k["WINDIR"], "SysWOW64"),
        k["PROGRAMFILES"],
        k["PROGRAMFILES86"],
        k["SYSTEMDRIVE"],
        k["HOME"],
        k["LOCAL"],
        k["ROAMING"],
        k["PROGRAMDATA"],
        os.path.join(k["HOME"], "Documents"),
        os.path.join(k["HOME"], "Desktop"),
        os.path.join(k["HOME"], "Downloads"),
        os.path.join(k["HOME"], "Pictures"),
        os.path.join(k["HOME"], "Videos"),
        os.path.join(k["HOME"], "Music"),
        os.path.join(k["HOME"], "OneDrive"),
    ]
    out = {_norm(p) for p in raw if p}
    out |= {_norm(d) for d in drives()}
    return out


PROTECTED = _protected_roots()


def is_protected(path: str) -> bool:
    """Нельзя удалять сам путь (но можно то, что глубже внутри)."""
    n = _norm(path)
    if not n or len(n) <= 3:
        return True
    return n in PROTECTED


def is_safe_target(path: str) -> tuple[bool, str]:
    """Финальная проверка перед удалением."""
    if not path:
        return False, "пустой путь"
    n = _norm(path)
    if len(n) <= 3:
        return False, "корень диска"
    if is_protected(path):
        return False, "системный/пользовательский корневой каталог"
    depth = n.replace("/", "\\").count("\\")
    if depth < 2:
        return False, "слишком высокий уровень в дереве"
    return True, ""


# --------------------------------------------------------------- удаление ---


def _on_rm_error(func: Callable, path: str, exc_info) -> None:
    """Снимаем read-only и пробуем ещё раз."""
    try:
        os.chmod(path, stat.S_IWRITE)
        func(path)
    except Exception:
        pass


def delete_path(path: str, to_recycle_bin: bool = True) -> tuple[bool, str]:
    """Удалить файл/каталог. Возвращает (успех, сообщение об ошибке)."""
    if not os.path.exists(path):
        return True, ""
    if to_recycle_bin and IS_WIN:
        ok, err = _recycle(path)
        if ok:
            return True, ""
        # корзина не справилась (сетевой диск, слишком большой объект) — не насилуем
        return False, err
    try:
        if os.path.isdir(path) and not os.path.islink(path):
            shutil.rmtree(path, onerror=_on_rm_error)
        else:
            try:
                os.chmod(path, stat.S_IWRITE)
            except OSError:
                pass
            os.remove(path)
        if os.path.exists(path):
            return False, "путь остался на месте"
        return True, ""
    except Exception as e:  # noqa: BLE001
        return False, str(e)


def _recycle(path: str) -> tuple[bool, str]:
    """SHFileOperationW с флагом FOF_ALLOWUNDO."""
    try:
        from ctypes import wintypes

        class SHFILEOPSTRUCTW(ctypes.Structure):
            _fields_ = [
                ("hwnd", wintypes.HWND),
                ("wFunc", wintypes.UINT),
                ("pFrom", wintypes.LPCWSTR),
                ("pTo", wintypes.LPCWSTR),
                ("fFlags", ctypes.c_uint16),
                ("fAnyOperationsAborted", wintypes.BOOL),
                ("hNameMappings", ctypes.c_void_p),
                ("lpszProgressTitle", wintypes.LPCWSTR),
            ]

        FO_DELETE = 3
        FOF_ALLOWUNDO = 0x0040
        FOF_NOCONFIRMATION = 0x0010
        FOF_NOERRORUI = 0x0400
        FOF_SILENT = 0x0004
        FOF_NOCONFIRMMKDIR = 0x0200

        op = SHFILEOPSTRUCTW()
        op.wFunc = FO_DELETE
        op.pFrom = os.path.abspath(path) + "\0\0"
        op.fFlags = (FOF_ALLOWUNDO | FOF_NOCONFIRMATION | FOF_NOERRORUI
                     | FOF_SILENT | FOF_NOCONFIRMMKDIR)
        res = ctypes.windll.shell32.SHFileOperationW(ctypes.byref(op))
        if res == 0 and not op.fAnyOperationsAborted:
            return True, ""
        return False, f"SHFileOperation код {res}"
    except Exception as e:  # noqa: BLE001
        return False, str(e)


def clear_contents(path: str, to_recycle_bin: bool = True) -> tuple[int, int, list[str]]:
    """Очистить содержимое каталога, сам каталог оставить.

    Возвращает (удалено, пропущено, ошибки).
    """
    done = skipped = 0
    errs: list[str] = []
    try:
        entries = list(os.scandir(path))
    except OSError as e:
        return 0, 0, [f"{path}: {e}"]
    for entry in entries:
        ok, err = delete_path(entry.path, to_recycle_bin)
        if ok:
            done += 1
        else:
            skipped += 1
            if err and len(errs) < 20:
                errs.append(f"{entry.name}: {err}")
    return done, skipped, errs


# ------------------------------------------------------------- процессы -----


def running_processes() -> set[str]:
    """Имена запущенных процессов в нижнем регистре."""
    names: set[str] = set()
    if not IS_WIN:
        return names
    try:
        out = subprocess.run(
            ["tasklist", "/FO", "CSV", "/NH"],
            capture_output=True, text=True, timeout=20,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
            errors="ignore",
        ).stdout
        for line in out.splitlines():
            line = line.strip()
            if line.startswith('"'):
                names.add(line.split('","')[0].strip('"').lower())
    except Exception:
        pass
    return names


def is_admin() -> bool:
    if not IS_WIN:
        return os.geteuid() == 0 if hasattr(os, "geteuid") else False
    try:
        return bool(ctypes.windll.shell32.IsUserAnAdmin())
    except Exception:
        return False


def open_in_explorer(path: str) -> None:
    try:
        if IS_WIN:
            if os.path.isdir(path):
                os.startfile(path)  # noqa: S606
            else:
                subprocess.Popen(["explorer", "/select,", os.path.abspath(path)])
        else:
            subprocess.Popen(["xdg-open", os.path.dirname(path)])
    except Exception:
        pass


def iter_globs(patterns: Iterable[str]) -> list[str]:
    import glob as _g
    out: list[str] = []
    for p in patterns:
        try:
            out.extend(_g.glob(expand(p), recursive=True))
        except Exception:
            continue
    return out