wii_disk 0.1.2

// SPDX-License-Identifier: LGPL-2.1-or-later OR GPL-2.0-or-later OR MPL-2.0
// SPDX-FileCopyrightText: 2026 Gabriel Marcano <gabemarcano@yahoo.com>

use crate::error::Error;
use crate::utils::from_latin1_or_shift_jis;
use crate::utils::trim;

use std::io;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;

use byteorder::BigEndian;
use byteorder::ReadBytesExt;

use aes::Aes128;
use aes::cipher::Block;
use aes::cipher::BlockDecryptMut;
use aes::cipher::KeyIvInit;
use cbc::Decryptor as CbcDecryptor;

/// Represents the metadata held by the Wii disk internal header.
#[derive(Debug)]
pub struct Header {
    /// Console ID. The first character of the game code field on disk.
    pub console_id: String,
    /// Unique game code. The second and third characters of the game code field.
    pub game_code: String,
    /// Country code. The last character of the game code field.
    pub country_code: String,
    /// Make code.
    pub maker_code: String,
    /// The ID of the disk.
    pub disk_id: u8,
    /// The version of the disk.
    pub version: u8,
    /// Whether or not audio streaming is enabled.
    pub audio_streaming: u8,
    /// The size of the audio stream buffer. Apparently 0 means 10?
    pub stream_buffer_size: u8,
    /// Wii magic word.
    pub magic: u32,
    /// The name of the game. On disk it is padded with null characters.
    pub game_name: String,

    /// Disables hash verification, makes all reads fail on retail units.
    pub disable_hash_verification: u8,
    /// Disables disk encryption, and h3 hash table loading and verification. Makes all reads fail
    /// on retail units.
    pub disable_disc_encryption: u8,
}

pub struct RegionSettings {
    pub region: u32,
    pub japan_taiwan: u8,
    pub usa: u8,
    pub germany: u8,
    pub pegi: u8,
    pub finland: u8,
    pub portugal: u8,
    pub britain: u8,
    pub australia: u8,
    pub korea: u8,
}

pub struct PartitionTables {
    pub partitions: Vec<Partition>,
}

pub struct Partition {
    // From Partition table entry
    pub offset: u64,
    pub type_: u32,
    // From Partition itself
    pub ticket: Ticket,
    pub tmd_size: u32,
    pub tmd_offset: u64,
    pub cert_chain_size: u32,
    pub cert_chain_offset: u64,
    pub h3_table_offset: u64,
    pub data_offset: u64,
    pub data_size: u64,
}
#[derive(Default, Copy, Debug, Clone)]
pub struct CcLimit {
    pub limit_type: u32,
    pub maximum_usage: u32,
}

#[derive(Debug, Clone)]
pub struct Ticket {
    pub signature_type: u32,
    pub signature: [u8; 0x100],
    pub signature_issuer: String, //[u8; 0x40],
    pub ecdh_data: [u8; 0x3C],
    pub format_version: u8,
    pub title_key: [u8; 0x10],
    pub id: [u8; 0x08],
    pub console_id: [u8; 4],
    pub title_id: [u8; 0x08],
    pub title_version: u16,
    pub permitted_titles_mask: u32,
    pub permit_mask: u32,
    pub title_export_allowed: u8,
    pub common_key_index: u8,
    pub content_access_permissions: [u8; 0x40],
    pub cc_limits: [CcLimit; 8],
}

pub trait TicketReader {
    /// Reads the Wii ticket.
    ///
    /// # Errors
    ///
    /// [`Error::Io`] if there are any underlying issues with IO.
    fn read_wii_ticket(&mut self) -> Result<Ticket, Error>;
}

