diff --git a/bmap-rs/Cargo.toml b/bmap-rs/Cargo.toml index 8a07985..cb2613e 100644 --- a/bmap-rs/Cargo.toml +++ b/bmap-rs/Cargo.toml @@ -12,4 +12,6 @@ anyhow = "1.0.66" nix = "0.26.1" flate2 = "1.0.24" clap = { version = "4.0.18", features = ["derive"] } -indicatif = "0.17.1" \ No newline at end of file +indicatif = "0.17.1" +sha2 = { version = "0.10.6", features = [ "asm" ] } +hex = "0.4.3" diff --git a/bmap-rs/src/main.rs b/bmap-rs/src/main.rs index 0f80185..b268681 100644 --- a/bmap-rs/src/main.rs +++ b/bmap-rs/src/main.rs @@ -1,9 +1,10 @@ -use anyhow::{anyhow, bail, Context, Result}; +use anyhow::{anyhow, bail, ensure, Context, Result}; use bmap::{Bmap, Discarder, SeekForward}; use clap::Parser; use flate2::read::GzDecoder; use indicatif::{ProgressBar, ProgressState, ProgressStyle}; use nix::unistd::ftruncate; +use sha2::{Digest, Sha256}; use std::ffi::OsStr; use std::fmt::Write; use std::fs::File; @@ -89,6 +90,24 @@ fn setup_input(path: &Path) -> Result { } } +fn bmap_integrity(checksum: String, xml: String) -> Result<()> { + //Unset the checksum + let mut bmap_hash = Sha256::new(); + let default = "0".repeat(64); + let before_checksum = xml.replace(&checksum, &default); + + //Compare given and created checksum + bmap_hash.update(before_checksum); + let digest = bmap_hash.finalize_reset(); + let new_checksum = hex::encode(digest.as_slice()); + ensure!( + checksum == new_checksum, + "Bmap file doesn't match its checksum. It could be corrupted or compromised." + ); + println!("Bmap integrity checked!"); + Ok(()) +} + fn copy(c: Copy) -> Result<()> { if !c.image.exists() { bail!("Image file doesn't exist") @@ -102,6 +121,7 @@ fn copy(c: Copy) -> Result<()> { b.read_to_string(&mut xml)?; let bmap = Bmap::from_xml(&xml)?; + bmap_integrity(bmap.bmap_file_checksum(), xml)?; let output = std::fs::OpenOptions::new() .write(true) .create(true)