psp-preservation-tool/app/nopaystation.py

"""NoPayStation (.pkg) browsing and extraction.

NoPayStation publishes per-platform TSV lists (``https://nopaystation.com/tsv/``)
with one row per downloadable PSN item. Each row links to an encrypted PSN
``.pkg`` (a ``\\x7FPKG`` container, not a ZIP) that is decrypted by the bundled
pure-Python decryptor (``app/pkg.py``) into a ``PSP/GAME/<TITLE_ID>/``-style
folder — no Wine or ``pkg2zip`` required.
"""
from __future__ import annotations

import logging
import os
import shutil
import subprocess
import urllib.parse
import urllib.request
from datetime import datetime
from pathlib import Path
from typing import List, Optional

from . import config

log = logging.getLogger("nopaystation")

_TSV_NAMES = {
    "psp": "PSP_GAMES.tsv",
    "psx": "PSX_GAMES.tsv",
    "theme": "PSP_THEMES.tsv",
    "dlc": "PSP_DLCS.tsv",
}

# NoPayStation TSVs are tab-separated with a header row whose exact columns
# differ per platform (PSP has "Type"/"RAP" columns, PSX has "Original Name").
# Map header names to canonical keys instead of relying on fixed positions.
_HEADER_ALIASES = {
    "title id": "title_id",
    "region": "region",
    "name": "name",
    "pkg direct link": "pkg_url",
    "content id": "content_id",
    "last modification date": "date",
    "file size": "size",
    "sha256": "sha256",
    "zrif": "zrif",
    "rap": "rap",
    "download .rap file": "rap_url",
    "type": "content_kind",
    "original name": "original_name",
}


def tsv_url(content_type: str, base: str = "") -> str:
    """TSV URL for a content type (``psp`` / ``psx`` / ``theme`` / ``dlc``)."""
    base = (base or config.NPS_TSV_BASE).rstrip("/")
    return f"{base}/{_TSV_NAMES.get(content_type, _TSV_NAMES['psp'])}"


def _tsv_dir() -> Path:
    d = config.DATA_DIR / "nps"
    d.mkdir(parents=True, exist_ok=True)
    return d


def _parse_date(value: str) -> str:
    value = (value or "").strip()
    if not value:
        return ""
    for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d"):
        try:
            return datetime.strptime(value, fmt).strftime("%Y-%m-%d")
        except ValueError:
            continue
    return value[:10]


def parse_tsv(text: str, content_type: str) -> List[dict]:
    """Parse a NoPayStation TSV into store item dicts.

    The header row drives column mapping (PSP and PSX TSVs differ). Each item
    carries the fields the Store page expects (``name``, ``size``,
    ``release_date``, ``mtime_ns``) plus ``title_id``, ``pkg_url``, ``sha256``,
    ``content_type`` and the cover ``system`` ("games" for PSP, "ps1" for PSX).
    """
    system = {"psp": "games", "psx": "ps1", "theme": "theme", "dlc": "dlc"}.get(content_type, "games")
    items: List[dict] = []
    header = None
    for line in text.splitlines():
        line = line.rstrip("\r\n")
        if not line.strip():
            continue
        cols = line.split("\t")
        if header is None:
            header = [c.strip().lower() for c in cols]
            continue
        row = {}
        for i, key in enumerate(header):
            canon = _HEADER_ALIASES.get(key)
            if canon and i < len(cols):
                row[canon] = cols[i].strip()
        name = row.get("name") or ""
        pkg_url = row.get("pkg_url") or ""
        # NoPayStation marks delisted packages with the literal "MISSING" in the
        # URL column — skip them so they can't be queued for download.
        if not name or not pkg_url or not pkg_url.lower().startswith(("http://", "https://")):
            continue
        title_id = row.get("title_id") or ""
        try:
            size = int(float(row.get("size") or 0))
        except (TypeError, ValueError):
            size = 0
        date = _parse_date(row.get("date") or "")
        try:
            mtime_ns = int(datetime.strptime(date, "%Y-%m-%d").timestamp() * 1_000_000_000) if date else 0
        except ValueError:
            mtime_ns = 0
        items.append({
            "name": name,
            "title_id": title_id,
            "region": row.get("region") or "",
            "pkg_url": pkg_url,
            "zrif": row.get("zrif") or "",
            "content_id": row.get("content_id") or "",
            "size": size,
            "sha256": (row.get("sha256") or "").strip().lower(),
            "release_date": date,
            "mtime_ns": mtime_ns,
            "content_type": content_type,
            "system": system,
            "content_kind": row.get("content_kind") or "",
        })
    return items


