oxyromon 0.22.0

use super::database::*;
use super::import_dats::reimport_orphan_romfiles;
use super::iso9660::{FileEntry, Iso9660Filesystem, SECTOR_SIZE, SectorReader};
use super::mimetype::*;
use super::model::*;
use super::progress::*;
use super::prompt::*;
use super::util::*;
use anyhow::{Result, bail};
use clap::value_parser;
use clap::{Arg, ArgAction, ArgMatches, Command};
use flate2::read::GzDecoder;
use indicatif::ProgressBar;
use sqlx::sqlite::SqliteConnection;
use std::collections::hash_map::{Entry, HashMap};
use std::io;
use std::io::prelude::*;
use std::path::{Path, PathBuf};
use std::str;
use strsim::jaro_winkler;

const IRD_VERSION: u8 = 9;

pub fn subcommand() -> Command {
    Command::new("import-irds")
        .about("Parse and import PlayStation 3 IRD files into oxyromon")
        .arg(
            Arg::new("IRDS")
                .help("Set the IRD files to import")
                .required(true)
                .num_args(1..)
                .index(1)
                .value_parser(value_parser!(PathBuf)),
        )
        .arg(
            Arg::new("INFO")
                .short('i')
                .long("info")
                .help("Show the IRD information and exit")
                .required(false)
                .action(ArgAction::SetTrue),
        )
        .arg(
            Arg::new("FORCE")
                .short('f')
                .long("force")
                .help("Force import of already imported IRD files")
                .required(false)
                .action(ArgAction::SetTrue),
        )
}

pub async fn main(
    connection: &mut SqliteConnection,
    matches: &ArgMatches,
    progress_bar: &ProgressBar,
) -> Result<()> {
    let ird_paths: Vec<&PathBuf> = matches.get_many::<PathBuf>("IRDS").unwrap().collect();
    let system = prompt_for_system_like(connection, None, "%PlayStation 3%").await?;
    let mut games = find_wanted_games_by_system_id(connection, system.id).await;

    for ird_path in ird_paths {
        let (irdfile, header) = parse_ird(ird_path).await?;

        print_header(progress_bar, &format!("IRD: {}", irdfile.game_name));
        print_info(progress_bar, &format!("IRD Version: {}", irdfile.version));
        print_info(progress_bar, &format!("Game ID: {}", irdfile.game_id));
        print_info(
            progress_bar,
            &format!("Update Version: {}", irdfile.update_version),
        );
        print_info(
            progress_bar,
            &format!("Game Version: {}", irdfile.game_version),
        );
        print_info(
            progress_bar,
            &format!("App Version: {}", irdfile.app_version),
        );
        print_info(progress_bar, &format!("Regions: {}", irdfile.regions_count));
        print_info(progress_bar, &format!("Files: {}", irdfile.files_count));

        if irdfile.version != IRD_VERSION {
            print_warning(progress_bar, "IRD version unsupported");
            continue;
        }

        if !matches.get_flag("INFO") {
            games.sort_by(|a, b| {
                jaro_winkler(&b.name.to_lowercase(), &irdfile.game_name.to_lowercase())
                    .partial_cmp(&jaro_winkler(
                        &a.name.to_lowercase(),
                        &irdfile.game_name.to_lowercase(),
                    ))
                    .unwrap()
            });
            if let Some(game) = prompt_for_game(&games, None)? {
                if game.jbfolder && !matches.get_flag("FORCE") {
                    print_skip(progress_bar, "IRD already exists");
                    continue;
                }
                import_ird(connection, progress_bar, game, &irdfile, header).await?;
            }
        }
        print_separator(progress_bar);
    }

    compute_system_completion(connection, progress_bar, &system).await?;

    Ok(())
}

