synta-python 0.2.5

//! Python bindings for HOTP (RFC 4226) and TOTP (RFC 6238) OTP generators.

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use synta_certificate::{constant_time_eq, default_hmac_provider, HmacProvider};

// ── Shared computation ────────────────────────────────────────────────────────

/// HOTP computation (RFC 4226 §5.3).
///
/// `counter` is an 8-byte big-endian integer.  `digits` controls the
/// output width (typically 6 or 8).  `algorithm` is one of `"sha1"`,
/// `"sha256"`, `"sha384"`, or `"sha512"`.
fn hotp_compute(key: &[u8], counter: u64, digits: u32, algorithm: &str) -> PyResult<String> {
    let hs = default_hmac_provider()
        .hmac_compute(algorithm, key, &counter.to_be_bytes())
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;

    // Dynamic truncation (RFC 4226 §5.3).
    let offset = (hs[hs.len() - 1] & 0x0f) as usize;
    let p = u32::from_be_bytes([hs[offset], hs[offset + 1], hs[offset + 2], hs[offset + 3]])
        & 0x7fff_ffff;
    let modulus = 10u32.checked_pow(digits).unwrap_or(u32::MAX);
    let otp = p % modulus;
    Ok(format!("{:0>width$}", otp, width = digits as usize))
}

fn validate_algorithm(algorithm: &str) -> PyResult<()> {
    match algorithm {
        "sha1" | "sha256" | "sha384" | "sha512" => Ok(()),
        other => Err(PyValueError::new_err(format!(
            "unsupported OTP algorithm: {other}; expected sha1, sha256, sha384, or sha512"
        ))),
    }
}

// ── HOTP ──────────────────────────────────────────────────────────────────────

/// HMAC-based One-Time Password generator (RFC 4226).
///
/// Generates and verifies counter-based OTPs using HMAC-SHA1 (default),
/// HMAC-SHA256, HMAC-SHA384, or HMAC-SHA512.
///
/// ```python,ignore
/// import synta.crypto
///
/// # RFC 4226 test vector — key is the ASCII string "12345678901234567890"
/// key = bytes.fromhex("3132333435363738393031323334353637383930")
/// hotp = synta.crypto.HOTP(key, digits=6, algorithm="sha1")
/// print(hotp.generate(0))    # "755224"
/// print(hotp.generate(1))    # "287082"
/// # Verify with look-ahead window of 5:
/// next_counter = hotp.verify("287082", counter=0, look_ahead=5)
/// assert next_counter == 1
/// ```
#[pyclass(frozen, name = "HOTP")]
pub struct PyHOTP {
    key: Vec<u8>,
    digits: u32,
    algorithm: String,
}

#[pymethods]
impl PyHOTP {
    /// Create a new HOTP instance.
    ///
    /// ``key`` is the shared secret (raw bytes).  ``digits`` is the OTP
    /// length (default 6).  ``algorithm`` is the HMAC hash (default
    /// ``"sha1"``); one of ``"sha1"``, ``"sha256"``, ``"sha384"``,
    /// ``"sha512"``.
    #[new]
    #[pyo3(signature = (key, digits = 6, algorithm = "sha1"))]
    fn new(key: &[u8], digits: u32, algorithm: &str) -> PyResult<Self> {
        validate_algorithm(algorithm)?;
        if !(1..=9).contains(&digits) {
            return Err(PyValueError::new_err(
                "digits must be between 1 and 9 inclusive",
            ));
        }
        Ok(Self {
            key: key.to_vec(),
            digits,
            algorithm: algorithm.to_string(),
        })
    }

    /// Generate the OTP for the given ``counter`` value.
    ///
    /// Returns a zero-padded decimal string of length ``digits``.
    fn generate(&self, counter: u64) -> PyResult<String> {
        hotp_compute(&self.key, counter, self.digits, &self.algorithm)
    }

    /// Verify ``code`` against a window of counters starting at ``counter``.
    ///
    /// Checks ``counter`` through ``counter + look_ahead`` (inclusive) using
    /// constant-time comparison.  Returns the matching counter value on
    /// success, or ``None`` if no counter in the window matches.
    #[pyo3(signature = (code, counter, look_ahead = 0))]
    fn verify(&self, code: &str, counter: u64, look_ahead: u64) -> PyResult<Option<u64>> {
        for i in 0..=look_ahead {
            let expected = hotp_compute(&self.key, counter + i, self.digits, &self.algorithm)?;
            if constant_time_eq(code.as_bytes(), expected.as_bytes()) {
                return Ok(Some(counter + i));
            }
        }
        Ok(None)
    }

    fn __repr__(&self) -> String {
        format!(
            "HOTP(digits={}, algorithm={:?})",
            self.digits, self.algorithm
        )
    }
}

// ── TOTP ──────────────────────────────────────────────────────────────────────

/// Time-based One-Time Password generator (RFC 6238).
///
/// The counter is derived from the Unix timestamp divided by the time step
/// (``period`` seconds, default 30).  Supports HMAC-SHA1, HMAC-SHA256,
/// HMAC-SHA384, and HMAC-SHA512.
///
/// ```python,ignore
/// import synta.crypto, time
///
/// key = bytes.fromhex("3132333435363738393031323334353637383930")
/// totp = synta.crypto.TOTP(key, digits=8, algorithm="sha1", period=30)
/// otp = totp.now()                        # current time step
/// otp = totp.generate(int(time.time()))   # explicit Unix timestamp
/// ```
#[pyclass(frozen, name = "TOTP")]
pub struct PyTOTP {
    key: Vec<u8>,
    digits: u32,
    algorithm: String,
    period: u64,
}

#[pymethods]
impl PyTOTP {
    /// Create a new TOTP instance.
    ///
    /// ``key`` is the shared secret (raw bytes).  ``digits`` is the OTP
    /// length (default 6).  ``algorithm`` is the HMAC hash (default
    /// ``"sha1"``); one of ``"sha1"``, ``"sha256"``, ``"sha384"``,
    /// ``"sha512"``.  ``period`` is the time step in seconds (default 30).
    #[new]
    #[pyo3(signature = (key, digits = 6, algorithm = "sha1", period = 30))]
    fn new(key: &[u8], digits: u32, algorithm: &str, period: u64) -> PyResult<Self> {
        validate_algorithm(algorithm)?;
        if !(1..=9).contains(&digits) {
            return Err(PyValueError::new_err(
                "digits must be between 1 and 9 inclusive",
            ));
        }
        if period == 0 {
            return Err(PyValueError::new_err("TOTP period must be positive"));
        }
        Ok(Self {
            key: key.to_vec(),
            digits,
            algorithm: algorithm.to_string(),
            period,
        })
    }

    /// Generate the OTP for the given Unix ``timestamp`` (seconds).
    fn generate(&self, timestamp: u64) -> PyResult<String> {
        let counter = timestamp / self.period;
        hotp_compute(&self.key, counter, self.digits, &self.algorithm)
    }

    /// Generate the OTP for the current system time.
    fn now(&self) -> PyResult<String> {
        let ts = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_err(|_| PyValueError::new_err("system clock error"))?
            .as_secs();
        self.generate(ts)
    }

    fn __repr__(&self) -> String {
        format!(
            "TOTP(digits={}, algorithm={:?}, period={})",
            self.digits, self.algorithm, self.period
        )
    }
}