synta-python 0.1.4

//! Symmetric-crypto primitives for Python: HMAC, PBKDF2, AES-CBC, 3DES-CBC,
//! PKCS#7 padding helpers, and the Fernet authenticated-encryption scheme.
//!
//! All operations are delegated to `synta-certificate`'s backend traits so
//! that no direct `openssl::*` imports are needed here.

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use synta_certificate::{
    default_block_cipher_provider, default_hmac_provider, default_pbkdf2_provider,
    default_secure_random, BlockCipherProvider, HmacProvider, Pbkdf2Provider, SecureRandom,
};

// ── URL-safe base64 helpers (used by Fernet) ─────────────────────────────────

/// Encode `data` as URL-safe base64 (RFC 4648 §5: `-` for `+`, `_` for `/`).
fn urlsafe_b64encode(data: &[u8]) -> String {
    synta_certificate::encode_base64(data)
        .replace('+', "-")
        .replace('/', "_")
}

/// Decode URL-safe base64 bytes.  Translates `-`/`_` back to `+`/`/` before
/// decoding.
fn urlsafe_b64decode(data: &[u8]) -> PyResult<Vec<u8>> {
    let s = std::str::from_utf8(data)
        .map_err(|_| PyValueError::new_err("Fernet key/token must be valid ASCII"))?;
    let standard = s.replace('-', "+").replace('_', "/");
    synta_certificate::decode_base64(standard.as_bytes())
        .ok_or_else(|| PyValueError::new_err("base64 decode error"))
}

// ── HMAC ─────────────────────────────────────────────────────────────────────

/// Compute an HMAC digest and return the raw MAC bytes.
///
/// ``algorithm`` must be one of ``"sha1"``, ``"sha224"``, ``"sha256"``,
/// ``"sha384"``, ``"sha512"``, or ``"md5"``.
///
/// ```python,ignore
/// import synta
///
/// key  = b"secret"
/// data = b"hello world"
/// mac  = synta.hmac_digest("sha256", key, data)
/// print(mac.hex())
/// ```
#[pyfunction]
pub fn hmac_digest<'py>(
    py: Python<'py>,
    algorithm: &str,
    key: &[u8],
    data: &[u8],
) -> PyResult<Bound<'py, PyBytes>> {
    let mac = default_hmac_provider()
        .hmac_compute(algorithm, key, data)
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
    Ok(PyBytes::new(py, &mac))
}

/// Verify an HMAC digest in constant time.
///
/// Raises :exc:`ValueError` if the computed HMAC does not match ``expected``.
/// Uses constant-time comparison to prevent timing side-channels.
///
/// ```python,ignore
/// import synta
///
/// synta.hmac_verify("sha256", key, data, expected_mac)
/// # raises ValueError on mismatch
/// ```
#[pyfunction]
pub fn hmac_verify(algorithm: &str, key: &[u8], data: &[u8], expected: &[u8]) -> PyResult<()> {
    default_hmac_provider()
        .hmac_verify(algorithm, key, data, expected)
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
    Ok(())
}

// ── PBKDF2 ───────────────────────────────────────────────────────────────────

/// Derive a key using PBKDF2-HMAC.
///
/// Returns ``length`` raw key bytes derived from ``password`` and ``salt``
/// using the given HMAC ``algorithm`` and the specified ``iterations`` count.
///
/// ```python,ignore
/// import synta
///
/// key = synta.pbkdf2_hmac("sha256", b"password", b"salt", 100_000, 32)
/// ```
#[pyfunction]
pub fn pbkdf2_hmac<'py>(
    py: Python<'py>,
    algorithm: &str,
    password: &[u8],
    salt: &[u8],
    iterations: u32,
    length: usize,
) -> PyResult<Bound<'py, PyBytes>> {
    let out = default_pbkdf2_provider()
        .pbkdf2_hmac(algorithm, password, salt, iterations as usize, length)
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
    Ok(PyBytes::new(py, &out))
}

// ── AES-CBC ──────────────────────────────────────────────────────────────────

