1use anyhow::Result;
4use pyo3::conversion::IntoPyObject;
5use pyo3::prelude::*;
6use pyo3::types::{IntoPyDict, PyDict, PyTuple};
7use serde_json::{Map, Value};
8
9#[pyfunction]
16pub fn create_cart(data: &[u8]) -> Result<Vec<u8>> {
17 malwaredb_client::encode_to_cart(data)
18}
19
20#[pyfunction]
27pub fn decode_cart(py: Python<'_>, data: &[u8]) -> Result<Py<PyTuple>> {
28 let (binary, header, footer) = malwaredb_client::decode_from_cart(data)?;
29
30 let header = if let Some(header) = header {
33 convert_json_map_to_dict(py, header)?
34 } else {
35 PyDict::new(py).into()
36 };
37
38 let footer = if let Some(footer) = footer {
39 convert_json_map_to_dict(py, footer)?
40 } else {
41 PyDict::new(py).into()
42 };
43
44 Ok(PyTuple::new(
45 py,
46 &[
47 binary.into_pyobject(py)?,
48 header.into_pyobject(py)?.into_any(),
49 footer.into_pyobject(py)?.into_any(),
50 ],
51 )?
52 .unbind())
53}
54
55#[cfg(not(feature = "rust_lib"))]
61pub fn register_cart_module(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
62 let cart_module = PyModule::new(parent_module.py(), "cart")?;
63 cart_module.add_function(wrap_pyfunction!(create_cart, &cart_module)?)?;
64 cart_module.add_function(wrap_pyfunction!(decode_cart, &cart_module)?)?;
65 parent_module.add_submodule(&cart_module)
66}
67
68#[inline]
70fn convert_json_map_to_dict(py: Python<'_>, map: Map<String, Value>) -> Result<Py<PyDict>> {
71 let mut vec_: Vec<(String, Py<PyAny>)> = Vec::with_capacity(map.len());
72
73 for (key, value) in map {
74 if let Some(val) = value.as_f64() {
77 vec_.push((key, Py::<PyAny>::from(val.into_pyobject(py)?)));
78 continue;
79 }
80 if let Some(val) = value.as_str() {
81 vec_.push((key, Py::<PyAny>::from(val.into_pyobject(py)?)));
82 continue;
83 }
84 if let Some(val) = value.as_u64() {
85 vec_.push((key, Py::<PyAny>::from(val.into_pyobject(py)?)));
86 continue;
87 }
88 if let Some(val) = value.as_i64() {
89 vec_.push((key, Py::<PyAny>::from(val.into_pyobject(py)?)));
90 continue;
91 }
92 println!("Parsing JSON to dict, found key {key} with unexpected or unknown value type: {value:?}");
93 }
94
95 Ok(vec_.into_py_dict(py)?.unbind())
96}