pyaxml_rs/

lib.rs

#[cfg(feature = "python")]
use pyo3::prelude::*;
#[cfg(feature = "python")]
use pyo3::types::PyBytes;

pub mod arsc;
pub mod axml;
mod error;
pub mod json_serde;
#[allow(clippy::all, dead_code)]
mod proto;
mod proto_compute;
mod proto_conv;
mod public;
mod resource_map;
mod string_pool;
mod typed_value;
mod xml_element;

/// Re-export the error type so downstream crates can name it.
pub use error::AxmlError;

#[cfg(feature = "python")]
use arsc::Arsc;
#[cfg(feature = "python")]
use axml::Axml;

/// Python-exposed AXML class.
#[cfg(feature = "python")]
#[pyclass(name = "AXML")]
pub struct PyAxml {
    inner: Axml,
}

#[cfg(feature = "python")]
#[pymethods]
impl PyAxml {
    #[new]
    pub fn new() -> Self {
        PyAxml { inner: Axml::new() }
    }

    /// Parse binary AXML data. Returns an AXML instance.
    #[staticmethod]
    pub fn from_axml(data: &[u8]) -> PyResult<Self> {
        let inner = Axml::from_axml(data)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        Ok(PyAxml { inner })
    }

    /// Serialize AXML to binary bytes.
    pub fn pack<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        let data = self.inner.pack();
        Ok(PyBytes::new(py, &data))
    }

    /// Convert binary AXML to an XML string.
    pub fn to_xml(&self) -> PyResult<String> {
        self.inner
            .to_xml()
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
    }

    /// Serialize AXML to protobuf text format.
    pub fn to_proto_text(&self) -> String {
        self.inner.to_proto_text()
    }

    /// Serialize AXML to indented protobuf text format.
    pub fn to_proto_text_pretty(&self) -> String {
        self.inner.to_proto_text_pretty()
    }

    /// Recalculate chunk sizes on next `pack()`.  Call before `pack()` for a valid binary.
    pub fn compute(&mut self) {
        self.inner.compute();
    }

    /// Apply new string-pool bytes without updating the chunk total-size header.
    ///
    /// `new_strings` is a list of raw string-data bytes (as stored in
    /// `StringBlock.data` in the proto: no length prefix, no null terminator).
    /// `pack()` after this call returns an intentionally "corrupted" binary where
    /// string content changed but the chunk-size field is stale.
    /// Call `compute()` before `pack()` to get a correctly-sized binary instead.
    pub fn apply_stringblocks_no_compute(&mut self, new_strings: Vec<Vec<u8>>) {
        self.inner.apply_stringblocks_no_compute(new_strings);
    }

    /// Build AXML from an XML string.
    pub fn from_xml(&mut self, xml: &str) -> PyResult<()> {
        self.inner
            .from_xml(xml)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
    }

    /// Return the stored AXML total-file-size value (bytes 4..8 of the chunk header).
    /// `pack()` writes this verbatim; `compute()` sets it to the correct value.
    pub fn get_total_size(&self) -> u32 {
        self.inner.total_size
    }

    /// Set the AXML total-file-size header field.
    /// Use this to produce intentionally malformed output without calling `compute()`.
    pub fn set_total_size(&mut self, size: u32) {
        self.inner.total_size = size;
    }

    /// Number of strings in the string pool.
    pub fn string_count(&self) -> usize {
        self.inner.stringblocks.inner.strings.len()
    }

    /// Decode a string from the string pool by index.
    pub fn get_string(&self, idx: u32) -> PyResult<String> {
        self.inner
            .stringblocks
            .inner
            .decode(idx)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
    }

    /// Serialize AXML to binary protobuf bytes (proto encoding, not Android binary XML).
    pub fn to_proto_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        Ok(PyBytes::new(py, &self.inner.to_proto_bytes()))
    }

    /// Deserialize AXML from binary protobuf bytes produced by `to_proto_bytes`.
    #[staticmethod]
    pub fn from_proto_bytes(data: &[u8]) -> PyResult<Self> {
        let inner = Axml::from_proto_bytes(data)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        Ok(PyAxml { inner })
    }

    /// Write binary proto file (proto encoding, not Android binary XML).
    pub fn export_proto_file(&self, path: &str) -> PyResult<()> {
        self.inner
            .export_proto_file(path)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
    }

    /// Load AXML from a binary proto file written by `export_proto_file`.
    #[staticmethod]
    pub fn from_proto_file(path: &str) -> PyResult<Self> {
        let inner = Axml::from_proto_file(path)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        Ok(PyAxml { inner })
    }

    /// Write the string pool as a binary proto file.
    pub fn export_stringblocks_proto_file(&self, path: &str) -> PyResult<()> {
        self.inner
            .export_stringblocks_proto_file(path)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
    }

    /// Replace the string pool from a binary proto file.
    pub fn import_stringblocks_proto_file(&mut self, path: &str) -> PyResult<()> {
        self.inner
            .import_stringblocks_proto_file(path)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
    }

    pub fn __repr__(&self) -> String {
        format!(
            "AXML(strings={}, elements={})",
            self.inner.stringblocks.inner.strings.len(),
            self.inner.elements.len()
        )
    }
}

