rs_peekaboo 0.2.4

Rust-native cross-platform computer-use CLI and library
Documentation
use crate::{PeekabooError, Result, Snapshot};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};

static SNAPSHOT_ID_COUNTER: AtomicU32 = AtomicU32::new(0);

pub fn cache_dir() -> Result<PathBuf> {
    let base = dirs::cache_dir()
        .or_else(|| dirs::home_dir().map(|home| home.join(".cache")))
        .ok_or(PeekabooError::MissingArgument("cache dir"))?;
    Ok(base.join("rs_peekaboo"))
}

pub fn snapshot_dir() -> Result<PathBuf> {
    Ok(cache_dir()?.join("snapshots"))
}

pub fn validate_snapshot_id(id: &str) -> Result<()> {
    if id.is_empty()
        || !id
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
    {
        return Err(PeekabooError::InvalidSnapshotId(id.to_string()));
    }
    Ok(())
}

fn snapshot_path(id: &str) -> Result<PathBuf> {
    validate_snapshot_id(id)?;
    let dir = snapshot_dir()?;
    let path = dir.join(format!("{id}.json"));
    ensure_within_dir(&dir, &path)?;
    Ok(path)
}

fn ensure_within_dir(dir: &Path, path: &Path) -> Result<()> {
    fs::create_dir_all(dir)?;
    let dir = fs::canonicalize(dir)?;
    let parent = path.parent().ok_or_else(|| {
        PeekabooError::InvalidSnapshotId(path.to_string_lossy().into_owned())
    })?;
    if parent != dir {
        return Err(PeekabooError::InvalidSnapshotId(
            path.to_string_lossy().into_owned(),
        ));
    }
    Ok(())
}

pub fn new_snapshot_id() -> String {
    let millis = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| duration.as_millis())
        .unwrap_or_default();
    let n = SNAPSHOT_ID_COUNTER.fetch_add(1, Ordering::Relaxed);
    format!("snap-{millis}-{n}")
}

pub fn save_snapshot(snapshot: &Snapshot) -> Result<()> {
    validate_snapshot_id(&snapshot.snapshot_id)?;
    let path = snapshot_path(&snapshot.snapshot_id)?;
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }
    fs::write(path, serde_json::to_vec_pretty(snapshot)?)?;
    Ok(())
}

pub fn load_snapshot(id: &str) -> Result<Snapshot> {
    let path = snapshot_path(id)?;
    let data = fs::read(path)?;
    Ok(serde_json::from_slice(&data)?)
}

pub fn clean_snapshots(all: bool, snapshot: Option<&str>) -> Result<usize> {
    let dir = snapshot_dir()?;
    if let Some(snapshot) = snapshot {
        let path = snapshot_path(snapshot)?;
        if path.exists() {
            fs::remove_file(path)?;
            return Ok(1);
        }
        return Ok(0);
    }
    if !all || !dir.exists() {
        return Ok(0);
    }
    let mut removed = 0;
    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        if entry.path().extension().is_some_and(|ext| ext == "json") {
            fs::remove_file(entry.path())?;
            removed += 1;
        }
    }
    Ok(removed)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn validate_snapshot_id_should_reject_path_traversal() {
        assert!(validate_snapshot_id("../etc/passwd").is_err());
        assert!(validate_snapshot_id("snap-ok").is_ok());
        assert!(validate_snapshot_id("snap_123").is_ok());
    }

    #[test]
    fn new_snapshot_id_should_be_unique_and_valid() {
        let id1 = new_snapshot_id();
        let id2 = new_snapshot_id();
        assert_ne!(id1, id2);
        assert!(validate_snapshot_id(&id1).is_ok());
        assert!(validate_snapshot_id(&id2).is_ok());
        assert!(id1.starts_with("snap-"));
        assert!(id2.starts_with("snap-"));
    }
}