Cutieguwu/kramer
1
0
Fork 0
Code Issues Pull Requests 4 Actions Packages Projects Releases Wiki Activity

Huge refactor. Introduce anyhow for error handling.

Browse Source
This commit is contained in:
Cutieguwu
2025-12-29 15:31:01 -05:00
parent e98383d9e5
commit 4b5460f754
13 changed files with 411 additions and 746 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/cli.rs Normal file
View File

@@ -0,0 +1,32 @@
use std::path::PathBuf;
use crate::FB_SECTOR_SIZE;
use clap::Parser;
#[derive(Parser, Debug)]
pub struct Args {
/// Path to source file or block device
#[arg(short, long, value_hint = clap::ValueHint::DirPath)]
pub input: PathBuf,
/// Path to output file. Defaults to {input}.iso
#[arg(short, long, value_hint = clap::ValueHint::DirPath)]
pub output: Option<PathBuf>,
/// Path to rescue map. Defaults to {input}.map
#[arg(short, long, value_hint = clap::ValueHint::DirPath)]
pub map: Option<PathBuf>,
/// Max number of consecutive sectors to test as a group
#[arg(short, long, default_value_t = 128)]
pub cluster_length: u16,
/// Number of brute force read passes
#[arg(short, long, default_value_t = 2)]
pub brute_passes: usize,
/// Sector size
#[arg(short, long, default_value_t = FB_SECTOR_SIZE)]
pub sector_size: u16,
}

21
src/io.rs
View File

@@ -1,4 +1,5 @@
use std::io::{self, Seek, SeekFrom};
use std::path::{Path, PathBuf};
/// Get length of data stream.
/// Physical length of data stream in bytes
@@ -13,3 +14,23 @@ 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) -> Option<PathBuf>
where
P: AsRef<Path>,
{
if let Some(f) = file_path {
return f.as_ref().to_path_buf().into();
}
PathBuf::from(format!(
"{:?}.{}",
source_name.as_ref().to_str()?,
extension
))
.as_path()
.to_owned()
.into()
}

121
src/main.rs
View File

@@ -1,87 +1,49 @@
mod cli;
mod io;
mod mapping;
mod recovery;
use std::fs::{File, OpenOptions};
use std::path::{Path, PathBuf};
use clap::Parser;
use mapping::MapFile;
use cli::Args;
use mapping::prelude::*;
use recovery::Recover;
use anyhow::{self, Context};
use clap::Parser;
const FB_SECTOR_SIZE: u16 = 2048;
const FB_PAD_VALUE: u8 = 0;
#[derive(Parser, Debug)]
struct Args {
/// Path to source file or block device
#[arg(short, long, value_hint = clap::ValueHint::DirPath)]
input: PathBuf,
/// Path to output file. Defaults to {input}.iso
#[arg(short, long, value_hint = clap::ValueHint::DirPath)]
output: Option<PathBuf>,
/// Path to rescue map. Defaults to {input}.map
#[arg(short, long, value_hint = clap::ValueHint::DirPath)]
map: Option<PathBuf>,
/// Max number of consecutive sectors to test as a group
#[arg(short, long, default_value_t = 128)]
cluster_length: u16,
/// Number of brute force read passes
#[arg(short, long, default_value_t = 2)]
brute_passes: usize,
/// Sector size
#[arg(short, long, default_value_t = FB_SECTOR_SIZE)]
sector_size: u16,
}
fn main() {
fn main() -> anyhow::Result<()> {
let config = Args::parse();
// Live with it, prefer to use expect() here.
// I'm lazy and don't want to mess around with comparing error types.
// Thus, any error in I/O here should be treated as fatal.
let mut input: File = OpenOptions::new()
.read(true)
.open(&config.input.as_path())
.expect("Failed to open input file");
let mut input: File = {
let input_path = &config.input.as_path();
OpenOptions::new()
.read(true)
.open(&input_path)
.with_context(move || format!("Failed to open input file: {}", input_path.display()))?
};
let output: File = {
// Keep this clean, make a short-lived binding.
let path = get_path(&config.output, &config.input.to_str().unwrap(), "iso");
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)
.expect("Failed to open/create output file")
.open(&path)
.with_context(|| format!("Failed to open/create output file at: {}", path.display()))?
};
// Check if output file is shorter than input.
// If so, autoextend the output file.
//
// Is this actually needed? I don't think I need to preallocate the space.
/*
{
let input_len =
get_stream_length(&mut input).expect("Failed to get the length of the input data.");
let output_len =
get_stream_length(&mut output).expect("Failed to get the length of the output file.");
if output_len < input_len {
output
.set_len(input_len)
.expect("Failed to autofill output file.")
}
}
*/
/*
let map: MapFile = {
let path = get_path(
@@ -107,49 +69,16 @@ fn main() {
};
*/
let buf_capacity =
crate::io::get_stream_length(&mut input).expect("Failed to get buffer capacity from input");
let mut recover_tool = {
let buf_capacity = crate::io::get_stream_length(&mut input)
.context("Failed to get buffer capacity from input")?;
let mut recover_tool = Recover::new(config, input, output, MapFile::default(), buf_capacity);
Recover::new(&config, input, output, MapFile::default(), buf_capacity)
};
recover_tool.run();
//todo!("Recovery, Map saving, and closure of all files.");
/*
let mut buf: Vec<u8> = Vec::with_capacity(
get_stream_length(&mut input).expect("Failed to get the length of the input data.")
as usize,
);
println!(
"Read {} bytes",
input
.read_to_end(&mut buf)
.expect("Failed to read complete input stream.")
);
println!("Wrote {} bytes", {
output
.write_all(&buf)
.expect("Failed to write complete output stream.");
&buf.len()
});
*/
}
/// Generates a file path if one not provided.
/// source_name for fallback name.
fn get_path<P>(file_path: &Option<P>, source_name: &str, extension: &str) -> PathBuf
where
P: AsRef<Path>,
{
if let Some(f) = file_path {
f.as_ref().to_path_buf()
} else {
PathBuf::from(format!("{:?}.{}", source_name, extension))
.as_path()
.to_owned()
}
Ok(())
}
#[cfg(test)]

