diff --git a/src/cli.rs b/src/cli.rs index bf8fd6a..4b86604 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,10 +1,13 @@ use std::path::PathBuf; +use std::sync::LazyLock; use crate::FB_SECTOR_SIZE; -use clap::Parser; +use clap::{ArgAction, Parser}; -#[derive(Parser, Debug)] +pub static CONFIG: LazyLock = LazyLock::new(|| Args::parse()); + +#[derive(Parser, Debug, Clone)] pub struct Args { /// Path to source file or block device #[arg(short, long, value_hint = clap::ValueHint::DirPath)] @@ -20,7 +23,7 @@ pub struct Args { /// Max number of consecutive sectors to test as a group #[arg(short, long, default_value_t = 128)] - pub cluster_length: u16, + pub cluster_length: usize, /// Number of brute force read passes #[arg(short, long, default_value_t = 2)] @@ -28,5 +31,12 @@ pub struct Args { /// Sector size #[arg(short, long, default_value_t = FB_SECTOR_SIZE)] - pub sector_size: u16, + pub sector_size: usize, + + // Behaviour is backwards. + // ArgAction::SetFalse by default evaluates to true, + // ArgAction::SetTrue by default evaluates to false. + /// Whether to reopen the file on a read error or not. + #[arg(short, long, action=ArgAction::SetTrue)] + pub reopen_on_error: bool, } diff --git a/src/io.rs b/src/io.rs index a7a2d48..56e765a 100644 --- a/src/io.rs +++ b/src/io.rs @@ -1,5 +1,7 @@ +use std::fs::{File, OpenOptions}; use std::io::{self, Seek, SeekFrom}; -use std::path::{Path, PathBuf}; + +use crate::cli::CONFIG; use anyhow::Context; @@ -17,28 +19,49 @@ pub fn get_stream_length(stream: &mut S) -> io::Result { len } -/// Generates a file path if one not provided. -/// source_name for fallback name. -pub fn get_path

