synta-python 0.2.5

//! Python wrapper for ASN.1 encoder

use pyo3::prelude::*;

use synta::Encoder;

use super::error::SyntaErr;

use super::types::{
    PyBitString, PyBmpString, PyBoolean, PyGeneralString, PyGeneralizedTime, PyIA5String,
    PyInteger, PyNull, PyNumericString, PyObjectIdentifier, PyOctetString, PyPrintableString,
    PyReal, PyTeletexString, PyUniversalString, PyUtcTime, PyUtf8String, PyVisibleString,
};
use super::PyEncoding;

/// ASN.1 Encoder
///
/// Encodes ASN.1 data in DER/BER format.
///
/// Example:
///     >>> encoder = Encoder(Encoding.DER)
///     >>> encoder.encode_integer(42)
///     >>> output = encoder.finish()
#[pyclass(name = "Encoder")]
pub struct PyEncoder {
    encoder: Encoder,
}

#[pymethods]
impl PyEncoder {
    /// Create a new encoder
    ///
    /// Args:
    ///     encoding: The encoding type (DER, BER, or CER)
    #[new]
    fn new(encoding: PyEncoding) -> Self {
        Self {
            encoder: Encoder::new(encoding.into()),
        }
    }

    /// Encode an INTEGER
    ///
    /// Accepts any Python ``int`` regardless of magnitude.  Values that fit in
    /// ``i64`` or ``i128`` are encoded via the native typed constructors;
    /// larger values (e.g. 20-byte X.509 serial numbers up to 160 bits) are
    /// converted through Python's ``int.to_bytes()`` using signed two's
    /// complement, then passed directly to the encoder.
    ///
    /// Args:
    ///     value: The integer value to encode (any Python int)
    fn encode_integer(
        &mut self,
        py: Python,
        value: &Bound<'_, pyo3::types::PyAny>,
    ) -> PyResult<()> {
        // Fast path: fits in i64 (covers most ASN.1 integer fields)
        if let Ok(v) = value.extract::<i64>() {
            self.encoder
                .encode(&synta::Integer::from_i64(v))
                .map_err(SyntaErr)?;
            return Ok(());
        }
        // Fits in i128 (covers 8–16 byte random serial numbers)
        if let Ok(v) = value.extract::<i128>() {
            self.encoder
                .encode(&synta::Integer::from_i128(v))
                .map_err(SyntaErr)?;
            return Ok(());
        }
        // Arbitrarily large: convert via Python int methods.
        // ceil((bit_length + 1) / 8) == (bit_length + 8) / 8 bytes are
        // sufficient for signed two's complement encoding of any integer.
        let bit_len: usize = value
            .call_method0(pyo3::intern!(py, "bit_length"))?
            .extract()?;
        let byte_len = ((bit_len + 8) / 8).max(1);
        let kwargs = pyo3::types::PyDict::new(py);
        kwargs.set_item(pyo3::intern!(py, "signed"), true)?;
        let raw: Vec<u8> = value
            .call_method(
                pyo3::intern!(py, "to_bytes"),
                (byte_len, pyo3::intern!(py, "big")),
                Some(&kwargs),
            )?
            .extract()?;
        self.encoder
            .encode(&synta::Integer::from_bytes(&raw))
            .map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode an INTEGER from a PyInteger object
    ///
    /// Args:
    ///     value: The PyInteger object to encode
    fn encode_integer_object(&mut self, value: &PyInteger) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode an OCTET STRING
    ///
    /// Args:
    ///     data: The bytes to encode
    fn encode_octet_string(&mut self, data: &[u8]) -> PyResult<()> {
        let octet_string = synta::OctetStringRef::new(data);
        self.encoder.encode(&octet_string).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode an OCTET STRING from a PyOctetString object
    ///
    /// Args:
    ///     value: The PyOctetString object to encode
    fn encode_octet_string_object(&mut self, value: &PyOctetString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode an OBJECT IDENTIFIER
    ///
    /// Args:
    ///     value: The PyObjectIdentifier to encode
    fn encode_oid(&mut self, value: &PyObjectIdentifier) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode an OBJECT IDENTIFIER from a PyObjectIdentifier object (alias for encode_oid)
    fn encode_oid_object(&mut self, value: &PyObjectIdentifier) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a BIT STRING
    ///
    /// Args:
    ///     value: The PyBitString to encode
    fn encode_bit_string(&mut self, value: &PyBitString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a BIT STRING from a PyBitString object (alias for encode_bit_string)
    fn encode_bit_string_object(&mut self, value: &PyBitString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a BOOLEAN
    ///
    /// Args:
    ///     value: The boolean value to encode
    fn encode_boolean(&mut self, value: bool) -> PyResult<()> {
        let boolean = synta::Boolean::new(value);
        self.encoder.encode(&boolean).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a BOOLEAN from a PyBoolean object
    ///
    /// Args:
    ///     value: The PyBoolean object to encode
    fn encode_boolean_object(&mut self, value: &PyBoolean) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a UTCTime
    ///
    /// Args:
    ///     value: The PyUtcTime to encode
    fn encode_utc_time(&mut self, value: &PyUtcTime) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a UTCTime from a PyUtcTime object (alias for encode_utc_time)
    fn encode_utc_time_object(&mut self, value: &PyUtcTime) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a GeneralizedTime
    ///
    /// Args:
    ///     value: The PyGeneralizedTime to encode
    fn encode_generalized_time(&mut self, value: &PyGeneralizedTime) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a GeneralizedTime from a PyGeneralizedTime object (alias for encode_generalized_time)
    fn encode_generalized_time_object(&mut self, value: &PyGeneralizedTime) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a REAL
    ///
    /// Args:
    ///     value: The float value to encode
    fn encode_real(&mut self, value: f64) -> PyResult<()> {
        let real = synta::Real::new(value);
        self.encoder.encode(&real).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a REAL from a PyReal object
    ///
    /// Args:
    ///     value: The PyReal object to encode
    fn encode_real_object(&mut self, value: &PyReal) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a NULL
    fn encode_null(&mut self) -> PyResult<()> {
        self.encoder.encode(&synta::Null).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a NULL from a PyNull object
    fn encode_null_object(&mut self, _value: &PyNull) -> PyResult<()> {
        self.encoder.encode(&synta::Null).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a UTF8String
    ///
    /// Args:
    ///     value: The string to encode
    fn encode_utf8_string(&mut self, value: &str) -> PyResult<()> {
        let s = synta::Utf8String::new(value.to_string());
        self.encoder.encode(&s).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a UTF8String from a PyUtf8String object
    fn encode_utf8_string_object(&mut self, value: &PyUtf8String) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a PrintableString
    ///
    /// Args:
    ///     value: The string to encode (must contain only PrintableString-valid characters)
    fn encode_printable_string(&mut self, value: &str) -> PyResult<()> {
        let s = synta::PrintableString::new(value.to_string()).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!("Invalid PrintableString: {:?}", e))
        })?;
        self.encoder.encode(&s).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a PrintableString from a PyPrintableString object
    fn encode_printable_string_object(&mut self, value: &PyPrintableString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode an IA5String
    ///
    /// Args:
    ///     value: The string to encode (must be ASCII only)
    fn encode_ia5_string(&mut self, value: &str) -> PyResult<()> {
        let s = synta::IA5String::new(value.to_string()).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!("Invalid IA5String: {:?}", e))
        })?;
        self.encoder.encode(&s).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode an IA5String from a PyIA5String object
    fn encode_ia5_string_object(&mut self, value: &PyIA5String) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a NumericString (tag 18)
    fn encode_numeric_string(&mut self, value: &str) -> PyResult<()> {
        let s = synta::NumericString::new(value.to_string()).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!("Invalid NumericString: {:?}", e))
        })?;
        self.encoder.encode(&s).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a NumericString from a PyNumericString object
    fn encode_numeric_string_object(&mut self, value: &PyNumericString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a TeletexString / T61String (tag 20) from raw bytes
    fn encode_teletex_string(&mut self, data: &[u8]) -> PyResult<()> {
        let s = synta::TeletexString::new(data.to_vec());
        self.encoder.encode(&s).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a TeletexString from a PyTeletexString object
    fn encode_teletex_string_object(&mut self, value: &PyTeletexString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a VisibleString (tag 26)
    fn encode_visible_string(&mut self, value: &str) -> PyResult<()> {
        let s = synta::VisibleString::new(value.to_string()).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!("Invalid VisibleString: {:?}", e))
        })?;
        self.encoder.encode(&s).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a VisibleString from a PyVisibleString object
    fn encode_visible_string_object(&mut self, value: &PyVisibleString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a GeneralString (tag 27) from raw bytes
    fn encode_general_string(&mut self, data: &[u8]) -> PyResult<()> {
        let s = synta::GeneralString::new(data.to_vec());
        self.encoder.encode(&s).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a GeneralString from a PyGeneralString object
    fn encode_general_string_object(&mut self, value: &PyGeneralString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a UniversalString (tag 28)
    fn encode_universal_string(&mut self, value: &str) -> PyResult<()> {
        let s = synta::UniversalString::new(value.to_string());
        self.encoder.encode(&s).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a UniversalString from a PyUniversalString object
    fn encode_universal_string_object(&mut self, value: &PyUniversalString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a BMPString (tag 30)
    ///
    /// Raises ValueError if the string contains code points outside the BMP (> U+FFFF).
    fn encode_bmp_string(&mut self, value: &str) -> PyResult<()> {
        let s = synta::BmpString::new(value.to_string()).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!("Invalid BMPString: {:?}", e))
        })?;
        self.encoder.encode(&s).map_err(SyntaErr)?;
        Ok(())
    }

    /// Encode a BMPString from a PyBmpString object
    fn encode_bmp_string_object(&mut self, value: &PyBmpString) -> PyResult<()> {
        self.encoder.encode(&value.inner).map_err(SyntaErr)?;
        Ok(())
    }

    /// Wrap pre-encoded bytes in a SEQUENCE TLV (tag 0x30) and append to the output.
    ///
    /// Use this to build SEQUENCE structures: encode the inner elements into a
    /// separate Encoder, call ``finish()`` to get the inner bytes, then pass those
    /// bytes to ``encode_sequence()`` on the outer Encoder.
    ///
    /// Example::
    ///
    ///     inner = synta.Encoder(synta.Encoding.DER)
    ///     inner.encode_integer(42)
    ///     inner.encode_boolean(True)
    ///     outer = synta.Encoder(synta.Encoding.DER)
    ///     outer.encode_sequence(inner.finish())
    ///     result = outer.finish()
    fn encode_sequence(&mut self, inner_bytes: &[u8]) -> PyResult<()> {
        use synta::tag::TAG_SEQUENCE;
        self.encoder
            .write_tag(synta::Tag::universal_constructed(TAG_SEQUENCE))
            .map_err(SyntaErr)?;
        self.encoder
            .write_length(inner_bytes.len())
            .map_err(SyntaErr)?;
        self.encoder.write_bytes(inner_bytes);
        Ok(())
    }

    /// Wrap pre-encoded bytes in a SET TLV (tag 0x31) and append to the output.
    fn encode_set(&mut self, inner_bytes: &[u8]) -> PyResult<()> {
        use synta::tag::TAG_SET;
        self.encoder
            .write_tag(synta::Tag::universal_constructed(TAG_SET))
            .map_err(SyntaErr)?;
        self.encoder
            .write_length(inner_bytes.len())
            .map_err(SyntaErr)?;
        self.encoder.write_bytes(inner_bytes);
        Ok(())
    }

    /// Wrap pre-encoded bytes with an explicit context tag ``[tag_num]`` and append
    /// to the output.
    ///
    /// ``tag_class`` must be one of ``"Context"`` (default for ``[N]`` notation),
    /// ``"Application"``, or ``"Private"``.
    ///
    /// Example — encode ``[1] EXPLICIT SEQUENCE { ... }``::
    ///
    ///     enc.encode_explicit_tag(1, "Context", sequence_bytes)
    fn encode_explicit_tag(
        &mut self,
        tag_num: u32,
        tag_class: &str,
        inner_bytes: &[u8],
    ) -> PyResult<()> {
        use pyo3::exceptions::PyValueError;
        use synta::tag::TagClass;

        let class = match tag_class {
            "Context" => TagClass::ContextSpecific,
            "Application" => TagClass::Application,
            "Private" => TagClass::Private,
            other => {
                return Err(PyValueError::new_err(format!(
                    "Unknown tag class '{}'; expected 'Context', 'Application', or 'Private'",
                    other
                )))
            }
        };
        // Explicit tags are always constructed (they wrap another TLV)
        let tag = synta::Tag::new(class, true, tag_num);
        self.encoder.write_tag(tag).map_err(SyntaErr)?;
        self.encoder
            .write_length(inner_bytes.len())
            .map_err(SyntaErr)?;
        self.encoder.write_bytes(inner_bytes);
        Ok(())
    }

    /// Wrap pre-encoded value bytes with an implicit tag and append to the output.
    ///
    /// For implicit tagging, the original type tag is replaced with the given
    /// tag; the value bytes (the content, *not* the full original TLV) are
    /// written as-is.  Set ``is_constructed=True`` when the underlying type is
    /// a SEQUENCE, SET, or other constructed type.
    ///
    /// ``tag_class`` must be one of ``"Context"`` (default for ``[N]`` notation),
    /// ``"Application"``, or ``"Private"``.
    ///
    /// Example — encode ``[1] IMPLICIT INTEGER`` with value bytes ``b'\\x2a'``:
    ///
    ///     enc.encode_implicit_tag(1, "Context", False, b'\\x2a')
    fn encode_implicit_tag(
        &mut self,
        tag_num: u32,
        tag_class: &str,
        is_constructed: bool,
        value_bytes: &[u8],
    ) -> PyResult<()> {
        use pyo3::exceptions::PyValueError;
        use synta::tag::TagClass;

        let class = match tag_class {
            "Context" => TagClass::ContextSpecific,
            "Application" => TagClass::Application,
            "Private" => TagClass::Private,
            other => {
                return Err(PyValueError::new_err(format!(
                    "Unknown tag class '{other}'; expected 'Context', 'Application', or 'Private'",
                )))
            }
        };
        let tag = synta::Tag::new(class, is_constructed, tag_num);
        self.encoder.write_tag(tag).map_err(SyntaErr)?;
        self.encoder
            .write_length(value_bytes.len())
            .map_err(SyntaErr)?;
        self.encoder.write_bytes(value_bytes);
        Ok(())
    }

    /// Finish encoding and return the encoded bytes
    ///
    /// Returns:
    ///     The encoded ASN.1 data as bytes
    fn finish<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, pyo3::types::PyBytes>> {
        // Capture encoding before the move so the replacement uses the same rules.
        let encoding = self.encoder.encoding();
        let encoder = std::mem::replace(&mut self.encoder, Encoder::new(encoding));
        let bytes = encoder.finish().map_err(SyntaErr)?;
        Ok(pyo3::types::PyBytes::new(py, &bytes))
    }

    fn __repr__(&self) -> String {
        format!("Encoder(encoding={:?})", self.encoder.encoding())
    }
}