/// Encrypt data using AES-CBC.
///
/// Key length determines the AES variant: 16 → AES-128, 24 → AES-192,
/// 32 → AES-256.  ``iv`` must be 16 bytes.
///
/// When ``pad=True`` (default) PKCS#7 padding is applied automatically.
/// When ``pad=False`` the caller is responsible for correct block alignment.
///
/// ```python,ignore
/// import synta
///
/// ct = synta.aes_cbc_encrypt(key_32, iv_16, plaintext)
/// ct_no_pad = synta.aes_cbc_encrypt(key_16, iv_16, aligned_data, pad=False)
/// ```
#[pyfunction]
#[pyo3(signature = (key, iv, plaintext, pad = true))]
pub fn aes_cbc_encrypt<'py>(
    py: Python<'py>,
    key: &[u8],
    iv: &[u8],
    plaintext: &[u8],
    pad: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    let ct = default_block_cipher_provider()
        .aes_cbc_encrypt(key, iv, plaintext, pad)
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
    Ok(PyBytes::new(py, &ct))
}

/// Decrypt data using AES-CBC.
///
/// Key length determines the AES variant: 16 → AES-128, 24 → AES-192,
/// 32 → AES-256.  ``iv`` must be 16 bytes.
///
/// When ``unpad=True`` (default) PKCS#7 padding is stripped automatically.
/// When ``unpad=False`` the raw decrypted bytes are returned.
///
/// ```python,ignore
/// import synta
///
/// pt = synta.aes_cbc_decrypt(key_32, iv_16, ciphertext)
/// pt_raw = synta.aes_cbc_decrypt(key_16, iv_16, ciphertext, unpad=False)
/// ```
#[pyfunction]
#[pyo3(signature = (key, iv, ciphertext, unpad = true))]
pub fn aes_cbc_decrypt<'py>(
    py: Python<'py>,
    key: &[u8],
    iv: &[u8],
    ciphertext: &[u8],
    unpad: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    let pt = default_block_cipher_provider()
        .aes_cbc_decrypt(key, iv, ciphertext, unpad)
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
    Ok(PyBytes::new(py, &pt))
}

// ── 3DES-EDE-CBC ─────────────────────────────────────────────────────────────

/// Encrypt data using 3DES-EDE-CBC.
///
/// ``key`` must be 16 bytes (two-key 3DES) or 24 bytes (three-key 3DES).
/// ``iv`` must be 8 bytes.
///
/// When ``pad=True`` (default) PKCS#7 padding is applied automatically.
/// When ``pad=False`` the caller is responsible for correct block alignment.
///
/// ```python,ignore
/// import synta
///
/// ct = synta.des3_cbc_encrypt(key_24, iv_8, plaintext)
/// ```
#[pyfunction]
#[pyo3(signature = (key, iv, plaintext, pad = true))]
pub fn des3_cbc_encrypt<'py>(
    py: Python<'py>,
    key: &[u8],
    iv: &[u8],
    plaintext: &[u8],
    pad: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    let ct = default_block_cipher_provider()
        .des3_cbc_encrypt(key, iv, plaintext, pad)
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
    Ok(PyBytes::new(py, &ct))
}

/// Decrypt data using 3DES-EDE-CBC.
///
/// ``key`` must be 16 bytes (two-key 3DES) or 24 bytes (three-key 3DES).
/// ``iv`` must be 8 bytes.
///
/// When ``unpad=True`` (default) PKCS#7 padding is stripped automatically.
/// When ``unpad=False`` raw decrypted bytes are returned.
///
/// ```python,ignore
/// import synta
///
/// pt = synta.des3_cbc_decrypt(key_24, iv_8, ciphertext)
/// ```
#[pyfunction]
#[pyo3(signature = (key, iv, ciphertext, unpad = true))]
pub fn des3_cbc_decrypt<'py>(
    py: Python<'py>,
    key: &[u8],
    iv: &[u8],
    ciphertext: &[u8],
    unpad: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    let pt = default_block_cipher_provider()
        .des3_cbc_decrypt(key, iv, ciphertext, unpad)
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
    Ok(PyBytes::new(py, &pt))
}

// ── PKCS#7 padding ───────────────────────────────────────────────────────────

