use sha2::{Digest, Sha256};
use crate::fs::VfsSnapshot;
use crate::interpreter::{ShellState, ShellStateOptions};
const SNAPSHOT_VERSION: u32 = 1;
const INTEGRITY_TAG: &[u8; 8] = b"BKSNAP01";
const DIGEST_LEN: usize = 32;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Snapshot {
pub version: u32,
pub shell: ShellState,
pub vfs: Option<VfsSnapshot>,
pub session_commands: u64,
pub session_exec_calls: u64,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct SnapshotOptions {
pub exclude_filesystem: bool,
pub exclude_functions: bool,
}
impl Snapshot {
pub fn to_bytes(&self) -> crate::Result<Vec<u8>> {
let json = serde_json::to_vec(self).map_err(|e| crate::Error::Internal(e.to_string()))?;
let digest = Self::compute_digest(&json);
let mut out = Vec::with_capacity(DIGEST_LEN + json.len());
out.extend_from_slice(&digest);
out.extend_from_slice(&json);
Ok(out)
}
pub fn from_bytes(data: &[u8]) -> crate::Result<Self> {
if data.len() < DIGEST_LEN {
return Err(crate::Error::Internal(
"snapshot too short: missing integrity digest".to_string(),
));
}
let (stored_digest, json) = data.split_at(DIGEST_LEN);
let expected = Self::compute_digest(json);
if stored_digest != expected.as_slice() {
return Err(crate::Error::Internal(
"snapshot integrity check failed: data may have been tampered with".to_string(),
));
}
let snap: Self =
serde_json::from_slice(json).map_err(|e| crate::Error::Internal(e.to_string()))?;
if snap.version != SNAPSHOT_VERSION {
return Err(crate::Error::Internal(format!(
"unsupported snapshot version {} (expected {})",
snap.version, SNAPSHOT_VERSION
)));
}
Ok(snap)
}
pub fn to_bytes_keyed(&self, key: &[u8]) -> crate::Result<Vec<u8>> {
let json = serde_json::to_vec(self).map_err(|e| crate::Error::Internal(e.to_string()))?;
let digest = Self::compute_hmac(key, &json);
let mut out = Vec::with_capacity(DIGEST_LEN + json.len());
out.extend_from_slice(&digest);
out.extend_from_slice(&json);
Ok(out)
}
pub fn from_bytes_keyed(data: &[u8], key: &[u8]) -> crate::Result<Self> {
use hmac::{Hmac, KeyInit, Mac};
type HmacSha256 = Hmac<Sha256>;
if data.len() < DIGEST_LEN {
return Err(crate::Error::Internal(
"snapshot too short: missing integrity digest".to_string(),
));
}
let (stored_digest, json) = data.split_at(DIGEST_LEN);
let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length");
mac.update(json);
if mac.verify_slice(stored_digest).is_err() {
return Err(crate::Error::Internal(
"snapshot integrity check failed: HMAC mismatch (wrong key or tampered data)"
.to_string(),
));
}
let snap: Self =
serde_json::from_slice(json).map_err(|e| crate::Error::Internal(e.to_string()))?;
if snap.version != SNAPSHOT_VERSION {
return Err(crate::Error::Internal(format!(
"unsupported snapshot version {} (expected {})",
snap.version, SNAPSHOT_VERSION
)));
}
Ok(snap)
}
fn compute_digest(payload: &[u8]) -> [u8; DIGEST_LEN] {
let mut hasher = Sha256::new();
hasher.update(INTEGRITY_TAG);
hasher.update(payload);
let result = hasher.finalize();
let mut out = [0u8; DIGEST_LEN];
out.copy_from_slice(&result);
out
}
fn compute_hmac(key: &[u8], payload: &[u8]) -> [u8; DIGEST_LEN] {
use hmac::{Hmac, KeyInit, Mac};
type HmacSha256 = Hmac<Sha256>;
let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length");
mac.update(payload);
let result = mac.finalize();
let mut out = [0u8; DIGEST_LEN];
out.copy_from_slice(&result.into_bytes());
out
}
}
impl crate::Bash {
fn build_snapshot(&self, options: SnapshotOptions) -> Snapshot {
let shell = self
.interpreter
.shell_state_with_options(ShellStateOptions {
include_functions: !options.exclude_functions,
});
let vfs = if options.exclude_filesystem {
None
} else {
self.fs.vfs_snapshot()
};
let counters = self.interpreter.counters();
Snapshot {
version: SNAPSHOT_VERSION,
shell,
vfs,
session_commands: counters.session_commands,
session_exec_calls: counters.session_exec_calls,
}
}
pub fn snapshot(&self) -> crate::Result<Vec<u8>> {
self.snapshot_with_options(SnapshotOptions::default())
}
pub fn snapshot_with_options(&self, options: SnapshotOptions) -> crate::Result<Vec<u8>> {
self.build_snapshot(options).to_bytes()
}
pub fn from_snapshot(data: &[u8]) -> crate::Result<Self> {
let snap = Snapshot::from_bytes(data)?;
let mut bash = Self::new();
bash.restore_snapshot_inner(&snap)?;
Ok(bash)
}
pub fn restore_snapshot(&mut self, data: &[u8]) -> crate::Result<()> {
let snap = Snapshot::from_bytes(data)?;
self.restore_snapshot_inner(&snap)
}
fn restore_snapshot_inner(&mut self, snap: &Snapshot) -> crate::Result<()> {
self.interpreter
.validate_shell_state_restore_limits(&snap.shell)?;
self.interpreter.restore_shell_state(&snap.shell);
if let Some(ref vfs) = snap.vfs {
self.fs.vfs_restore(vfs);
}
Ok(())
}
pub fn snapshot_to_bytes_keyed(&self, key: &[u8]) -> crate::Result<Vec<u8>> {
self.snapshot_to_bytes_keyed_with_options(key, SnapshotOptions::default())
}
pub fn snapshot_to_bytes_keyed_with_options(
&self,
key: &[u8],
options: SnapshotOptions,
) -> crate::Result<Vec<u8>> {
self.build_snapshot(options).to_bytes_keyed(key)
}
pub fn from_snapshot_keyed(data: &[u8], key: &[u8]) -> crate::Result<Self> {
let snap = Snapshot::from_bytes_keyed(data, key)?;
let mut bash = Self::new();
bash.restore_snapshot_inner(&snap)?;
Ok(bash)
}
pub fn restore_snapshot_keyed(&mut self, data: &[u8], key: &[u8]) -> crate::Result<()> {
let snap = Snapshot::from_bytes_keyed(data, key)?;
self.restore_snapshot_inner(&snap)
}
}