irradiate/src/main.rs

// Irradiate v0.2
// Evelene Ultraviolet
// Ellie Photodiode
// Ariane Mechanism
// Amelia Rainfall
// July 2019

use std::io::prelude::*;
use std::io::stdout;
use std::io::SeekFrom;
use std::fs::OpenOptions;
use std::time::{Duration,Instant};
use std::process::exit;

extern crate rand;
use rand::Rng;

extern crate lapp;

// sets a bit in-place. depending on the bit's initial state, may or may not produce an error.
fn set_bit(byte:&mut u8, index:usize, state:bool) {
	let norm_index = (index % 8) as u8;
	match state {
		true  => *byte |= (0x01 << norm_index) ^ 0x00,
		false => *byte &= (0x01 << norm_index) ^ 0xff,
	};
}

// flips a bit in-place, ensuring exactly one error will be produced.
fn flip_bit(byte:&mut u8, index:usize) {
	let norm_index = (index % 8) as u8;
	*byte ^= 0x01 << norm_index;
}

fn main() {
	let args = lapp::parse_args("
		Irradiate file corruptor v0.2
		<filename> (string) name of file on which to operate
		-n, --count (default 10) integer number of bits to affect in the file
		-0, --zero never set affected bits to 1, only set them to 0.
		-1, --one never set affected bits to 0, only set them to 1.
		-f, --flip invert affected bits, ensuring that the file contains an exact number of errors
		-h, --holdoff (default 0) number of bytes at the start of the file to leave unaffected, so as not to corrupt the header and render the file unopenable
	");
	let filename = args.get_string("filename");
	let reps = args.get_integer("count") as u64;
	let zero_only = args.get_bool("zero");
	let one_only = args.get_bool("one");
	let flip = args.get_bool("flip");
	let holdoff = args.get_integer("holdoff") as u64;

	// make sure only one of the mode flags was given; gripe otherwise
	if (zero_only as u8) + (one_only as u8) + (flip as u8) > 1 {
		eprintln!("irradiate: the --flip, --zero, and --one flags (-f, -0, and -1) are mutually exclusive and cannot be combined");
		exit(1);
	}

	// set up the progress meter by determining the numbers of affected bits that correspond to each percentage value.
	let mut milestones:[u64;99] = [0;99];
	let mut current_milestone:usize = 0;
	for i in 1..100 {
		milestones[i-1] = reps*(i as u64)/100;
	}
	
	let mut rng = rand::thread_rng();
	let mut buffer:[u8;1] = [0x00;1];
	let mut sout = stdout();

	// we can't use lapp's nifty builtin file-fetcher because we need to open the file in read-write mode, 
	// allowing us to modify it in-place.
	let mut infile = match OpenOptions::new()
		.read(true)
		.write(true)
		.append(false)
		.truncate(false)
		.create(false)
		.create_new(false)
		.open(&filename) 
	{
		Ok(file) => file,
		Err(why) => {
			eprintln!("irradiate: could not open {} - {}", filename, why);
			exit(1);
		},
	};

	// seek to the end of the file and figure out where that is - a better trick for determining length that,
	// unlike fs metadata, actually works on block devices.
	let infile_len:u64 = match infile.seek(SeekFrom::End(0)) {
		Ok(n) => match n {
			0 => {
				eprintln!("irradiate: this file appears to be empty");
				exit(1);
			},
			_ => n,
		},
		Err(why) => {
			eprintln!("irradiate: could not determine the size of {} - {}", filename, why);
			exit(1);
		},
	};

	let mut total_seektime:Duration = Duration::new(0,0);
	let mut total_readtime:Duration = Duration::new(0,0);
	let mut total_writtime:Duration = Duration::new(0,0);

	let mut seekstart:Instant;
	let mut readstart:Instant;
	let mut writstart:Instant;

	print!("\rirradiating [  0%]");
	for rep in 0..reps {
		// go to a randomly-selected location, between the holdoff position and the file end.
		let position = rng.gen_range(holdoff, infile_len);
		seekstart = Instant::now();
		match infile.seek(SeekFrom::Start(position)) {
			Ok(_) => (),
			Err(why) => {
				eprintln!("irradiate: seek error - {}", why);
				exit(1);
			},
		};
		total_seektime += seekstart.elapsed();
		// read data into the buffer 
		readstart = Instant::now();
		match infile.read(&mut buffer) {
			Ok(n) => match n {
				0 => break,
				_ => (),
			},
			Err(why) => {
				eprintln!("irradiate: read error - {}", why);
				exit(1);
			},
		};
		total_readtime += readstart.elapsed();
		// go back to the same position from which we read the data
		seekstart = Instant::now();
		match infile.seek(SeekFrom::Start(position)) {
			Ok(_) => (),
			Err(why) => {
				eprintln!("irradiate: seek error - {}", why);
				exit(1);
			},
		};
		total_seektime += seekstart.elapsed();
		// set a random bit of this byte
		let bit = rng.gen_range(0,8);
		if one_only {
			set_bit(&mut buffer[0], bit, true);
		} else if zero_only {
			set_bit(&mut buffer[0], bit, false);
		} else if flip {
			flip_bit(&mut buffer[0], bit);
		} else {
			set_bit(&mut buffer[0], bit, rng.gen::<bool>());
		}
		// write the modified data back to the file
		writstart = Instant::now();
		match infile.write_all(&buffer) {
			Ok(_) => (),
			Err(why) => {
				eprintln!("irradiate: write error - {}", why);
				exit(1);
			},
		};
		let _ = infile.flush();
		total_writtime += writstart.elapsed();
		// update the progress meter
		while milestones[current_milestone] <= rep && current_milestone < 98 {
			print!("\rirradiating [{:3}%]", current_milestone);
			current_milestone += 1;
		}
		// flush stdout, to make sure the progress meter actually updates
		let _ = sout.flush();
	}
	print!("\rirradiating [100%]");
	println!();
	println!("irradiating done.");
	println!("time seeking: {} s", (total_seektime.as_micros() as f32)/1000000.0);
	println!("time reading: {} s", (total_readtime.as_micros() as f32)/1000000.0);
	println!("time writing: {} s", (total_writtime.as_micros() as f32)/1000000.0);
}