7
src/mapping/cluster.rs
View File

@@ -54,3 +54,10 @@ impl Cluster {
self
}
}
#[cfg(test)]
mod tests {
use super::*;
// Test for Cluster::subdivide()
}

5
src/mapping/domain.rs
View File

@@ -20,3 +20,8 @@ impl Domain {
self.end - self.start
}
}
#[cfg(test)]
mod tests {
use super::*;
}

183
src/mapping/map.rs
View File

@@ -1,7 +1,9 @@
use std::fs::File;
use std::io::Write;
use super::{Cluster, Domain, Stage};
use anyhow;
use ron::de::from_reader;
use ron::error::SpannedError;
use serde::{Deserialize, Serialize};
@@ -197,11 +199,178 @@ impl MapFile {
}
}
pub fn write_map_to_file(file: File, map: &MapFile) -> ron::error::Result<String> {
ron::ser::to_string_pretty(
map,
ron::ser::PrettyConfig::new()
.new_line("\n".to_string())
.struct_names(true),
)
pub fn write_map_to_file(file: &mut File, map: &MapFile) -> anyhow::Result<usize> {
let written_bytes = file.write(
ron::ser::to_string_pretty(
map,
ron::ser::PrettyConfig::new()
.new_line("\n".to_string())
.struct_names(true),
)?
.as_bytes(),
)?;
Ok(written_bytes)
}
#[cfg(test)]
mod tests {
use super::*;
// Test for MapFile::update()
// Test for MapFile::get_stage()
#[test]
fn test_get_stage() {
use std::vec;
let mut mf = MapFile::default();
let mut mf_stage = mf.get_stage();
// If this fails here, there's something SERIOUSLY wrong.
assert!(
mf_stage == Stage::Untested,
"Determined stage to be {:?}, when {:?} was expeccted.",
mf_stage,
Stage::Untested
);
let stages = vec![
Stage::Damaged,
Stage::ForIsolation(1),
Stage::ForIsolation(0),
Stage::Untested,
];
mf.map = vec![];
for stage in stages {
mf.map.push(*Cluster::default().set_stage(stage));
mf_stage = mf.get_stage();
assert!(
stage == mf_stage,
"Expected stage to be {:?}, determined {:?} instead.",
stage,
mf_stage
)
}
}
// Test for MapFile::get_clusters()
#[test]
fn test_get_clusters() {
let mut mf = MapFile::default();
mf.map = vec![
*Cluster::default().set_stage(Stage::Damaged),
*Cluster::default().set_stage(Stage::ForIsolation(0)),
*Cluster::default().set_stage(Stage::ForIsolation(1)),
Cluster::default(),
Cluster::default(),
*Cluster::default().set_stage(Stage::ForIsolation(1)),
*Cluster::default().set_stage(Stage::ForIsolation(0)),
*Cluster::default().set_stage(Stage::Damaged),
];
let stages = vec![
Stage::Damaged,
Stage::ForIsolation(1),
Stage::ForIsolation(0),
Stage::Untested,
];
for stage in stages {
let expected = vec![
*Cluster::default().set_stage(stage),
*Cluster::default().set_stage(stage),
];
let received = mf.get_clusters(stage);
assert!(
expected == received,
"Expected clusters {:?}, got {:?}.",
expected,
received
)
}
}
// Test for MapFile::defrag()
#[test]
fn test_defrag() {
let mut mf = MapFile {
sector_size: 1,
domain: Domain { start: 0, end: 8 },
map: vec![
Cluster {
domain: Domain { start: 0, end: 1 },
stage: Stage::Untested,
},
Cluster {
domain: Domain { start: 1, end: 2 },
stage: Stage::Untested,
},
Cluster {
domain: Domain { start: 2, end: 3 },
stage: Stage::Untested,
},
Cluster {
domain: Domain { start: 3, end: 4 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 4, end: 5 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 5, end: 6 },
stage: Stage::ForIsolation(1),
},
Cluster {
domain: Domain { start: 6, end: 7 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 7, end: 8 },
stage: Stage::Damaged,
},
],
};
let expected = vec![
Cluster {
domain: Domain { start: 0, end: 3 },
stage: Stage::Untested,
},
Cluster {
domain: Domain { start: 3, end: 5 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 5, end: 6 },
stage: Stage::ForIsolation(1),
},
Cluster {
domain: Domain { start: 6, end: 7 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 7, end: 8 },
stage: Stage::Damaged,
},
];
mf.defrag();
let received = mf.map;
assert!(
expected == received,
"Expected {:?} after defragging, got {:?}.",
expected,
received
)
}
}

185
src/mapping/mod.rs
View File

@@ -1,187 +1,12 @@
#![allow(unused_imports)]
pub mod cluster;
pub mod domain;
pub mod map;
use serde::{Deserialize, Serialize};
pub mod prelude;
pub mod stage;
pub use cluster::Cluster;
pub use domain::Domain;
pub use map::MapFile;
#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, PartialOrd)]
pub enum Stage {
Intact,
Untested,
ForIsolation(u8),
Damaged,
}
impl Default for Stage {
fn default() -> Self {
Stage::Untested
}
}
#[cfg(test)]
mod tests {
use super::*;
// Test for Cluster::subdivide()
// Test for MapFile::update()
// Test for MapFile::get_stage()
#[test]
fn test_get_stage() {
use std::vec;
let mut mf = MapFile::default();
let mut mf_stage = mf.get_stage();
// If this fails here, there's something SERIOUSLY wrong.
assert!(
mf_stage == Stage::Untested,
"Determined stage to be {:?}, when {:?} was expeccted.",
mf_stage,
Stage::Untested
);
let stages = vec![
Stage::Damaged,
Stage::ForIsolation(1),
Stage::ForIsolation(0),
Stage::Untested,
];
mf.map = vec![];
for stage in stages {
mf.map.push(*Cluster::default().set_stage(stage));
mf_stage = mf.get_stage();
assert!(
stage == mf_stage,
"Expected stage to be {:?}, determined {:?} instead.",
stage,
mf_stage
)
}
}
// Test for MapFile::get_clusters()
#[test]
fn test_get_clusters() {
let mut mf = MapFile::default();
mf.map = vec![
*Cluster::default().set_stage(Stage::Damaged),
*Cluster::default().set_stage(Stage::ForIsolation(0)),
*Cluster::default().set_stage(Stage::ForIsolation(1)),
Cluster::default(),
Cluster::default(),
*Cluster::default().set_stage(Stage::ForIsolation(1)),
*Cluster::default().set_stage(Stage::ForIsolation(0)),
*Cluster::default().set_stage(Stage::Damaged),
];
let stages = vec![
Stage::Damaged,
Stage::ForIsolation(1),
Stage::ForIsolation(0),
Stage::Untested,
];
for stage in stages {
let expected = vec![
*Cluster::default().set_stage(stage),
*Cluster::default().set_stage(stage),
];
let received = mf.get_clusters(stage);
assert!(
expected == received,
"Expected clusters {:?}, got {:?}.",
expected,
received
)
}
}
// Test for MapFile::defrag()
#[test]
fn test_defrag() {
let mut mf = MapFile {
sector_size: 1,
domain: Domain { start: 0, end: 8 },
map: vec![
Cluster {
domain: Domain { start: 0, end: 1 },
stage: Stage::Untested,
},
Cluster {
domain: Domain { start: 1, end: 2 },
stage: Stage::Untested,
},
Cluster {
domain: Domain { start: 2, end: 3 },
stage: Stage::Untested,
},
Cluster {
domain: Domain { start: 3, end: 4 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 4, end: 5 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 5, end: 6 },
stage: Stage::ForIsolation(1),
},
Cluster {
domain: Domain { start: 6, end: 7 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 7, end: 8 },
stage: Stage::Damaged,
},
],
};
let expected = vec![
Cluster {
domain: Domain { start: 0, end: 3 },
stage: Stage::Untested,
},
Cluster {
domain: Domain { start: 3, end: 5 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 5, end: 6 },
stage: Stage::ForIsolation(1),
},
Cluster {
domain: Domain { start: 6, end: 7 },
stage: Stage::ForIsolation(0),
},
Cluster {
domain: Domain { start: 7, end: 8 },
stage: Stage::Damaged,
},
];
mf.defrag();
let received = mf.map;
assert!(
expected == received,
"Expected {:?} after defragging, got {:?}.",
expected,
received
)
}
}
pub use stage::Stage;

6
src/mapping/prelude.rs Normal file
View File

@@ -0,0 +1,6 @@
#![allow(unused_imports)]
pub use super::cluster::Cluster;
pub use super::domain::Domain;
pub use super::map::{write_map_to_file, MapFile};
pub use super::stage::Stage;

20
src/mapping/stage.rs Normal file
View File

@@ -0,0 +1,20 @@
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, PartialOrd)]
pub enum Stage {
Intact,
Untested,
ForIsolation(u8),
Damaged,
}
impl Default for Stage {
fn default() -> Self {
Stage::Untested
}
}
#[cfg(test)]
mod tests {
use super::*;
}

