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;
}

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.