synta-python 0.2.5

//! Python wrapper for ASN.1 decoder

use pyo3::prelude::*;
use pyo3::types::PyBytes;

use synta::traits::Decode;
use synta::{Decoder, Encoding, TagClass};

use super::error::SyntaErr;

fn parse_tag_class(s: &str) -> PyResult<TagClass> {
    match s {
        "Universal" => Ok(TagClass::Universal),
        "Context" => Ok(TagClass::ContextSpecific),
        "Application" => Ok(TagClass::Application),
        "Private" => Ok(TagClass::Private),
        other => Err(pyo3::exceptions::PyValueError::new_err(format!(
            "Unknown tag class '{other}'; expected 'Universal', 'Context', 'Application', or 'Private'",
        ))),
    }
}

fn tag_class_str(class: TagClass) -> &'static str {
    match class {
        TagClass::Universal => "Universal",
        TagClass::Application => "Application",
        TagClass::ContextSpecific => "Context",
        TagClass::Private => "Private",
    }
}

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 Decoder
///
/// Decodes ASN.1 DER/BER encoded data.
///
/// Example:
///     >>> decoder = Decoder(b'\\x02\\x01\\x2A')          # DER is the default
///     >>> decoder = Decoder(b'\\x02\\x01\\x2A', Encoding.DER)  # explicit
///     >>> integer = decoder.decode_integer()
///     >>> print(integer.to_int())
///     42
#[pyclass(name = "Decoder")]
pub struct PyDecoder {
    // Store owned data to ensure lifetime validity
    data: Vec<u8>,
    // Position in the data (to track progress)
    position: usize,
    // Encoding type
    encoding: Encoding,
}

impl PyDecoder {
    /// Advance the position by `by` bytes, asserting the invariant in debug builds.
    fn advance(&mut self, by: usize) {
        self.position += by;
        debug_assert!(
            self.position <= self.data.len(),
            "decoder position {} overshot data length {}",
            self.position,
            self.data.len()
        );
    }

    /// Decode a single owned value of type `T`, then advance the parent position.
    ///
    /// This encapsulates the repeated pattern:
    ///   1. create a sub-`Decoder` over the remaining slice
    ///   2. call `T::decode` to decode a value
    ///   3. advance `self.position` by however many bytes were consumed
    ///
    /// `T` must implement `Decode<'a>` for *any* lifetime `'a` (i.e. it must own
    /// its decoded representation, not borrow from the decoder's input buffer).
    fn decode_one<T>(&mut self) -> PyResult<T>
    where
        T: for<'a> Decode<'a>,
    {
        let mut decoder = Decoder::new(&self.data[self.position..], self.encoding);
        let result = T::decode(&mut decoder).map_err(SyntaErr)?;
        self.advance(decoder.position());
        Ok(result)
    }
}

#[pymethods]
impl PyDecoder {
    /// Create a new decoder
    ///
    /// Args:
    ///     data: The ASN.1 encoded data to decode
    ///     encoding: The encoding type (DER, BER, or CER); defaults to DER
    #[new]
    #[pyo3(signature = (data, encoding=None))]
    fn new(data: Vec<u8>, encoding: Option<PyEncoding>) -> Self {
        Self {
            data,
            position: 0,
            encoding: encoding.unwrap_or(PyEncoding::DER).into(),
        }
    }

    /// Decode an INTEGER
    fn decode_integer(&mut self) -> PyResult<PyInteger> {
        self.decode_one::<synta::Integer>()
            .map(|inner| PyInteger { inner })
    }

    /// Decode an OCTET STRING
    fn decode_octet_string(&mut self) -> PyResult<PyOctetString> {
        self.decode_one::<synta::OctetString>()
            .map(|inner| PyOctetString { inner })
    }

    /// Decode an OBJECT IDENTIFIER
    fn decode_oid(&mut self) -> PyResult<PyObjectIdentifier> {
        self.decode_one::<synta::ObjectIdentifier>()
            .map(PyObjectIdentifier::from_oid)
    }

    /// Decode a BIT STRING
    fn decode_bit_string(&mut self) -> PyResult<PyBitString> {
        self.decode_one::<synta::BitString>()
            .map(|inner| PyBitString { inner })
    }

