psp-preservation-tool/app/psp.py

"""PSP device detection and memory-stick inventory.

The PSP mounts as a plain USB mass-storage device whose root contains a
``PSP/`` folder (and, on CFW, an ``ISO/`` folder).  This module finds mounted
PSP sticks, lists what is installed on them, and reports free space.
"""
import os
import shutil
import struct
import subprocess
from pathlib import Path

from . import matcher


def _candidate_bases():
    bases = ["/media", "/run/media", "/mnt"]
    user = os.environ.get("USER")
    if user:
        bases += [f"/media/{user}", f"/run/media/{user}"]
    return list(dict.fromkeys(bases))


def is_psp_root(path):
    """True if ``path`` looks like the root of a PSP memory stick.

    A real stick has a ``PSP`` folder containing the firmware-created ``GAME``
    or ``SYSTEM`` subfolders, and CFW sticks also have an ``ISO`` folder.  A
    plain folder named ``psp`` that only holds saves (e.g. a cloud-save
    collection) is *not* a memory stick, so it is rejected.
    """
    p = Path(path)
    if not p.is_dir():
        return False
    psp_dir = p / "PSP"
    if psp_dir.is_dir():
        return (psp_dir / "GAME").is_dir() or (psp_dir / "SYSTEM").is_dir()
    return (p / "ISO").is_dir()


def detect_mounts():
    """Return a list of mounted filesystem roots that look like PSP sticks."""
    found = []
    seen = set()
    for base in _candidate_bases():
        base_path = Path(base)
        if not base_path.is_dir():
            continue
        try:
            children = list(base_path.iterdir())
        except OSError:
            continue
        for child in sorted(children):
            if not child.is_dir():
                continue
            try:
                key = str(child.resolve())
            except OSError:
                key = str(child)
            if key in seen:
                continue
            seen.add(key)
            if is_psp_root(child):
                found.append(str(child))
    return found


def free_space(root):
    """Return ``{total, used, free}`` bytes, or None on failure."""
    try:
        usage = shutil.disk_usage(root)
    except OSError:
        return None
    return {"total": usage.total, "used": usage.used, "free": usage.free}


def _file_item(path, kind):
    try:
        st = path.stat()
        size = st.st_size
        mtime = int(st.st_mtime)
    except OSError:
        size = 0
        mtime = 0
    return {"name": path.name, "path": str(path), "size": size, "mtime": mtime, "kind": kind}


DISC_SUFFIXES = {".iso", ".cso"}


def _iso_items(root):
    """Split installed disc images into games (ISO/), videos (ISO/VIDEO/), and
    homebrew ISOs (ISO/Homebrew/, used by Category Lite). Handles both ``.iso``
    and compressed ``.cso`` images."""
    games, videos, homebrew = [], [], []
    iso_dir = root / "ISO"
    if iso_dir.is_dir():
        for path in sorted(iso_dir.rglob("*")):
            if not path.is_file() or path.suffix.lower() not in DISC_SUFFIXES:
                continue
            rel = path.relative_to(iso_dir)
            first = rel.parts[0].lower() if rel.parts else ""
            if first in ("video", "videos"):
                videos.append(_file_item(path, "video"))
            elif "homebrew" in first:
                # Category Lite may use "Homebrew", "CAT_Homebrew", etc.
                homebrew.append(_file_item(path, "homebrew_iso"))
            else:
                item = _file_item(path, "game")
                item["category"] = rel.parts[0] if len(rel.parts) > 1 else ""
                games.append(item)

    # UMD video ISOs may also live in PSP/VIDEO on some CFW setups.
    psp_video = root / "PSP" / "VIDEO"
    if psp_video.is_dir():
        for path in sorted(psp_video.rglob("*")):
            if path.is_file() and path.suffix.lower() in DISC_SUFFIXES:
                videos.append(_file_item(path, "video"))
    return games, videos, homebrew


def _read_pbp_param_sfo(pbp_path):
    """Extract the PARAM.SFO blob embedded in an EBOOT.PBP, or None."""
    try:
        with open(pbp_path, "rb") as fh:
            header = fh.read(0x28)
    except OSError:
        return None
    if len(header) < 0x28 or header[:4] != b"\x00PBP":
        return None
    offsets = struct.unpack_from("<8I", header, 0x08)
    start, end = offsets[0], offsets[1]
    if end <= start:
        return None
    try:
        with open(pbp_path, "rb") as fh:
            fh.seek(start)
            return fh.read(end - start)
    except OSError:
        return None


def _sfo_strings(data):
    """Parse a PARAM.SFO blob into a ``{key: value}`` dict of strings/ints."""
    out = {}
    if len(data) < 20 or data[:4] != b"\x00PSF":
        return out
    _version, key_table, data_table, count = struct.unpack_from("<IIII", data, 0x04)
    if count <= 0:
        return out
    size = len(data)
    # The index table sits right after the header, before the key table.
    index_start = key_table - count * 16
    if index_start < 20:
        index_start = 20
    if index_start + count * 16 > size:
        return out
    for i in range(count):
        off = index_start + i * 16
        key_off, fmt, length, _max_len, data_off = struct.unpack_from("<HHIII", data, off)
        if key_table + key_off >= size:
            continue
        key = data[key_table + key_off:].split(b"\x00", 1)[0].decode("utf-8", "replace")
        value_off = data_table + data_off
        if value_off + length > size:
            continue
        raw = data[value_off:value_off + length]
        if fmt in (0x0004, 0x0404) and length == 4:  # uint32
            out[key] = struct.unpack_from("<I", data, value_off)[0]
        else:  # UTF-8 string
            out[key] = raw.split(b"\x00", 1)[0].decode("utf-8", "replace").strip()
    return out


