ZAK decryptor/extractor for Wonder Boy Anniversary Collection

zak.cpp

#include <cstdio>
#include <cstdlib>
#include <span>
#include <string_view>
#include <filesystem>

#include <lz4.h>

/* Buffers. */
namespace {
	using u8 = std::uint8_t;
	using u16 = std::uint16_t;
	using u32 = std::uint32_t;
	using u64 = std::uint64_t;

	template<std::size_t Size>
	union Buffer {
		static constexpr size_t SizeInBytes			= Size / sizeof(u8);
		static constexpr size_t SizeInShorts		= Size / sizeof(u16);
		static constexpr size_t SizeInWords			= Size / sizeof(u32);
		static constexpr size_t SizeInDoubleWords	= Size / sizeof(u64);

		u8  m8  [SizeInBytes];
		u16 m16 [SizeInShorts];
		u32 m32 [SizeInWords];
		u64 m64 [SizeInDoubleWords];
	};

	using Buffer8   = Buffer<8>;
	using Buffer12  = Buffer<12>;
	using Buffer16  = Buffer<16>;
	using Buffer32  = Buffer<32>;
	using Buffer64  = Buffer<64>;
}

/* Decryption. */
namespace {
	struct ChaCha20 {

		using Block = Buffer64;

		static constexpr Buffer16 s_Constant = {.m8{
			'e','x','p','a','n','d',' ','3','2','-','b','y','t','e',' ','k'
		}};

		Block m_State;
		Block m_KeyStream;
		size_t m_KeyStreamIndex;
		Buffer32 m_Key;
		Buffer12 m_Nonce;
		std::uint32_t m_InitialCounter;

		constexpr ChaCha20(const Buffer32 &key, const Buffer12& nonce, const std::uint32_t initialCounter) :
			m_State({}),
			m_KeyStream({}),
			m_KeyStreamIndex(Block::SizeInBytes),
			m_Key(key),
			m_Nonce(nonce),
            m_InitialCounter(initialCounter)
		{}

		constexpr void InitState() {
			m_State.m32[0] = s_Constant.m32[0];
			m_State.m32[1] = s_Constant.m32[1];
			m_State.m32[2] = s_Constant.m32[2];
			m_State.m32[3] = s_Constant.m32[3];

			m_State.m32[4] = m_Key.m32[0];
			m_State.m32[5] = m_Key.m32[1];
			m_State.m32[6] = m_Key.m32[2];
			m_State.m32[7] = m_Key.m32[3];
			m_State.m32[8] = m_Key.m32[4];
			m_State.m32[9] = m_Key.m32[5];
			m_State.m32[10] = m_Key.m32[6];
			m_State.m32[11] = m_Key.m32[7];

			m_State.m32[12] = m_InitialCounter;

			m_State.m32[13] = m_Nonce.m32[0];
			m_State.m32[14] = m_Nonce.m32[1];
			m_State.m32[15] = m_Nonce.m32[2];
		}

		static constexpr std::uint32_t Rotl(const std::uint32_t x, const int n) {
			return (x << n) | (x >> (32 - n));
		}

		static constexpr void QuarterRound(Block& block, const std::uint32_t x, const std::uint32_t y, const std::uint32_t z, const std::uint32_t w) {
			block.m32[x] += block.m32[y]; block.m32[w] = Rotl(block.m32[w] ^ block.m32[x], 16);
			block.m32[z] += block.m32[w]; block.m32[y] = Rotl(block.m32[y] ^ block.m32[z], 12);
			block.m32[x] += block.m32[y]; block.m32[w] = Rotl(block.m32[w] ^ block.m32[x], 8);
			block.m32[z] += block.m32[w]; block.m32[y] = Rotl(block.m32[y] ^ block.m32[z], 7);
		}

		constexpr void IncrementCounter() {
			auto* counter = &m_State.m32[12];
			counter[0]++;
			if (counter[0] == 0) {
				counter[1]++;
				/* TODO: assert counter[1] doesn't roll over. Don't decrypt ~1180 exabytes, please. */
			}
		}