/// Apply PKCS#7 padding to ``data``.
///
/// ``block_size`` is the cipher block size in bytes (16 for AES, 8 for 3DES).
/// Raises :exc:`ValueError` if ``block_size`` is 0 or greater than 255.
///
/// ```python,ignore
/// import synta
///
/// padded = synta.pkcs7_pad(b"hello", 8)   # → b"hello\x03\x03\x03"
/// ```
#[pyfunction]
pub fn pkcs7_pad<'py>(
    py: Python<'py>,
    data: &[u8],
    block_size: usize,
) -> PyResult<Bound<'py, PyBytes>> {
    if block_size == 0 || block_size > 255 {
        return Err(PyValueError::new_err(
            "block_size must be between 1 and 255",
        ));
    }
    let pad_len = block_size - (data.len() % block_size);
    let mut padded = Vec::with_capacity(data.len() + pad_len);
    padded.extend_from_slice(data);
    padded.extend(std::iter::repeat_n(pad_len as u8, pad_len));
    Ok(PyBytes::new(py, &padded))
}

/// Remove PKCS#7 padding from ``data``.
///
/// Raises :exc:`ValueError` if the padding is missing or inconsistent.
///
/// ```python,ignore
/// import synta
///
/// unpadded = synta.pkcs7_unpad(b"hello\x03\x03\x03", 8)   # → b"hello"
/// ```
#[pyfunction]
pub fn pkcs7_unpad<'py>(
    py: Python<'py>,
    data: &[u8],
    block_size: usize,
) -> PyResult<Bound<'py, PyBytes>> {
    if data.is_empty() {
        return Err(PyValueError::new_err("data is empty"));
    }
    let pad_byte = *data.last().unwrap() as usize;
    if pad_byte == 0 || pad_byte > block_size || pad_byte > data.len() {
        return Err(PyValueError::new_err("invalid PKCS#7 padding"));
    }
    let pad_start = data.len() - pad_byte;
    if data[pad_start..].iter().any(|&b| b as usize != pad_byte) {
        return Err(PyValueError::new_err("inconsistent PKCS#7 padding bytes"));
    }
    Ok(PyBytes::new(py, &data[..pad_start]))
}

// ── Fernet ───────────────────────────────────────────────────────────────────

/// Fernet symmetric authenticated encryption.
///
/// Implements the `Fernet specification <https://github.com/fernet/spec/blob/master/Spec.md>`_
/// using AES-128-CBC encryption and HMAC-SHA256 authentication.  Tokens
/// produced by this class are byte-for-byte compatible with those produced
/// by ``cryptography.fernet.Fernet``.
///
/// **Key format:** a URL-safe base64-encoded string of 32 raw bytes, where
/// the first 16 bytes are the signing key and the last 16 bytes are the
/// encryption key.
///
/// **Token format (before base64url encoding):**
///
/// * 1 byte: version (``0x80``)
/// * 8 bytes: timestamp (big-endian unsigned 64-bit Unix seconds)
/// * 16 bytes: AES-CBC initialization vector
/// * N bytes: AES-128-CBC ciphertext (PKCS#7-padded plaintext)
/// * 32 bytes: HMAC-SHA256 of all preceding bytes
///
/// ```python,ignore
/// import synta
///
/// key   = synta.Fernet.generate_key()
/// fern  = synta.Fernet(key)
/// token = fern.encrypt(b"secret data")
/// plain = fern.decrypt(token)
/// assert plain == b"secret data"
///
/// # TTL check (seconds):
/// plain = fern.decrypt(token, ttl=300)
/// ```
#[pyclass(frozen, name = "Fernet")]
pub struct PyFernet {
    signing_key: [u8; 16],
    encryption_key: [u8; 16],
}

#[pymethods]
impl PyFernet {
    /// Create a :class:`Fernet` instance from a URL-safe base64-encoded 32-byte key.
    ///
    /// ```python,ignore
    /// key  = synta.Fernet.generate_key()
    /// fern = synta.Fernet(key)
    /// ```
    #[new]
    fn new(key: &[u8]) -> PyResult<Self> {
        let raw = urlsafe_b64decode(key)?;
        if raw.len() != 32 {
            return Err(PyValueError::new_err(
                "Fernet key must decode to exactly 32 bytes",
            ));
        }
        let mut signing_key = [0u8; 16];
        let mut encryption_key = [0u8; 16];
        signing_key.copy_from_slice(&raw[0..16]);
        encryption_key.copy_from_slice(&raw[16..32]);
        Ok(Self {
            signing_key,
            encryption_key,
        })
    }

