Tool to unpack/repack XNA game framework XNB images (5 version only; Color/Bgr565 formats only; skips non-image data)

xna5_tool.py

#!/usr/bin/env python3
"""
XNA Texture Extractor/Repacker
Extracts XNB files to PNG and repacks PNG files to XNB format.
"""

import os, sys, struct, argparse, logging, re
from pathlib import Path
from enum import IntEnum
from PIL import Image
import numpy as np

XNB_VERSION = 5
PLATFORM_PC = b'd'
XNB_MAGIC = b'XNB'

class SurfaceFormat(IntEnum):
    Color = 0
    Bgr565 = 1
    Bgra5551 = 2
    Bgra4444 = 3
    Dxt1 = 4
    Dxt3 = 5
    Dxt5 = 6


def read_uleb128(data, pos):
    """Read a 7-bit encoded integer from data."""
    result = 0
    shift = 0
    while pos < len(data):
        byte = data[pos]
        pos += 1
        result |= (byte & 0x7F) << shift
        if (byte & 0x80) == 0:
            break
        shift += 7
    return result, pos


def write_uleb128(value):
    """Write a 7-bit encoded integer."""
    if value == 0:
        return b'\0'

    result = bytearray()
    while value != 0:
        byte = value & 0x7F
        value >>= 7
        if value != 0:
            byte |= 0x80
        result.append(byte)
    return bytes(result)


def read_string(data, pos):
    """Read a string from XNB data."""
    length, pos = read_uleb128(data, pos)
    string_data = data[pos:pos + length]
    return string_data.decode('utf-8'), pos + length


def write_string(text):
    """Write a string in XNB format."""
    if not text:
        return b'\0'

    encoded = text.encode('utf-8')
    return write_uleb128(len(encoded)) + encoded


def extract_xnb(file_path):
    """Extract XNB file to PNG image(s)."""
    try:
        with open(file_path, 'rb') as f:
            data = f.read()

        pos = 0

        if data[pos:pos+3] != XNB_MAGIC:
            return extract_headerless_texture(file_path, data)

        pos += 3
        platform = chr(data[pos])
        pos += 1
        version = data[pos]
        pos += 1
        flags = data[pos]
        pos += 1

        if flags & 0x80 or flags & 0x40:
            logging.error(f"Compressed XNB files not supported: {file_path}")
            return False

        # Skip file size
        file_size = struct.unpack('<I', data[pos:pos+4])[0]
        pos += 4

        reader_count, pos = read_uleb128(data, pos)

        for _ in range(reader_count):
            reader_name, pos = read_string(data, pos)
            reader_version = struct.unpack('<i', data[pos:pos+4])[0]
            pos += 4

        # Skip shared resources
        shared_count, pos = read_uleb128(data, pos)

        # Read primary asset
        asset_index, pos = read_uleb128(data, pos)

        # Read texture data
        surface_format = struct.unpack('<i', data[pos:pos+4])[0]
        pos += 4
        width = struct.unpack('<i', data[pos:pos+4])[0]
        pos += 4
        height = struct.unpack('<i', data[pos:pos+4])[0]
        pos += 4
        mip_count = struct.unpack('<i', data[pos:pos+4])[0]
        pos += 4

        # Extract each mip level
        output_base = Path(file_path).with_suffix('')
        extracted_count = 0

        for level in range(mip_count):
            data_size = struct.unpack('<I', data[pos:pos+4])[0]
            pos += 4

            texture_data = data[pos:pos+data_size]
            pos += data_size

            # Convert texture to image
            mip_width = width >> level
            mip_height = height >> level
            image = convert_texture_to_image(texture_data, mip_width, mip_height, surface_format)
            if image:
                output_path = f"{output_base}-{level}.png" if mip_count > 1 else f"{output_base}.png"
                image.save(output_path)
                logging.info(f"Extracted: {output_path}")
                extracted_count += 1

        return extracted_count > 0

    except Exception as e:
        logging.error(f"Failed to extract {file_path}: {e}")
        return False