pub async fn parse_ird<P: AsRef<Path>>(path: &P) -> Result<(Irdfile, Vec<u8>)> {
    let mimetype = get_mimetype(path).await?;

    if mimetype.is_none() {
        bail!("Not an IRD file");
    }

    let mut reader = match mimetype.unwrap().extension() {
        GZ_EXTENSION => Box::new(GzDecoder::new(get_reader_sync(&path)?)) as Box<dyn Read>,
        IRD_EXTENSION => Box::new(get_reader_sync(&path)?) as Box<dyn Read>,
        _ => bail!("Not an IRD file"),
    };

    // read magic
    let mut magic = [0u8; 4];
    reader.read_exact(&mut magic).unwrap();

    // read version
    let mut version = [0u8];
    reader.read_exact(&mut version).unwrap();

    // read game id
    let mut game_id = [0u8; 9];
    reader.read_exact(&mut game_id).unwrap();

    // read game name
    let mut game_name_length = [0u8];
    reader.read_exact(&mut game_name_length).unwrap();
    let mut game_name = vec![0u8; game_name_length[0] as usize];
    reader.read_exact(&mut game_name).unwrap();

    // read update version
    let mut update_version = [0u8; 4];
    reader.read_exact(&mut update_version).unwrap();

    // read game version
    let mut game_version = [0u8; 5];
    reader.read_exact(&mut game_version).unwrap();

    // read app version
    let mut app_version = [0u8; 5];
    reader.read_exact(&mut app_version).unwrap();

    // read header
    let mut gzipped_header_length = [0u8; 4];
    reader.read_exact(&mut gzipped_header_length).unwrap();
    let gzipped_header_length = u32::from_le_bytes(gzipped_header_length);
    let mut gzipped_header = vec![0u8; gzipped_header_length as usize];
    reader.read_exact(&mut gzipped_header).unwrap();
    let mut gzipped_header_decoder = GzDecoder::new(io::Cursor::new(gzipped_header));
    let mut header: Vec<u8> = vec![];
    gzipped_header_decoder.read_to_end(&mut header).unwrap();

    // read footer
    let mut footer_length = [0u8; 4];
    reader.read_exact(&mut footer_length).unwrap();
    let footer_length = u32::from_le_bytes(footer_length);
    let mut footer = vec![0u8; footer_length as usize];
    reader.read_exact(&mut footer).unwrap();

    // read region hashes
    let mut regions_count = [0u8];
    reader.read_exact(&mut regions_count).unwrap();
    let regions_count = regions_count[0] as usize;
    let mut regions_hashes: Vec<String> = Vec::with_capacity(regions_count);
    while regions_hashes.len() < regions_count {
        let mut region_hash = [0u8; 16];
        reader.read_exact(&mut region_hash).unwrap();
        #[allow(clippy::format_collect)]
        regions_hashes.push(region_hash.iter().map(|b| format!("{:02x}", b)).collect());
    }

    // read file hashes
    let mut files_count = [0u8; 4];
    reader.read_exact(&mut files_count).unwrap();
    let files_count = u32::from_le_bytes(files_count) as usize;
    let mut files_hashes: HashMap<u64, String> = HashMap::with_capacity(files_count);
    while files_hashes.len() < files_count {
        let mut sector = [0u8; 8];
        reader.read_exact(&mut sector).unwrap();
        let sector = u64::from_le_bytes(sector);
        let mut file_hash = [0u8; 16];
        reader.read_exact(&mut file_hash).unwrap();
        files_hashes.insert(
            sector,
            #[allow(clippy::format_collect)]
            file_hash.iter().map(|b| format!("{:02x}", b)).collect(),
        );
    }

    Ok((
        Irdfile {
            version: version[0],
            game_id: str::from_utf8(&game_id).unwrap().trim_end().to_string(),
            game_name: str::from_utf8(&game_name).unwrap().trim_end().to_string(),
            update_version: str::from_utf8(&update_version)
                .unwrap()
                .trim_end()
                .to_string(),
            game_version: str::from_utf8(&game_version)
                .unwrap()
                .trim_end()
                .to_string(),
            app_version: str::from_utf8(&app_version).unwrap().trim_end().to_string(),
            regions_count,
            regions_hashes,
            files_count,
            files_hashes,
        },
        header,
    ))
}

/// A [`SectorReader`] backed by the ISO header embedded in an IRD file.
///
/// IRD files only carry the beginning of the disc, enough for the volume
/// descriptors and the directory extents, so reads past its end are zero filled
/// rather than refused. A short sector would make the default `read_bytes` spin
/// forever, and an error would abort the whole walk.
struct HeaderSectorReader {
    header: Vec<u8>,
}

