use std::{
borrow::Cow,
fs,
fs::File,
io::{BufRead, Write},
path::{Path, PathBuf},
};
use argp::FromArgs;
use itertools::Itertools;
use nod::{
Disc, Fst, Node, OpenOptions, PartitionBase, PartitionKind, PartitionMeta, ResultContext,
};
use size::{Base, Size};
use zerocopy::IntoBytes;
use crate::util::{display, has_extension};
#[derive(FromArgs, Debug)]
#[argp(subcommand, name = "extract")]
pub struct Args {
#[argp(positional)]
file: PathBuf,
#[argp(positional)]
out: Option<PathBuf>,
#[argp(switch, short = 'q')]
quiet: bool,
#[argp(switch, short = 'h')]
validate: bool,
#[argp(option, short = 'p')]
partition: Option<String>,
}
pub fn run(args: Args) -> nod::Result<()> {
let output_dir: PathBuf;
if let Some(dir) = args.out {
output_dir = dir;
} else if has_extension(&args.file, "nfs") {
if let Some(parent) = args.file.parent() {
output_dir = parent.with_file_name("extracted");
} else {
output_dir = args.file.with_extension("");
}
} else {
output_dir = args.file.with_extension("");
}
let disc = Disc::new_with_options(&args.file, &OpenOptions {
rebuild_encryption: false,
validate_hashes: args.validate,
})?;
let header = disc.header();
let is_wii = header.is_wii();
if let Some(partition) = args.partition {
if partition.eq_ignore_ascii_case("all") {
for info in disc.partitions() {
let mut out_dir = output_dir.clone();
out_dir.push(info.kind.dir_name().as_ref());
let mut partition = disc.open_partition(info.index)?;
extract_partition(&disc, partition.as_mut(), &out_dir, is_wii, args.quiet)?;
}
} else if partition.eq_ignore_ascii_case("data") {
let mut partition = disc.open_partition_kind(PartitionKind::Data)?;
extract_partition(&disc, partition.as_mut(), &output_dir, is_wii, args.quiet)?;
} else if partition.eq_ignore_ascii_case("update") {
let mut partition = disc.open_partition_kind(PartitionKind::Update)?;
extract_partition(&disc, partition.as_mut(), &output_dir, is_wii, args.quiet)?;
} else if partition.eq_ignore_ascii_case("channel") {
let mut partition = disc.open_partition_kind(PartitionKind::Channel)?;
extract_partition(&disc, partition.as_mut(), &output_dir, is_wii, args.quiet)?;
} else {
let idx = partition.parse::<usize>().map_err(|_| "Invalid partition index")?;
let mut partition = disc.open_partition(idx)?;
extract_partition(&disc, partition.as_mut(), &output_dir, is_wii, args.quiet)?;
}
} else {
let mut partition = disc.open_partition_kind(PartitionKind::Data)?;
extract_partition(&disc, partition.as_mut(), &output_dir, is_wii, args.quiet)?;
}
Ok(())
}
fn extract_partition(
disc: &Disc,
partition: &mut dyn PartitionBase,
out_dir: &Path,
is_wii: bool,
quiet: bool,
) -> nod::Result<()> {
let meta = partition.meta()?;
extract_sys_files(disc, meta.as_ref(), out_dir, quiet)?;
let files_dir = out_dir.join("files");
fs::create_dir_all(&files_dir)
.with_context(|| format!("Creating directory {}", display(&files_dir)))?;
let fst = Fst::new(&meta.raw_fst)?;
let mut path_segments = Vec::<(Cow<str>, usize)>::new();
for (idx, node, name) in fst.iter() {
let mut new_size = 0;
for (_, end) in path_segments.iter() {
if *end == idx {
break;
}
new_size += 1;
}
path_segments.truncate(new_size);
let end = if node.is_dir() { node.length() as usize } else { idx + 1 };
path_segments.push((name?, end));
let path = path_segments.iter().map(|(name, _)| name.as_ref()).join("/");
if node.is_dir() {
fs::create_dir_all(files_dir.join(&path))
.with_context(|| format!("Creating directory {}", path))?;
} else {
extract_node(node, partition, &files_dir, &path, is_wii, quiet)?;
}
}
Ok(())
}
fn extract_sys_files(
disc: &Disc,
data: &PartitionMeta,
out_dir: &Path,
quiet: bool,
) -> nod::Result<()> {
let sys_dir = out_dir.join("sys");
fs::create_dir_all(&sys_dir)
.with_context(|| format!("Creating directory {}", display(&sys_dir)))?;
extract_file(data.raw_boot.as_ref(), &sys_dir.join("boot.bin"), quiet)?;
extract_file(data.raw_bi2.as_ref(), &sys_dir.join("bi2.bin"), quiet)?;
extract_file(data.raw_apploader.as_ref(), &sys_dir.join("apploader.img"), quiet)?;
extract_file(data.raw_fst.as_ref(), &sys_dir.join("fst.bin"), quiet)?;
extract_file(data.raw_dol.as_ref(), &sys_dir.join("main.dol"), quiet)?;
let disc_header = disc.header();
if disc_header.is_wii() {
let disc_dir = out_dir.join("disc");
fs::create_dir_all(&disc_dir)
.with_context(|| format!("Creating directory {}", display(&disc_dir)))?;
extract_file(&disc_header.as_bytes()[..0x100], &disc_dir.join("header.bin"), quiet)?;
if let Some(region) = disc.region() {
extract_file(region, &disc_dir.join("region.bin"), quiet)?;
}
if let Some(ticket) = data.raw_ticket.as_deref() {
extract_file(ticket, &out_dir.join("ticket.bin"), quiet)?;
}
if let Some(tmd) = data.raw_tmd.as_deref() {
extract_file(tmd, &out_dir.join("tmd.bin"), quiet)?;
}
if let Some(cert_chain) = data.raw_cert_chain.as_deref() {
extract_file(cert_chain, &out_dir.join("cert.bin"), quiet)?;
}
if let Some(h3_table) = data.raw_h3_table.as_deref() {
extract_file(h3_table, &out_dir.join("h3.bin"), quiet)?;
}
}
Ok(())
}
fn extract_file(bytes: &[u8], out_path: &Path, quiet: bool) -> nod::Result<()> {
if !quiet {
println!(
"Extracting {} (size: {})",
display(out_path),
Size::from_bytes(bytes.len()).format().with_base(Base::Base10)
);
}
fs::write(out_path, bytes).with_context(|| format!("Writing file {}", display(out_path)))?;
Ok(())
}
fn extract_node(
node: Node,
partition: &mut dyn PartitionBase,
base_path: &Path,
name: &str,
is_wii: bool,
quiet: bool,
) -> nod::Result<()> {
let file_path = base_path.join(name);
if !quiet {
println!(
"Extracting {} (size: {})",
display(&file_path),
Size::from_bytes(node.length()).format().with_base(Base::Base10)
);
}
let mut file = File::create(&file_path)
.with_context(|| format!("Creating file {}", display(&file_path)))?;
let mut r = partition.open_file(node).with_context(|| {
format!(
"Opening file {} on disc for reading (offset {}, size {})",
name,
node.offset(is_wii),
node.length()
)
})?;
loop {
let buf =
r.fill_buf().with_context(|| format!("Extracting file {}", display(&file_path)))?;
let len = buf.len();
if len == 0 {
break;
}
file.write_all(buf).with_context(|| format!("Writing file {}", display(&file_path)))?;
r.consume(len);
}
file.flush().with_context(|| format!("Flushing file {}", display(&file_path)))?;
Ok(())
}