zcmodkit 0.1.0__py3-none-any.whl

zcmodkit/formats/oodle.py

@@ -0,0 +1,72 @@
+"""Oodle decompression, for reading compressed entries out of the game's paks."""
+
+from __future__ import annotations
+
+import ctypes
+import os
+from pathlib import Path
+
+_SEARCH_NAMES = ("oodle-data-shared.dll", "oo2core_9_win64.dll", "oo2core_8_win64.dll")
+_lib = None
+_fn = None
+
+
+def _candidates() -> list[Path]:
+    """The usual places an Oodle DLL turns up."""
+    out: list[Path] = []
+    env = os.environ.get("ZCMODKIT_OODLE")
+    if env:
+        out.append(Path(env))
+    for base in (Path.cwd(), Path(__file__).parent, Path.home() / "Downloads"):
+        out.extend(base / name for name in _SEARCH_NAMES)
+    # FModel keeps one in its Output/.data folder
+    out.extend((Path.home() / "Downloads").glob("*/Output/.data/oodle-data-shared.dll"))
+    return out
+
+
+def load(path: str | os.PathLike | None = None):
+    """Load the Oodle DLL and hand back its decompress function."""
+    global _lib, _fn
+    if _fn is not None and path is None:
+        return _fn
+    tried = [Path(path)] if path else _candidates()
+    for p in tried:
+        if not p.is_file():
+            continue
+        lib = ctypes.CDLL(str(p))
+        fn = lib.OodleLZ_Decompress
+        fn.restype = ctypes.c_ssize_t
+        fn.argtypes = [
+            ctypes.c_void_p,
+            ctypes.c_ssize_t,
+            ctypes.c_void_p,
+            ctypes.c_ssize_t,
+            ctypes.c_int,
+            ctypes.c_int,
+            ctypes.c_int,
+            ctypes.c_void_p,
+            ctypes.c_ssize_t,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_ssize_t,
+            ctypes.c_int,
+        ]
+        _lib, _fn = lib, fn
+        return fn
+    raise FileNotFoundError(
+        "Could not find an Oodle DLL. Set ZCMODKIT_OODLE to the full path of "
+        "oodle-data-shared.dll (FModel ships one in its Output/.data folder)."
+    )
+
+
+def decompress(
+    data: bytes, out_size: int, dll: str | os.PathLike | None = None
+) -> bytes:
+    """Decompress one Oodle block."""
+    fn = load(dll)
+    dst = ctypes.create_string_buffer(out_size)
+    n = fn(data, len(data), dst, out_size, 1, 0, 0, None, 0, None, None, None, 0, 3)
+    if n != out_size:
+        raise RuntimeError(f"Oodle decompress failed: got {n}, expected {out_size}")
+    return dst.raw[:out_size]

zcmodkit/formats/pak.py