		constexpr void NextBlock() {
			m_KeyStream = m_State;

			/* Perform 20 rounds. */
			for (size_t i = 0; i < 10; i++) {
				QuarterRound(m_KeyStream, 0, 4, 8, 12);
				QuarterRound(m_KeyStream, 1, 5, 9, 13);
				QuarterRound(m_KeyStream, 2, 6, 10, 14);
				QuarterRound(m_KeyStream, 3, 7, 11, 15);

				QuarterRound(m_KeyStream, 0, 5, 10, 15);
				QuarterRound(m_KeyStream, 1, 6, 11, 12);
				QuarterRound(m_KeyStream, 2, 7, 8, 13);
				QuarterRound(m_KeyStream, 3, 4, 9, 14);
			}
			for (size_t i = 0; i < Block::SizeInWords; i++) {
				m_KeyStream.m32[i] += m_State.m32[i];
			}
			IncrementCounter();
		}

		constexpr void Crypt(std::span<std::byte> input) {
			InitState();
			for (auto& b : input)
			{
				/* If we've run out of key stream, calculate the next block of it. */
				if ( m_KeyStreamIndex >= Block::SizeInBytes )
				{
					NextBlock();
					m_KeyStreamIndex = 0;
				}
				/* Xor input with keystream and continue. */
				b ^= static_cast<std::byte>(this->m_KeyStream.m8[m_KeyStreamIndex]);
				m_KeyStreamIndex++;
			}
		}
	};

	struct Header {
		std::uint32_t m_MaxDecompressedSize;
		std::uint32_t m_CompressedSize;
		std::uint8_t m_Data[1];
	};
}

/* Utils. */
namespace {

	class UniqueArray {
		UniqueArray(std::unique_ptr<std::byte[]> data, const size_t size) :
			m_Data(std::move(data)),
			m_Size(size)
		{}

		std::unique_ptr<std::byte[]> m_Data;
		size_t m_Size;

	public:
		/* Disable copying. */
		UniqueArray() = delete;
		UniqueArray(const UniqueArray&) = delete;
		UniqueArray operator=(const UniqueArray&) = delete;

		/* Allow moving. */
		UniqueArray(UniqueArray&&) = default;

		[[nodiscard]] std::byte* Data() const {
			return m_Data.get();
		}

		[[nodiscard]] size_t Size() const {
			return m_Size;
		}

		[[nodiscard]] std::span<const std::byte> SpanUnsafe() const {
			return { m_Data.get(), m_Size };
		}

		[[nodiscard]] std::span<std::byte> SpanUnsafe() {
			return { m_Data.get(), m_Size };
		}

		static UniqueArray Create(const size_t size) {
			return {
				std::make_unique<std::byte[]>(size),
				size
			};
		}
	};

#ifndef _MSC_VER
	/* Microsoft provide fopen_s. */
	int fopen_s(FILE** out, const char* path, const char* mode) {
		*out = fopen(path, mode);
		return errno;
	}
#endif

	UniqueArray ReadFile(const char *path) {
		FILE* f = nullptr;
		const auto openErr = fopen_s(&f, path, "rb");
		if(f == nullptr) {
			printf("Failed to open \"%s\" (%d)\n", path, openErr);
			exit(1);
		}

		fseek(f, 0, SEEK_END);
		const size_t size = ftell(f);
		fseek(f, 0, SEEK_SET);

		auto data = UniqueArray::Create(size);
		size_t res = fread(data.Data(), size, 1, f);
		fclose(f);

		return data;
	}

	void WriteFile(const std::string& path, const std::span<const std::byte> data) {
		FILE* f = nullptr;
		const auto openErr = fopen_s(&f, path.c_str(), "w+");
		if(f == nullptr) {
			printf("Failed to open \"%s\" (%d)\n", path.c_str(), openErr);
			exit(1);
		}

		fwrite(data.data(), 1, data.size_bytes(), f);
		fclose(f);
	}
}

static UniqueArray Decompress(const UniqueArray &input) {
	const auto* outerCompressed = reinterpret_cast<Header*>(input.Data());

	auto decompress = [](const Header* header) {
		printf("m_CompressedSize = %x\nm_MaxDecompressedSize = %x\n", header->m_CompressedSize, header->m_MaxDecompressedSize);
		/* Allocate space for decompressed data. */
		auto decompressed = UniqueArray::Create(header->m_MaxDecompressedSize);
		/* Decompress. */
		const int actualDecompSizeOrErr = LZ4_decompress_safe(
			reinterpret_cast<const char*>(&header->m_Data),
			reinterpret_cast<char*>(decompressed.Data()),
			static_cast<int>(header->m_CompressedSize),
			static_cast<int>(header->m_MaxDecompressedSize)
		);
		/* Ensure some sanity. */
		if (actualDecompSizeOrErr != header->m_MaxDecompressedSize) {
			printf("Decompress fail! %x\n", actualDecompSizeOrErr);
			abort();
		}
		return decompressed;
	};
	const auto outerDecompressed = decompress(outerCompressed);
	const auto innerHeader = reinterpret_cast<Header*>(outerDecompressed.Data());
	return decompress(innerHeader);
}