def extract_headerless_texture(file_path, data):
    """Extract headerless texture data (fallback method)."""
    try:
        # Skip unknown header bytes and try to read texture info
        pos = 10

        surface_format = struct.unpack('<i', data[pos:pos+4])[0]
        pos += 4
        width = struct.unpack('<i', data[pos:pos+4])[0]
        pos += 4
        height = struct.unpack('<i', data[pos:pos+4])[0]
        pos += 4
        mip_count = struct.unpack('<i', data[pos:pos+4])[0]
        pos += 4

        output_base = Path(file_path).with_suffix('')

        for level in range(mip_count):
            data_size = struct.unpack('<I', data[pos:pos+4])[0]
            pos += 4

            texture_data = data[pos:pos+data_size]
            pos += data_size

            mip_width = width >> level
            mip_height = height >> level

            image = convert_texture_to_image(texture_data, mip_width, mip_height, surface_format)
            if image:
                output_path = f"{output_base}-{level}.png" if mip_count > 1 else f"{output_base}.png"
                image.save(output_path)
                logging.info(f"Extracted (headerless): {output_path}")

        return True

    except Exception as e:
        logging.error(f"Failed to extract headerless texture {file_path}: {e}")
        return False


def convert_texture_to_image(texture_data, width, height, surface_format):
    """Convert raw texture data to PIL Image."""
    try:
        if surface_format == SurfaceFormat.Color:
            # XNA Color format is RGBA (not BGRA as I initially thought)
            expected_size = width * height * 4
            if len(texture_data) != expected_size:
                logging.warning(f"Texture size mismatch: expected {expected_size}, got {len(texture_data)}")
                return None

            # The data is already in RGBA format
            pixels = np.frombuffer(texture_data, dtype=np.uint8).copy()
            pixels = pixels.reshape((height, width, 4))

            return Image.fromarray(pixels, 'RGBA')

        elif surface_format == SurfaceFormat.Bgr565:
            # BGR565 format
            expected_size = width * height * 2
            if len(texture_data) != expected_size:
                return None

            # Convert BGR565 to RGBA
            pixels = np.frombuffer(texture_data, dtype=np.uint16).reshape((height, width))

            # Extract color components (BGR565 layout)
            b = ((pixels >> 0) & 0x1F) * 255 // 31
            g = ((pixels >> 5) & 0x3F) * 255 // 63
            r = ((pixels >> 11) & 0x1F) * 255 // 31
            a = np.full((height, width), 255)

            rgba = np.stack([r, g, b, a], axis=2).astype(np.uint8)
            return Image.fromarray(rgba, 'RGBA')

        else:
            logging.warning(f"Unsupported surface format: {surface_format}")
            # Try to interpret as RGBA anyway
            expected_size = width * height * 4
            if len(texture_data) == expected_size:
                pixels = np.frombuffer(texture_data, dtype=np.uint8).copy()
                pixels = pixels.reshape((height, width, 4))
                return Image.fromarray(pixels, 'RGBA')
            return None

    except Exception as e:
        logging.error(f"Failed to convert texture: {e}")
        return None