@@ -0,0 +1,248 @@
+"""Small reader and writer for UE4 and UE5 .pak archives, version 11, no encryption."""
+
+from __future__ import annotations
+
+import hashlib
+import io
+import struct
+from dataclasses import dataclass
+from pathlib import Path
+
+from . import oodle
+
+PAK_MAGIC = 0x5A6F12E1
+NUL = bytes([0])
+
+
+def _rstr(st: io.BytesIO) -> str:
+    """Read an FString."""
+    (n,) = struct.unpack("<i", st.read(4))
+    if n == 0:
+        return ""
+    if n < 0:
+        return st.read(-n * 2).decode("utf-16-le").rstrip(chr(0))
+    return st.read(n).decode("latin-1").rstrip(chr(0))
+
+
+def _wstr(s: str) -> bytes:
+    """Write an FString. ASCII, null terminated."""
+    raw = s.encode("latin-1") + NUL
+    return struct.pack("<i", len(raw)) + raw
+
+
+@dataclass
+class Entry:
+    """One file inside a pak."""
+
+    path: str
+    offset: int
+    size: int
+    uncompressed_size: int
+    compression_index: int
+    block_count: int
+
+
+class PakReader:
+    """Reads a .pak index and pulls files back out."""
+
+    def __init__(self, path: str | Path):
+        self.path = Path(path)
+        self._f = open(self.path, "rb")  # noqa: SIM115, closed by close()/__exit__
+        self.entries: dict[str, Entry] = {}
+        self._read_index()
+
+    def close(self) -> None:
+        self._f.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.close()
+
+    def _read_index(self) -> None:
+        f = self._f
+        f.seek(0, 2)
+        size = f.tell()
+        tail_len = min(4096, size)
+        f.seek(size - tail_len)
+        tail = f.read(tail_len)
+        h = tail.rfind(struct.pack("<I", PAK_MAGIC))
+        if h < 0:
+            raise ValueError(f"{self.path.name}: not a pak file (no magic)")
+        (self.version,) = struct.unpack_from("<i", tail, h + 4)
+        idx_off, idx_size = struct.unpack_from("<qq", tail, h + 8)
+
+        f.seek(idx_off)
+        idx = io.BytesIO(f.read(idx_size))
+        self.mount_point = _rstr(idx)
+        idx.read(4 + 8)  # NumEntries, PathHashSeed
+        if struct.unpack("<i", idx.read(4))[0]:
+            idx.read(36)  # path hash index offset/size/hash
+        fdi_off = fdi_size = None
+        if struct.unpack("<i", idx.read(4))[0]:
+            fdi_off, fdi_size = struct.unpack("<qq", idx.read(16))
+            idx.read(20)
+        (enc_size,) = struct.unpack("<i", idx.read(4))
+        encoded = idx.read(enc_size)
+        if fdi_off is None:
+            raise ValueError("pak has no full directory index; unsupported")
+
+        f.seek(fdi_off)
+        fdi = io.BytesIO(f.read(fdi_size))
+        (ndirs,) = struct.unpack("<i", fdi.read(4))
+        for _ in range(ndirs):
+            d = _rstr(fdi)
+            (nfiles,) = struct.unpack("<i", fdi.read(4))
+            for _ in range(nfiles):
+                name = _rstr(fdi)
+                (eo,) = struct.unpack("<i", fdi.read(4))
+                self.entries[d + name] = self._decode(d + name, encoded, eo)
+
+    @staticmethod
+    def _decode(path: str, buf: bytes, at: int) -> Entry:
+        """Unpack one bit-packed index entry."""
+        (v,) = struct.unpack_from("<I", buf, at)
+        p = at + 4
+        cmi = (v >> 23) & 0x3F
+        fmt = "<I" if v & (1 << 31) else "<q"
+        (offset,) = struct.unpack_from(fmt, buf, p)
+        p += 4 if v & (1 << 31) else 8
+        fmt = "<I" if v & (1 << 30) else "<q"
+        (usize,) = struct.unpack_from(fmt, buf, p)
+        p += 4 if v & (1 << 30) else 8
+        if cmi:
+            fmt = "<I" if v & (1 << 29) else "<q"
+            (size,) = struct.unpack_from(fmt, buf, p)
+        else:
+            size = usize
+        return Entry(path, offset, size, usize, cmi, (v >> 6) & 0xFFFF)
+
+    def read(self, path: str) -> bytes:
+        """Pull one file out of the pak, decompressing it if it needs it."""
+        e = self.entries.get(path)
+        if e is None:
+            raise KeyError(path)
+        f = self._f
+        f.seek(e.offset)
+        head = f.read(64 + e.block_count * 16 + 32)
+        # header: Offset, Size, UncompressedSize, CompressionMethodIndex,
+        #         Hash, [blocks], Flags, BlockSize
+        p = 24
+        (cmi,) = struct.unpack_from("<i", head, p)
+        p += 4
+        p += 20
+        blocks = []
+        if cmi:
+            (nb,) = struct.unpack_from("<i", head, p)
+            p += 4
+            for _ in range(nb):
+                s, t = struct.unpack_from("<qq", head, p)
+                p += 16
+                blocks.append((s, t))
+        p += 1
+        (block_size,) = struct.unpack_from("<i", head, p)
+        p += 4
+
+        if not cmi:
+            f.seek(e.offset + p)
+            return f.read(e.size)
+
+        out = bytearray()
+        base = e.offset if blocks and blocks[0][0] < e.offset else 0
+        for s, t in blocks:
+            f.seek(base + s)
+            comp = f.read(t - s)
+            want = min(block_size, e.uncompressed_size - len(out))
+            out += oodle.decompress(comp, want)
+        return bytes(out)
+
+
+class PakWriter:
+    """Builds a version 11 pak, uncompressed and unencrypted."""
+
+    VERSION = 11
+    MOUNT_POINT = "../../../"
+
+    def __init__(self):
+        self._files: dict[str, bytes] = {}
+
+    def add(self, path: str, data: bytes) -> None:
+        """Queue up a file at an in-pak path like 'SWZeroCompany/Content/...'."""
+        self._files[path.replace("\\", "/")] = data
+
+    @staticmethod
+    def _header(offset: int, size: int, sha: bytes) -> bytes:
+        # UnrealPak puts 0 in the inline header Offset field, not the real one.
+        del offset
+        return (
+            struct.pack("<qqq", 0, size, size)
+            + struct.pack("<i", 0)
+            + sha
+            + bytes([0])
+            + struct.pack("<i", 0)
+        )
+
+    @staticmethod
+    def _encoded(offset: int, size: int) -> bytes:
+        flags = 0
+        if offset <= 0xFFFFFFFF:
+            flags |= 1 << 31
+        if size <= 0xFFFFFFFF:
+            flags |= 1 << 30
+        out = struct.pack("<I", flags)
+        out += struct.pack("<I" if offset <= 0xFFFFFFFF else "<q", offset)
+        out += struct.pack("<I" if size <= 0xFFFFFFFF else "<q", size)
+        return out
+
+    def write(self, dest: str | Path) -> Path:
+        """Write the pak out and give back its path."""
+        dest = Path(dest)
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        body = bytearray()
+        encoded = bytearray()
+        offsets: dict[str, int] = {}
+        for path, data in self._files.items():
+            offset = len(body)
+            sha = hashlib.sha1(data).digest()
+            body += self._header(offset, len(data), sha) + data
+            offsets[path] = len(encoded)
+            encoded += self._encoded(offset, len(data))
+
+        # full directory index: {directory: {filename: encoded entry offset}}
+        dirs: dict[str, dict[str, int]] = {}
+        for path in self._files:
+            d, _, name = path.rpartition("/")
+            dirs.setdefault(d + "/", {})[name] = offsets[path]
+        fdi = bytearray(struct.pack("<i", len(dirs)))
+        for d, files in dirs.items():
+            fdi += _wstr(d) + struct.pack("<i", len(files))
+            for name, eo in files.items():
+                fdi += _wstr(name) + struct.pack("<i", eo)
+
+        index_offset = len(body)
+        fdi_offset = index_offset  # patched below once primary size is known
+        primary = bytearray()
+        primary += _wstr(self.MOUNT_POINT)
+        primary += struct.pack("<i", len(self._files))
+        primary += struct.pack("<Q", 0)  # path hash seed
+        primary += struct.pack("<i", 0)  # no path hash index
+        primary += struct.pack("<i", 1)  # has full directory index
+        fdi_ptr = len(primary)
+        primary += struct.pack("<qq", 0, len(fdi)) + bytes(20)
+        primary += struct.pack("<i", len(encoded)) + bytes(encoded)
+        primary += struct.pack("<i", 0)  # no non-encoded entries
+
+        fdi_offset = index_offset + len(primary)
+        struct.pack_into("<q", primary, fdi_ptr, fdi_offset)
+        struct.pack_into("<qq", primary, fdi_ptr, fdi_offset, len(fdi))
+        primary[fdi_ptr + 16 : fdi_ptr + 36] = hashlib.sha1(bytes(fdi)).digest()
+
+        out = bytes(body) + bytes(primary) + bytes(fdi)
+        footer = bytes(16) + bytes([0])
+        footer += struct.pack("<I", PAK_MAGIC) + struct.pack("<i", self.VERSION)
+        footer += struct.pack("<qq", index_offset, len(primary))
+        footer += hashlib.sha1(bytes(primary)).digest()
+        footer += bytes(32 * 5)
+        dest.write_bytes(out + footer)
+        return dest