/* Zak. */
namespace zak {

	struct File;

	struct Header {
		static constexpr std::uint32_t sMagic = 0x24B415A;

		std::uint32_t m_Magic;
        /* Unknown? */
		char padding[4];
		std::uint32_t m_FileCount;
		std::uint32_t m_DataStart;

		[[nodiscard]] const void* End() const  {
			return reinterpret_cast<const std::byte*>(this) + sizeof(Header);
		}

		[[nodiscard]] const std::byte* Data() const {
			return reinterpret_cast<const std::byte*>(this) + m_DataStart;
		}

		[[nodiscard]] const File* Files() const {
			return static_cast<const File*>(End());
		}
	};

	struct File {
		u32 Unk;
		u32 m_StrLen;
		u32 m_Offset;
		u32 m_Size;

		[[nodiscard]] const void* End() const {
			return reinterpret_cast<const std::byte*>(this) + sizeof(File) + m_StrLen;
		}

		[[nodiscard]] const File* Next() const {
			return static_cast<const File*>(End());
		}

		[[nodiscard]] std::string_view GetFileName() const {
			return { reinterpret_cast<const char*>(this) + sizeof(File), static_cast<size_t>(m_StrLen) };
		}
	};
}

static constexpr Buffer32 s_Key = { .m8 {
	0x62, 0x1F, 0x1C, 0x38, 0xF1, 0x63, 0x30, 0x16, 0xBC, 0x51, 0x49, 0x47, 0xBE, 0xC1, 0x58, 0xDB,
	0xF2, 0xC0, 0x8C, 0x6F, 0x45, 0xB1, 0xCF, 0xEC, 0x04, 0x9A, 0xA1, 0x33, 0xBB, 0xCF, 0x90, 0xC5
}};

/* Nonce is derived from the size of the file, for some reason. */
static constexpr Buffer12 ComputeNonce(const size_t size) {
	return {.m32 {
		static_cast<u32>(size),
		static_cast<u32>(size + 1),
		static_cast<u32>(~size),
	}};
}

int main(int argc, char** argv) {
	if(argc != 3) {
		printf("%s: [input] [output]\n", argv[0]);
		return 0;
	}

	const char* inputPath = argv[1];
	const auto outputPath = std::filesystem::path(argv[2]);

	auto read = [&] {
		printf("Reading...\n");
		auto input = ReadFile(inputPath);

		printf("Decrypting...\n");
		ChaCha20 ctx(s_Key, ComputeNonce(input.Size()), 0);
		ctx.Crypt(input.SpanUnsafe());

		printf("Writing compressed zak...\n");
		WriteFile((outputPath / "compressed_decrypted.zak").string(), input.SpanUnsafe());

		printf("Decompressing...\n");
		return Decompress(input);
	};
	auto zakData = read();
	printf("Writing decompressed zak...\n");
	WriteFile((outputPath / "decompressed_decrypted.zak").string(), zakData.SpanUnsafe());

	std::filesystem::create_directories(outputPath);
	const auto header = reinterpret_cast<const zak::Header*>(zakData.Data());
	if(header->m_Magic == zak::Header::sMagic) {
		printf("Extracting...\n");
		auto file = header->Files();
		for(size_t i = 0; i < header->m_FileCount; i++) {
			auto str = std::string(file->GetFileName());
			printf("Writing %s...\n", str.c_str());

			auto path = outputPath / str;
            std::filesystem::create_directories(path.parent_path());

			const auto fileData = std::span { header->Data() + file->m_Offset, file->m_Size };
			WriteFile(path.string(), fileData);

			file = file->Next();
		}
	} else {
		printf("Decrypted file seems corrupt. Skipping extraction...\n");
	}

    return 0;
}

Compile with g++ --std=gnu++20 zak.cpp -O3 -llz4 -o zak
Depends on liblz4-dev

Might you be able to compile this for Windows/macOS? I geep ketting the "zak.cpp:5:10: fatal error: 'lz4.h' file not found" error message. I've got lz4 installed via HomeBrew on my mac and can't seem to fix the issue on Windows. Thanks so much!

Ensure you have the -dev and non-dev variant of the package as well. I don't have a Mac to compile with, and don't really mess with building on Windows natively much (despite daily driving it).