def repack_png(file_path, surface_format=SurfaceFormat.Color):
    """Repack PNG file to XNB format."""
    try:
        image = Image.open(file_path)
        if image.mode == 'RGBA':
            pixels = np.array(image)
        elif image.mode == 'RGB':
            pixels = np.array(image)
            alpha = np.full((pixels.shape[0], pixels.shape[1], 1), 0xFF, dtype=pixels.dtype)
            pixels = np.concatenate([pixels, alpha], axis=2)
        else:
            image = image.convert('RGBA')
            pixels = np.array(image)

        width, height = image.size

        # Set RGB channels to 0 (black) where alpha is 0 (fully transparent)
        white_color = np.array([0xFF, 0xFF, 0xFF])
        transparent_mask = (pixels[:, :, 3] == 0) #& np.all(pixels[:, :, :3] == white_color, axis=2)
        pixels[transparent_mask] = [0, 0, 0, 0]

        if surface_format == SurfaceFormat.Color:
            texture_data = pixels.tobytes()
        else:
            logging.error(f"Repacking for surface format {surface_format} not implemented")
            return False

        # Build XNB file
        output = bytearray()

        # Write header
        output.extend(XNB_MAGIC)
        output.extend(PLATFORM_PC)  # Platform
        output.append(XNB_VERSION)     # Version
        output.append(0)     # Flags (no compression)

        # Reserve space for file size
        file_size_pos = len(output)
        output.extend(struct.pack('<I', 0))

        # Write type readers
        output.extend(write_uleb128(1))  # One type reader
        output.extend(write_string("Microsoft.Xna.Framework.Content.Texture2DReader"))
        output.extend(struct.pack('<i', 0))  # Version

        # Write shared resources
        output.extend(write_uleb128(0))  # No shared resources

        # Write primary asset
        output.extend(write_uleb128(1))  # Type reader index

        # Write texture data
        output.extend(struct.pack('<i', surface_format))
        output.extend(struct.pack('<i', width))
        output.extend(struct.pack('<i', height))
        output.extend(struct.pack('<i', 1))  # Mip count

        # Write texture bytes
        output.extend(struct.pack('<I', len(texture_data)))
        output.extend(texture_data)

        # Update file size
        file_size = len(output)
        output[file_size_pos:file_size_pos+4] = struct.pack('<I', file_size)

        # Write to file
        output_path = Path(file_path).with_suffix('.xnb')
        with open(output_path, 'wb') as f:
            f.write(output)

        logging.info(f"Repacked: {file_path} -> {output_path}")
        return True

    except Exception as e:
        logging.error(f"Failed to repack {file_path}: {e}")
        return False


def process_directory(directory, operation, surface_format=SurfaceFormat.Color):
    """Process all files in directory and subdirectories."""
    directory = Path(directory)
    success_count = 0
    error_count = 0

    if operation == 'extract':
        files = list(directory.rglob("*.xnb"))
        logging.info(f"Found {len(files)} XNB files to extract")

        for file_path in files:
            if extract_xnb(file_path):
                success_count += 1
            else:
                error_count += 1

    else:  # repack
        files = list(directory.rglob("*.png"))
        for f in files:
            print(f.stem)
        # Filter out extracted mip levels
        files = [f for f in files if not any(f.stem.endswith(f'-{i}') for i in range(10))]
        logging.info(f"Found {len(files)} PNG files to repack")

        for file_path in files:
            if repack_png(file_path, surface_format):
                success_count += 1
            else:
                error_count += 1

    return success_count, error_count


def main():
    parser = argparse.ArgumentParser(
        description="Extract XNA texture files to PNG or repack PNG files to XNA format"
    )

    mode_group = parser.add_mutually_exclusive_group(required=True)
    mode_group.add_argument('-e', '--extract', action='store_true',
                           help='Extract all XNB files to PNG format')
    mode_group.add_argument('-r', '--repack', action='store_true',
                           help='Repack all PNG files to XNB format')

    parser.add_argument('directory', help='Directory to process (includes subdirectories)')
    parser.add_argument('-v', '--verbose', action='store_true', help='Enable verbose logging')
    parser.add_argument('-f', '--format', type=int, default=0,
                       help='Surface format for repacking (0=Color, 1=Bgr565)')

    args = parser.parse_args()

    # Setup logging
    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=log_level, format='%(levelname)s: %(message)s')

    # Validate directory
    if not os.path.isdir(args.directory):
        logging.error(f"Directory not found: {args.directory}")
        sys.exit(1)

    # Process files
    operation = 'extract' if args.extract else 'repack'
    surface_format = SurfaceFormat(args.format)

    logging.info(f"{'Extracting' if args.extract else 'Repacking'} files in: {args.directory}")
    success, errors = process_directory(args.directory, operation, surface_format)

    # Print summary
    print(f"\nOperation complete:")
    print(f"  Successful: {success} files")
    if errors > 0:
        print(f"  Errors: {errors} files")


if __name__ == "__main__":
    main()