    /// Decode a BOOLEAN
    fn decode_boolean(&mut self) -> PyResult<PyBoolean> {
        self.decode_one::<synta::Boolean>()
            .map(|inner| PyBoolean { inner })
    }

    /// Decode a UTCTime
    fn decode_utc_time(&mut self) -> PyResult<PyUtcTime> {
        self.decode_one::<synta::UtcTime>()
            .map(|inner| PyUtcTime { inner })
    }

    /// Decode a GeneralizedTime
    fn decode_generalized_time(&mut self) -> PyResult<PyGeneralizedTime> {
        self.decode_one::<synta::GeneralizedTime>()
            .map(|inner| PyGeneralizedTime { inner })
    }

    /// Decode a NULL
    fn decode_null(&mut self) -> PyResult<PyNull> {
        self.decode_one::<synta::Null>()?;
        Ok(PyNull)
    }

    /// Decode a REAL
    fn decode_real(&mut self) -> PyResult<PyReal> {
        self.decode_one::<synta::Real>()
            .map(|inner| PyReal { inner })
    }

    /// Decode a UTF8String
    fn decode_utf8_string(&mut self) -> PyResult<PyUtf8String> {
        self.decode_one::<synta::Utf8String>()
            .map(|inner| PyUtf8String { inner })
    }

    /// Decode a PrintableString
    fn decode_printable_string(&mut self) -> PyResult<PyPrintableString> {
        self.decode_one::<synta::PrintableString>()
            .map(|inner| PyPrintableString { inner })
    }

    /// Decode an IA5String
    fn decode_ia5_string(&mut self) -> PyResult<PyIA5String> {
        self.decode_one::<synta::IA5String>()
            .map(|inner| PyIA5String { inner })
    }

    /// Decode a NumericString (tag 18)
    fn decode_numeric_string(&mut self) -> PyResult<PyNumericString> {
        self.decode_one::<synta::NumericString>()
            .map(|inner| PyNumericString { inner })
    }

    /// Decode a TeletexString / T61String (tag 20)
    fn decode_teletex_string(&mut self) -> PyResult<PyTeletexString> {
        self.decode_one::<synta::TeletexString>()
            .map(|inner| PyTeletexString { inner })
    }

    /// Decode a VisibleString (tag 26)
    fn decode_visible_string(&mut self) -> PyResult<PyVisibleString> {
        self.decode_one::<synta::VisibleString>()
            .map(|inner| PyVisibleString { inner })
    }

    /// Decode a GeneralString (tag 27)
    fn decode_general_string(&mut self) -> PyResult<PyGeneralString> {
        self.decode_one::<synta::GeneralString>()
            .map(|inner| PyGeneralString { inner })
    }

    /// Decode a UniversalString (tag 28)
    fn decode_universal_string(&mut self) -> PyResult<PyUniversalString> {
        self.decode_one::<synta::UniversalString>()
            .map(|inner| PyUniversalString { inner })
    }

    /// Decode a BMPString (tag 30)
    fn decode_bmp_string(&mut self) -> PyResult<PyBmpString> {
        self.decode_one::<synta::BmpString>()
            .map(|inner| PyBmpString { inner })
    }