def _pbp_title(pbp_path):
    """Return the TITLE from an EBOOT.PBP's PARAM.SFO, or None."""
    try:
        sfo = _read_pbp_param_sfo(pbp_path)
        if sfo:
            return _sfo_strings(sfo).get("TITLE") or None
    except Exception:  # noqa: BLE001
        return None
    return None


def _homebrew_items(root):
    """PSP/GAME/**/EBOOT.PBP (homebrew, PSN, and Category Lite folders).

    Category Lite nests homebrew one level deeper
    (e.g. PSP/GAME/CAT_Homebrew/App/EBOOT.PBP), so search recursively.  The
    display name is the title from the PBP's embedded PARAM.SFO when available,
    falling back to the folder name."""
    items = []
    game_dir = root / "PSP" / "GAME"
    if game_dir.is_dir():
        for eboot in sorted(game_dir.rglob("EBOOT.PBP")):
            if not eboot.is_file():
                continue
            try:
                st = eboot.stat()
                size = st.st_size
                mtime = int(st.st_mtime)
            except OSError:
                size = 0
                mtime = 0
            name = _pbp_title(eboot) or eboot.parent.name
            items.append({"name": name, "id": eboot.parent.name, "path": str(eboot), "size": size, "mtime": mtime, "kind": "homebrew"})
    return items


def _pkg_items(root):
    """.pkg installer files at the root of the stick (CFW package installers)."""
    items = []
    try:
        entries = list(root.iterdir())
    except OSError:
        return items
    for path in sorted(entries):
        if path.is_file() and path.suffix.lower() == ".pkg":
            items.append(_file_item(path, "pkg"))
    return items


def list_installed(root):
    """Return ``{games, videos, homebrew, pkg}`` inventories for a PSP root."""
    root = Path(root)
    games, videos, homebrew_isos = _iso_items(root)
    return {
        "games": games,
        "videos": videos,
        "homebrew": _homebrew_items(root) + homebrew_isos,
        "pkg": _pkg_items(root),
    }


def build_index(inventory):
    """Build a lookup index for matching library discs against the PSP.

    Returns ``{"raw", "names", "homebrew_ids", "homebrew_names"}`` where
    ``raw`` holds lowercase ISO stems (exact filename match), ``names`` holds
    fuzzy-normalized stems, and the homebrew keys index ``PSP/GAME/**/EBOOT.PBP``
    (homebrew + PS1 Classics, including Category Lite subfolders) by folder id
    and embedded title.
    """
    raw, names = set(), set()
    for group in ("games", "videos"):
        for it in (inventory or {}).get(group, []):
            stem = Path(it["name"]).stem
            raw.add(stem.lower())
            norm = matcher.normalize_name(stem)
            if norm:
                names.add(norm)

    hb_ids, hb_names = set(), set()
    for it in (inventory or {}).get("homebrew", []):
        ident = str(it.get("id") or "").strip().lower()
        if ident:
            hb_ids.add(ident)
        stem = Path(it.get("name") or "").stem.strip().lower()
        if stem:
            hb_names.add(stem)
        norm = matcher.normalize_name(Path(it.get("name") or "").stem)
        if norm:
            hb_names.add(norm)
    return {"raw": raw, "names": names,
            "homebrew_ids": hb_ids, "homebrew_names": hb_names}


def find_umd_iso(root):
    """Return the path of the ``UMD9660.ISO`` at a mount root, or None."""
    try:
        entries = list(Path(root).iterdir())
    except OSError:
        return None
    for path in entries:
        if path.is_file() and path.name.lower() == "umd9660.iso":
            return path
    return None


def umd_mounts():
    """Return mounted UMD volumes (USB UMD mode).

    In USB UMD mode the PSP exposes a single volume whose name is the disc's
    catalogue number (e.g. ``ULUS-10268``) and whose root holds one ISO named
    ``UMD9660.ISO``."""
    found, seen = [], set()
    for base in _candidate_bases():
        base_path = Path(base)
        if not base_path.is_dir():
            continue
        try:
            children = list(base_path.iterdir())
        except OSError:
            continue
        for child in sorted(children):
            if not child.is_dir():
                continue
            try:
                key = str(child.resolve())
            except OSError:
                key = str(child)
            if key in seen:
                continue
            seen.add(key)
            if find_umd_iso(child) is not None:
                found.append(str(child))
    return found


def block_device_for_mount(root):
    """Return the block device backing a mount point, or None."""
    try:
        proc = subprocess.run(
            ["findmnt", "-n", "-o", "SOURCE", "--target", str(root)],
            capture_output=True, text=True, timeout=10,
        )
    except (OSError, subprocess.SubprocessError):
        return None
    if proc.returncode != 0:
        return None
    lines = [ln.strip() for ln in proc.stdout.splitlines() if ln.strip()]
    return lines[0] if lines else None


def eject(root):
    """Unmount a mounted PSP stick so it is safe to unplug.

    Returns ``(ok, message)``.
    """
    root = Path(root)
    device = block_device_for_mount(root)
    if not device:
        return False, f"No block device found for {root}"
    try:
        proc = subprocess.run(
            ["udisksctl", "unmount", "-b", device],
            capture_output=True, text=True, timeout=30,
        )
    except OSError as exc:
        return False, f"udisksctl failed: {exc}"
    if proc.returncode == 0:
        return True, proc.stdout.strip() or f"Unmounted {root}"
    err = (proc.stderr or proc.stdout or "").strip()
    return False, err or f"Unmount failed (exit code {proc.returncode})"