#[cfg(feature = "python")]
impl Default for PyAxml {
    fn default() -> Self {
        Self::new()
    }
}

/// Python-exposed ARSC class.
#[cfg(feature = "python")]
#[pyclass(name = "ARSC")]
pub struct PyArsc {
    inner: Arsc,
}

#[cfg(feature = "python")]
#[pymethods]
impl PyArsc {
    /// Parse a binary resources.arsc file (plain file or APK entry).
    #[staticmethod]
    pub fn from_axml(data: &[u8]) -> PyResult<Self> {
        let inner = Arsc::from_axml(data)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        Ok(PyArsc { inner })
    }

    /// List resource entries grouped by locale.
    /// Pass `language=None` for all locales, `language=Some("default")` for no-locale
    /// entries, or a language tag like `Some("en")` or `Some("fr-FR")`.
    #[pyo3(signature = (language=None))]
    pub fn list_packages(&self, language: Option<&str>) -> String {
        self.inner.list_packages(language)
    }

    /// Return the stored ARSC total-file-size header field (bytes 4..8).
    /// `pack()` writes this verbatim; `compute()` recalculates it.
    pub fn get_total_size(&self) -> u32 {
        self.inner.total_size
    }

    /// Set the ARSC total-file-size header field.
    pub fn set_total_size(&mut self, size: u32) {
        self.inner.total_size = size;
    }

    /// Return the stored package-count header field (bytes 8..12).
    /// `pack()` writes this verbatim; `compute()` sets it to the actual package count.
    pub fn get_header_package_count(&self) -> u32 {
        self.inner.package_count
    }

    /// Set the package-count header field.
    pub fn set_header_package_count(&mut self, count: u32) {
        self.inner.package_count = count;
    }

    /// Return the stored chunk-size header field for package at `idx` (bytes 4..8 of the package chunk).
    /// Returns 0 if `idx` is out of range.
    pub fn get_pkg_chunk_size(&self, idx: usize) -> u32 {
        self.inner
            .packages
            .get(idx)
            .map(|p| p.chunk_size)
            .unwrap_or(0)
    }

    /// Set the chunk-size header field for package at `idx`.
    /// No-op if `idx` is out of range.
    pub fn set_pkg_chunk_size(&mut self, idx: usize, size: u32) {
        if let Some(pkg) = self.inner.packages.get_mut(idx) {
            pkg.chunk_size = size;
        }
    }

    /// Number of packages in this ARSC file.
    pub fn package_count(&self) -> usize {
        self.inner.packages.len()
    }

    /// Recalculate string-pool chunk sizes on next `pack()`.
    ///
    /// `recursive=true` (default): recomputes global pool + all package type/key pools.
    /// `recursive=false`: recomputes only the global string pool.
    #[pyo3(signature = (recursive = true))]
    pub fn compute(&mut self, recursive: bool) {
        self.inner.compute(recursive);
    }

    /// Serialize ARSC back to binary `resources.arsc` format.
    pub fn pack<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        Ok(PyBytes::new(py, &self.inner.pack()))
    }

    /// Add a new resource entry (type_name / name → file_path string value).
    ///
    /// Finds or creates `type_name` in the first package's type string pool,
    /// adds `name` to the key string pool, adds `file_path` to the global
    /// string pool, then appends TypeSpec + TypeType chunks for the default
    /// locale.  Returns the new resource ID (0x7fTTEEEE).
    pub fn add_resource(&mut self, type_name: &str, name: &str, file_path: &str) -> u32 {
        self.inner.add_resource(type_name, name, file_path)
    }

    /// Serialize ARSC to protobuf text format.
    pub fn to_proto_text(&self) -> String {
        self.inner.to_proto_text()
    }

    /// Serialize ARSC to indented protobuf text format.
    pub fn to_proto_text_pretty(&self) -> String {
        self.inner.to_proto_text_pretty()
    }

    /// Serialize ARSC to binary protobuf bytes (proto encoding, not Android binary ARSC format).
    pub fn to_proto_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        Ok(PyBytes::new(py, &self.inner.to_proto_bytes()))
    }

    /// Deserialize ARSC from binary protobuf bytes produced by `to_proto_bytes`.
    #[staticmethod]
    pub fn from_proto_bytes(data: &[u8]) -> PyResult<Self> {
        let inner = Arsc::from_proto_bytes(data)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        Ok(PyArsc { inner })
    }

    /// Write binary proto file.
    pub fn export_proto_file(&self, path: &str) -> PyResult<()> {
        self.inner
            .export_proto_file(path)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
    }

    /// Load ARSC from a binary proto file written by `export_proto_file`.
    #[staticmethod]
    pub fn from_proto_file(path: &str) -> PyResult<Self> {
        let inner = Arsc::from_proto_file(path)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        Ok(PyArsc { inner })
    }

    /// Write package `idx` as a binary proto file.
    pub fn export_pkg_proto_file(&self, idx: usize, path: &str) -> PyResult<()> {
        self.inner
            .export_pkg_proto_file(idx, path)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
    }

    /// Replace package `idx` from a binary proto file.
    pub fn import_pkg_proto_file(&mut self, idx: usize, path: &str) -> PyResult<()> {
        self.inner
            .import_pkg_proto_file(idx, path)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
    }

    pub fn __repr__(&self) -> String {
        format!("ARSC(packages={})", self.inner.packages.len())
    }
}

