malwaredb-client-py 0.3.4

// SPDX-License-Identifier: Apache-2.0

use anyhow::Result;
use pyo3::conversion::IntoPyObject;
use pyo3::prelude::*;
use pyo3::types::{IntoPyDict, PyDict, PyTuple};
use serde_json::{Map, Value};

/// Encase a binary in a `CaRT` file
///
/// # Errors
///
/// There should not be any errors, but the underlying library can't guarantee that. However, any data
/// passed to it is correct.
#[pyfunction]
pub fn create_cart(data: &[u8]) -> Result<Vec<u8>> {
    malwaredb_client::encode_to_cart(data)
}

/// Decode a `CaRT` file into the original binary and metadata
///
/// # Errors
///
/// There should not be any errors, but the underlying library can't guarantee that. However, any data
/// passed to it is correct.
#[pyfunction]
pub fn decode_cart(py: Python<'_>, data: &[u8]) -> Result<Py<PyTuple>> {
    let (binary, header, footer) = malwaredb_client::decode_from_cart(data)?;

    // Todo: Figure out how to do this without so many possible errors.

    let header = if let Some(header) = header {
        convert_json_map_to_dict(py, header)?
    } else {
        PyDict::new(py).into()
    };

    let footer = if let Some(footer) = footer {
        convert_json_map_to_dict(py, footer)?
    } else {
        PyDict::new(py).into()
    };

    Ok(PyTuple::new(
        py,
        &[
            binary.into_pyobject(py)?,
            header.into_pyobject(py)?.into_any(),
            footer.into_pyobject(py)?.into_any(),
        ],
    )?
    .unbind())
}

/// Registers this `CaRT` module to an exiting module
///
/// # Errors
///
/// Should not have errors but not yet known
#[cfg(not(feature = "rust_lib"))]
pub fn register_cart_module(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
    let cart_module = PyModule::new(parent_module.py(), "cart")?;
    cart_module.add_function(wrap_pyfunction!(create_cart, &cart_module)?)?;
    cart_module.add_function(wrap_pyfunction!(decode_cart, &cart_module)?)?;
    parent_module.add_submodule(&cart_module)
}

/// Convenience function to handle converting from the hash map to a dictionary.
#[inline]
fn convert_json_map_to_dict(py: Python<'_>, map: Map<String, Value>) -> Result<Py<PyDict>> {
    let mut vec_: Vec<(String, Py<PyAny>)> = Vec::with_capacity(map.len());

    for (key, value) in map {
        // Todo: Figure out how to use PyOxide's `serde` feature, since it seems to offer nothing.
        // Alternatively, figure out how to do this conversion manually.
        if let Some(val) = value.as_f64() {
            vec_.push((key, Py::<PyAny>::from(val.into_pyobject(py)?)));
            continue;
        }
        if let Some(val) = value.as_str() {
            vec_.push((key, Py::<PyAny>::from(val.into_pyobject(py)?)));
            continue;
        }
        if let Some(val) = value.as_u64() {
            vec_.push((key, Py::<PyAny>::from(val.into_pyobject(py)?)));
            continue;
        }
        if let Some(val) = value.as_i64() {
            vec_.push((key, Py::<PyAny>::from(val.into_pyobject(py)?)));
            continue;
        }
        println!("Parsing JSON to dict, found key {key} with unexpected or unknown value type: {value:?}");
    }

    Ok(vec_.into_py_dict(py)?.unbind())
}