impl<T: Seek + Read> TicketReader for T {
    fn read_wii_ticket(&mut self) -> Result<Ticket, Error> {
        let signature_type = self.read_u32::<BigEndian>()?;
        let mut signature = [0u8; 0x100];
        self.read_exact(&mut signature)?;
        self.seek(SeekFrom::Current(0x3C))?;
        let mut signature_issuer = [0u8; 0x40];
        self.read_exact(&mut signature_issuer)?;
        let signature_issuer = trim(str::from_utf8(&signature_issuer)?).to_string();
        let mut ecdh_data = [0u8; 0x3C];
        self.read_exact(&mut ecdh_data)?;
        let format_version = self.read_u8()?;
        self.seek(SeekFrom::Current(2))?;
        let mut title_key = [0u8; 0x10];
        self.read_exact(&mut title_key)?;
        self.seek(SeekFrom::Current(1))?;
        let mut id = [0u8; 8];
        self.read_exact(&mut id)?;
        let mut console_id = [0u8; 4];
        self.read_exact(&mut console_id)?;
        let mut title_id = [0u8; 0x08];
        self.read_exact(&mut title_id)?;
        self.seek(SeekFrom::Current(2))?;
        let title_version = self.read_u16::<BigEndian>()?;
        let permitted_titles_mask = self.read_u32::<BigEndian>()?;
        let permit_mask = self.read_u32::<BigEndian>()?;
        let title_export_allowed = self.read_u8()?;
        let common_key_index = self.read_u8()?;
        self.seek(SeekFrom::Current(0x30))?;
        let mut content_access_permissions = [0u8; 0x40];
        self.read_exact(&mut content_access_permissions)?;
        self.seek(SeekFrom::Current(0x2))?;
        let mut cc_limits = [CcLimit::default(); 8];
        for limit in &mut cc_limits {
            *limit = CcLimit {
                limit_type: self.read_u32::<BigEndian>()?,
                maximum_usage: self.read_u32::<BigEndian>()?,
            };
        }

        Ok(Ticket {
            signature_type,
            signature,
            signature_issuer,
            ecdh_data,
            format_version,
            title_key,
            id,
            console_id,
            title_id,
            title_version,
            permitted_titles_mask,
            permit_mask,
            title_export_allowed,
            common_key_index,
            content_access_permissions,
            cc_limits,
        })
    }
}

/// Represents the metadata of a Wii disc.
pub struct Metadata {
    pub header: Header,
    pub partition_tables: PartitionTables,
}

pub trait MetadataRead {
    /// Parses the Wii ISO metadata from the object provided, returning a [`Header`] object with
    /// the Wii ISO metadata.
    ///
    /// # Errors
    ///
    /// Returns [`Error::Parse`] if the header cannot be found or if a field in the header contains
    /// an unexpected value.
    /// Returns [`Error::Io`] if an IO error took place while reading from the file.
    fn read_wii_header(&mut self) -> Result<Header, Error>;

    /// Extracts partition table information.
    ///
    /// # Errors
    /// [`Error::Io`] if an IO error took place while reading from the underlying IO object.
    fn read_wii_partitions(&mut self) -> Result<PartitionTables, Error>;
}

impl Metadata {
    /// Returns a new Metadata instance.
    ///
    /// # Errors
    ///
    /// See [`MetadataRead::read_wii_header`] and [`MetadataRead::read_wii_partitions`] for details on possible
    /// errors.
    pub fn try_from<T: Read + Seek>(io: &mut T) -> Result<Self, Error> {
        let header = io.read_wii_header()?;
        let partition_tables = io.read_wii_partitions()?;
        Ok(Self {
            header,
            partition_tables,
        })
    }
}

impl<T: Read + Seek> MetadataRead for T {
    fn read_wii_header(&mut self) -> Result<Header, Error> {
        // Check file header, and 2 magic bytes
        self.seek(SeekFrom::Start(0))?;
        let mut game_code = [0u8; 4];
        self.read_exact(&mut game_code)?;
        let game_code = from_latin1_or_shift_jis(&game_code)?;
        let maker_code = self.read_u16::<BigEndian>()?.to_be_bytes();
        let maker_code = from_latin1_or_shift_jis(&maker_code)?;
        let disk_id = self.read_u8()?;
        let version = self.read_u8()?;
        let audio_streaming = self.read_u8()?;
        let stream_buffer_size = self.read_u8()?;
        self.seek(SeekFrom::Current(14))?;
        let magic = self.read_u32::<BigEndian>()?;
        let mut game_name = [0u8; 0x40];
        self.read_exact(&mut game_name)?;
        let game_name = trim(&from_latin1_or_shift_jis(&game_name)?).to_string();

        if game_code.len() != 4 {
            return Err(Error::Parse("bad game code string, wrong size".into()));
        }

        if !game_code.is_ascii() {
            return Err(Error::Parse("bad game code string, not ASCII".into()));
        }

        let disable_hash_verification = self.read_u8()?;
        let disable_disc_encryption = self.read_u8()?;

        Ok(Header {
            console_id: game_code[0..1].to_string(),
            game_code: game_code[1..3].to_string(),
            country_code: game_code[3..4].to_string(),
            maker_code,
            disk_id,
            version,
            audio_streaming,
            stream_buffer_size,
            magic,
            game_name,
            disable_hash_verification,
            disable_disc_encryption,
        })
    }