def fetch_items(content_type: str, url: str = "", timeout: int = 60) -> List[dict]:
    """Download (and cache) a TSV, returning parsed item dicts."""
    url = url or tsv_url(content_type)
    name = _TSV_NAMES.get(content_type, "PSP_GAMES.tsv")
    cache = _tsv_dir() / name
    if cache.is_file():
        try:
            return parse_tsv(cache.read_text(encoding="utf-8", errors="replace"), content_type)
        except OSError:
            pass
    req = urllib.request.Request(url, headers={"User-Agent": "psp-preserve/0.1"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        data = resp.read()
    text = data.decode("utf-8", errors="replace")
    cache.write_text(text, encoding="utf-8")
    return parse_tsv(text, content_type)


def find_pkg2zip(configured_path: str = "") -> Optional[Path]:
    """Locate pkg2zip: configured path, then the bundled binary, then PATH."""
    if configured_path:
        p = Path(configured_path)
        return p if p.is_file() else None
    bundled = config.DATA_DIR / "bin" / "pkg2zip.exe"
    if bundled.is_file():
        return bundled
    found = shutil.which("pkg2zip")
    return Path(found) if found else None


def _pkg2zip_cmd(pkg2zip: Path) -> List[str]:
    """Command prefix for pkg2zip, wrapping the bundled .exe in Wine on Linux."""
    if pkg2zip.suffix.lower() == ".exe" and os.name != "nt":
        wine = shutil.which("wine") or shutil.which("wine64")
        if not wine:
            raise RuntimeError(
                "The bundled pkg2zip is a Windows binary and Wine is not installed — "
                "install Wine (or point Settings → Store at a native pkg2zip)."
            )
        return [wine, str(pkg2zip)]
    return [str(pkg2zip)]


def _safe_folder(value: str) -> str:
    import re
    cleaned = re.sub(r"[^A-Za-z0-9]", "", value or "").upper()
    return cleaned or "PKG"


def _ensure_keys_bin(folder, kind):
    """PS1 Classics need KEYS.BIN next to EBOOT.PBP to boot on CFW (npdrm)."""
    if kind != "psx":
        return
    src = config.DATA_DIR / "loader" / "KEYS.BIN"
    dest = Path(folder) / "KEYS.BIN"
    if src.is_file() and not dest.exists():
        try:
            shutil.copy2(src, dest)
        except OSError:
            pass


def extract_pkg(pkg_path, title_id: str, dest_dir, pkg2zip: Path = None,
                kind_hint: str = None, progress_cb=None, cancel_event=None) -> Path:
    """Decrypt/extract a ``.pkg`` into ``dest_dir/<TITLE_ID>/``.

    Uses the bundled pure-Python decryptor (no Wine / pkg2zip needed); falls
    back to pkg2zip only if that fails and a binary is available. Returns the
    destination game folder path.
    """
    from . import pkg as pkgmod
    from .store import TransferCancelled

    pkg_path = Path(pkg_path)
    dest = Path(dest_dir) / _safe_folder(title_id)
    if dest.is_dir() and any(dest.iterdir()):
        return dest

    # 1. Pure-Python decryptor (primary).
    try:
        _kind, _tid, folder = pkgmod.extract(pkg_path, dest_dir, kind_hint,
                                             progress_cb, cancel_event)
        _ensure_keys_bin(folder, _kind)
        return folder
    except TransferCancelled:
        raise
    except Exception as exc:  # noqa: BLE001
        py_exc = exc
        log.warning("python pkg decrypt failed (%s); trying pkg2zip fallback", exc)

    # 2. pkg2zip fallback (needs a binary; usually Wine + the bundled .exe).
    if pkg2zip is None:
        pkg2zip = find_pkg2zip()
    if pkg2zip is None:
        raise RuntimeError(
            f"could not decrypt {pkg_path.name} (no pkg2zip available): {py_exc}"
        )

    folder_name = _safe_folder(title_id)
    tmp = Path(dest_dir).with_name(".nps-tmp-" + folder_name)
    if tmp.exists():
        shutil.rmtree(tmp, ignore_errors=True)
    tmp.mkdir(parents=True, exist_ok=True)

    if progress_cb is not None:
        progress_cb("Extracting", 0.0, "")
    try:
        cmd = _pkg2zip_cmd(Path(pkg2zip)) + ["-x", str(pkg_path)]
        env = dict(os.environ)
        env["WINEDEBUG"] = "-all"  # silence Wine's console chatter
        proc = subprocess.Popen(
            cmd,
            cwd=str(tmp), stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            text=True, bufsize=1, env=env,
        )
        assert proc.stdout is not None
        for _line in proc.stdout:
            if cancel_event is not None and cancel_event.is_set():
                proc.kill()
                raise TransferCancelled()
        proc.wait()
        if proc.returncode != 0:
            raise RuntimeError(f"pkg2zip failed (exit {proc.returncode}) for {pkg_path.name}")
    finally:
        if cancel_event is not None and cancel_event.is_set():
            shutil.rmtree(tmp, ignore_errors=True)
            raise TransferCancelled()

    eboots = list(tmp.rglob("EBOOT.PBP"))
    if not eboots:
        shutil.rmtree(tmp, ignore_errors=True)
        raise RuntimeError(f"pkg2zip produced no EBOOT.PBP from {pkg_path.name}")
    src_folder = eboots[0].parent

    dest.mkdir(parents=True, exist_ok=True)
    for p in sorted(src_folder.rglob("*")):
        if not p.is_file():
            continue
        rel = p.relative_to(src_folder)
        target = dest / rel
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(p, target)
    if progress_cb is not None:
        progress_cb("Extracting", 1.0, "")
    shutil.rmtree(tmp, ignore_errors=True)
    return dest