malwaredb/
cart.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use anyhow::Result;
4use pyo3::conversion::IntoPyObject;
5use pyo3::prelude::*;
6use pyo3::types::{IntoPyDict, PyDict, PyTuple};
7use serde_json::{Map, Value};
8
9/// Encase a binary in a `CaRT` file
10///
11/// # Errors
12///
13/// There should not be any errors, but the underlying library can't guarantee that. However, any data
14/// passed to it is correct.
15#[pyfunction]
16pub fn create_cart(data: &[u8]) -> Result<Vec<u8>> {
17    malwaredb_client::encode_to_cart(data)
18}
19
20/// Decode a `CaRT` file into the original binary and metadata
21///
22/// # Errors
23///
24/// There should not be any errors, but the underlying library can't guarantee that. However, any data
25/// passed to it is correct.
26#[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    // Todo: Figure out how to do this without so many possible errors.
31
32    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/// Registers this `CaRT` module to an exiting module
56///
57/// # Errors
58///
59/// Should not have errors but not yet known
60#[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/// Convenience function to handle converting from the hash map to a dictionary.
69#[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        // Todo: Figure out how to use PyOxide's `serde` feature, since it seems to offer nothing.
75        // Alternatively, figure out how to do this conversion manually.
76        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}