( - file_path: &Option

, - source_name: &P, - extension: &str, -) -> anyhow::Result -where - P: AsRef, -{ - if let Some(f) = file_path { - return Ok(f.as_ref().to_path_buf()); - } - - Ok(PathBuf::from(format!( - "{}.{}", - source_name - .as_ref() - .to_str() - .context("source_name path was not UTF-8 valid.")?, - extension - )) - .as_path() - .to_owned()) +pub fn load_input() -> anyhow::Result { + OpenOptions::new() + .read(true) + .open(&CONFIG.input) + .with_context(|| format!("Failed to open input file: {}", &CONFIG.input.display())) +} + +pub fn load_output() -> anyhow::Result { + OpenOptions::new() + .read(true) + .write(true) + .create(true) + .open(crate::path::OUTPUT_PATH.clone()) + .with_context(|| { + format!( + "Failed to open/create output file at: {}", + crate::path::OUTPUT_PATH.display() + ) + }) +} + +pub fn load_map_read() -> anyhow::Result { + OpenOptions::new() + .read(true) + .open(crate::path::MAP_PATH.clone()) + .with_context(|| { + format!( + "Failed to open/create mapping file at: {}", + crate::path::MAP_PATH.display() + ) + }) +} + +pub fn load_map_write() -> anyhow::Result { + OpenOptions::new() + .write(true) + .create(true) + .truncate(true) // Wipe old map. Should really make a backup first. + .open(crate::path::MAP_PATH.clone()) + .with_context(|| { + format!( + "Failed to open map file at: {}", + crate::path::MAP_PATH.display() + ) + }) } diff --git a/src/main.rs b/src/main.rs index 66d1de9..d444ec6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,79 +1,19 @@ mod cli; mod io; mod mapping; +mod path; mod recovery; -use std::fs::{File, OpenOptions}; - -use cli::Args; -use mapping::prelude::*; use recovery::Recover; -use anyhow::{self, Context}; -use clap::Parser; +use anyhow; -const FB_SECTOR_SIZE: u16 = 2048; -const FB_PAD_VALUE: u8 = 0; +const FB_SECTOR_SIZE: usize = 2048; +const FB_NULL_VALUE: u8 = 0; fn main() -> anyhow::Result<()> { - let config = Args::parse(); - - let mut input: File = { - let input_path = &config.input.as_path(); - - OpenOptions::new() - .read(true) - .open(&input_path) - .with_context(|| format!("Failed to open input file: {}", input_path.display()))? - }; - - let output: File = { - // Keep this clean, make a short-lived binding. - let path = crate::io::get_path(&config.output, &config.input, "iso") - .context("Failed to generate output path.")?; - - OpenOptions::new() - .read(true) - .write(true) - .create(true) - .open(&path) - .with_context(|| format!("Failed to open/create output file at: {}", path.display()))? - }; - - let map: MapFile = { - let path = crate::io::get_path(&config.map, &config.input, "map") - .context("Failed to generate map path.")?; - - MapFile::try_from( - OpenOptions::new() - .read(true) - .write(true) - .create(true) - .open(&path) - .with_context(|| { - format!("Failed to open/create mapping file at: {}", path.display()) - })?, - ) - .unwrap_or(MapFile::new(config.sector_size)) - }; - - let mut recover_tool = Recover::new(&config, &mut input, output, map)?; - + let mut recover_tool = Recover::new()?; recover_tool.run()?; Ok(()) } - -#[cfg(test)] -#[allow(unused)] -mod tests { - use super::*; - - // Test for get_path - // Need to determine how to package files to test with, or at least - // how to test with PathBuf present. - // Test must also check unwrapping of file name, not just generation. - - // Test for get_stream_length - // Need to determine how to test with Seek-able objects. -} diff --git a/src/mapping/map.rs b/src/mapping/map.rs index b6f379c..d8cba29 100644 --- a/src/mapping/map.rs +++ b/src/mapping/map.rs @@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize}; #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] pub struct MapFile { - pub sector_size: u16, + pub sector_size: usize, pub domain: Domain, pub map: Vec, } @@ -37,11 +37,11 @@ impl Default for MapFile { } impl MapFile { - pub fn new(sector_size: u16) -> Self { + pub fn new(sector_size: usize) -> Self { MapFile::default().set_sector_size(sector_size).to_owned() } - pub fn set_sector_size(&mut self, sector_size: u16) -> &mut Self { + pub fn set_sector_size(&mut self, sector_size: usize) -> &mut Self { self.sector_size = sector_size; self } diff --git a/src/path.rs b/src/path.rs new file mode 100644 index 0000000..b130d75 --- /dev/null +++ b/src/path.rs @@ -0,0 +1,40 @@ +use std::path::{Path, PathBuf}; +use std::sync::LazyLock; + +use crate::cli::CONFIG; + +use anyhow::{self, Context}; + +/// Generates a file path if one not provided. +/// root_path for fallback name. +pub fn get_path

(path: &Option

, root_path: &P, extension: &str) -> anyhow::Result +where + P: AsRef, +{ + if let Some(f) = path { + return Ok(f.as_ref().to_path_buf()); + } + + Ok(PathBuf::from(format!( + "{}.{}", + root_path + .as_ref() + .to_str() + .context("source_name path was not UTF-8 valid.")?, + extension + )) + .as_path() + .to_owned()) +} + +pub static MAP_PATH: LazyLock = LazyLock::new(|| { + get_path(&CONFIG.map, &CONFIG.input, "map") + .context("Failed to generate map path.") + .unwrap() +}); + +pub static OUTPUT_PATH: LazyLock = LazyLock::new(|| { + get_path(&CONFIG.output, &CONFIG.input, "iso") + .context("Failed to generate output path.") + .unwrap() +}); diff --git a/src/recovery.rs b/src/recovery.rs index 3e2c050..fadee86 100644 --- a/src/recovery.rs +++ b/src/recovery.rs @@ -1,42 +1,39 @@ -use std::fs::{File, OpenOptions}; -use std::io::{BufWriter, Read, Seek, Write}; +use std::fs::File; +use std::io::{BufWriter, Read, Seek, SeekFrom, Write}; use std::usize; use anyhow::Context; -use crate::cli::Args; +use crate::cli::CONFIG; use crate::mapping::prelude::*; #[derive(Debug)] #[allow(dead_code)] -pub struct Recover<'a> { - config: &'a Args, - stream_len: u64, - - input: &'a mut File, +pub struct Recover { + input: File, output: BufWriter, map: MapFile, } -impl<'a> Recover<'a> { - pub fn new( - config: &'a Args, - input: &'a mut File, - output: File, - map: MapFile, - ) -> anyhow::Result { - let stream_len = - crate::io::get_stream_length(input).context("Failed to get input stream length.")?; +impl Recover { + pub fn new() -> anyhow::Result { + let input: File = crate::io::load_input()?; + + let output: File = crate::io::load_output()?; + + let map: MapFile = { + crate::io::load_map_read()? + .try_into() + .unwrap_or(MapFile::new(CONFIG.sector_size)) + }; let mut r = Recover { - config, - stream_len, input, - output: BufWriter::with_capacity(stream_len as usize, output), + output: BufWriter::with_capacity(map.domain.end as usize, output), map, }; - r.restore(); + r.restore()?; Ok(r) } @@ -70,8 +67,6 @@ impl<'a> Recover<'a> { self.output .rewind() .context("Failed to reset output seek position")?; - - dbg!(); } // Temporary. @@ -81,8 +76,13 @@ impl<'a> Recover<'a> { /// Restore current progress based on MapFile. /// Also updates MapFile if needed, such as to extend the MapFile domain. - pub fn restore(&mut self) { - self.map.extend(self.stream_len as usize); + pub fn restore(&mut self) -> anyhow::Result<()> { + self.map.extend( + crate::io::get_stream_length(&mut self.input) + .context("Failed to get input stream length.")? as usize, + ); + + Ok(()) } /// Attempt to copy all untested blocks. @@ -92,11 +92,13 @@ impl<'a> Recover<'a> { let mut read_position: usize; let mut cluster: Cluster; let mut buf_capacity = self.get_buf_capacity() as usize; - let mut buf = vec![crate::FB_PAD_VALUE; buf_capacity]; + let mut buf = vec![crate::FB_NULL_VALUE; buf_capacity]; read_position = untested.domain.start; while read_position < untested.domain.end { + dbg!(read_position); + cluster = Cluster { domain: Domain { start: read_position, @@ -116,8 +118,13 @@ impl<'a> Recover<'a> { // That padding should have a cli arg to control it. println!("Hit error: {:?}", err); + if CONFIG.reopen_on_error { + self.reload_input() + .context("Failed to reload input file after previous error.")?; + } + self.input - .seek_relative(buf_capacity as i64) + .seek_relative((read_position + buf_capacity) as i64) .context("Failed to seek input by buf_capacity to skip previous error")?; // I don't remember what level was for. @@ -135,21 +142,7 @@ impl<'a> Recover<'a> { } } - self.map - .write_to({ - let map_path = crate::io::get_path(&self.config.map, &self.config.input, "map") - .context("Failed to generate map path.")?; - - &mut OpenOptions::new() - .write(true) - .create(true) - .truncate(true) // Wipe old map. Should really make a backup first. - .open(&map_path) - .with_context(|| { - format!("Failed to open map file at: {}", map_path.display()) - })? - }) - .context("Failed to write map file")?; + self.map.write_to(&mut crate::io::load_map_write()?)?; Ok(()) } @@ -157,7 +150,16 @@ impl<'a> Recover<'a> { /// Set buffer capacity as cluster length in bytes. /// Varies depending on the recovery stage. fn get_buf_capacity(&mut self) -> u64 { - self.config.sector_size as u64 * self.config.cluster_length as u64 + CONFIG.sector_size as u64 * CONFIG.cluster_length as u64 + } + + /// Reloads the input and restores the seek position. + fn reload_input(&mut self) -> anyhow::Result<()> { + let seek_pos = self.input.stream_position()?; + + self.input = crate::io::load_input()?; + self.input.seek(SeekFrom::Start(seek_pos))?; + Ok(()) } }