78
src/recovery.rs
View File

@@ -1,16 +1,17 @@
use std::fs::{File, OpenOptions};
use std::io::{BufWriter, Read, Seek, SeekFrom, Write};
use std::ptr::read;
use std::io::{BufWriter, Read, Seek, Write};
use crate::mapping::{Cluster, Domain, MapFile, Stage};
use crate::Args;
use anyhow::Context;
use crate::cli::Args;
use crate::mapping::prelude::*;
#[derive(Debug)]
#[allow(dead_code)]
pub struct Recover {
pub struct Recover<'a> {
/// Buffer capacity in bytes.
buf_capacity: u64,
config: Args,
config: &'a Args,
input: File,
output: BufWriter<File>,
@@ -18,8 +19,14 @@ pub struct Recover {
stage: Stage,
}
impl Recover {
pub fn new(config: Args, input: File, output: File, map: MapFile, buf_capacity: u64) -> Self {
impl<'a> Recover<'a> {
pub fn new(
config: &'a Args,
input: File,
output: File,
map: MapFile,
buf_capacity: u64,
) -> Self {
let stage = map.get_stage();
let mut r = Recover {
@@ -64,27 +71,25 @@ impl Recover {
}
/// Attempt to copy all untested blocks.
fn copy_untested(&mut self) -> &mut Self {
fn copy_untested(&mut self) -> anyhow::Result<()> {
let mut buf = vec![crate::FB_PAD_VALUE; self.buf_capacity as usize];
// Purely caching.
let mut read_position = 0_u64;
let last_read_position = crate::io::get_stream_length(&mut self.input)
.expect("Failed to get length of input stream")
.context("Failed to get length of input stream")?
- self.buf_capacity;
while read_position < last_read_position {
dbg!(read_position);
if let Err(err) = self.input.read_exact(&mut buf) {
println!("Hit error: {:?}", err);
self.input
.seek_relative(self.buf_capacity as i64)
.expect("Failed to seek input by buf_capacity to skip previous error");
.context("Failed to seek input by buf_capacity to skip previous error")?;
} else {
self.output
.write_all(buf.as_slice())
.expect("Failed to write data to output file");
.context("Failed to write data to output file")?;
self.map.update(Cluster {
domain: Domain {
@@ -99,41 +104,26 @@ impl Recover {
}
crate::mapping::map::write_map_to_file(
OpenOptions::new()
.create(true)
.write(true)
.open(crate::get_path(
&self.config.map,
self.config.input.to_str().unwrap(),
"map",
))
.expect("Failed to open map file"),
{
let map_path = crate::io::get_path(&self.config.map, &self.config.input, "map")
.context("Failed to generate map path.")?;
&mut OpenOptions::new()
.create(true)
.write(true)
.open(&map_path)
.with_context(|| {
format!("Failed to open map file at: {}", map_path.display())
})?
},
&self.map,
)
.expect("Failed to write map file");
.context("Failed to write map file")?;
/*
let mut untested: Vec<Cluster> = vec![];
for cluster in self.map.get_clusters(Stage::Untested).iter_mut() {
untested.append(&mut cluster.subdivide(self.map.sector_size as usize));
}
todo!("Read and save data.");
*/
self
Ok(())
}
/// Attempt to copy blocks via isolation at pass level.
fn copy_isolate(&mut self, level: u8) -> &mut Self {
todo!();
self
}
/// Set buffer capacities as cluster length in bytes.
/// Set buffer capacity as cluster length in bytes.
/// Varies depending on the recovery stage.
fn set_buf_capacity(&mut self) -> &mut Self {
self.buf_capacity = self.config.sector_size as u64 * self.config.cluster_length as u64;