grapl_graph_descriptions_py/
asset_node.rs

1use grapl_graph_descriptions::graph_description::Asset as InnerAsset;
2use grapl_graph_descriptions::graph_description::AssetBuilder;
3use pyo3::create_exception;
4use pyo3::prelude::*;
5
6create_exception!(asset_node, AssetBuilderError, pyo3::exceptions::ValueError);
7
8#[pyclass]
9#[derive(Clone)]
10pub struct AssetNode {
11    pub(crate) inner_node: InnerAsset,
12}
13
14impl<'source> pyo3::FromPyObject<'source> for AssetNode {
15    fn extract(ob: &'source pyo3::types::PyAny) -> pyo3::PyResult<Self> {
16        Ok(
17            pyo3::PyTryFrom::try_from(ob).map(|x: &Self| x.clone())?
18        )
19    }
20}
21
22#[pyclass]
23#[derive(Clone, Default)]
24pub struct AssetNodeBuilder {
25    builder: AssetBuilder,
26}
27
28#[pymethods]
29impl AssetNodeBuilder {
30    #[new]
31    fn new(
32        obj: &PyRawObject,
33    ) {
34        obj.init(
35            Self::default()
36        )
37    }
38
39    pub fn with_asset_id(&mut self, asset_id: Option<String>) -> Self {
40        self.builder.asset_id(asset_id);
41        self.clone()
42    }
43
44    pub fn with_hostname(&mut self, hostname: Option<String>) -> Self {
45        self.builder.hostname(hostname);
46        self.clone()
47    }
48
49    pub fn with_first_timestamp_ms(&mut self, timestamp: u64) -> Self {
50        self.builder.first_seen_timestamp(timestamp);
51        self.clone()
52    }
53
54    pub fn with_last_seen_timestamp_ms(&mut self, timestamp: u64) -> Self {
55        self.builder.last_seen_timestamp(timestamp);
56        self.clone()
57    }
58
59    pub fn build(&self) -> PyResult<AssetNode> {
60
61        let built_node = match self.builder.build() {
62            Ok(built_node) => built_node,
63            Err(e) => {
64                return Err(
65                    PyErr::new::<AssetBuilderError, _>(format!("{}", e))
66                )
67            }
68        };
69
70        Ok(
71            AssetNode {
72                inner_node: built_node
73            }
74        )
75    }
76}