    /// Decode the next ASN.1 element as a Python ``str``.
    ///
    /// Accepts any of the standard ASN.1 string types and converts to a
    /// native Python ``str`` using the appropriate encoding for each tag:
    ///
    /// * UTF8String (12), PrintableString (19), IA5String (22),
    ///   NumericString (18), VisibleString (26), GeneralString (27) —
    ///   decoded as UTF-8 (lossily; all are ASCII or UTF-8 subsets in practice)
    /// * TeletexString / T61String (20) — each byte decoded as a Latin-1
    ///   code point (U+0000–U+00FF)
    /// * BMPString (30) — decoded from UCS-2 big-endian (BMP only)
    /// * UniversalString (28) — decoded from UCS-4 big-endian
    ///
    /// Raises ``ValueError`` for any other tag, and ``EOFError`` when no data
    /// remains.
    ///
    /// This is the single-call replacement for the duck-typing pattern that
    /// ``decode_any()`` requires:
    ///
    /// ```python
    /// # Before — three-way probe after decode_any():
    /// val = decoder.decode_any()
    /// if hasattr(val, 'as_str'):
    ///     s = val.as_str()
    /// elif hasattr(val, 'to_bytes'):   # TeletexString / BmpString
    ///     s = val.to_bytes().decode('latin-1')
    /// else:
    ///     raise ValueError(f"expected a string type, got {type(val)}")
    ///
    /// # After — one call, correct encoding for all nine string types:
    /// s = decoder.decode_any_str()
    /// ```
    fn decode_any_str<'py>(
        &mut self,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, pyo3::types::PyString>> {
        use pyo3::exceptions::{PyEOFError, PyValueError};
        use pyo3::types::PyString;
        use synta::tag::*;

        if self.position >= self.data.len() {
            return Err(PyEOFError::new_err("No more data"));
        }

        let (tag_num, val_bytes, advance) = {
            let mut decoder = Decoder::new(&self.data[self.position..], self.encoding);
            let tag = decoder.read_tag().map_err(SyntaErr)?;
            let length = decoder.read_length().map_err(SyntaErr)?;
            let len = length.definite().map_err(SyntaErr)?;
            let val = decoder.read_bytes(len).map_err(SyntaErr)?.to_vec();
            (tag.number(), val, decoder.position())
        };

        let s: String = match tag_num {
            TAG_UTF8_STRING | TAG_PRINTABLE_STRING | TAG_IA5_STRING | TAG_VISIBLE_STRING
            | TAG_NUMERIC_STRING | TAG_GENERAL_STRING => {
                String::from_utf8_lossy(&val_bytes).into_owned()
            }

            // TeletexString / T61String: treat each byte as a Latin-1 codepoint.
            TAG_TELETEX_STRING => val_bytes.iter().map(|&b| b as char).collect(),

            // BMPString: UCS-2 big-endian (Basic Multilingual Plane only).
            TAG_BMP_STRING => val_bytes
                .chunks_exact(2)
                .filter_map(|pair| char::from_u32(u16::from_be_bytes([pair[0], pair[1]]) as u32))
                .collect(),

            // UniversalString: UCS-4 big-endian.
            TAG_UNIVERSAL_STRING => val_bytes
                .chunks_exact(4)
                .filter_map(|quad| {
                    char::from_u32(u32::from_be_bytes([quad[0], quad[1], quad[2], quad[3]]))
                })
                .collect(),

            other => {
                return Err(PyValueError::new_err(format!(
                    "tag {other} is not an ASN.1 string type; \
                     expected UTF8String (12), PrintableString (19), IA5String (22), \
                     NumericString (18), TeletexString (20), VisibleString (26), \
                     GeneralString (27), BMPString (30), or UniversalString (28)"
                )));
            }
        };

        self.advance(advance);
        Ok(PyString::new(py, &s))
    }

    /// Enter a SEQUENCE and return a child Decoder positioned over its contents.
    ///
    /// Reads and validates the SEQUENCE tag and length from the current position,
    /// advances this decoder past the whole SEQUENCE, and returns a new Decoder
    /// whose data is exactly the content bytes of the SEQUENCE.  Call typed
    /// ``decode_*`` methods on the child decoder to read elements one by one;
    /// use ``is_empty()`` to detect the end of the SEQUENCE.
    fn decode_sequence(&mut self) -> PyResult<PyDecoder> {
        use pyo3::exceptions::PyValueError;
        use synta::tag::TAG_SEQUENCE;
        use synta::TagClass;

        let (data, advance) = {
            let mut decoder = Decoder::new(&self.data[self.position..], self.encoding);
            let tag = decoder.read_tag().map_err(SyntaErr)?;
            if tag.class() != TagClass::Universal
                || tag.number() != TAG_SEQUENCE
                || !tag.is_constructed()
            {
                return Err(PyValueError::new_err(format!(
                    "Expected SEQUENCE (0x30), got tag {:?}",
                    tag
                )));
            }
            let length = decoder.read_length().map_err(SyntaErr)?;
            let len = length.definite().map_err(SyntaErr)?;
            let data = decoder.read_bytes(len).map_err(SyntaErr)?.to_vec();
            (data, decoder.position())
            // decoder drops here, releasing the borrow of self.data
        };
        self.advance(advance);
        Ok(PyDecoder {
            data,
            position: 0,
            encoding: self.encoding,
        })
    }

