UnityExtract64

unity_extractor.py

#!/usr/bin/env python3
"""Source-based Unity NSISBI extractor for legacy and Unity 6000 installers.

Supported layouts discovered from Unity installers:

* Legacy 2020-2022: 4-byte flagged raw-LZMA1 chunks. The decoded NSIS
  instruction table starts at 0x1b74, uses 36-byte records with opcode in
  field 0, and field 3 is the compressed data-chunk offset. Each data chunk
  expands directly to one file; there is no extra record-size word.
* Unity 6000: 3-byte raw-LZMA1 chunks, an eight-byte declared-size header
  prefix, and a 36-byte Unity command table whose data stream contains
  8-byte record headers.

All progress/diagnostics belong on stderr; stdout is JSON. Extraction uses
same-directory temporary files and atomic replacement; listing and verification
never create a temporary workspace.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import logging
import lzma
import os
import re
import struct
import sys
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO

MAGIC = b"NullsoftInst"
SIG = 0xDEADBEEF
FLAG_EXTERNAL = 0x20
FLAG_STUB = 0x40
MAX_CHUNK_OUTPUT = 256 * 1024 * 1024
MAX_RECORD_SIZE = 1 << 42
LOGGER = logging.getLogger("unity-nsisbi")

@dataclass
class Header:
    offset: int
    flags: int
    header_size: int
    following: int
    low: int
    high: int
    chunk_width: int

    @property
    def external(self) -> bool:
        return bool(self.flags & FLAG_EXTERNAL)

@dataclass
class FilePlan:
    offset: int
    paths: list[str]


def u32(data: bytes, at: int) -> int:
    return struct.unpack_from("<I", data, at)[0]


def u16(data: bytes, at: int) -> int:
    return struct.unpack_from("<H", data, at)[0]


def u64(data: bytes, at: int) -> int:
    return struct.unpack_from("<Q", data, at)[0]


def find_header(f: BinaryIO) -> Header:
    f.seek(0); carry = b""; absolute = 0
    while True:
        block = f.read(1024 * 1024)
        if not block: break
        hay = carry + block; at = hay.find(MAGIC)
        if at >= 0:
            marker = absolute - len(carry) + at; off = marker - 8
            f.seek(off); raw = f.read(36)
            if len(raw) == 36 and u32(raw, 4) == SIG and raw[8:20] == MAGIC:
                f.seek(off + 36); probe = f.read(9); f.seek(off + 36)
                width = 4 if len(probe) >= 9 and (u32(probe, 0) & 0x80000000) else 3
                header = Header(off, u32(raw, 0), u32(raw, 20), u32(raw, 24), u32(raw, 28), u32(raw, 32), width)
                LOGGER.debug(
                    "header: offset=0x%x flags=0x%x declared_size=%d following=%d low=0x%x high=0x%x chunk_width=%d",
                    header.offset, header.flags, header.header_size, header.following,
                    header.low, header.high, header.chunk_width,
                )
                return header
        carry = hay[-(len(MAGIC) - 1):]; absolute += len(block)
    raise ValueError("NSISBI first header not found")


def lzma_filter(props: bytes) -> dict[str, int]:
    if len(props) != 5: raise ValueError("truncated LZMA properties")
    p = props[0]
    if p >= 225: raise ValueError(f"invalid LZMA property 0x{p:02x}")
    rem = p // 9
    dictionary = int.from_bytes(props[1:5], "little")
    if dictionary < 4096 or dictionary > 1 << 30: raise ValueError("invalid LZMA dictionary")
    return {"id": lzma.FILTER_LZMA1, "dict_size": dictionary, "lc": p % 9, "lp": rem % 5, "pb": rem // 5}


def decode_chunk(f: BinaryIO, width: int, expected: int | None = None) -> tuple[bytes, int]:
    start = f.tell()
    raw = f.read(width)
    if len(raw) != width: raise EOFError("truncated chunk length")
    word = int.from_bytes(raw, "little")
    length = (word & 0x7fffffff) if width == 4 else word
    if length < 5: raise ValueError(f"invalid chunk length {length}")
    props = f.read(5); compressed = f.read(length - 5)
    if len(props) != 5 or len(compressed) != length - 5: raise EOFError("truncated chunk")
    decoder = lzma.LZMADecompressor(format=lzma.FORMAT_RAW, filters=[lzma_filter(props)])
    output = decoder.decompress(compressed)
    if not decoder.eof or decoder.unused_data: raise ValueError("LZMA chunk did not terminate cleanly")
    if len(output) > MAX_CHUNK_OUTPUT: raise ValueError("chunk output exceeds safety limit")
    if expected is not None and len(output) != expected: raise ValueError(f"chunk produced {len(output)}, expected {expected}")
    LOGGER.debug(
        "chunk: file_offset=0x%x width=%d stored_length=%d compressed=%d decoded=%d next=0x%x",
        start, width, length, length - 5, len(output), f.tell(),
    )
    return output, width + length


def decode_header(f: BinaryIO, h: Header) -> tuple[bytes, int]:
    f.seek(h.offset + 36); result = bytearray(); chunk_index = 0
    # Unity 6000 stores an eight-byte logical-size prefix in the decoded
    # header. The NSISBI field excludes that prefix, so its final chunk ends
    # eight decoded bytes after h.header_size. Older layouts do not have it.
    first, _ = decode_chunk(f, h.chunk_width)
    result.extend(first); chunk_index += 1
    modern_prefix = len(first) >= 8 and int.from_bytes(first[:8], "little") == h.header_size
    target = h.header_size + 8 if modern_prefix else h.header_size
    LOGGER.debug(
        "header stream: target=%d (declared=%d, modern_size_prefix=%s)",
        target, h.header_size, modern_prefix,
    )
    while len(result) < target:
        chunk, _ = decode_chunk(f, h.chunk_width)
        result.extend(chunk); chunk_index += 1
    if len(result) != target:
        raise ValueError(
            f"header stream decoded {len(result)} bytes, expected {target} "
            f"(declared {h.header_size}{' + 8-byte Unity 6000 prefix' if modern_prefix else ''})"
        )
    LOGGER.debug("header decoded: chunks=%d bytes=%d data_start=0x%x", chunk_index, len(result), f.tell())
    return bytes(result), f.tell()


def decode_string(data: bytes, base: int, end: int, reference: int, legacy: bool) -> str:
    # References are UTF-16 code-unit offsets from the first string unit.
    # Legacy strings begin directly at the block boundary; Unity 6000 has an
    # eight-byte block header before its first string unit.
    at = base + (0 if legacy else 8) + reference * 2
    if not base <= at < end - 1: return ""
    out: list[str] = []
    while at + 1 < end:
        unit = u16(data, at)
        at += 2
        if unit == 0: break
        if unit == 3 and at + 1 < end:
            encoded = u16(data, at)
            at += 2
            var_id = (encoded & 0x7f) | (((encoded >> 8) & 0x7f) << 7)
            variables = {
                20: "$CMDLINE", 21: "$INSTDIR", 22: "$OUTDIR",
                23: "$EXEDIR", 24: "$LANGUAGE", 25: "$TEMP",
                26: "$PLUGINSDIR", 27: "$EXEFILE", 28: "$EXEPATH",
                29: "$HWNDPARENT", 31: "$INSTDIR",
            }
            if var_id < 10: out.append(f"${var_id}")
            elif var_id < 20: out.append(f"$R{var_id - 10}")
            else: out.append(variables.get(var_id, f"$VAR{var_id}"))
            continue
        if unit < 32 and unit not in (9, 10, 13): return ""
        out.append(chr(unit))
    return "".join(out)


WINDOWS_RESERVED_NAMES = {
    "CON", "PRN", "AUX", "NUL",
    *(f"COM{index}" for index in range(1, 10)),
    *(f"LPT{index}" for index in range(1, 10)),
}


def safe_name(name: str) -> str:
    name = name.replace("\\", "/")
    name = re.sub(r"^\$(?:INSTDIR|OUTDIR|EXEDIR|TEMP|PLUGINSDIR|EXEFILE|EXEPATH|VAR\d+)/?", "", name, flags=re.I)
    name = re.sub(r"^[A-Za-z]:/?", "", name).lstrip("/")
    parts = []
    for part in name.split("/"):
        if part in ("", "."): continue
        if part == "..":
            if parts: parts.pop()
            continue
        parts.append(part)
    sanitized: list[str] = []
    for part in parts:
        value = re.sub(r"[^A-Za-z0-9._-]+", "_", part)
        stem = value.rstrip(" .").split(".", 1)[0].upper()
        if stem in WINDOWS_RESERVED_NAMES:
            value = f"_{value}"
        sanitized.append(value)
    return "/".join(sanitized) or "unnamed.bin"


def destination_key(relative: str) -> str:
    """Return a filesystem-insensitive key for collision detection."""
    return relative.replace("\\", "/").casefold()


SPECIAL_ROOT_RE = re.compile(
    r"^\$(INSTDIR|OUTDIR|EXEDIR|TEMP|PLUGINSDIR|EXEFILE|EXEPATH|VAR\d+)(?:[\\\\/]|$)",
    re.IGNORECASE,
)


def join_path(folder: str, filename: str) -> str:
    # NSIS can put an absolute special-folder token in the filename itself.
    # Preserve PLUGINSDIR as a visible tree root; the other installer-root
    # tokens intentionally resolve relative to the requested output root.
    special = SPECIAL_ROOT_RE.match(filename)
    if special:
        token = special.group(1).upper()
        remainder = filename[special.end():]
        if token == "PLUGINSDIR":
            return safe_name(f"PLUGINSDIR/{remainder}")
        return safe_name(remainder)
    filename = safe_name(filename)
    folder = safe_name(folder) if folder else ""
    return safe_name(f"{folder}/{filename}" if folder else filename)


def legacy_plans(header: bytes) -> tuple[list[FilePlan], dict[str, object]]:
    entries = u32(header, 0x14); count = u32(header, 0x18); entries_end = u32(header, 0x1c)
    strings = entries_end
    strings_end = u32(header, 0x24)
    if not (entries == 0x1b74 and entries_end == entries + count * 36 and entries_end <= strings < strings_end <= len(header)):
        raise ValueError("unsupported legacy NSISBI table boundaries")
    folder = "PLUGINSDIR"; by_offset: dict[int, list[str]] = {}; commands = 0; mapped = 0
    for i in range(count):
        fields = struct.unpack_from("<9I", header, entries + i * 36); opcode = fields[0]
        if opcode == 11:
            value = decode_string(header, strings, strings_end, fields[1], True)
            if value: folder = value
        elif opcode == 20:
            commands += 1
            filename = decode_string(header, strings, strings_end, fields[2], True)
            if not filename: continue
            destination = join_path(folder, filename); mapped += 1
            paths = by_offset.setdefault(fields[3], [])
            if destination not in paths: paths.append(destination)
    plans = [FilePlan(offset, paths) for offset, paths in by_offset.items()]
    meta = {"format": "legacy-36-byte-opcode0-4byte-chunks", "entry_offset": entries, "entry_count": count, "entry_end": entries_end, "string_offset": strings, "string_end": strings_end, "extract_command_count": commands, "mapped_extract_command_count": mapped, "named_data_offsets": len(by_offset), "mapped_destination_count": sum(len(x) for x in by_offset.values())}
    if commands < 32 or mapped * 100 < commands * 90: raise ValueError("legacy table did not pass mapping sanity checks")
    return plans, meta


def modern_plans(header: bytes) -> tuple[list[FilePlan], dict[str, object]]:
    entries = u32(header, 0x1c); strings = u32(header, 0x24); strings_end = u32(header, 0x2c)
    if not (0 < entries < strings < strings_end <= len(header) and (strings - entries) % 36 == 0): raise ValueError("unsupported Unity 6000 table boundaries")
    folder = "PLUGINSDIR"; by_offset: dict[int, list[str]] = {}; commands = mapped = 0
    for i in range((strings - entries) // 36):
        fields = struct.unpack_from("<9I", header, entries + i * 36); opcode = fields[2]
        if opcode == 11:
            value = decode_string(header, strings, strings_end, fields[3], False)
            if value: folder = value
        elif opcode == 20:
            commands += 1; filename = decode_string(header, strings, strings_end, fields[4], False)
            if not filename: continue
            destination = join_path(folder, filename); mapped += 1
            paths = by_offset.setdefault(fields[5], [])
            if destination not in paths: paths.append(destination)
    meta = {"format": "unity6000-36-byte-opcode2-3byte-records", "entry_offset": entries, "string_offset": strings, "string_end": strings_end, "extract_command_count": commands, "mapped_extract_command_count": mapped, "named_data_offsets": len(by_offset), "mapped_destination_count": sum(len(x) for x in by_offset.values())}
    if commands < 32 or mapped * 100 < commands * 90: raise ValueError("Unity 6000 table did not pass mapping sanity checks")
    return [FilePlan(k, v) for k, v in by_offset.items()], meta


def validate_destinations(plans: list[FilePlan]) -> None:
    seen: dict[str, str] = {}
    for plan in plans:
        for relative in plan.paths:
            key = destination_key(relative)
            previous = seen.setdefault(key, relative)
            if previous != relative:
                raise ValueError(
                    f"destination paths collide on a case-insensitive filesystem: "
                    f"{previous!r} and {relative!r}"
                )


def tree(paths: list[str]) -> dict[str, object]:
    root = {"directories": {}, "files": []}
    for path in paths:
        node = root; parts = [x for x in path.split("/") if x]
        for part in parts[:-1]: node = node["directories"].setdefault(part, {"directories": {}, "files": []})
        if parts: node["files"].append(parts[-1]) if parts[-1] not in node["files"] else None
    return root


def magic(data: bytes) -> str:
    if data.startswith(b"MZ"): return "MZ/PE"
    if data.startswith(b"PK"): return "ZIP"
    if data.startswith(b"UnityFS"): return "UnityFS"
    if data.startswith(b"\x7fELF"): return "ELF"
    if data.startswith(b"\x89PNG"): return "PNG"
    return data[:8].hex() if data else "empty"


def destination_path(root: Path, relative: str) -> Path:
    """Resolve a planned path without following an output-tree symlink."""
    parts = [part for part in relative.replace("\\", "/").split("/") if part]
    if not parts:
        raise ValueError("empty extraction path")
    candidate = root.joinpath(*parts)
    current = root
    for part in parts:
        current = current / part
        if current.is_symlink():
            raise ValueError(f"refusing symlinked extraction path: {relative}")
    root_resolved = root.resolve()
    resolved = candidate.resolve(strict=False)
    try:
        resolved.relative_to(root_resolved)
    except ValueError as exc:
        raise ValueError(f"extraction path escapes output directory: {relative}") from exc
    return candidate


class ModernPartialRecord(EOFError):
    def __init__(self, offset: int, expected: int, available: int) -> None:
        super().__init__(
            f"modern record at virtual offset {offset} is truncated: "
            f"expected {expected} bytes, have {available}"
        )
        self.offset = offset
        self.expected = expected
        self.available = available


class ModernDataStream:
    """Sequentially expose Unity 6000's virtual record stream.

    Modern command offsets address the decompressed stream, not installer file
    offsets. Each record is encoded as an eight-byte little-endian payload
    size followed by that many bytes of payload. Compressed chunks are decoded
    only as needed, keeping memory bounded to the current chunks and record.
    """

    def __init__(self, file: BinaryIO, width: int, data_start: int) -> None:
        self.file = file
        self.width = width
        self.data_start = data_start
        self.buffer = bytearray()
        self.position = 0
        self.file.seek(data_start)

    def _fill(self, required: int) -> None:
        while len(self.buffer) < required:
            chunk, _ = decode_chunk(self.file, self.width)
            self.buffer.extend(chunk)

    def read(self, size: int) -> bytes:
        if size < 0 or size > MAX_RECORD_SIZE:
            raise ValueError(f"invalid modern record read size {size}")
        try:
            self._fill(size)
        except EOFError as exc:
            raise EOFError(
                f"modern data stream truncated at virtual offset {self.position}"
            ) from exc
        value = bytes(self.buffer[:size])
        del self.buffer[:size]
        self.position += size
        return value

    def iter_record(self) -> tuple[int, bool, int, object]:
        """Read the next sequential Unity 6000 record.

        The command-table offsets are lookup keys. They are not a promise that
        every command offset is a top-level record boundary (package entries
        can refer into a containing payload), so records must be consumed in
        stream order rather than sought by plan offset.
        """
        offset = self.position
        raw = int.from_bytes(self.read(8), "little")
        compressed = bool(raw & (1 << 63))
        size = raw & ((1 << 63) - 1)
        if size > MAX_RECORD_SIZE:
            raise ValueError(f"modern record at virtual offset {offset} has unreasonable size {size}")

        def chunks():
            remaining = size
            while remaining:
                part = min(remaining, 8 * 1024 * 1024)
                try:
                    yield self.read(part)
                except EOFError as exc:
                    available = len(self.buffer)
                    raise ModernPartialRecord(offset, size, size - remaining + available) from exc
                remaining -= part

        return offset, compressed, size, chunks()


def replace_with_retry(source: Path, destination: Path) -> None:
    for attempt in range(6):
        try:
            os.replace(source, destination)
            return
        except PermissionError:
            if attempt == 5:
                raise
            time.sleep(0.25 * (attempt + 1))


def atomic_write(path: Path, payload: bytes) -> None:
    """Write a payload through a same-directory temporary file."""
    temporary: Path | None = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="wb", prefix=f".{path.name}.", suffix=".part",
            dir=path.parent, delete=False,
        ) as stream:
            temporary = Path(stream.name)
            stream.write(payload)
            stream.flush()
            try:
                os.fsync(stream.fileno())
            except OSError as exc:
                LOGGER.debug("filesystem sync unavailable for %s: %s", path, exc)
        for attempt in range(6):
            try:
                os.replace(temporary, path)
                temporary = None
                break
            except PermissionError:
                if attempt == 5:
                    raise
                time.sleep(0.25 * (attempt + 1))
    finally:
        if temporary is not None:
            try:
                temporary.unlink()
            except FileNotFoundError:
                pass


def run(installer: Path, output: Path | None, list_only: bool, verify: bool, limit: int | None) -> dict[str, object]:
    with installer.open("rb") as f:
        h = find_header(f); header, data_start = decode_header(f, h)
        modern = len(header) >= 8 and int.from_bytes(header[:8], "little") == h.header_size
        LOGGER.debug("layout: %s", "Unity 6000" if modern else "legacy 2020-2022")
        plans, table = modern_plans(header) if modern else legacy_plans(header)
        LOGGER.debug(
            "table: format=%s extract_commands=%s mapped=%s named_offsets=%s",
            table.get("format"), table.get("extract_command_count"),
            table.get("mapped_extract_command_count"), table.get("named_data_offsets"),
        )
        validate_destinations(plans)
        paths = sorted({p for plan in plans for p in plan.paths})
        result: dict[str, object] = {
            "installer": str(installer),
            "header_offset": hex(h.offset),
            "flags": hex(h.flags),
            "header_size": h.header_size,
            "chunk_width": h.chunk_width,
            "header_prefix_size": 8 if modern else 0,
            "data_start": hex(data_start),
            "external_data": h.external,
            "table": table,
        }
        if list_only:
            listed = paths if limit is None else paths[:limit]
            result.update({"tree": tree(paths), "paths": listed, "tree_file_count": len(paths), "listed_tree_file_count": len(listed), "data_stream_scanned": False, "records": []})
            return result
        if not verify and output is None: raise ValueError("output required unless --list or --verify")
        output and output.mkdir(parents=True, exist_ok=True)
        records: list[dict[str, object]] = []; unmapped = 0; partial_record = None
        if modern:
            stream = ModernDataStream(f, h.chunk_width, data_start)
            paths_by_offset = {plan.offset: plan.paths for plan in plans}
            index = 0
            while limit is None or len(records) < limit:
                try:
                    offset, compressed, payload_size, payload_chunks = stream.iter_record()
                    relatives = paths_by_offset.get(offset)
                    if not relatives:
                        for _ in payload_chunks:
                            pass
                        unmapped += 1
                        index += 1
                        continue
                    digest = hashlib.sha256(); prefix = bytearray(); actual = []
                    streams = []
                    if output is not None:
                        for rel in relatives:
                            path = destination_path(output, rel)
                            path.parent.mkdir(parents=True, exist_ok=True)
                            path = destination_path(output, rel)
                            temporary = tempfile.NamedTemporaryFile(
                                mode="wb", prefix=f".{path.name}.", suffix=".part",
                                dir=path.parent, delete=False,
                            )
                            streams.append((path, Path(temporary.name), temporary))
                    try:
                        for part in payload_chunks:
                            digest.update(part)
                            if len(prefix) < 8:
                                prefix.extend(part[:8 - len(prefix)])
                            for _, _, target in streams:
                                target.write(part)
                        for path, temporary, target in streams:
                            target.flush()
                            try:
                                os.fsync(target.fileno())
                            except OSError as exc:
                                LOGGER.debug("filesystem sync unavailable for %s: %s", path, exc)
                            target.close()
                            replace_with_retry(temporary, path)
                            actual.append(str(path.relative_to(output)).replace("\\", "/"))
                        streams.clear()
                    finally:
                        for _, temporary, target in streams:
                            try:
                                target.close()
                            finally:
                                try:
                                    temporary.unlink()
                                except FileNotFoundError:
                                    pass
                    records.append({"index": index, "offset": offset, "stored_size": payload_size, "compressed": compressed, "paths": relatives, "path": relatives[0], "magic": magic(bytes(prefix)), "sha256": digest.hexdigest(), "written_paths": actual})
                    index += 1
                except ModernPartialRecord as exc:
                    partial_record = {"offset": exc.offset, "expected_bytes": exc.expected, "available_bytes": exc.available}
                    if not h.external:
                        raise
                    break
                except EOFError:
                    if h.external:
                        partial_record = {"offset": stream.position, "expected_bytes": None, "available_bytes": 0}
                        break
                    raise
        else:
            for index, plan in enumerate(plans):
                if limit is not None and len(records) >= limit: break
                f.seek(data_start + plan.offset)
                payload, _ = decode_chunk(f, h.chunk_width)
                sha = hashlib.sha256(payload).hexdigest(); actual = []
                if output is not None:
                    for rel in plan.paths:
                        path = destination_path(output, rel)
                        path.parent.mkdir(parents=True, exist_ok=True)
                        path = destination_path(output, rel)
                        atomic_write(path, payload)
                        actual.append(rel)
                records.append({"index": index, "offset": plan.offset, "stored_size": len(payload), "paths": plan.paths, "path": plan.paths[0], "magic": magic(payload), "sha256": sha, "written_paths": actual})
        result.update({"data_stream_scanned": True, "record_count": len(records), "records": records, "unmapped_record_offsets": unmapped, "partial_record": partial_record, "external_data_required": bool(h.external and partial_record)})
        if output is not None:
            atomic_write(
                output / "manifest.json",
                (json.dumps(result, indent=2) + "\n").encode("utf-8"),
            )
        return result


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Extract Unity NSISBI legacy and Unity 6000 installers")
    p.add_argument("installer", type=Path); p.add_argument("output", type=Path, nargs="?"); p.add_argument("--list", action="store_true"); p.add_argument("--verify", action="store_true"); p.add_argument("--limit", type=int)
    p.add_argument("-v", "--verbose", action="store_true", help="write parser diagnostics to stderr")
    a = p.parse_args(argv)
    if a.verbose:
        logging.basicConfig(level=logging.DEBUG, format="[nsisbi] %(message)s", stream=sys.stderr)
    if a.list and a.verify: p.error("--list and --verify cannot be combined")
    if a.verify and a.output: p.error("--verify does not take an output directory")
    try: print(json.dumps(run(a.installer, a.output, a.list, a.verify, a.limit), indent=2)); return 0
    except (OSError, ValueError, EOFError, lzma.LZMAError) as exc: print(f"error: {exc}", file=sys.stderr); return 2

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