impl SectorReader for HeaderSectorReader {
    fn read_sector(&mut self, lba: u64) -> Result<Vec<u8>> {
        let mut sector = vec![0u8; SECTOR_SIZE as usize];
        let offset = usize::try_from(lba.saturating_mul(SECTOR_SIZE)).unwrap_or(usize::MAX);
        if let Some(available) = self.header.get(offset..) {
            let length = available.len().min(SECTOR_SIZE as usize);
            sector[..length].copy_from_slice(&available[..length]);
        }
        Ok(sector)
    }
}

/// Collect every file below `directory`, keyed by its path relative to the root.
///
/// Files larger than 4 GiB span several extents, each of which gets its own
/// directory record, hence the vector of entries per path.
fn walk_directory(
    filesystem: &mut Iso9660Filesystem,
    directory: &FileEntry,
) -> Result<HashMap<String, Vec<FileEntry>>> {
    let mut files: HashMap<String, Vec<FileEntry>> = HashMap::new();
    let mut directories: Vec<FileEntry> = vec![];
    for entry in filesystem.list_directory(directory)? {
        if entry.is_directory() {
            directories.push(entry);
            continue;
        }
        // Entry paths are absolute, and ISO 9660 spells an extensionless file
        // "NAME." while DAT files and IRD hashes know it as "NAME"
        let path = entry.path.trim_start_matches('/');
        let path = path.strip_suffix('.').unwrap_or(path).to_string();
        match files.entry(path.clone()) {
            Entry::Vacant(e) => {
                e.insert(Vec::from([entry]));
            }
            _ => {
                files.get_mut(&path).unwrap().push(entry);
            }
        }
    }
    for directory in directories {
        files.extend(walk_directory(filesystem, &directory)?);
    }
    Ok(files)
}

pub async fn import_ird(
    connection: &mut SqliteConnection,
    progress_bar: &ProgressBar,
    game: &Game,
    irdfile: &Irdfile,
    header: Vec<u8>,
) -> Result<()> {
    let roms = find_roms_by_game_id_no_parents(connection, game.id).await;
    let parent_rom = prompt_for_rom(&roms, None)?;
    if parent_rom.is_none() {
        return Ok(());
    }

    let mut transaction = begin_transaction(connection).await;

    // parse ISO header
    let mut filesystem = Iso9660Filesystem::new(Box::new(HeaderSectorReader { header }))?;
    let root = filesystem.root();
    let files = walk_directory(&mut filesystem, &root)?;

    if files.len() != irdfile.files_count {
        bail!(
            "Files count differ, expected {}, actual {}",
            irdfile.files_count,
            files.len()
        );
    }

    // convert files into roms
    let mut orphan_romfile_ids: Vec<i64> = vec![];
    for file in files {
        let size = file.1.iter().map(|file| file.size).sum::<u64>() as i64;
        let location = file.1.first().unwrap().location;
        match find_rom_by_name_and_game_id(&mut transaction, &file.0, game.id).await {
            Some(rom) => {
                update_rom(
                    &mut transaction,
                    rom.id,
                    &file.0,
                    size,
                    irdfile.files_hashes.get(&location).unwrap(),
                    game.id,
                    parent_rom.as_ref().map(|rom| rom.id),
                    false,
                )
                .await;
                if (size != rom.size
                    || irdfile.files_hashes.get(&location).unwrap() != rom.md5.as_ref().unwrap())
                    && let Some(romfile_id) = rom.romfile_id
                {
                    orphan_romfile_ids.push(romfile_id);
                    update_rom_romfile(&mut transaction, rom.id, None).await;
                }
                rom.id
            }
            None => {
                create_rom(
                    &mut transaction,
                    &file.0,
                    size,
                    irdfile.files_hashes.get(&location).unwrap(),
                    game.id,
                    parent_rom.as_ref().map(|rom| rom.id),
                    false,
                )
                .await
            }
        };
    }

    // mark game as jbfolder
    update_game_jbfolder(&mut transaction, game.id, true).await;

    // reimport orphan romfiles
    if !orphan_romfile_ids.is_empty() {
        print_subheader(progress_bar, "Processing orphan romfiles");
        reimport_orphan_romfiles(
            &mut transaction,
            progress_bar,
            game.system_id,
            orphan_romfile_ids,
        )
        .await?;
    }

    commit_transaction(transaction).await;

    Ok(())
}

#[cfg(test)]
mod test_ird;