    /// Consume an explicit context tag ``[tag_num]`` and return a child Decoder
    /// positioned over the tagged content.
    ///
    /// Reads the outer explicit tag+length from the current position, checks that the
    /// tag number matches ``tag_num`` and that the class is Context-specific, then
    /// returns a new Decoder whose data is the content of the tagged element.
    fn decode_explicit_tag(&mut self, tag_num: u32) -> PyResult<PyDecoder> {
        use pyo3::exceptions::PyValueError;
        use synta::TagClass;

        let (data, advance) = {
            let mut decoder = Decoder::new(&self.data[self.position..], self.encoding);
            let tag = decoder.read_tag().map_err(SyntaErr)?;
            if tag.class() != TagClass::ContextSpecific || tag.number() != tag_num {
                return Err(PyValueError::new_err(format!(
                    "Expected explicit context tag [{}], got {:?}",
                    tag_num, tag
                )));
            }
            let length = decoder.read_length().map_err(SyntaErr)?;
            let len = length.definite().map_err(SyntaErr)?;
            let data = decoder.read_bytes(len).map_err(SyntaErr)?.to_vec();
            (data, decoder.position())
            // decoder drops here, releasing the borrow of self.data
        };
        self.advance(advance);
        Ok(PyDecoder {
            data,
            position: 0,
            encoding: self.encoding,
        })
    }

    /// Enter a SET and return a child Decoder positioned over its contents.
    ///
    /// Symmetric counterpart to ``decode_sequence()`` for SET (tag 0x31).
    fn decode_set(&mut self) -> PyResult<PyDecoder> {
        use pyo3::exceptions::PyValueError;
        use synta::tag::TAG_SET;

        let (data, advance) = {
            let mut decoder = Decoder::new(&self.data[self.position..], self.encoding);
            let tag = decoder.read_tag().map_err(SyntaErr)?;
            if tag.class() != TagClass::Universal
                || tag.number() != TAG_SET
                || !tag.is_constructed()
            {
                return Err(PyValueError::new_err(format!(
                    "Expected SET (0x31), got tag {:?}",
                    tag
                )));
            }
            let length = decoder.read_length().map_err(SyntaErr)?;
            let len = length.definite().map_err(SyntaErr)?;
            let data = decoder.read_bytes(len).map_err(SyntaErr)?.to_vec();
            (data, decoder.position())
        };
        self.advance(advance);
        Ok(PyDecoder {
            data,
            position: 0,
            encoding: self.encoding,
        })
    }

    /// Consume an implicit tag ``[tag_num]`` and return a child Decoder over the
    /// value content bytes.
    ///
    /// For implicit tags the original type tag is replaced by a context-specific
    /// (or other class) tag.  This method validates the tag, strips the T+L, and
    /// returns a Decoder over the raw value bytes.  The caller must supply the
    /// expected ``tag_class`` (``"Context"``, ``"Application"``, ``"Private"``,
    /// or ``"Universal"``).
    ///
    /// To decode the content as a typed value after stripping, pass the returned
    /// bytes to a new ``Decoder``:
    ///
    /// ```python
    /// raw = decoder.decode_implicit_tag(0, "Context")
    /// inner = synta.Decoder(raw, synta.Encoding.DER)
    /// value = inner.decode_integer()
    /// ```
    fn decode_implicit_tag(&mut self, tag_num: u32, tag_class: &str) -> PyResult<PyDecoder> {
        use pyo3::exceptions::PyValueError;

        let class = parse_tag_class(tag_class)?;

        let (data, advance) = {
            let mut decoder = Decoder::new(&self.data[self.position..], self.encoding);
            let tag = decoder.read_tag().map_err(SyntaErr)?;
            if tag.class() != class || tag.number() != tag_num {
                return Err(PyValueError::new_err(format!(
                    "Expected implicit tag [{tag_num} {tag_class}], got {:?}",
                    tag
                )));
            }
            let length = decoder.read_length().map_err(SyntaErr)?;
            let len = length.definite().map_err(SyntaErr)?;
            let data = decoder.read_bytes(len).map_err(SyntaErr)?.to_vec();
            (data, decoder.position())
        };
        self.advance(advance);
        Ok(PyDecoder {
            data,
            position: 0,
            encoding: self.encoding,
        })
    }

