use super::config::*;
use super::crc32::*;
use super::database::*;
use super::generate_playlists::DISC_REGEX;
use super::mimetype::*;
use super::model::*;
use super::progress::*;
use super::prompt::prompt_for_games;
use super::util::*;
use anyhow::{Context, Result, bail};
use clap::ArgMatches;
use core::fmt;
use digest::Digest;
use digest_io::IoWrapper;
use indexmap::IndexMap;
use indicatif::ProgressBar;
use md5::Md5;
use num_traits::FromPrimitive;
use rayon::prelude::*;
use sha1::Sha1;
use sqlx::SqliteConnection;
use std::collections::HashMap;
use std::io;
use std::io::prelude::*;
use std::path::{Path, PathBuf};
use std::{fs::File, str::FromStr};
fn to_hex(bytes: &[u8]) -> String {
bytes.iter().map(|b| format!("{:02x}", b)).collect()
}
pub async fn select_games(
connection: &mut SqliteConnection,
matches: &ArgMatches,
system_id: i64,
) -> Result<Vec<Game>> {
let games = match matches.get_many::<String>("GAME") {
Some(game_names) => {
let mut games: Vec<Game> = vec![];
for game_name in game_names {
games.append(
&mut find_full_games_by_name_and_system_id(
connection, game_name, system_id,
)
.await,
);
}
games.dedup_by_key(|game| game.id);
prompt_for_games(games, cfg!(test))?
}
None => find_full_games_by_system_id(connection, system_id).await,
};
Ok(games)
}
pub async fn load_rom_data(
connection: &mut SqliteConnection,
games: Vec<Game>,
original_roms: bool,
) -> (IndexMap<i64, Vec<Rom>>, HashMap<i64, Game>, HashMap<i64, Romfile>) {
let game_ids: Vec<i64> = games.par_iter().map(|game| game.id).collect();
let roms = if original_roms {
find_original_roms_with_romfile_by_game_ids(connection, &game_ids).await
} else {
find_roms_with_romfile_by_game_ids(connection, &game_ids).await
};
let romfiles = find_romfiles_by_ids(
connection,
roms.par_iter()
.map(|rom| rom.romfile_id.unwrap())
.collect::<Vec<i64>>()
.as_slice(),
)
.await;
let mut roms_by_game_id: IndexMap<i64, Vec<Rom>> = IndexMap::new();
roms.into_iter().for_each(|rom| {
let group = roms_by_game_id.entry(rom.game_id).or_default();
group.push(rom);
});
let games_by_id: HashMap<i64, Game> =
games.into_par_iter().map(|game| (game.id, game)).collect();
let romfiles_by_id: HashMap<i64, Romfile> = romfiles
.into_par_iter()
.map(|romfile| (romfile.id, romfile))
.collect();
(roms_by_game_id, games_by_id, romfiles_by_id)
}
pub fn partition_games_by_extensions(
roms_by_game_id: IndexMap<i64, Vec<Rom>>,
romfiles_by_id: &HashMap<i64, Romfile>,
extensions: &[&str],
) -> (IndexMap<i64, Vec<Rom>>, IndexMap<i64, Vec<Rom>>) {
roms_by_game_id.into_iter().partition(|(_, roms)| {
roms.par_iter().any(|rom| {
let path = &romfiles_by_id.get(&rom.romfile_id.unwrap()).unwrap().path;
extensions.iter().any(|extension| path.ends_with(extension))
})
})
}
pub fn partition_games_by_all_extensions(
roms_by_game_id: IndexMap<i64, Vec<Rom>>,
romfiles_by_id: &HashMap<i64, Romfile>,
extensions: &[&str],
) -> (IndexMap<i64, Vec<Rom>>, IndexMap<i64, Vec<Rom>>) {
roms_by_game_id.into_iter().partition(|(_, roms)| {
extensions.iter().all(|extension| {
roms.par_iter().any(|rom| {
romfiles_by_id
.get(&rom.romfile_id.unwrap())
.unwrap()
.path
.ends_with(extension)
})
})
})
}
#[derive(Clone)]
pub struct CommonRomfile {
pub path: PathBuf,
pub system_id: Option<i64>,
}
impl CommonRomfile {
pub fn with_system(mut self, system_id: Option<i64>) -> Self {
self.system_id = system_id;
self
}
}
pub struct IsoRomfile {
pub romfile: CommonRomfile,
}
pub struct CueBinRomfile {
pub cue_romfile: CommonRomfile,
pub bin_romfiles: Vec<CommonRomfile>,
}
pub trait FromPath<T> {
fn from_path<P: AsRef<Path>>(path: &P) -> Result<T>;
}
pub trait CommonFile {
async fn get_sorted_path(
&self,
connection: &mut SqliteConnection,
system: &System,
game: &Game,
rom: &Rom,
subfolders: &Option<SubfolderScheme>,
extension: &Option<&str>,
) -> Result<PathBuf>;
async fn get_relative_path(&self, connection: &mut SqliteConnection) -> Result<&Path>;
async fn rename<P: AsRef<Path>>(
&self,
progress_bar: &ProgressBar,
new_path: &P,
quiet: bool,
) -> Result<CommonRomfile>;
async fn copy<P: AsRef<Path>>(
&self,
progress_bar: &ProgressBar,
destination_directory: &P,
quiet: bool,
) -> Result<CommonRomfile>;
async fn delete(&self, progress_bar: &ProgressBar, quiet: bool) -> Result<()>;
}
pub trait GetRomfile {
fn romfile(&self) -> &CommonRomfile;
}
impl GetRomfile for CommonRomfile {
fn romfile(&self) -> &CommonRomfile {
self
}
}
impl GetRomfile for IsoRomfile {
fn romfile(&self) -> &CommonRomfile {
&self.romfile
}
}
pub trait AsCommon {
async fn as_common(&self, connection: &mut SqliteConnection) -> Result<CommonRomfile>;
}
pub trait ToCommon {
async fn to_common<P: AsRef<Path>>(
&self,
progress_bar: &ProgressBar,
destination_directory: &P,
) -> Result<CommonRomfile>;
}
pub trait AsIso {
fn as_iso(self) -> Result<IsoRomfile>;
}
pub trait AsCueBin {
fn as_cue_bin(self, bin_romfiles: Vec<CommonRomfile>) -> Result<CueBinRomfile>;
}
pub trait ToIso {
async fn to_iso<P: AsRef<Path>>(
&self,
progress_bar: &ProgressBar,
destination_directory: &P,
) -> Result<IsoRomfile>;
}
pub trait ToCueBin {
async fn to_cue_bin<P: AsRef<Path>>(
&self,
progress_bar: &ProgressBar,
destination_directory: &P,
cue_romfile: Option<CommonRomfile>,
bin_roms: &[&Rom],
quiet: bool,
) -> Result<CueBinRomfile>;
}
#[allow(dead_code)]
pub trait Patch {
async fn patch<P: AsRef<Path>>(
&self,
progress_bar: &ProgressBar,
romfile: &CommonRomfile,
destination_directory: &P,
) -> Result<CommonRomfile>;
}
pub trait Playlist {
async fn get_playlist_path(
&self,
connection: &mut SqliteConnection,
system: &System,
subfolders: &Option<SubfolderScheme>,
) -> Result<PathBuf>;
}
pub trait Size {
async fn get_size(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
) -> Result<u64>;
}
pub trait HashAndSize {
async fn get_hash_and_size(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
position: usize,
total: usize,
hash_algorithm: &HashAlgorithm,
) -> Result<(String, u64)>;
}
pub trait HeaderedHashAndSize {
async fn get_file_and_header_size(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
header: &Header,
) -> Result<(File, u64)>;
async fn get_headered_size(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
header: &Header,
) -> Result<u64>;
async fn get_headered_hash_and_size(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
header: &Header,
position: usize,
total: usize,
hash_algorithm: &HashAlgorithm,
) -> Result<(String, u64)>;
}
pub trait Check {
async fn check(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
header: &Option<Header>,
roms: &[&Rom],
) -> Result<()>;
}
pub trait Persist {
async fn create(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
romfile_type: RomfileType,
) -> Result<i64>;
async fn update(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
id: i64,
) -> Result<()>;
}
impl FromPath<CommonRomfile> for CommonRomfile {
fn from_path<P: AsRef<Path>>(path: &P) -> Result<CommonRomfile> {
Ok(CommonRomfile {
path: path.as_ref().to_path_buf(),
system_id: None,
})
}
}
impl fmt::Display for CommonRomfile {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.path.as_os_str().to_str().unwrap())
}
}
impl CommonFile for CommonRomfile {
async fn get_relative_path(&self, connection: &mut SqliteConnection) -> Result<&Path> {
let rom_directory = match get_directory(connection, "ROM_DIRECTORY", self.system_id).await {
Some(dir) => dir,
None => get_rom_directory(connection).await,
};
let relative_path = self.path.strip_prefix(rom_directory).with_context(|| {
format!(
"Failed to convert \"{}\"to relative path",
self.path.as_os_str().to_str().unwrap()
)
})?;
Ok(relative_path)
}
async fn rename<P: AsRef<Path>>(
&self,
progress_bar: &ProgressBar,
new_path: &P,
quiet: bool,
) -> Result<CommonRomfile> {
if self.path != new_path.as_ref() {
rename_file(progress_bar, &self.path, new_path, quiet).await?;
}
CommonRomfile::from_path(new_path)
}
async fn copy<P: AsRef<Path>>(
&self,
progress_bar: &ProgressBar,
destination_directory: &P,
quiet: bool,
) -> Result<CommonRomfile> {
let new_path = destination_directory
.as_ref()
.join(self.path.file_name().unwrap());
copy_file(progress_bar, &self.path, &new_path, quiet).await?;
CommonRomfile::from_path(&new_path)
}
async fn delete(&self, progress_bar: &ProgressBar, quiet: bool) -> Result<()> {
remove_file(progress_bar, &self.path, quiet).await?;
Ok(())
}
async fn get_sorted_path(
&self,
connection: &mut SqliteConnection,
system: &System,
game: &Game,
rom: &Rom,
subfolder_scheme: &Option<SubfolderScheme>,
extension: &Option<&str>,
) -> Result<PathBuf> {
let extension = extension.or_else(|| self.path.extension()?.to_str());
let mut sorted_path = match game.sorting {
s if s == Sorting::OneRegion as i64 => {
get_one_region_directory(connection, system).await?
}
s if s == Sorting::Ignored as i64 => {
get_trash_directory(connection, Some(system)).await?
}
_ => get_system_directory(connection, system).await?,
};
let subfolder_scheme = match subfolder_scheme {
Some(scheme) => scheme,
None => match Sorting::from_i64(game.sorting) {
Some(Sorting::OneRegion) => &SubfolderScheme::from_str(
&get_string(connection, "REGIONS_ONE_SUBFOLDERS", Some(system.id))
.await
.unwrap(),
)
.unwrap(),
Some(Sorting::AllRegions) => &SubfolderScheme::from_str(
&get_string(connection, "REGIONS_ALL_SUBFOLDERS", Some(system.id))
.await
.unwrap(),
)
.unwrap(),
Some(Sorting::Ignored) | None => &SubfolderScheme::None,
},
};
if subfolder_scheme == &SubfolderScheme::Alpha {
sorted_path.push(compute_alpha_subfolder(&game.name));
}
if (system.arcade && !extension.is_some_and(|ext| ARCHIVE_EXTENSIONS.contains(&ext)))
|| game.jbfolder
{
sorted_path.push(&game.name);
}
let filename = match extension {
Some(ext) if NON_ORIGINAL_EXTENSIONS.contains(&ext) => {
if system.arcade && !ARCHIVE_EXTENSIONS.contains(&ext) {
format!("{}.{}", rom.name, ext)
} else {
format!("{}.{}", game.name, ext)
}
}
_ => match &system.custom_extension {
Some(custom_ext) => format!("{}.{}", game.name, custom_ext),
None => rom.name.clone(),
},
};
sorted_path.push(filename);
Ok(sorted_path)
}
}
impl Size for CommonRomfile {
async fn get_size(
&self,
_connection: &mut SqliteConnection,
_progress_bar: &ProgressBar,
) -> Result<u64> {
Ok(self.path.metadata().unwrap().len())
}
}
impl HashAndSize for CommonRomfile {
async fn get_hash_and_size(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
position: usize,
total: usize,
hash_algorithm: &HashAlgorithm,
) -> Result<(String, u64)> {
progress_bar.reset();
progress_bar.set_message(format!(
"Computing {} ({}/{})",
hash_algorithm, position, total
));
let mut file = open_file_sync(&self.path)?;
let (hash, _) = hash_reader(&mut file, progress_bar, hash_algorithm)?;
let size = self.get_size(connection, progress_bar).await?;
progress_bar.set_message("");
Ok((hash, size))
}
}
pub fn hash_reader<R: Read + ?Sized>(
reader: &mut R,
progress_bar: &ProgressBar,
hash_algorithm: &HashAlgorithm,
) -> Result<(String, u64)> {
let size;
let hash = match hash_algorithm {
HashAlgorithm::Crc => {
let mut digest = Crc32::new();
size = io::copy(reader, &mut progress_bar.wrap_write(&mut digest))
.context("Failed to copy data")?;
to_hex(&digest.finalize())
}
HashAlgorithm::Md5 => {
let mut digest = IoWrapper(Md5::new());
size = io::copy(reader, &mut progress_bar.wrap_write(&mut digest))
.context("Failed to copy data")?;
to_hex(&digest.0.finalize())
}
HashAlgorithm::Sha1 => {
let mut digest = IoWrapper(Sha1::new());
size = io::copy(reader, &mut progress_bar.wrap_write(&mut digest))
.context("Failed to copy data")?;
to_hex(&digest.0.finalize())
}
};
Ok((hash, size))
}
pub fn get_hash_algorithm(rom: &Rom) -> Result<HashAlgorithm> {
if rom.crc.is_some() {
Ok(HashAlgorithm::Crc)
} else if rom.md5.is_some() {
Ok(HashAlgorithm::Md5)
} else if rom.sha1.is_some() {
Ok(HashAlgorithm::Sha1)
} else {
bail!("Not possible")
}
}
pub fn compare_hash_and_size(
rom: &Rom,
hash: &str,
size: u64,
hash_algorithm: &HashAlgorithm,
) -> Result<()> {
if rom.size > 0 && size != rom.size as u64 {
bail!("Size mismatch");
};
let expected = match hash_algorithm {
HashAlgorithm::Crc => rom.crc.as_ref(),
HashAlgorithm::Md5 => rom.md5.as_ref(),
HashAlgorithm::Sha1 => rom.sha1.as_ref(),
};
if Some(&hash.to_string()) != expected {
bail!("Checksum mismatch");
}
Ok(())
}
impl HeaderedHashAndSize for CommonRomfile {
async fn get_file_and_header_size(
&self,
connection: &mut SqliteConnection,
_progress_bar: &ProgressBar,
header: &Header,
) -> Result<(File, u64)> {
let mut file = open_file_sync(&self.path)?;
let mut header_size: u64 = 0;
let rules = find_rules_by_header_id(connection, header.id).await;
let mut buffer: Vec<u8> = Vec::with_capacity(header.size as usize);
(&mut file)
.take(header.size as u64)
.read_to_end(&mut buffer)
.context("Failed to read into buffer")?;
let mut matches: Vec<bool> = vec![];
for rule in rules {
let start_byte = rule.start_byte as usize;
let hex_value = to_hex(&buffer[start_byte..]);
matches.push(hex_value.starts_with(&rule.hex_value.to_lowercase()));
}
if matches.iter().all(|&m| m) {
header_size = header.size as u64;
} else {
file.rewind().context("Failed to rewind file")?;
}
Ok((file, header_size))
}
async fn get_headered_size(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
header: &Header,
) -> Result<u64> {
let (file, header_size) = self
.get_file_and_header_size(connection, progress_bar, header)
.await?;
Ok(file.metadata().unwrap().len() - header_size)
}
async fn get_headered_hash_and_size(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
header: &Header,
position: usize,
total: usize,
hash_algorithm: &HashAlgorithm,
) -> Result<(String, u64)> {
let size = self
.get_headered_size(connection, progress_bar, header)
.await?;
progress_bar.reset();
progress_bar.set_message(format!(
"Computing {} ({}/{})",
hash_algorithm, position, total
));
progress_bar.set_style(get_bytes_progress_style());
progress_bar.set_length(size);
let mut file = self
.get_file_and_header_size(connection, progress_bar, header)
.await?
.0;
let hash = match hash_algorithm {
HashAlgorithm::Crc => {
let mut digest = Crc32::new();
io::copy(&mut file, &mut progress_bar.wrap_write(&mut digest))
.context("Failed to copy data")?;
to_hex(&digest.finalize())
}
HashAlgorithm::Md5 => {
let mut digest = IoWrapper(Md5::new());
io::copy(&mut file, &mut progress_bar.wrap_write(&mut digest))
.context("Failed to copy data")?;
to_hex(&digest.0.finalize())
}
HashAlgorithm::Sha1 => {
let mut digest = IoWrapper(Sha1::new());
io::copy(&mut file, &mut progress_bar.wrap_write(&mut digest))
.context("Failed to copy data")?;
to_hex(&digest.0.finalize())
}
};
progress_bar.set_message("");
progress_bar.set_style(get_none_progress_style());
Ok((hash, size))
}
}
impl Check for CommonRomfile {
async fn check(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
header: &Option<Header>,
roms: &[&Rom],
) -> Result<()> {
print_action(progress_bar, &format!("Checking \"{}\"", self));
let rom = roms[0];
let hash_algorithm = get_hash_algorithm(rom)?;
let (hash, size) = match header {
Some(header) => {
self.get_headered_hash_and_size(
connection,
progress_bar,
header,
1,
1,
&hash_algorithm,
)
.await?
}
None => {
let (hash, size) = self
.get_hash_and_size(connection, progress_bar, 1, 1, &hash_algorithm)
.await?;
(hash, size)
}
};
compare_hash_and_size(rom, &hash, size, &hash_algorithm)
}
}
impl Persist for CommonRomfile {
async fn create(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
romfile_type: RomfileType,
) -> Result<i64> {
let path = &self.get_relative_path(connection).await?;
let size = self.get_size(connection, progress_bar).await?;
Ok(create_romfile(
connection,
path.as_os_str().to_str().unwrap(),
size,
romfile_type,
)
.await)
}
async fn update(
&self,
connection: &mut SqliteConnection,
progress_bar: &ProgressBar,
id: i64,
) -> Result<()> {
let path = &self.get_relative_path(connection).await?;
let size = self.get_size(connection, progress_bar).await?;
update_romfile(connection, id, path.as_os_str().to_str().unwrap(), size).await;
Ok(())
}
}
impl AsCommon for Romfile {
async fn as_common(&self, connection: &mut SqliteConnection) -> Result<CommonRomfile> {
let rom_directory = get_rom_directory(connection).await;
CommonRomfile::from_path(&rom_directory.join(&self.path))
}
}
impl AsIso for CommonRomfile {
fn as_iso(self) -> Result<IsoRomfile> {
if self
.path
.extension()
.unwrap()
.to_str()
.unwrap()
.to_lowercase()
!= ISO_EXTENSION
{
bail!("Not a valid iso");
}
Ok(IsoRomfile { romfile: self })
}
}
impl AsCueBin for CommonRomfile {
fn as_cue_bin(self, bin_romfiles: Vec<CommonRomfile>) -> Result<CueBinRomfile> {
if self
.path
.extension()
.unwrap()
.to_str()
.unwrap()
.to_lowercase()
!= CUE_EXTENSION
{
bail!("Not a valid cue");
}
for bin_romfile in &bin_romfiles {
if bin_romfile
.path
.extension()
.unwrap()
.to_str()
.unwrap()
.to_lowercase()
!= BIN_EXTENSION
{
bail!("Not a valid bin");
}
}
Ok(CueBinRomfile {
cue_romfile: self,
bin_romfiles,
})
}
}
impl Playlist for Game {
async fn get_playlist_path(
&self,
connection: &mut SqliteConnection,
system: &System,
subfolder_scheme: &Option<SubfolderScheme>,
) -> Result<PathBuf> {
let mut playlist_path = if self.sorting == Sorting::OneRegion as i64 {
get_one_region_directory(connection, system).await?
} else {
get_system_directory(connection, system).await?
};
let subfolder_scheme = match subfolder_scheme {
Some(scheme) => scheme,
None => match Sorting::from_i64(self.sorting) {
Some(Sorting::OneRegion) => &SubfolderScheme::from_str(
&get_string(connection, "REGIONS_ONE_SUBFOLDERS", Some(system.id))
.await
.unwrap(),
)
.unwrap(),
Some(Sorting::AllRegions) => &SubfolderScheme::from_str(
&get_string(connection, "REGIONS_ALL_SUBFOLDERS", Some(system.id))
.await
.unwrap(),
)
.unwrap(),
Some(Sorting::Ignored) | None => &SubfolderScheme::None,
},
};
if subfolder_scheme == &SubfolderScheme::Alpha {
playlist_path = playlist_path.join(compute_alpha_subfolder(&self.name));
}
playlist_path = playlist_path.join(format!(
"{}.{}",
DISC_REGEX.replace(&self.name, ""),
M3U_EXTENSION
));
Ok(playlist_path)
}
}
#[cfg(test)]
mod test_path_archive_custom_extension;
#[cfg(test)]
mod test_path_archive_multiple_files;
#[cfg(test)]
mod test_path_archive_single_file;
#[cfg(test)]
mod test_path_chd_multiple_tracks;
#[cfg(test)]
mod test_path_chd_single_track;
#[cfg(test)]
mod test_path_cso;
#[cfg(test)]
mod test_path_custom_extension;
#[cfg(test)]
mod test_path_original;
#[cfg(test)]
mod test_path_playlist;
#[cfg(test)]
mod test_path_playlist_subfolder_alpha;
#[cfg(test)]
mod test_path_rvz;
#[cfg(test)]
mod test_path_subfolder_alpha_letter;
#[cfg(test)]
mod test_path_subfolder_alpha_other;