# Copyright (c) 2018 Stephen Bunn <stephen@bunn.io>
# MIT License <https://choosealicense.com/licenses/mit/>
import warnings
from typing import Tuple, Generator
from pathlib import PureWindowsPath
from construct import (
Array,
Bytes,
Const,
Int8ul,
Struct,
Switch,
VarInt,
Default,
Int16ul,
Int32ul,
Int64ul,
Container,
FlagsEnum,
Compressed,
GreedyBytes,
PaddedString,
PascalString,
)
from .. import __version__
from ._common import ArchiveFile, BaseArchive
from ..contrib.dds import (
DDS_HEADER,
MAKEFOURCC,
DDS_HEADER_DX10,
DXGIFormats,
D3D10ResourceDimension,
)
[docs]class BTDXArchive(BaseArchive):
"""Archive type for BTDX files (aka. BA2).
BTDX files (utilize the extension ``.ba2``) are Bethesda's second framework revision
for archives.
These files have virtually the same goal as :class:`.BSAArchive` but for optimized
loading of archived textures directly into the engine instead of simply compressing
the files.
This is done by splitting the BTDX archive into 2 different types:
- ``GNRL``: Storage for general files that are simply compressed
- ``DX10``: Storage for Microsoft DirectDraw textures in an optimized format
The extraction for ``GNRL`` files is simple.
But the extraction for ``DX10`` requires rebuliding the DDS headers for each of the
texture chunks.
For this reason the :mod:`.bethesda_structs.contrib.dds` module was added.
Note:
BTDX archives to not read the file data on initialization.
Header's, records and names are read in and files are built during
:func:`~BTDXArchive.iter_files`.
**Reference**:
- `BAE <https://github.com/jonwd7/bae>`_
"""
header_struct = Struct(
"magic" / Bytes(4),
"version" / Int32ul,
"type" / PaddedString(4, "utf8"),
"file_count" / Int32ul,
"names_offset" / Int64ul,
)
"""The structure of BTDX headers.
Returns:
:class:`~construct.core.Struct`: The structure of BTDX headers
"""
file_struct = Struct(
"hash" / Int32ul,
"ext" / PaddedString(4, "utf8"),
"directory_hash" / Int32ul,
"_unknown_0" / Int32ul,
"offset" / Int64ul,
"packed_size" / Int32ul,
"unpacked_size" / Int32ul,
"_unknown_1" / Int32ul,
)
"""The structure of GNRL files.
Returns:
:class:`~construct.core.Struct`: The structure of GNRL files
"""
tex_header_struct = Struct(
"hash" / Int32ul,
"ext" / PaddedString(4, "utf8"),
"directory_hash" / Int32ul,
"_unknown_0" / Int8ul,
"chunks_count" / Int8ul,
"chunk_header_size" / Int16ul,
"height" / Int16ul,
"width" / Int16ul,
"mips_count" / Int8ul,
"format" / Int8ul,
"_unknown_1" / Int16ul,
)
"""The structure of DX10 file headers.
Returns:
:class:`~construct.core.Struct`: The structure of DX10 file headers.
"""
tex_chunk_struct = Struct(
"offset" / Int64ul,
"packed_size" / Int32ul,
"unpacked_size" / Int32ul,
"start_mip" / Int16ul,
"end_mip" / Int16ul,
"_unknown_0" / Int32ul,
)
"""The structure of DX10 chunks.
Returns:
:class:`~construct.core.Struct`: The structure of DX10 chunks
"""
tex_struct = Struct(
"header" / tex_header_struct,
"chunks" / Array(lambda this: this.header.chunks_count, tex_chunk_struct),
)
"""The structure of DX10 tex files.
Returns:
:class:`~construct.core.Struct`: The structure of DX10 tex files
"""
archive_struct = Struct(
"header" / header_struct,
"files"
/ Array(
lambda this: this.header.file_count,
Switch(
lambda this: this.header.type, {"GNRL": file_struct, "DX10": tex_struct}
),
),
)
"""The **partial** structure of BTDX archives.
Returns:
:class:`~construct.core.Struct`: The **partial** structure of BTDX archives
"""
[docs] @classmethod
def can_handle(cls, filepath: str) -> bool:
"""Determines if a given file can be handled by the current archive.
Args:
filepath (str): The filepath to check if can be handled
"""
header = cls.header_struct.parse_file(filepath)
return header.magic == b"BTDX" and header.version >= 1
def _build_dds_headers(self, file_container: Container) -> Tuple[bytes, bytes]:
"""Builds DDS and DX10 secion headers for a given `file_container`.
Args:
file_container (Container): File container to build headers for
Returns:
Tuple[bytes, bytes]: A tuple of `DDS_HEADER` and `DX10_HEADER` (maybe None)
"""
header_data = {
"dwFlags": {
"DDSD_CAPS": True,
"DDSD_HEIGHT": True,
"DDSD_WIDTH": True,
"DDSD_PIXELFORMAT": True,
"DDSD_MIPMAPCOUNT": True,
"DDSD_LINEARSIZE": True,
},
"dwHeight": file_container.header.height,
"dwWidth": file_container.header.width,
"dwMipMapCount": file_container.header.mips_count,
"dwCaps": {
"DDSCAPS_COMPLEX": True,
"DDSCAPS_TEXTURE": True,
"DDSCAPS_MIPMAP": True,
},
}
# NOTE: I'm unsure what this field "is", but BAE has logic to build complete
# cubemaps for the DDS_HEADER if set to 2049
if file_container.header._unknown_1 == 2049:
header_data.update(
{
"dwCaps2": {
"DDSCAPS2_CUBEMAP": True,
"DDSCAPS2_CUBEMAP_POSITIVEX": True,
"DDSCAPS2_CUBEMAP_NEGATIVEX": True,
"DDSCAPS2_CUBEMAP_POSITIVEY": True,
"DDSCAPS2_CUBEMAP_NEGATIVEY": True,
"DDSCAPS2_CUBEMAP_POSITIVEZ": True,
"DDSCAPS2_CUBEMAP_NEGATIVEZ": True,
}
}
)
# TODO: find a cleaner more obvious way of implementing this logic
dx10_header_data = {}
pixel_data = {}
if file_container.header.format == DXGIFormats.DXGI_FORMAT_BC1_UNORM:
pixel_data.update(
dict(
dwFlags=dict(DDPF_FOURCC=True),
dwFourCC=MAKEFOURCC("D", "X", "T", "1"),
)
)
header_data.update(
dict(
dwPitchOrLinearSize=(
(file_container.header.width * file_container.header.height)
// 2
)
)
)
elif file_container.header.format == DXGIFormats.DXGI_FORMAT_BC2_UNORM:
pixel_data.update(
dict(
dwFlags=dict(DDPF_FOURCC=True),
dwFourCC=MAKEFOURCC("D", "X", "T", "3"),
)
)
header_data.update(
dict(
dwPitchOrLinearSize=(
file_container.header.width * file_container.header.height
)
)
)
elif file_container.header.format == DXGIFormats.DXGI_FORMAT_BC3_UNORM:
pixel_data.update(
dict(
dwFlags=dict(DDPF_FOURCC=True),
dwFourCC=MAKEFOURCC("D", "X", "T", "5"),
)
)
header_data.update(
dict(
dwPitchOrLinearSize=(
file_container.header.width * file_container.header.height
)
)
)
elif file_container.header.format == DXGIFormats.DXGI_FORMAT_BC5_UNORM:
pixel_data.update(
dict(
dwFlags=dict(DDPF_FOURCC=True),
dwFourCC=MAKEFOURCC("A", "T", "I", "2"),
)
)
header_data.update(
dict(
dwPitchOrLinearSize=(
file_container.header.width * file_container.header.height
)
)
)
elif file_container.header.format in (
DXGIFormats.DXGI_FORMAT_BC7_UNORM,
DXGIFormats.DXGI_FORMAT_BC7_UNORM_SRGB,
):
# FIXME: There may be a header differnce between BC7_UNORM and
# BC7_UNORM_SRGB, but I haven't noticed any
# (someone with more experience will have to let me know)
pixel_data.update(
dict(
dwFlags=dict(DDPF_FOURCC=True),
dwFourCC=MAKEFOURCC("D", "X", "1", "0"),
)
)
header_data.update(
dict(
dwPitchOrLinearSize=(
file_container.header.width * file_container.header.height
)
)
)
dx10_header_data.update(
{
"dxgiFormat": file_container.header.format,
"resourceDimension": (
D3D10ResourceDimension.D3D10_RESOURCE_DIMENSION_TEXTURE2D.value
),
"miscFlag": 0,
"arraySize": 1,
"miscFlags2": 0,
}
)
elif file_container.header.format == DXGIFormats.DXGI_FORMAT_B8G8R8A8_UNORM:
pixel_data.update(
dict(
dwFlags=dict(DDPF_ALPHA=True, DDPF_RBG=True),
dwRGBBitCount=32,
dwABitMask=0xFF000000,
dwRBitMask=0x00FF0000,
dwGBitMask=0x0000FF00,
dwBBitMask=0x000000FF,
)
)
header_data.update(
dict(
dwPitchOrLinearSize=(
(file_container.header.width * file_container.header.height) * 4
)
)
)
elif file_container.header.format == DXGIFormats.DXGI_FORMAT_R8_UNORM:
pixel_data.update(
dict(
dwFlags=dict(DDPF_RGB=True), dwRGBBitCount=8, dwRBitMask=0x000000FF
)
)
header_data.update(
dict(
dwPitchOrLinearSize=(
file_container.header.width * file_container.header.height
)
)
)
else:
warnings.warn(
(
f"unsupported DXGI format "
f"{DXGIFormats(file_container.header.format).name}, "
f"please create an issue on {__version__.__repo__} if you see this"
),
UserWarning,
)
return
header_data.update({"ddspf": pixel_data})
dx10_header = None
if len(dx10_header_data) > 0:
dx10_header = DDS_HEADER_DX10.build(dx10_header_data)
return (DDS_HEADER.build(header_data), dx10_header)
def _iter_gnrl_files(self) -> Generator[ArchiveFile, None, None]:
"""Iterates over the parsed data for GNRL fiels and yields instances of
`ArchiveFile`.
Raises:
ValueError: If a filename cannot be determined for a specific file record
Yields:
:class:`.ArchiveFile`: A file contained within the archive
"""
filename_offset = 0
for file_container in self.container.files:
filepath_content = self.content[
(self.container.header.names_offset + filename_offset) :
]
filepath = PascalString(VarInt, "utf8").parse(filepath_content)
# filename offset increased by length of parsed string accounting for
# prefix and suffix bytes
filename_offset += len(filepath) + 2
file_data = self.content[
file_container.offset : (
file_container.offset + file_container.unpacked_size
)
]
if file_container.packed_size > 0:
file_data = Compressed(GreedyBytes, "zlib").parse(file_data)
yield ArchiveFile(filepath=PureWindowsPath(filepath[1:]), data=file_data)
def _iter_dx10_files(self) -> Generator[ArchiveFile, None, None]:
"""Iterates over the parsed data for DX10 archives and yields instances of
`ArchiveFile`.
Raises:
ValueError: If a filename cannot be determined for a specific file record
Yields:
:class:`.ArchiveFile`: A file contained within the archive
"""
filename_offset = 0
for file_container in self.container.files:
filepath_content = self.content[
(self.container.header.names_offset + filename_offset) :
]
filepath = PascalString(Int16ul, "utf8").parse(filepath_content)
filename_offset += len(filepath) + 2
(dds_header, dx10_header) = self._build_dds_headers(file_container)
if dds_header:
dds_content = b"DDS "
dds_content += dds_header
if dx10_header:
dds_content += dx10_header
for tex_chunk in file_container.chunks:
if tex_chunk.packed_size > 0:
dds_content += Compressed(GreedyBytes, "zlib").parse(
self.content[
tex_chunk.offset : (
tex_chunk.offset + tex_chunk.packed_size
)
]
)
else:
dds_content += self.content[
tex_chunk.offset : (
tex_chunk.offset + tex_chunk.unpacked_size
)
]
yield ArchiveFile(filepath=PureWindowsPath(filepath), data=dds_content)
[docs] def iter_files(self) -> Generator[ArchiveFile, None, None]:
"""Iterates over the parsed data and yields instances of `ArchiveFile`
Raises:
ValueError: If a filename cannot be determined for a specific file record
Yields:
:class:`.ArchiveFile`: A file contained within the archive
"""
iter_method = {"GNRL": self._iter_gnrl_files, "DX10": self._iter_dx10_files}[
self.container.header.type
]
for archive_file in iter_method():
yield archive_file