Tons of cleanup. More useless git notices.
This commit is contained in:
18
src/cli.rs
18
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<Args> = 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,
|
||||
}
|
||||
|
||||
73
src/io.rs
73
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<S: Seek>(stream: &mut S) -> io::Result<u64> {
|
||||
len
|
||||
}
|
||||
|
||||
/// Generates a file path if one not provided.
|
||||
/// source_name for fallback name.
|
||||
pub fn get_path<P>(
|
||||
file_path: &Option<P>,
|
||||
source_name: &P,
|
||||
extension: &str,
|
||||
) -> anyhow::Result<PathBuf>
|
||||
where
|
||||
P: AsRef<Path>,
|
||||
{
|
||||
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<File> {
|
||||
OpenOptions::new()
|
||||
.read(true)
|
||||
.open(&CONFIG.input)
|
||||
.with_context(|| format!("Failed to open input file: {}", &CONFIG.input.display()))
|
||||
}
|
||||
|
||||
pub fn load_output() -> anyhow::Result<File> {
|
||||
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<File> {
|
||||
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<File> {
|
||||
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()
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
70
src/main.rs
70
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.
|
||||
}
|
||||
|
||||
@@ -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<Cluster>,
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
40
src/path.rs
Normal file
40
src/path.rs
Normal file
@@ -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<P>(path: &Option<P>, root_path: &P, extension: &str) -> anyhow::Result<PathBuf>
|
||||
where
|
||||
P: AsRef<Path>,
|
||||
{
|
||||
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<PathBuf> = LazyLock::new(|| {
|
||||
get_path(&CONFIG.map, &CONFIG.input, "map")
|
||||
.context("Failed to generate map path.")
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
pub static OUTPUT_PATH: LazyLock<PathBuf> = LazyLock::new(|| {
|
||||
get_path(&CONFIG.output, &CONFIG.input, "iso")
|
||||
.context("Failed to generate output path.")
|
||||
.unwrap()
|
||||
});
|
||||
@@ -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<File>,
|
||||
map: MapFile,
|
||||
}
|
||||
|
||||
impl<'a> Recover<'a> {
|
||||
pub fn new(
|
||||
config: &'a Args,
|
||||
input: &'a mut File,
|
||||
output: File,
|
||||
map: MapFile,
|
||||
) -> anyhow::Result<Self> {
|
||||
let stream_len =
|
||||
crate::io::get_stream_length(input).context("Failed to get input stream length.")?;
|
||||
impl Recover {
|
||||
pub fn new() -> anyhow::Result<Self> {
|
||||
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(())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user