    /// Inspect the next tag without consuming any bytes.
    ///
    /// Returns ``(tag_number, tag_class, is_constructed)`` where ``tag_class``
    /// is one of ``"Universal"``, ``"Context"``, ``"Application"``, or
    /// ``"Private"``.  Raises ``EOFError`` if no data remains.
    ///
    /// Use this to implement CHOICE dispatch or to test for optional fields:
    ///
    /// ```python
    /// tag_num, tag_class, _ = decoder.peek_tag()
    /// if tag_class == "Context" and tag_num == 0:
    ///     version = decoder.decode_explicit_tag(0)
    /// ```
    fn peek_tag(&self) -> PyResult<(u32, String, bool)> {
        use pyo3::exceptions::PyEOFError;
        if self.position >= self.data.len() {
            return Err(PyEOFError::new_err("No more data to peek"));
        }
        let mut decoder = Decoder::new(&self.data[self.position..], self.encoding);
        let tag = decoder.read_tag().map_err(SyntaErr)?;
        Ok((
            tag.number(),
            tag_class_str(tag.class()).to_string(),
            tag.is_constructed(),
        ))
    }

    /// Read the complete next TLV (tag + length + value) as raw bytes and
    /// advance past it.
    ///
    /// Useful for capturing an element whose type is unknown or whose decoding
    /// should be deferred:
    ///
    /// ```python
    /// raw_tlv = decoder.decode_raw_tlv()
    /// inner = synta.Decoder(raw_tlv, synta.Encoding.DER)
    /// ```
    fn decode_raw_tlv<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        use pyo3::exceptions::PyEOFError;
        if self.position >= self.data.len() {
            return Err(PyEOFError::new_err("No more data"));
        }
        let total = {
            let mut decoder = Decoder::new(&self.data[self.position..], self.encoding);
            decoder.read_tag().map_err(SyntaErr)?;
            let length = decoder.read_length().map_err(SyntaErr)?;
            let content_len = length.definite().map_err(SyntaErr)?;
            decoder.read_bytes(content_len).map_err(SyntaErr)?;
            decoder.position()
        };
        let end = self.position + total;
        let bytes = PyBytes::new(py, &self.data[self.position..end]);
        self.advance(total);
        Ok(bytes)
    }

    /// Return all remaining bytes from the current position as a Python bytes object.
    ///
    /// After ``decode_implicit_tag()`` strips the context tag and length, the
    /// child ``Decoder``'s data contains only the raw value bytes (no TLV
    /// header).  For **constructed** implicit types (e.g.
    /// ``[0] IMPLICIT SEQUENCE``) the child can be decoded element-by-element
    /// using the ``decode_*`` methods.  For **primitive** implicit types (e.g.
    /// ``[2] IMPLICIT IA5String`` for dNSName, ``[7] IMPLICIT OCTET STRING``
    /// for iPAddress) the child data is just the bare value bytes —
    /// ``remaining_bytes()`` retrieves them directly without a new
    /// encode/decode round-trip.
    ///
    /// Example — decode a dNSName GeneralName::
    ///
    ///     tag_num, tag_class, _ = decoder.peek_tag()
    ///     if tag_class == "Context" and tag_num == 2:
    ///         child = decoder.decode_implicit_tag(2, "Context")
    ///         dns_name = child.remaining_bytes().decode("ascii")
    fn remaining_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.data[self.position..])
    }

    /// Decode the next ASN.1 element as a dynamically-typed Python object.
    ///
    /// Sequences and Sets are returned as Python lists.
    /// Tagged elements are returned as ``TaggedElement`` objects; use
    /// ``isinstance(x, synta.TaggedElement)`` to detect them.
    fn decode_any(&mut self, py: Python) -> PyResult<Py<PyAny>> {
        let (py_obj, advance) = {
            let mut decoder = Decoder::new(&self.data[self.position..], self.encoding);
            let element = synta::Element::decode(&mut decoder).map_err(SyntaErr)?;
            let advance = decoder.position();
            let py_obj = super::types::element_to_pyobject(element, py)?;
            (py_obj, advance)
        };
        self.advance(advance);
        Ok(py_obj)
    }

    /// Check if there is more data to decode
    fn is_empty(&self) -> bool {
        self.position == self.data.len()
    }

    /// Get the current position in the data
    fn position(&self) -> usize {
        self.position
    }

    /// Get the remaining bytes
    fn remaining(&self) -> usize {
        self.data.len().saturating_sub(self.position)
    }

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