anti-sec 0.1.2

use std::fs::File;
use std::io::{self, Read};
use std::path::{Path, PathBuf};

use walkdir::WalkDir;

use hmac::{Hmac, Mac};
use md5::Md5;
use sha1::Sha1;
use sha2::Digest;
use sha2::Sha256;

/// Supported hash algorithms
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HashAlgorithm {
    Md5,
    Sha1,
    Sha256,
    HmacMd5,
    HmacSha1,
    HmacSha256,
}

/// Compute the hash of an arbitrary byte slice.
pub fn hash_bytes(data: &[u8], algorithm: HashAlgorithm, key: Option<&[u8]>) -> io::Result<String> {
    let result = match algorithm {
        HashAlgorithm::Md5 => format!("{:x}", Md5::digest(data)),
        HashAlgorithm::Sha1 => format!("{:x}", Sha1::digest(data)),
        HashAlgorithm::Sha256 => format!("{:x}", Sha256::digest(data)),
        HashAlgorithm::HmacMd5 => {
            let key =
                key.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "missing HMAC key"))?;
            type HmacMd5 = Hmac<Md5>;
            let mut mac = HmacMd5::new_from_slice(key)
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid key"))?;
            mac.update(data);
            to_hex(&mac.finalize().into_bytes())
        }
        HashAlgorithm::HmacSha1 => {
            let key =
                key.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "missing HMAC key"))?;
            type HmacSha1 = Hmac<Sha1>;
            let mut mac = HmacSha1::new_from_slice(key)
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid key"))?;
            mac.update(data);
            to_hex(&mac.finalize().into_bytes())
        }
        HashAlgorithm::HmacSha256 => {
            let key =
                key.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "missing HMAC key"))?;
            type HmacSha256 = Hmac<Sha256>;
            let mut mac = HmacSha256::new_from_slice(key)
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid key"))?;
            mac.update(data);
            to_hex(&mac.finalize().into_bytes())
        }
    };

    Ok(result)
}

/// Hash data from a reader.
pub fn hash_reader<R: Read>(
    mut reader: R,
    algorithm: HashAlgorithm,
    key: Option<&[u8]>,
) -> io::Result<String> {
    let mut buffer = Vec::new();
    reader.read_to_end(&mut buffer)?;
    hash_bytes(&buffer, algorithm, key)
}

/// Compute the hash of a file using the selected algorithm. For HMAC
/// algorithms, a secret key must be provided via the `key` argument.
pub fn hash_file(path: &Path, algorithm: HashAlgorithm, key: Option<&[u8]>) -> io::Result<String> {
    let file = File::open(path)?;
    hash_reader(file, algorithm, key)
}

/// Recursively hash a directory by hashing each contained file and then
/// hashing the concatenation of those hashes. The search depth can be
/// limited with `max_depth`.
pub fn hash_directory(
    path: &Path,
    algorithm: HashAlgorithm,
    key: Option<&[u8]>,
    max_depth: usize,
) -> io::Result<String> {
    let mut files: Vec<PathBuf> = WalkDir::new(path)
        .max_depth(max_depth)
        .into_iter()
        .filter_map(|e| e.ok())
        .filter(|e| e.file_type().is_file())
        .map(|e| e.into_path())
        .collect();

    files.sort();

    let mut combined = String::new();
    for f in files {
        let h = hash_file(&f, algorithm, key)?;
        combined.push_str(&h);
    }

    hash_bytes(combined.as_bytes(), algorithm, key)
}

fn to_hex(bytes: &[u8]) -> String {
    let mut out = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        use std::fmt::Write as _;
        write!(&mut out, "{:02x}", b).unwrap();
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::{self, File};
    use std::io::Write;
    use tempfile::tempdir;

    fn create_temp_file(content: &[u8]) -> std::path::PathBuf {
        use std::time::{SystemTime, UNIX_EPOCH};
        let mut path = std::env::temp_dir();
        let ts = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        path.push(format!("anti_sec_test_{}", ts));
        let mut f = File::create(&path).unwrap();
        f.write_all(content).unwrap();
        path
    }

    #[test]
    fn test_md5_hash() {
        let path = create_temp_file(b"hello");
        let digest = hash_file(&path, HashAlgorithm::Md5, None).unwrap();
        fs::remove_file(&path).unwrap();
        assert_eq!(digest, "5d41402abc4b2a76b9719d911017c592");
    }

    #[test]
    fn test_sha1_hash() {
        let path = create_temp_file(b"hello");
        let digest = hash_file(&path, HashAlgorithm::Sha1, None).unwrap();
        fs::remove_file(&path).unwrap();
        assert_eq!(digest, "aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d");
    }

    #[test]
    fn test_sha256_hash() {
        let path = create_temp_file(b"hello");
        let digest = hash_file(&path, HashAlgorithm::Sha256, None).unwrap();
        fs::remove_file(&path).unwrap();
        assert_eq!(
            digest,
            "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824"
        );
    }

    #[test]
    fn test_hmac_sha256_hash() {
        let path = create_temp_file(b"hello");
        let digest = hash_file(&path, HashAlgorithm::HmacSha256, Some(b"secret")).unwrap();
        fs::remove_file(&path).unwrap();
        assert_eq!(
            digest,
            "88aab3ede8d3adf94d26ab90d3bafd4a2083070c3bcce9c014ee04a443847c0b"
        );
    }

    #[test]
    fn test_hash_reader() {
        let data = b"reader";
        let digest = hash_reader(&data[..], HashAlgorithm::Sha256, None).unwrap();
        assert_eq!(
            digest,
            "3d0941964aa3ebdcb00ccef58b1bb399f9f898465e9886d5aec7f31090a0fb30"
        );
    }

    #[test]
    fn test_hash_directory() {
        let dir = tempdir().unwrap();
        let file_a = dir.path().join("a.txt");
        let file_b = dir.path().join("b.txt");
        fs::write(&file_a, b"foo").unwrap();
        fs::write(&file_b, b"bar").unwrap();
        let digest = hash_directory(dir.path(), HashAlgorithm::Sha256, None, 1).unwrap();
        // Compute expected: hash of files individually then concat and hash again
        let h1 = hash_file(&file_a, HashAlgorithm::Sha256, None).unwrap();
        let h2 = hash_file(&file_b, HashAlgorithm::Sha256, None).unwrap();
        let mut concat = String::new();
        let mut names = vec![(&file_a, h1), (&file_b, h2)];
        names.sort_by_key(|(p, _)| p.to_path_buf());
        for (_, h) in names {
            concat.push_str(&h);
        }
        let expected = hash_bytes(concat.as_bytes(), HashAlgorithm::Sha256, None).unwrap();
        assert_eq!(digest, expected);
    }
}