NightSpire/web-next/sql/db2/wdc_db2.py

"""Lector mínimo de DB2 WDC3/WDC4 (formato de wowdev.wiki).

Solo lo necesario para sacar columnas enteras: registros normales, bitpacked,
common data y pallet (indexed). No implementa strings ni arrays de pallet, que
no hacen falta para el mapeo item -> displayId.
"""
import struct

FIELD_COMPRESSION_NONE = 0
FIELD_COMPRESSION_BITPACKED = 1
FIELD_COMPRESSION_COMMON_DATA = 2
FIELD_COMPRESSION_BITPACKED_INDEXED = 3
FIELD_COMPRESSION_BITPACKED_INDEXED_ARRAY = 4
FIELD_COMPRESSION_BITPACKED_SIGNED = 5


def _bits(data, bit_offset, bit_size):
    """Lee bit_size bits a partir de bit_offset (little-endian)."""
    byte = bit_offset // 8
    shift = bit_offset % 8
    nbytes = (shift + bit_size + 7) // 8
    chunk = data[byte:byte + nbytes]
    if len(chunk) < nbytes:
        chunk = chunk + b'\x00' * (nbytes - len(chunk))
    val = int.from_bytes(chunk, 'little')
    return (val >> shift) & ((1 << bit_size) - 1)


