doot/crates/doot-lang/src/builtins/crypto.rs

use crate::evaluator::{EvalError, Value};
use std::path::PathBuf;

#[tracing::instrument(level = "trace", skip_all)]
pub fn hash_file(args: &[Value]) -> Result<Value, EvalError> {
    let path = match args.first() {
        Some(Value::Path(p)) => p.clone(),
        Some(Value::Str(s)) => PathBuf::from(s),
        _ => return Err(EvalError::TypeError("hash_file expects a path".to_string())),
    };

    let content = std::fs::read(&path)?;
    let hash = blake3::hash(&content);
    Ok(Value::Str(hash.to_hex().to_string()))
}

#[tracing::instrument(level = "trace", skip_all)]
pub fn hash_str(args: &[Value]) -> Result<Value, EvalError> {
    let s = match args.first() {
        Some(Value::Str(s)) => s,
        _ => {
            return Err(EvalError::TypeError(
                "hash_str expects a string".to_string(),
            ));
        }
    };

    let hash = blake3::hash(s.as_bytes());
    Ok(Value::Str(hash.to_hex().to_string()))
}

#[tracing::instrument(level = "trace", skip_all)]
pub fn encrypt_age(args: &[Value]) -> Result<Value, EvalError> {
    let content = match args.first() {
        Some(Value::Str(s)) => s,
        _ => {
            return Err(EvalError::TypeError(
                "encrypt_age expects content string".to_string(),
            ));
        }
    };

    let recipient = match args.get(1) {
        Some(Value::Str(s)) => s,
        _ => {
            return Err(EvalError::TypeError(
                "encrypt_age requires recipient public key".to_string(),
            ));
        }
    };

    let recipient = recipient
        .parse::<age::x25519::Recipient>()
        .map_err(|e| EvalError::TypeError(format!("invalid recipient: {}", e)))?;

    let encryptor = age::Encryptor::with_recipients(vec![Box::new(recipient)])
        .expect("failed to create encryptor");

    let mut encrypted = vec![];
    let mut writer = encryptor
        .wrap_output(&mut encrypted)
        .map_err(|e| EvalError::TypeError(format!("encryption error: {}", e)))?;

    use std::io::Write;
    writer
        .write_all(content.as_bytes())
        .map_err(|e| EvalError::TypeError(format!("encryption error: {}", e)))?;
    writer
        .finish()
        .map_err(|e| EvalError::TypeError(format!("encryption error: {}", e)))?;

    Ok(Value::Str(base64_encode(&encrypted)))
}

#[tracing::instrument(level = "trace", skip_all)]
pub fn decrypt_age(args: &[Value]) -> Result<Value, EvalError> {
    let encrypted = match args.first() {
        Some(Value::Str(s)) => s,
        _ => {
            return Err(EvalError::TypeError(
                "decrypt_age expects encrypted string".to_string(),
            ));
        }
    };

    let identity_str = match args.get(1) {
        Some(Value::Str(s)) => s,
        _ => {
            return Err(EvalError::TypeError(
                "decrypt_age requires identity".to_string(),
            ));
        }
    };

    let identity = identity_str
        .parse::<age::x25519::Identity>()
        .map_err(|e| EvalError::TypeError(format!("invalid identity: {}", e)))?;

    let encrypted_bytes = base64_decode(encrypted)
        .map_err(|e| EvalError::TypeError(format!("invalid base64: {}", e)))?;

    let decryptor = match age::Decryptor::new(&encrypted_bytes[..])
        .map_err(|e| EvalError::TypeError(format!("decryption error: {}", e)))?
    {
        age::Decryptor::Recipients(d) => d,
        _ => {
            return Err(EvalError::TypeError(
                "unexpected decryptor type".to_string(),
            ));
        }
    };

    let mut decrypted = vec![];
    let mut reader = decryptor
        .decrypt(std::iter::once(&identity as &dyn age::Identity))
        .map_err(|e| EvalError::TypeError(format!("decryption error: {}", e)))?;

    use std::io::Read;
    reader
        .read_to_end(&mut decrypted)
        .map_err(|e| EvalError::TypeError(format!("decryption error: {}", e)))?;

    Ok(Value::Str(String::from_utf8(decrypted).map_err(|e| {
        EvalError::TypeError(format!("invalid UTF-8: {}", e))
    })?))
}

pub fn base64_encode(data: &[u8]) -> String {
    const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut result = String::new();

    for chunk in data.chunks(3) {
        let b0 = chunk[0] as usize;
        let b1 = chunk.get(1).copied().unwrap_or(0) as usize;
        let b2 = chunk.get(2).copied().unwrap_or(0) as usize;

        result.push(ALPHABET[b0 >> 2] as char);
        result.push(ALPHABET[((b0 & 0x03) << 4) | (b1 >> 4)] as char);

        if chunk.len() > 1 {
            result.push(ALPHABET[((b1 & 0x0f) << 2) | (b2 >> 6)] as char);
        } else {
            result.push('=');
        }

        if chunk.len() > 2 {
            result.push(ALPHABET[b2 & 0x3f] as char);
        } else {
            result.push('=');
        }
    }

    result
}

pub fn base64_decode(s: &str) -> Result<Vec<u8>, String> {
    const DECODE: [i8; 256] = {
        let mut table = [-1i8; 256];
        let alphabet = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
        let mut i = 0;
        while i < 64 {
            table[alphabet[i] as usize] = i as i8;
            i += 1;
        }
        table
    };

    let s = s.trim_end_matches('=');
    let mut result = Vec::with_capacity(s.len() * 3 / 4);

    let chars: Vec<u8> = s.bytes().collect();
    for chunk in chars.chunks(4) {
        let mut buf = [0u8; 4];
        for (i, &c) in chunk.iter().enumerate() {
            let val = DECODE[c as usize];
            if val < 0 {
                return Err(format!("invalid base64 character: {}", c as char));
            }
            buf[i] = val as u8;
        }

        result.push((buf[0] << 2) | (buf[1] >> 4));
        if chunk.len() > 2 {
            result.push((buf[1] << 4) | (buf[2] >> 2));
        }
        if chunk.len() > 3 {
            result.push((buf[2] << 6) | buf[3]);
        }
    }

    Ok(result)
}