Wii Disc decryption

disc_image.hpp

#pragma once
#include <filesystem>
#include <limits>
#include <optional>
#include <type_traits>

#include "helpers.hpp"
#include "io_file.hpp"
#include "swap.hpp"

class DiscImage {
  protected:
	std::filesystem::path path = "";
	IOFile file;

  public:
	static constexpr u64 PARTITION_NONE = std::numeric_limits<u64>::max();

	virtual std::optional<u64> read(u64 offset, u64 size, u8* output, u64 partition = PARTITION_NONE) = 0;
	virtual u64 getGamePartition() { return PARTITION_NONE; }

	template <typename T>
	std::optional<T> readBE(u64 offset, u64 partition) {
		T value;

		std::optional<u64> size = read(offset, sizeof(T), (u8*)&value, partition);
		if (!size.has_value() || *size != sizeof(T)) {
			return std::nullopt;
		}

#ifdef COMMON_LITTLE_ENDIAN
		if constexpr (std::is_same<T, u32>()) {
			return Common::swap32(value);
		} else {
			Helpers::panic("DiscImage::ReadBE with unknown type");
			return T{};
		}
#else
		return value;
#endif
	}

	DiscImage(const std::filesystem::path& path) : path(path), file(path, "rb") {}
};

wii_disc.cpp

#include "loader/wii_disc.hpp"

#include <cryptopp/aes.h>
#include <cryptopp/modes.h>
#include <fmt/format.h>
#include <fmt/ranges.h>

#include <algorithm>
#include <cstring>
#include <tuple>

static constexpr WiiDisc::Key commonKey = {/* REDACTED */}

WiiDisc::WiiDisc(const std::filesystem::path& path) : DiscImage(path) {
	u32 partitionCount, partitionTableOffset;

	file.seek(0x40000);
	auto [success, bytes] = file.readBytes(&partitionCount, sizeof(u32));
	if (!success || bytes != sizeof(u32)) {
		Helpers::panic("Failed to read partition info from Wii disc!");
	}

	std::tie(success, bytes) = file.readBytes(&partitionTableOffset, sizeof(u32));
	if (!success || bytes != sizeof(u32)) {
		Helpers::panic("Failed to read partition info from Wii disc!");
	}

	// Byteswap because BE
	partitionCount = Common::nativeToBE32(partitionCount);
	// Table offset is encoded in 32-bit words, not bytes
	partitionTableOffset = Common::nativeToBE32(partitionTableOffset) << 2;
	fmt::println("Wii disc partition count: {}, partition table offset: {:08X}", partitionCount, partitionTableOffset);

	file.seek(partitionTableOffset);

	for (u32 i = 0; i < partitionCount; i++) {
		file.seek(partitionTableOffset + i * 8);

		u32 partitionOffset;
		u32 partitionType;

		std::tie(success, bytes) = file.readBytes(&partitionOffset, sizeof(u32));
		if (!success || bytes != sizeof(u32)) {
			Helpers::panic("Failed to read partition offset from Wii disc!");
		}

		std::tie(success, bytes) = file.readBytes(&partitionType, sizeof(u32));
		if (!success || bytes != sizeof(u32)) {
			Helpers::panic("Failed to read partition type from Wii disc!");
		}

		partitionOffset = Common::nativeToBE32(partitionOffset);
		partitionType = Common::nativeToBE32(partitionType);

		Partition partition{
			.offset = u64(partitionOffset) << 2,
			.type = partitionType,
		};

		// Retrieve the title ID/AES Initialization Vector and encrypted title key
		Key encryptedTitleKey;
		std::array<u8, 8> titleID, ticketID;

		file.seek(partition.offset + 0x1BF);
		std::tie(success, bytes) = file.readBytes(&encryptedTitleKey[0], 16);
		if (!success || bytes != 16) {
			Helpers::panic("Failed to read title key from Wii disc Ticket!");
		}

		file.seek(partition.offset + 0x1D0);
		std::tie(success, bytes) = file.readBytes(&ticketID[0], 8);
		if (!success || bytes != 8) {
			Helpers::panic("Failed to read ticket ID from Wii disc Ticket!");
		}

		file.seek(partition.offset + 0x1DC);
		std::tie(success, bytes) = file.readBytes(&titleID[0], 8);
		if (!success || bytes != 8) {
			Helpers::panic("Failed to read title ID from Wii disc Ticket!");
		}

		Key titleKey;
		// First 8 bytes of the IV are the title ID, last 8 bytes are 0
		Key iv;
		iv.fill(0);
		
		std::copy(titleID.begin(), titleID.end(), iv.begin());

		CryptoPP::CBC_Mode<CryptoPP::AES>::Decryption decrypter(commonKey.data(), commonKey.size(), iv.data());
		decrypter.ProcessData(titleKey.data(), encryptedTitleKey.data(), titleKey.size());

		partition.titleKey = titleKey;
		partition.titleID = titleID;
		partition.ticketID = ticketID;

		fmt::println("Partition {}: Offset = {:08X}, type: {}", i, partition.offset, partition.getTypeName());
		partitions.push_back(partition);

		fmt::println("Decrypted title key: {::#x}, title ID: {::#x}", titleKey, titleID);
		fmt::println("TicketID: {::#x}", ticketID);
	}

	for (auto& partition : partitions) {
		if (partition.type == Partition::Type::Data) {
			gamePartition = &partition;
			break;
		}
	}

	u8 tmp;

	file.seek(0x60);
	file.readBytes(&tmp, sizeof(u8));
	const bool containsHashes = (tmp == 0);

	file.readBytes(&tmp, sizeof(u8));
	const bool isEncrypted = (tmp == 0);

	if (!containsHashes || !isEncrypted) {
		Helpers::panic("Wii ISO doesn't contain hashes or is not encrypted");
	}
}