    /// Generate a fresh random Fernet key.
    ///
    /// Returns 32 cryptographically random bytes encoded as URL-safe base64.
    ///
    /// ```python,ignore
    /// key = synta.Fernet.generate_key()
    /// fern = synta.Fernet(key)
    /// ```
    #[staticmethod]
    fn generate_key<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        let mut raw = [0u8; 32];
        default_secure_random()
            .rand_bytes(&mut raw)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;
        Ok(PyBytes::new(py, urlsafe_b64encode(&raw).as_bytes()))
    }

    /// Encrypt ``data`` and return a Fernet token as URL-safe base64 bytes.
    ///
    /// The token includes a timestamp set to the current UTC time; use
    /// :meth:`decrypt` with a ``ttl`` argument to enforce expiry.
    ///
    /// ```python,ignore
    /// token = fern.encrypt(b"hello")
    /// ```
    fn encrypt<'py>(&self, py: Python<'py>, data: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
        let rng = default_secure_random();
        let cipher = default_block_cipher_provider();
        let hmac = default_hmac_provider();

        // Generate a random 16-byte IV.
        let mut iv = [0u8; 16];
        rng.rand_bytes(&mut iv)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;

        // Current Unix timestamp (seconds).
        let ts = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_err(|_| PyValueError::new_err("system clock error"))?
            .as_secs();

        // AES-128-CBC encrypt with PKCS#7 padding.
        let ciphertext = cipher
            .aes_cbc_encrypt(&self.encryption_key, &iv, data, true)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;

        // Assemble the pre-MAC bytes: version(1) || ts_be(8) || iv(16) || ciphertext.
        let mut pre_mac = Vec::with_capacity(1 + 8 + 16 + ciphertext.len());
        pre_mac.push(0x80u8);
        pre_mac.extend_from_slice(&ts.to_be_bytes());
        pre_mac.extend_from_slice(&iv);
        pre_mac.extend_from_slice(&ciphertext);

        // HMAC-SHA256 over the pre-MAC bytes.
        let mac = hmac
            .hmac_compute("sha256", &self.signing_key, &pre_mac)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;

        // Token = pre_mac || mac, base64url-encoded.
        let mut token_bytes = pre_mac;
        token_bytes.extend_from_slice(&mac);
        Ok(PyBytes::new(py, urlsafe_b64encode(&token_bytes).as_bytes()))
    }

    /// Decrypt a Fernet ``token``, optionally enforcing a TTL.
    ///
    /// ``token`` must be URL-safe base64-encoded bytes as returned by
    /// :meth:`encrypt`.
    ///
    /// If ``ttl`` is given (seconds), raises :exc:`ValueError` when the token
    /// is older than ``ttl`` seconds relative to the current time.
    ///
    /// Raises :exc:`ValueError` on malformed tokens, unknown version bytes,
    /// HMAC verification failure, or TTL expiry.
    ///
    /// ```python,ignore
    /// plain = fern.decrypt(token)
    /// plain = fern.decrypt(token, ttl=300)
    /// ```
    #[pyo3(signature = (token, ttl = None))]
    fn decrypt<'py>(
        &self,
        py: Python<'py>,
        token: &[u8],
        ttl: Option<u64>,
    ) -> PyResult<Bound<'py, PyBytes>> {
        let hmac = default_hmac_provider();
        let cipher = default_block_cipher_provider();

        // Decode URL-safe base64.
        let raw = urlsafe_b64decode(token)?;

        // Minimum token length: version(1) + ts(8) + iv(16) + one AES block(16) + hmac(32) = 73.
        if raw.len() < 1 + 8 + 16 + 16 + 32 {
            return Err(PyValueError::new_err("invalid Fernet token: too short"));
        }

        // Check the version byte.
        if raw[0] != 0x80 {
            return Err(PyValueError::new_err(
                "invalid Fernet token: unknown version byte",
            ));
        }

        // Extract fields.
        let ts = u64::from_be_bytes(
            raw[1..9]
                .try_into()
                .map_err(|_| PyValueError::new_err("internal error: Fernet timestamp slice"))?,
        );
        let iv = &raw[9..25];
        let ciphertext = &raw[25..raw.len() - 32];
        let token_mac = &raw[raw.len() - 32..];

        // Verify HMAC-SHA256 in constant time.
        let pre_mac = &raw[..raw.len() - 32];
        let computed_mac = hmac
            .hmac_compute("sha256", &self.signing_key, pre_mac)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;

        if !synta_certificate::constant_time_eq(&computed_mac, token_mac) {
            return Err(PyValueError::new_err(
                "invalid Fernet token: HMAC verification failed",
            ));
        }

        // TTL check (done after HMAC so we don't leak timing info about bad tokens).
        if let Some(ttl_secs) = ttl {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map_err(|_| PyValueError::new_err("system clock error"))?
                .as_secs();
            // Reject tokens issued in the future (60-second clock-skew tolerance).
            const MAX_CLOCK_SKEW: u64 = 60;
            if ts > now.saturating_add(MAX_CLOCK_SKEW) {
                return Err(PyValueError::new_err(
                    "Fernet token timestamp is in the future",
                ));
            }
            if now.saturating_sub(ts) > ttl_secs {
                return Err(PyValueError::new_err("Fernet token has expired"));
            }
        }

        // AES-128-CBC decrypt with PKCS#7 unpadding.
        let plaintext = cipher
            .aes_cbc_decrypt(&self.encryption_key, iv, ciphertext, true)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;

        Ok(PyBytes::new(py, &plaintext))
    }
}