Ah, ok. I'm currently trying to install/set up Windows Subsystem for Linux. I can get it compiled but, it failed to read the roms.zak file:
./zak
./zak: [input] [output]

./zak roms.zak roms
Reading...
Failed to open "roms.zak"

The hashes of my "roms.zak" file:
crc32: a3a4bcfc
MD5: 1f78695cf4f2a1eab637be4a31e6faa2
SHA-1: 1654ba08f072c8d93ef44ded4ec8a63ba0d24460

Am I missing something? Trying to install a normal Ubuntu VM instead of WSL. Thanks :)

I'd guess the file isn't in the working directory or you need to use ./roms.zak in place of roms.zak? Otherwise, adjust the code to print the error it gets (line 396 probably), and we can work from there.

Yep, that was it. Needed normal Linux set up

Thanks so much!

Could you please do a hash check on the resulting ROMs?

I would like to double-check that I am getting the correct files because most of them don't match known DATs:
I'll start (MD5):

• WBIMW_GEN.md = 389D3B2CD7F2D50EC4656817D9363A38
• WB3DT_MK3.sms = CA91C753BB4002A0C35BA4386A7E34EB
• WBIMW.sms = 0051D55923A28569E268A01D62479F3C
• WBIMW_palette.sms = A030B627E28CB5144D136C1AB65FD891
• WB3DT_GGE.gg = A155F79606FCF6E5EA49B92A43C09D0C
• WB3DT_GGJ.gg = 0205718D459439CBA81DE1204F9109C2
• WB3DT_SMS.sms = 81B438C9E69CA345FC06E4D0CDAB33E7
• WBIML.sms = EBDE5134B4760A0D454C64C146A3F175
• WB.sms = 0AA1932D8B6F33C4BD20C27882E0F0F0
• WB_GGE.gg = 0C367E6FF35D50A8413672DE3674FE11
• WB_GGJ.gg = 95174071D45215CF04FC146C5525ED8E
• WB_SG1000.bin = 76F29FEE0CD3AF921E488A276DACE636

The following match known DATs:
WB3ML_GEN.md = Wonder Boy III - Monster Lair (Japan, Europe) (En)
WBIMW_SMD.md = Wonder Boy V - Monster World III (Japan, Korea) (Ja)
MW4_SMD.md = Monster World IV (USA) (En,Ja) (Genesis Mini)

Finally, I also cannot make heads or tails of the Arcade ROMs.

I'm not aware of a hash being stored in the ZAK archive, so there would be nothing to check against. Assuming your files don't look like complete noise (suggesting broken decryption), I would guess the non-matching is due to changes by the publisher/developers of the collections.

No, I don’t think there is. I was talking about manually hash-checking the extracted files.

I modified your code for Windows and compiled using VS19.
For LZ4, I added lz4.h and lz4.c directly to the project folder, using the reference implementation from v1.10.0:
https://github.com/lz4/lz4/releases/tag/v1.10.0

The compiled executable runs fine -the same files are extracted- but I want to confirm that nothing broke during porting.

While testing, I noticed the "WBIMW_GEN.md" ROM freezes during the game demo. The JP version "WBIMW_SMD.md" (which matches a known DAT) runs correctly, so I’m trying to determine if the issue is with the extraction binary or the ROM itself.

Obviously you're welcome to review the code or include it in your GitHub!

EDIT: After digging, I found that the WBIMW_GEN.md ROM is almost identical to the one included with the Mega Drive Mini console. The only difference is at addresses 0x14F4–0x14F5.

The Mini’s ROM runs fine in an emulator and contains the bytes "16 1A" (ordinary data).

The extracted ROM, however, has "4A FC", which on the Motorola 68000 (the Mega Drive CPU) is the instruction TSTA.W which "tests the contents of address register A6" as a 16‑bit word.

This feels too deliberate to be a bad dump. I suspect that the ROM was modified, possibly to do "something" with the emulator it runs on.

All of this is to say, it appears the ported code works as expected and creates the correct roms, meaning we now have a working Windows executable!

Ah, I misunderstood your initial question (I thought you meant add a hash check functionality to the program).
Anyways, good to hear the ROMs are likely in-tact. Unfortunately, I'm not the right person to research what the differences actually mean.
I went ahead and revamped my implementation to be compatible with GCC and MSVC, so it should run on Windows/Linux.
In addition, I figured out the encryption algorithm is plain ChaCha20, so now it's implemented normally rather than being sourced via decompilation. I used your hashes to compare to ensure there was no regressions here but let me know if seems like I made a mistake.