mbEditorPro/mbEditorPro2.0/arcane/util/ArcWpak.py

#  • ▌ ▄ ·. ▄▄▄▄·     ▄▄▄ .·▄▄▄▄  ▪  ▄▄▄▄▄      ▄▄▄       ▄▄▄·▄▄▄
#  ·██ ▐███▪▐█ ▀█▪    ▀▄.▀·██▪ ██ ██ •██  ▪     ▀▄ █·    ▐█ ▄█▀▄ █·▪
#  ▐█ ▌▐▌▐█·▐█▀▀█▄    ▐▀▀▪▄▐█· ▐█▌▐█· ▐█.▪ ▄█▀▄ ▐▀▀▄      ██▀·▐▀▀▄  ▄█▀▄
#  ██ ██▌▐█▌██▄▪▐█    ▐█▄▄▌██. ██ ▐█▌ ▐█▌·▐█▌.▐▌▐█•█▌    ▐█▪·•▐█•█▌▐█▌.▐▌
#  ▀▀  █▪▀▀▀·▀▀▀▀      ▀▀▀ ▀▀▀▀▀• ▀▀▀ ▀▀▀  ▀█▄▀▪.▀  ▀    .▀   .▀  ▀ ▀█▄▀▪
#                Magicbane Emulator Project © 2013 - 2022
#                           www.magicbane.com

import os
import time
import zipfile

from .ArcFileCache import next_available_path

# A .wpak is a zip (client uses minizip, ArcWPakFile 0x54D3C0) whose entries are
# each Blowfish CFB-64 encrypted with a zero IV (BF_cfb64_encrypt 0xCD9000),
# except image entries.
WPAK_KEY = bytes([
    0x85, 0x71, 0x40, 0x3C, 0x14, 0x50, 0x0B, 0x52,
    0x73, 0x2D, 0x10, 0x08, 0x63, 0x59, 0x5B, 0xAA,
])  # hardcoded at 0x54DEC0

WPAK_COMPRESS_LEVEL = 9  # zipOpenNewFileInZip level in ArcWPakFile writer

# Map icon paks keep their .tga entries unencrypted (LoadMapZoneData 0xAACA60).
WPAK_PLAIN_EXTENSIONS = ('.tga',)

MASK32 = 0xFFFFFFFF


def pi_fraction_words(count):
    # Blowfish's initial P-array and S-boxes are the hex digits of pi's fraction.
    guard_bits = 64
    bits = count * 32 + guard_bits
    one = 1 << bits

    def arctan_inverse(x):
        x_squared = x * x
        term = one // x
        total = term
        sign = -1
        n = 1
        while term:
            term //= x_squared
            total += sign * (term // (2 * n + 1))
            sign = -sign
            n += 1
        return total

    pi = 4 * (4 * arctan_inverse(5) - arctan_inverse(239))
    fraction = (pi - (3 << bits)) >> guard_bits

    return [(fraction >> (32 * (count - 1 - i))) & MASK32 for i in range(count)]


class ArcBlowfish:
    p_array = None
    s_boxes = None

    def __init__(self, key):
        words = pi_fraction_words(18 + 4 * 256)
        self.p_array = words[:18]
        self.s_boxes = [words[18 + 256 * i: 18 + 256 * (i + 1)] for i in range(4)]
        self.set_key(key)

    def set_key(self, key):
        left = 0
        right = 0

        for i in range(18):
            key_word = 0
            for j in range(4):
                key_word = (key_word << 8) | key[(i * 4 + j) % len(key)]
            self.p_array[i] ^= key_word

        for i in range(0, 18, 2):
            left, right = self.encrypt_block(left, right)
            self.p_array[i] = left
            self.p_array[i + 1] = right

        for s_box in self.s_boxes:
            for i in range(0, 256, 2):
                left, right = self.encrypt_block(left, right)
                s_box[i] = left
                s_box[i + 1] = right

    def encrypt_block(self, left, right):
        p_array = self.p_array
        s0, s1, s2, s3 = self.s_boxes

        for i in range(16):
            left ^= p_array[i]
            feistel = ((s0[left >> 24] + s1[(left >> 16) & 0xFF]) & MASK32) ^ s2[(left >> 8) & 0xFF]
            right ^= (feistel + s3[left & 0xFF]) & MASK32
            left, right = right, left

        return right ^ p_array[17], left ^ p_array[16]

    def keystream_block(self, iv):
        left, right = self.encrypt_block(
            int.from_bytes(iv[:4], 'big'),
            int.from_bytes(iv[4:], 'big')
        )
        return left.to_bytes(4, 'big') + right.to_bytes(4, 'big')

    def cfb64_decrypt(self, data):
        output = bytearray(len(data))
        iv = bytes(8)

        for offset in range(0, len(data), 8):
            keystream = self.keystream_block(iv)
            cipher_block = data[offset:offset + 8]
            for i, cipher_byte in enumerate(cipher_block):
                output[offset + i] = cipher_byte ^ keystream[i]
            iv = cipher_block

        return bytes(output)

    def cfb64_encrypt(self, data):
        output = bytearray(len(data))
        iv = bytes(8)

        for offset in range(0, len(data), 8):
            keystream = self.keystream_block(iv)
            plain_block = data[offset:offset + 8]
            for i, plain_byte in enumerate(plain_block):
                output[offset + i] = plain_byte ^ keystream[i]
            iv = bytes(output[offset:offset + 8])

        return bytes(output)


WPAK_CIPHER = ArcBlowfish(WPAK_KEY)


def is_encrypted_entry(filename):
    return not filename.lower().endswith(WPAK_PLAIN_EXTENSIONS)


def load_wpak_file(filename):
    entries = []

    with zipfile.ZipFile(filename) as wpak_file:
        for info in wpak_file.infolist():
            if info.is_dir():
                continue
            data = wpak_file.read(info)
            if is_encrypted_entry(info.filename):
                data = WPAK_CIPHER.cfb64_decrypt(data)
            entries.append([info.filename, data])

    return entries


def load_wpak_directory(directory):
    entries = []

    # Client finds entries by name (unzLocateFile 0xCD7920), so order only needs to be stable.
    for filename in sorted(os.listdir(directory), key=str.upper):
        filepath = os.path.join(directory, filename)
        if not os.path.isfile(filepath):
            continue
        with open(filepath, 'rb') as fp:
            entries.append([filename, fp.read()])

    return entries


def build_wpak_bytes(entries, wpak_file):
    with zipfile.ZipFile(wpak_file, 'w') as archive:
        for filename, data in entries:
            info = zipfile.ZipInfo(filename, date_time=time.localtime()[:6])
            info.compress_type = zipfile.ZIP_DEFLATED
            if is_encrypted_entry(filename):
                data = WPAK_CIPHER.cfb64_encrypt(data)
            archive.writestr(info, data, compresslevel=WPAK_COMPRESS_LEVEL)


def save_wpak_file(filename, entries):
    filename = next_available_path(filename)
    print('Packing to ' + filename)

    with open(filename, 'wb') as wpak_file:
        build_wpak_bytes(entries, wpak_file)