class DB2:
    def __init__(self, path):
        self.d = d = open(path, 'rb').read()
        self.magic = d[:4].decode()
        if self.magic not in ('WDC3', 'WDC4'):
            raise ValueError(f'formato no soportado: {self.magic}')
        (self.record_count, self.field_count, self.record_size, self.string_table_size,
         self.table_hash, self.layout_hash, self.min_id, self.max_id, self.locale,
         self.flags, self.id_index, self.total_field_count, self.bitpacked_data_offset,
         self.lookup_column_count, self.field_storage_info_size, self.common_data_size,
         self.pallet_data_size, self.section_count) = struct.unpack_from('<IIIIIIIIIHHIIIIIII', d, 4)

        o = 4 + struct.calcsize('<IIIIIIIIIHHIIIIIII')
        # WDC4 mete aquí un bloque extra de "encrypted section" opcional; en la
        # práctica de estos ficheros la cabecera de sección sigue directamente.
        self.sections = []
        for _ in range(self.section_count):
            (tact, file_offset, rec_count, str_size, offset_records_end, id_list_size,
             rel_size, offset_map_id_count, copy_count) = struct.unpack_from('<QIIIIIIII', d, o)
            o += 40
            self.sections.append(dict(file_offset=file_offset, record_count=rec_count,
                                      string_table_size=str_size, offset_records_end=offset_records_end,
                                      id_list_size=id_list_size, relationship_data_size=rel_size,
                                      offset_map_id_count=offset_map_id_count, copy_table_count=copy_count))

        # field_structure
        self.fields = []
        for _ in range(self.field_count):
            size, pos = struct.unpack_from('<hH', d, o); o += 4
            self.fields.append((size, pos))

        # field_storage_info
        self.storage = []
        for _ in range(self.field_count):
            (off_bits, size_bits, extra, stype, c1, c2, c3) = struct.unpack_from('<HHIIIII', d, o)
            o += 24
            self.storage.append(dict(offset_bits=off_bits, size_bits=size_bits,
                                     extra=extra, type=stype, c1=c1, c2=c2, c3=c3))

        self.pallet = d[o:o + self.pallet_data_size]; o += self.pallet_data_size
        self.common_raw = d[o:o + self.common_data_size]; o += self.common_data_size

        # Base de cada campo dentro de pallet_data: suma acumulada de los
        # additional_data_size de los campos pallet anteriores (c1/c2 son los bits
        # de bitpacking, NO un offset).
        self.pallet_base = {}
        acc = 0
        for fi, st in enumerate(self.storage):
            if st['type'] in (FIELD_COMPRESSION_BITPACKED_INDEXED,
                              FIELD_COMPRESSION_BITPACKED_INDEXED_ARRAY):
                self.pallet_base[fi] = acc
                acc += st['extra']

        # common data: por campo, `extra`/8 pares (id, value)
        self.common = {}
        p = 0
        for fi, st in enumerate(self.storage):
            if st['type'] == FIELD_COMPRESSION_COMMON_DATA and st['extra']:
                m = {}
                for _ in range(st['extra'] // 8):
                    k, v = struct.unpack_from('<II', self.common_raw, p); p += 8
                    m[k] = v
                self.common[fi] = m

        self.rows = []
        self.row_pos = []  # posición en fichero de cada fila (para string_at)
        for s in self.sections:
            self._read_section(s)

    def _read_section(self, s):
        d = self.d
        o = s['file_offset']
        if self.flags & 0x1:  # sparse: offset map en vez de registros fijos
            raise NotImplementedError('tabla sparse (offset map) no soportada')
        rec_start = o
        rec_data = d[o:o + s['record_count'] * self.record_size]
        o += s['record_count'] * self.record_size
        o += s['string_table_size']
        ids = None
        if s['id_list_size']:
            ids = list(struct.unpack_from(f'<{s["id_list_size"] // 4}I', d, o))
            o += s['id_list_size']
        # copy table y relationship data: no se usan aquí

        for i in range(s['record_count']):
            rec = rec_data[i * self.record_size:(i + 1) * self.record_size]
            # Posición del registro en el fichero: la necesita string_at, porque
            # los offsets de los strings son relativos al propio campo.
            self.row_pos.append(rec_start + i * self.record_size)
            row = [self._field(rec, fi, i) for fi in range(self.field_count)]
            rid = ids[i] if ids is not None else row[self.id_index]
            # common data: valor por id, o el default (c1) si no está listado
            for fi, st in enumerate(self.storage):
                if st['type'] == FIELD_COMPRESSION_COMMON_DATA:
                    row[fi] = self.common.get(fi, {}).get(rid, st['c1'])
            self.rows.append((rid, row))

    def _field(self, rec, fi, rec_index):
        st = self.storage[fi]
        t = st['type']
        if t == FIELD_COMPRESSION_NONE:
            if st['size_bits'] == 0:
                return 0
            return _bits(rec, st['offset_bits'], st['size_bits'])
        if t in (FIELD_COMPRESSION_BITPACKED, FIELD_COMPRESSION_BITPACKED_SIGNED):
            v = _bits(rec, st['offset_bits'], st['size_bits'])
            if t == FIELD_COMPRESSION_BITPACKED_SIGNED:
                sign = 1 << (st['size_bits'] - 1)
                if v & sign:
                    v -= 1 << st['size_bits']
            return v
        if t == FIELD_COMPRESSION_COMMON_DATA:
            return None  # se resuelve en _read_section, que sí conoce el id
        if t in (FIELD_COMPRESSION_BITPACKED_INDEXED, FIELD_COMPRESSION_BITPACKED_INDEXED_ARRAY):
            idx = _bits(rec, st['offset_bits'], st['size_bits'])
            off = self.pallet_base[fi] + idx * 4
            if off + 4 <= len(self.pallet):
                return struct.unpack_from('<I', self.pallet, off)[0]
            return 0
        raise NotImplementedError(f'storage type {t}')

    def string_at(self, row_index, fi, value):
        """Resuelve un campo string de una tabla NO sparse.

        El valor guardado no es un offset absoluto ni un índice a la string
        table: es un offset RELATIVO a la posición del propio campo dentro del
        fichero. 0 = cadena vacía.
        """
        if not value:
            return ''
        pos = self.row_pos[row_index] + self.storage[fi]['offset_bits'] // 8 + value
        end = self.d.find(b'\x00', pos)
        return self.d[pos:end].decode('utf-8', 'replace')