Decryption program for Battleblock Theater .wbt files

wbt-decrypt.py

#!/usr/bin/env python
import array
import blowfish
import os
import sys
import struct
from io import BytesIO
from itertools import chain


class MersenneTwister:
	def __init__(self):
		self._state = [0 for i in range(624)]
		self._index = 0

	def seed(self, seed):
		self._state[0] = seed
		for i in range(1, 624):
			self._state[i] = (i + 0x6C078965 * (self._state[i - 1] ^ (self._state[i - 1] >> 30))) & 0xFFFFFFFF

	def _reseed(self):
		for i in range(624):
			y = (self._state[i] & 0x80000000) + (self._state[(i + 1) % 624] & 0x7FFFFFFF)
			self._state[i] = self._state[(i + 397) % 624] ^ (y >> 1)
			if y % 2 != 0:
				self._state[i] ^= 0x9908B0DF

	def next(self):
		if self._index == 0:
			self._reseed()
		y = self._state[self._index]
		y ^= (y >> 11)
		y ^= (y << 7) & 0x9D2C5680
		y ^= (y << 15) & 0xEFC60000
		y ^= (y >> 18)

		self._index = (self._index + 1) % 624
		return y


def get_rand_seed(basename):
	seed = 0x19570320
	for idx, char in enumerate(basename):
		m = ord(char) >> (idx & 3)
		seed = (seed * m) + ord(char) & 0xFFFFFFFF
	return seed


def get_file_key(basename):
	rng = MersenneTwister()
	seed = get_rand_seed(basename)
	rng.seed(seed)
	key = array.array("I", [rng.next() for i in range(4)])
	key.byteswap()
	return key


def get_basename(filename):
	basename = os.path.basename(filename).upper()
	basename, dot, ext = basename.rpartition('.')
	return basename


class Blowfish:
	def __init__(self, key):
		self._key = key
		self._key1 = blowfish.PI_P_ARRAY
		self._key2 = chain(*blowfish.PI_S_BOXES)

		self._key1_buf = array.array('I', self._key1)
		self._key2_buf = array.array('I', self._key2)

	def _descramble_keys(self):
		# Some shorthands.
		k1 = self._key1_buf
		k2 = self._key2_buf

		# Make sure that key1 has the filename data inside it.
		for i in range(len(k1)):
			k1[i] ^= self._key[i % len(self._key)]

		def descramble(a, b):
			v5 = a
			v4 = b

			for i in range(0, 16, 4):
				v1 = (v5 ^ k1[i + 0])
				v2 = (v4 ^ k1[i + 1] ^ k2[(v1 & 0xFF) + 0x300] + (k2[((v1 >> 8) & 0xFF) + 0x200] ^ (k2[((v1 >> 16) & 0xFF) + 0x100]) + k2[((v1 >> 24) & 0xFF)])) & 0xFFFFFFFF
				v3 = (v1 ^ k1[i + 2] ^ k2[(v2 & 0xFF) + 0x300] + (k2[((v2 >> 8) & 0xFF) + 0x200] ^ (k2[((v2 >> 16) & 0xFF) + 0x100]) + k2[((v2 >> 24) & 0xFF)])) & 0xFFFFFFFF
				v4 = (v2 ^ k1[i + 3] ^ k2[(v3 & 0xFF) + 0x300] + (k2[((v3 >> 8) & 0xFF) + 0x200] ^ (k2[((v3 >> 16) & 0xFF) + 0x100]) + k2[((v3 >> 24) & 0xFF)])) & 0xFFFFFFFF
				v5 = (v3 ^ k2[(v4 & 0xFF) + 0x300] + (k2[((v4 >> 8) & 0xFF) + 0x200] ^ (k2[((v4 >> 16) & 0xFF) + 0x100]) + k2[((v4 >> 24) & 0xFF)])) & 0xFFFFFFFF

			a = v4 ^ k1[17]
			b = v5 ^ k1[16]

			return a, b

		# Descramble the key buffers.
		a = 0
		b = 0
		for i in range(0, len(k1), 2):
			a, b = descramble(a, b)
			k1[i + 0] = a
			k1[i + 1] = b

		for i in range(0, len(k2), 2):
			a, b = descramble(a, b)
			k2[i + 0] = a
			k2[i + 1] = b

	def decrypt(self, data):
		self._descramble_keys()

		buf = array.array("I")
		buf.frombytes(data)

		def _decrypt(a, b):
			k1 = self._key1_buf
			k2 = self._key2_buf

			v2 = a
			v3 = b

			for i in range(16):
				v1 = v2 ^ (k1[17 - i])
				v2 = v3 ^ (k2[(v1 & 0xFF) + 0x300] + (k2[((v1 >> 8) & 0xFF) + 0x200] ^ (k2[((v1 >> 16) & 0xFF) + 0x100]) + k2[((v1 >> 24) & 0xFF)])) & 0xFFFFFFFF
				v3 = v1

			a = v1 ^ k1[0]
			b = v2 ^ k1[1]
			return a, b

		ret = BytesIO()
		for i in range(0, len(buf), 2):
			a = buf[i + 0]
			b = buf[i + 1]
			a, b = _decrypt(a, b)
			if i > 4:
				ret.write(struct.pack('II', a, b))

		ret.seek(0)
		return ret.read()


def decrypt_file(filename, outname):
	basename = get_basename(filename)
	outname = filename.replace(".wbt", ".wav")
	key = get_file_key(basename)
	cipher = Blowfish(key)

	with open(filename, "rb") as f:
		data = f.read()

	sys.stdout.write("Decrypting %r... " % (filename))
	sys.stdout.flush()
	decrypted_data = cipher.decrypt(data)

	with open(outname, "wb") as f:
		f.write(decrypted_data)

	sys.stdout.write("Done. %i bytes written to %r\n" % (len(data), outname))


def main():
	for filename in sys.argv[1:]:
		decrypt_file(filename, filename.replace(".wbt", ".wav"))


if __name__ == "__main__":
	main()