// ──────────────────────── compute bytes-in/bytes-out ─────────────────────────

#[cfg(feature = "python")]
#[pyfunction]
fn compute_string_block_proto<'py>(
    py: Python<'py>,
    data: &[u8],
    is_utf8: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    proto_compute::compute_string_block_bytes(data, is_utf8)
        .map(|b| PyBytes::new(py, &b))
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[cfg(feature = "python")]
#[pyfunction]
fn compute_string_blocks_proto<'py>(
    py: Python<'py>,
    data: &[u8],
    recursive: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    proto_compute::compute_string_blocks_bytes(data, recursive)
        .map(|b| PyBytes::new(py, &b))
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[cfg(feature = "python")]
#[pyfunction]
fn compute_resource_map_proto<'py>(py: Python<'py>, data: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
    proto_compute::compute_resource_map_bytes(data)
        .map(|b| PyBytes::new(py, &b))
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[cfg(feature = "python")]
#[pyfunction]
fn compute_xml_element_proto<'py>(py: Python<'py>, data: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
    proto_compute::compute_xml_element_bytes(data)
        .map(|b| PyBytes::new(py, &b))
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[cfg(feature = "python")]
#[pyfunction]
fn compute_arsc_res_table_config_proto<'py>(
    py: Python<'py>,
    data: &[u8],
) -> PyResult<Bound<'py, PyBytes>> {
    proto_compute::compute_arsc_res_table_config_bytes(data)
        .map(|b| PyBytes::new(py, &b))
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[cfg(feature = "python")]
#[pyfunction]
fn compute_arsc_res_type_proto<'py>(
    py: Python<'py>,
    data: &[u8],
    recursive: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    proto_compute::compute_arsc_res_type_bytes(data, recursive)
        .map(|b| PyBytes::new(py, &b))
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[cfg(feature = "python")]
#[pyfunction]
fn compute_arsc_res_table_package_proto<'py>(
    py: Python<'py>,
    data: &[u8],
    recursive: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    proto_compute::compute_arsc_res_table_package_bytes(data, recursive)
        .map(|b| PyBytes::new(py, &b))
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[cfg(feature = "python")]
#[pyfunction]
fn compute_axml_proto_bytes<'py>(
    py: Python<'py>,
    data: &[u8],
    recursive: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    proto_compute::compute_axml_bytes(data, recursive)
        .map(|b| PyBytes::new(py, &b))
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[cfg(feature = "python")]
#[pyfunction]
fn compute_arsc_proto_bytes<'py>(
    py: Python<'py>,
    data: &[u8],
    recursive: bool,
) -> PyResult<Bound<'py, PyBytes>> {
    proto_compute::compute_arsc_bytes(data, recursive)
        .map(|b| PyBytes::new(py, &b))
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

/// The pyaxml_rs Python extension module.
#[cfg(feature = "python")]
#[pymodule]
fn _pyaxml(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyAxml>()?;
    m.add_class::<PyArsc>()?;
    m.add_function(wrap_pyfunction!(compute_string_block_proto, m)?)?;
    m.add_function(wrap_pyfunction!(compute_string_blocks_proto, m)?)?;
    m.add_function(wrap_pyfunction!(compute_resource_map_proto, m)?)?;
    m.add_function(wrap_pyfunction!(compute_xml_element_proto, m)?)?;
    m.add_function(wrap_pyfunction!(compute_arsc_res_table_config_proto, m)?)?;
    m.add_function(wrap_pyfunction!(compute_arsc_res_type_proto, m)?)?;
    m.add_function(wrap_pyfunction!(compute_arsc_res_table_package_proto, m)?)?;
    m.add_function(wrap_pyfunction!(compute_axml_proto_bytes, m)?)?;
    m.add_function(wrap_pyfunction!(compute_arsc_proto_bytes, m)?)?;
    Ok(())
}