azathoth_transformers 0.1.0

Data transformers for the AzathothC2 project
Documentation
use crate::Transformer;
use crate::traits::TransformerError;
use alloc::boxed::Box;
use alloc::vec::Vec;

/// Appends a suffix to data.
///
/// ## Example:
/// ```
/// use azathoth_transformers::{Transformer, TransformerError, TransformerExt, basic::AppendTransformer};
/// fn main() -> Result<(), TransformerError> {
///     let data = "foo";
///     let suffix = "bar";
///     let mut at = AppendTransformer::new(suffix);
///     let result = at.apply(data)?;
///     assert_eq!(result.as_slice(), b"foobar");
///     Ok(())
/// }
#[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))]
#[derive(Clone, Debug)]
pub struct AppendTransformer {
    suffix: Vec<u8>,
}

impl AppendTransformer {
    /// Creates a new [`AppendTransformer`]
    #[must_use]
    pub fn new(suffix: impl Into<Vec<u8>>) -> Self {
        Self {
            suffix: suffix.into(),
        }
    }
}

impl Transformer for AppendTransformer {
    fn apply_bytes(&mut self, data: &[u8]) -> Result<Vec<u8>, TransformerError> {
        let mut out = Vec::with_capacity(data.len() + self.suffix.len());
        out.extend_from_slice(data);
        out.extend_from_slice(&self.suffix);
        Ok(out)
    }

    fn reverse_bytes(&mut self, data: &[u8]) -> Result<Vec<u8>, TransformerError> {
        if data.len() < self.suffix.len() {
            return Err("data shorter than suffix".into());
        }
        if !data.ends_with(&self.suffix) {
            return Err("data does not end with suffix".into());
        }
        Ok(data[..data.len() - self.suffix.len()].to_vec())
    }

    fn name(&self) -> Option<&'static str> {
        Some("append")
    }

    fn box_clone(&self) -> Box<dyn Transformer> {
        Box::new(self.clone())
    }
}

/// Prepends a prefix to data.
///
/// ## Example:
/// ```
///  use azathoth_transformers::{Transformer, TransformerError, TransformerExt, basic::PrependTransformer};
/// fn main() -> Result<(), TransformerError> {
///     let data = "baz";
///     let prefix = "bar ";
///     let mut at = PrependTransformer::new(prefix);
///     let result = at.apply(data)?;
///     assert_eq!(result.as_slice(), b"bar baz");
///     Ok(())
/// }
#[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))]
#[derive(Clone, Debug)]
pub struct PrependTransformer {
    prefix: Vec<u8>,
}

impl PrependTransformer {
    /// Creates a new [`PrependTransformer`] object
    #[must_use]
    pub fn new(prefix: impl Into<Vec<u8>>) -> Self {
        Self {
            prefix: prefix.into(),
        }
    }
}

impl Transformer for PrependTransformer {
    fn apply_bytes(&mut self, data: &[u8]) -> Result<Vec<u8>, TransformerError> {
        let mut out = Vec::with_capacity(self.prefix.len() + data.len());
        out.extend_from_slice(&self.prefix);
        out.extend_from_slice(data);
        Ok(out)
    }
    fn reverse_bytes(&mut self, data: &[u8]) -> Result<Vec<u8>, TransformerError> {
        if data.len() < self.prefix.len() {
            return Err("data shorter than prefix".into());
        }
        if !data.starts_with(&self.prefix) {
            return Err("data does not start with prefix".into());
        }
        Ok(data[self.prefix.len()..].to_vec())
    }

    fn name(&self) -> Option<&'static str> {
        Some("prepend")
    }

    fn box_clone(&self) -> Box<dyn Transformer> {
        Box::new(self.clone())
    }
}

/// Simple XOR operation transformer
///
/// ## Example:
/// ```
///  use azathoth_transformers::{Transformer, TransformerError, TransformerExt, basic::XorTransformer};
/// fn main() -> Result<(), TransformerError> {
///     let data = "hello!";
///     let key = "key";
///     let mut xor = XorTransformer::new(key)?;
///     let result = xor.apply(data)?;
///     assert_eq!(result.as_slice(), &[0x03, 0x00, 0x15, 0x07, 0x0a, 0x58]);
///     let reversed = xor.reverse(&result)?;
///     assert_eq!(reversed.as_slice(), b"hello!");
///     Ok(())
/// }
#[cfg_attr(feature = "std", derive(serde::Deserialize, serde::Serialize))]
#[derive(Clone, Debug)]
pub struct XorTransformer {
    key: Vec<u8>,
}

impl XorTransformer {
    /// Creates a new [`XorTransformer`] object
    ///
    /// # Errors
    /// This function returns an error if the given key is empty
    pub fn new(key: impl Into<Vec<u8>>) -> Result<Self, TransformerError> {
        let key = key.into();
        if key.is_empty() {
            return Err("empty key".into());
        }
        Ok(Self { key })
    }
}

impl Transformer for XorTransformer {
    fn apply_bytes(&mut self, data: &[u8]) -> Result<Vec<u8>, TransformerError> {
        let mut key_iter = self.key.iter().cycle();
        Ok(data.iter().map(|x| x ^ key_iter.next().unwrap()).collect())
    }

    fn reverse_bytes(&mut self, data: &[u8]) -> Result<Vec<u8>, TransformerError> {
        self.apply_bytes(data)
    }

    fn name(&self) -> Option<&'static str> {
        Some("xor")
    }

    fn box_clone(&self) -> Box<dyn Transformer> {
        Box::new(self.clone())
    }
}