    fn read_wii_partitions(&mut self) -> Result<PartitionTables, Error> {
        // Check file header, and 2 magic bytes
        let mut partitions = Vec::<Partition>::default();

        for i in 0..4 {
            self.seek(SeekFrom::Start(0x40000 + 8 * i))?;
            let total_partitions = self.read_u32::<BigEndian>()?;
            let partition_table_offset = u64::from(self.read_u32::<BigEndian>()?) << 2;
            if total_partitions == 0 {
                continue;
            }
            for j in 0..total_partitions {
                self.seek(SeekFrom::Start(partition_table_offset + u64::from(j) * 8))?;
                let offset = u64::from(self.read_u32::<BigEndian>()?) << 2;
                let type_ = self.read_u32::<BigEndian>()?;
                self.seek(SeekFrom::Start(offset))?;

                /*let region = self.read_u32::<BigEndian>()?;
                self.seek(SeekFrom::Current(12))?;
                let japan_taiwan = self.read_u8()?;
                let usa = self.read_u8()?;
                self.seek(SeekFrom::Current(1))?;
                let germany = self.read_u8()?;
                let pegi = self.read_u8()?;
                let findland = self.read_u8()?;
                let portugal = self.read_u8()?;
                let britain = self.read_u8()?;
                let australia = self.read_u8()?;
                let korea= self.read_u8()?;*/

                let ticket = self.read_wii_ticket()?;
                let tmd_size = self.read_u32::<BigEndian>()?;
                let tmd_offset = u64::from(self.read_u32::<BigEndian>()?) << 2;
                let cert_chain_size = self.read_u32::<BigEndian>()?;
                let cert_chain_offset = u64::from(self.read_u32::<BigEndian>()?) << 2;
                let h3_table_offset = u64::from(self.read_u32::<BigEndian>()?) << 2;
                let data_offset = u64::from(self.read_u32::<BigEndian>()?) << 2;
                let data_size = u64::from(self.read_u32::<BigEndian>()?) << 2;
                let partition = Partition {
                    offset,
                    type_,
                    ticket,
                    tmd_size,
                    tmd_offset,
                    cert_chain_size,
                    cert_chain_offset,
                    h3_table_offset,
                    data_offset,
                    data_size,
                };
                partitions.push(partition);
            }
        }
        Ok(PartitionTables { partitions })
    }
}

fn print_key(key: &[u8]) {
    for byte in key {
        print!("{byte:02X}");
    }
    println!();
}

const WII_COMMON_KEY: [u8; 16] = [
    0xeb, 0xe4, 0x2a, 0x22, 0x5e, 0x85, 0x93, 0xe4, 0x48, 0xd9, 0xc5, 0x45, 0x73, 0x81, 0xaa, 0xf7,
];

impl Partition {
    /// Reads a sector from the partition.
    ///
    /// # Errors
    ///
    /// [`Error::Io`] if there are any issues with the IO object, or for trying to access a sector
    /// outside of the partition.
    pub fn read_sector<T: Read + Seek>(&self, io: &mut T, sector: u32) -> io::Result<Vec<u8>> {
        // Find location of sector in partition
        let position = u64::from(sector) * 0x8000 + self.offset + self.data_offset;
        let max = self.offset + self.data_offset + self.data_size;
        if position >= max {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "sector out of partition range",
            ));
        }

        io.seek(SeekFrom::Start(position))?;
        let mut data = vec![0u8; 0x8000];
        io.read_exact(&mut data)?;
        // Data IV is in the encrypted hash table, copy it before the table gets decrypted...
        // This should never panic, as data is 0x8000 wide
        #[allow(clippy::missing_panics_doc)]
        let data_iv: [u8; 16] = data[0x3D0..0x3E0].try_into().unwrap();
        // Now, decrypt hashes in place, and then data
        let mut title_iv = [0u8; 16];
        title_iv[0..8].copy_from_slice(&self.ticket.title_id);
        // Copy title key so we can decrypt it in place
        let mut title_key = self.ticket.title_key;

        let mut title_aes = CbcDecryptor::<Aes128>::new(&WII_COMMON_KEY.into(), &title_iv.into());
        title_aes.decrypt_block_mut((&mut title_key).into());
        let title_key = title_key;

        // Decrypt hash
        let mut aes_hash = CbcDecryptor::<Aes128>::new(&title_key.into(), &[0u8; 16].into());
        for block in data[0..0x400].chunks_exact_mut(16) {
            let block = Block::<Aes128>::from_mut_slice(block);
            aes_hash.decrypt_block_mut(block);
        }

        // Decrypt data
        let mut data_aes = CbcDecryptor::<Aes128>::new(&title_key.into(), &data_iv.into());
        for block in data[0x400..0x8000].chunks_exact_mut(16) {
            let block = Block::<Aes128>::from_mut_slice(block);
            data_aes.decrypt_block_mut(block);
        }

        Ok(data)
    }
}