std::optional<u64> WiiDisc::read(u64 offset, u64 size, u8* output, u64 partition) {
	if (!file.isOpen()) {
		return std::nullopt;
	}

	std::array<u8, sectorSize> sectorBuffer;

	// The start of a partition contains ticket data, which must be skipped
	u64 startingOffset = gamePartition->offset + (isEncrypted ? ticketSizeEncrypted : ticketSizeUnencrypted);
	// The disc is also split into smaller sectors
	u64 currentSector = offset / dataPerSectorSize;
	u64 offsetInSector = offset % dataPerSectorSize;

	u64 bytesRead = 0;

	// Read sector by sector until enough bytes are consumed
	while (size > 0) {
		const u64 currentOffset = startingOffset + currentSector * sectorSize + offsetInSector + hashSize;
		const u64 offsetWithinSector = offsetInSector % sectorSize;
		// Read either as much of a sector as we can, or size if it's smaller
		const u64 bytesToConsume = std::min(dataPerSectorSize - offsetWithinSector, size);

		// Seek to the start of the sector
		if (!file.seek(startingOffset + currentSector * sectorSize)) {
			return bytesRead;
		}

		// Read one entire sector
		auto [success, bytes] = file.readBytes(sectorBuffer.data(), sectorSize);
		if (!success) {
			return bytesRead;
		}

		// Both the SHA-1 hashes (offset 0-0x3FF) and the user data (offset 0x0400-7FFF) are encrypted using the title key
		// However the IV for the SHA-1 hashes is all zeroes, whereas the IV for user data, comes from the ENCRYPTED
		// SHA-1 block, bytes 0x3D0-0x3DF
		Key iv;
		std::memcpy(iv.data(), &sectorBuffer[0x3D0], 0x10);

		// Decrypt sector in-place
		// For now we only decrypt the user bytes, not the SHA hashes.
		CryptoPP::CBC_Mode<CryptoPP::AES>::Decryption decrypter(gamePartition->titleKey.data(), gamePartition->titleKey.size(), iv.data());
		decrypter.ProcessData(&sectorBuffer[hashSize], &sectorBuffer[hashSize], dataPerSectorSize);

		// Copy sector to output
		std::memcpy(output, &sectorBuffer[hashSize + offsetWithinSector], bytesToConsume);
		output += bytesToConsume;

		// We're starting from a new sector, so the offset will be 0 by default
		offsetInSector = 0;
		currentSector += 1;
		size -= bytesToConsume;
		bytesRead += bytesToConsume;
	}

	fmt::println("Read {:X} bytes from DVD offset {:X} (Starting physical offset on disc: {:X})", bytesRead, offset, startingOffset);
	return bytesRead;
}

wii_disc.hpp

#pragma once
#include <array>

#include "loader/disc_image.hpp"

class WiiDisc : public DiscImage {
  public:
	using Key = std::array<u8, 16>;

  private:
	struct Partition {
		enum Type : u32 {
			Data = 0,
			Update = 1,
			ChannelInstaller = 2,
		};

		u64 offset = 0;
		u32 type = 0;
		std::array<u8, 8> titleID = {}; // Used as the IV for decrypting the title key (key == common key)
		std::array<u8, 8> ticketID = {}; // Used as the IV for decrypting disc images
		Key titleKey = {}; // Used as the key for decrypting disc images

		const char* getTypeName() {
			switch (type) {
				case 0: return "Data";
				case 1: return "Update";
				case 2: return "Channel installer";
				default: return "Unknown";
			}
		}
	};

	std::vector<Partition> partitions;
	Partition* gamePartition = nullptr;
	bool isEncrypted = true;

	// Each partition starts with a ticket which is used for decrypting it
	static constexpr u64 ticketSizeEncrypted = 0x20000;
	static constexpr u64 ticketSizeUnencrypted = 0x8000;
	// The partition is split into smaller sectors, which start with 0x400 bytes of SHA1 hash data, followed by 0x7C00 bytes of encrypted user data
	static constexpr u64 sectorSize = 0x8000;
	static constexpr u64 hashSize = 0x400;
	static constexpr u64 dataPerSectorSize = 0x7C00;

  public:
	WiiDisc(const std::filesystem::path& path);
	virtual std::optional<u64> read(u64 offset, u64 size, u8* output, u64 partition) override;
};