zcmodkit/formats/usmap.py

@@ -0,0 +1,328 @@
+"""Reading .usmap mappings files.
+
+Cooked packages store property values but not the names or types of the fields
+they belong to. A mappings file fills that in, which is the only reason editing
+a value by name is possible at all.
+"""
+
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import ClassVar
+
+MAGIC = 0x30C4
+
+#: The format revisions that change how we read things.
+VERSION_LONG_FNAME = 2
+VERSION_LARGE_ENUMS = 3
+
+
+class MappingsError(ValueError):
+    """A mappings file would not parse, or does not match this asset."""
+
+
+class PropertyType:
+    """EPropertyType, the way a mappings file writes it."""
+
+    BYTE = 0
+    BOOL = 1
+    INT = 2
+    FLOAT = 3
+    OBJECT = 4
+    NAME = 5
+    DELEGATE = 6
+    DOUBLE = 7
+    ARRAY = 8
+    STRUCT = 9
+    STR = 10
+    TEXT = 11
+    INTERFACE = 12
+    MULTICAST_DELEGATE = 13
+    WEAK_OBJECT = 14
+    LAZY_OBJECT = 15
+    ASSET_OBJECT = 16
+    SOFT_OBJECT = 17
+    UINT64 = 18
+    UINT32 = 19
+    UINT16 = 20
+    INT64 = 21
+    INT16 = 22
+    INT8 = 23
+    MAP = 24
+    SET = 25
+    ENUM = 26
+    FIELD_PATH = 27
+    OPTIONAL = 28
+
+    NAMES: ClassVar[dict[int, str]] = {
+        BYTE: "Byte",
+        BOOL: "Bool",
+        INT: "Int",
+        FLOAT: "Float",
+        OBJECT: "Object",
+        NAME: "Name",
+        DELEGATE: "Delegate",
+        DOUBLE: "Double",
+        ARRAY: "Array",
+        STRUCT: "Struct",
+        STR: "Str",
+        TEXT: "Text",
+        INTERFACE: "Interface",
+        MULTICAST_DELEGATE: "MulticastDelegate",
+        WEAK_OBJECT: "WeakObject",
+        LAZY_OBJECT: "LazyObject",
+        ASSET_OBJECT: "AssetObject",
+        SOFT_OBJECT: "SoftObject",
+        UINT64: "UInt64",
+        UINT32: "UInt32",
+        UINT16: "UInt16",
+        INT64: "Int64",
+        INT16: "Int16",
+        INT8: "Int8",
+        MAP: "Map",
+        SET: "Set",
+        ENUM: "Enum",
+        FIELD_PATH: "FieldPath",
+        OPTIONAL: "Optional",
+    }
+
+
+@dataclass
+class TypeInfo:
+    """A property's type, plus any inner types it wraps."""
+
+    kind: int
+    struct_name: str | None = None
+    enum_name: str | None = None
+    inner: list[TypeInfo] = field(default_factory=list)
+
+    @property
+    def name(self) -> str:
+        return PropertyType.NAMES.get(self.kind, f"Unknown({self.kind})")
+
+    def __repr__(self) -> str:
+        extra = self.struct_name or self.enum_name or ""
+        return f"{self.name}{'<' + extra + '>' if extra else ''}"
+
+
+@dataclass
+class PropertyInfo:
+    """One property of a schema that actually gets written out."""
+
+    name: str
+    schema_index: int
+    array_size: int
+    type: TypeInfo
+
+    def __repr__(self) -> str:
+        return f"<Property {self.name} {self.type!r} @{self.schema_index}>"
+
+
+@dataclass
+class Schema:
+    """A class or struct, and its properties in the order they get written."""
+
+    name: str
+    super_name: str | None
+    property_count: int
+    properties: list[PropertyInfo]
+
+    def __repr__(self) -> str:
+        return f"<Schema {self.name} properties={len(self.properties)}>"
+
+
+class Mappings:
+    """Schemas for every class and struct in one build of the game."""
+
+    def __init__(
+        self,
+        names: list[str],
+        enums: dict[str, dict[int, str]],
+        schemas: dict[str, Schema],
+        version: int,
+    ):
+        self.names = names
+        self.enums = enums
+        self.schemas = schemas
+        self.version = version
+
+    @classmethod
+    def load(cls, path: str | Path) -> Mappings:
+        """Parse a .usmap off disk."""
+        return cls.loads(Path(path).read_bytes())
+
+    @classmethod
+    def loads(cls, data: bytes) -> Mappings:
+        """Parse .usmap bytes."""
+        (magic,) = struct.unpack_from("<H", data, 0)
+        if magic != MAGIC:
+            raise MappingsError(f"not a .usmap file (magic 0x{magic:04x})")
+        version = data[2]
+        # bHasVersioning is written as a 32-bit bool.
+        (has_versioning,) = struct.unpack_from("<I", data, 3)
+        if has_versioning:
+            raise MappingsError("versioned mappings files are not supported")
+        compression = data[7]
+        compressed, _uncompressed = struct.unpack_from("<II", data, 8)
+        body = data[16:]
+        if compression != 0:
+            raise MappingsError(
+                f"compressed mappings are not supported (method {compression}); "
+                "export an uncompressed .usmap"
+            )
+        if compressed != len(body):
+            raise MappingsError(
+                f"mappings declare {compressed} bytes but {len(body)} follow the header"
+            )
+        return cls._read_body(body, version)
+
+    @staticmethod
+    def _read_body(b: bytes, version: int) -> Mappings:
+        at = 0
+
+        def u8() -> int:
+            nonlocal at
+            at += 1
+            return b[at - 1]
+
+        def u16() -> int:
+            nonlocal at
+            v = struct.unpack_from("<H", b, at)[0]
+            at += 2
+            return v
+
+        def u32() -> int:
+            nonlocal at
+            v = struct.unpack_from("<I", b, at)[0]
+            at += 4
+            return v
+
+        count = u32()
+        names: list[str] = []
+        for _ in range(count):
+            length = u16() if version >= VERSION_LONG_FNAME else u8()
+            names.append(b[at : at + length].decode("utf-8", "replace"))
+            at += length
+
+        def u64() -> int:
+            nonlocal at
+            v = struct.unpack_from("<Q", b, at)[0]
+            at += 8
+            return v
+
+        # Every entry has its numeric value first, then its name.
+        enums: dict[str, dict[int, str]] = {}
+        for _ in range(u32()):
+            enum_name = names[u32()]
+            entries = u16() if version >= VERSION_LARGE_ENUMS else u8()
+            enums[enum_name] = {u64(): names[u32()] for _ in range(entries)}
+
+        def read_type() -> TypeInfo:
+            kind = u8()
+            info = TypeInfo(kind)
+            if kind == PropertyType.ENUM:
+                info.inner.append(read_type())
+                info.enum_name = names[u32()]
+            elif kind == PropertyType.STRUCT:
+                info.struct_name = names[u32()]
+            elif kind in (PropertyType.ARRAY, PropertyType.SET, PropertyType.OPTIONAL):
+                info.inner.append(read_type())
+            elif kind == PropertyType.MAP:
+                info.inner.append(read_type())
+                info.inner.append(read_type())
+            return info
+
+        schemas: dict[str, Schema] = {}
+        for _ in range(u32()):
+            schema_name = names[u32()]
+            super_index = u32()
+            super_name = None if super_index == 0xFFFFFFFF else names[super_index]
+            prop_count = u16()
+            serializable = u16()
+            props: list[PropertyInfo] = []
+            for _ in range(serializable):
+                schema_index = u16()
+                array_size = u8()
+                prop_name = names[u32()]
+                props.append(
+                    PropertyInfo(prop_name, schema_index, array_size, read_type())
+                )
+            schemas[schema_name] = Schema(schema_name, super_name, prop_count, props)
+
+        # Later revisions tack on sections we do not need, module paths and
+        # so on. Stopping after the schemas is fine. Running off the end is not.
+        if at > len(b):
+            raise MappingsError("mappings body is truncated")
+        return Mappings(names, enums, schemas, version)
+
+    def schema(self, name: str) -> Schema:
+        """Look up a schema by class or struct name."""
+        try:
+            return self.schemas[name]
+        except KeyError:
+            raise MappingsError(f"no schema for {name!r} in these mappings") from None
+
+    def properties(self, name: str) -> list[PropertyInfo]:
+        """A schema's properties, inherited ones included, base class first.
+
+        A .usmap numbers each schema's properties from zero, but serialization
+        numbers them straight through the whole chain. So a class's own
+        properties start after however many its ancestors have between them,
+        and every inherited index gets shifted up to match. Without that, a
+        derived class collides with its own base and the walk reads the wrong
+        property.
+        """
+        chain, current = [], name
+        seen = set()
+        while current and current not in seen:
+            seen.add(current)
+            schema = self.schemas.get(current)
+            if schema is None:
+                break
+            chain.append(schema)
+            current = schema.super_name
+
+        out: list[PropertyInfo] = []
+        base = 0
+        for schema in reversed(chain):
+            out.extend(
+                PropertyInfo(
+                    prop.name, base + prop.schema_index, prop.array_size, prop.type
+                )
+                for prop in schema.properties
+            )
+            base += schema.property_count
+        return out
+
+    def __repr__(self) -> str:
+        return (
+            f"<Mappings version={self.version} names={len(self.names)} "
+            f"enums={len(self.enums)} schemas={len(self.schemas)}>"
+        )
+
+
+#: Which build of the game these bundled mappings came from.
+BUNDLED_BUILD = "SWZeroCompany-5.6.1-196320+++ProjectBruno+Stable-a1e7f571"
+
+_bundled: Mappings | None = None
+
+
+def bundled() -> Mappings:
+    """The mappings that ship with zcmodkit, parsed once and kept around.
+
+    They describe one build of the game. If a patch moves a struct around, the
+    parse stops landing exactly where it should and the editor raises instead
+    of writing a value to the wrong place.
+    """
+    global _bundled
+    if _bundled is None:
+        path = Path(__file__).parent.parent / "data" / "mappings.usmap"
+        if not path.is_file():
+            raise MappingsError(
+                f"bundled mappings are missing from the install ({path}). "
+                "Editing values needs them; text edits do not."
+            )
+        _bundled = Mappings.load(path)
+    return _bundled