// ── Submodule registration ────────────────────────────────────────────────────

/// Register the ``synta.crypto`` submodule.
///
/// Exposes all symmetric-crypto primitives under ``synta.crypto``:
///
/// * :class:`Fernet` — authenticated encryption (AES-128-CBC + HMAC-SHA256)
/// * :func:`hmac_digest` / :func:`hmac_verify` — HMAC generation and verification
/// * :func:`pbkdf2_hmac` — PBKDF2 key derivation
/// * :func:`aes_cbc_encrypt` / :func:`aes_cbc_decrypt` — AES-CBC
/// * :func:`des3_cbc_encrypt` / :func:`des3_cbc_decrypt` — 3DES-EDE-CBC
/// * :func:`pkcs7_pad` / :func:`pkcs7_unpad` — PKCS#7 block padding helpers
pub fn register_crypto_module(parent: &Bound<'_, PyModule>) -> PyResult<()> {
    use pyo3::prelude::*;
    let py = parent.py();
    let m = PyModule::new(py, "crypto")?;

    m.add_class::<PyFernet>()?;
    m.add_function(wrap_pyfunction!(hmac_digest, &m)?)?;
    m.add_function(wrap_pyfunction!(hmac_verify, &m)?)?;
    m.add_function(wrap_pyfunction!(pbkdf2_hmac, &m)?)?;
    m.add_function(wrap_pyfunction!(aes_cbc_encrypt, &m)?)?;
    m.add_function(wrap_pyfunction!(aes_cbc_decrypt, &m)?)?;
    m.add_function(wrap_pyfunction!(des3_cbc_encrypt, &m)?)?;
    m.add_function(wrap_pyfunction!(des3_cbc_decrypt, &m)?)?;
    m.add_function(wrap_pyfunction!(pkcs7_pad, &m)?)?;
    m.add_function(wrap_pyfunction!(pkcs7_unpad, &m)?)?;
    m.add_class::<crate::otp::PyHOTP>()?;
    m.add_class::<crate::otp::PyTOTP>()?;

    crate::install_submodule(
        parent,
        &m,
        "synta.crypto",
        Some(concat!(
            "Symmetric cryptography primitives: HMAC, PBKDF2, AES-CBC, 3DES-CBC, ",
            "PKCS#7 padding, and Fernet authenticated encryption (RFC-compatible ",
            "with the Python ``cryptography`` library).",
        )),
    )
}