1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

use grapl_graph_descriptions::graph_description::IpAddress as InnerIpAddress;
use grapl_graph_descriptions::graph_description::IpAddressBuilder;
use pyo3::create_exception;
use pyo3::prelude::*;

create_exception!(ip_address_node, IpAddressBuilderError, pyo3::exceptions::ValueError);

#[pyclass]
#[derive(Clone)]
pub struct IpAddressNode {
    pub(crate) inner_node: InnerIpAddress,
}

impl<'source> pyo3::FromPyObject<'source> for IpAddressNode {
    fn extract(ob: &'source pyo3::types::PyAny) -> pyo3::PyResult<Self> {
        Ok(
            pyo3::PyTryFrom::try_from(ob).map(|x: &Self| x.clone())?
        )
    }
}

#[pyclass]
#[derive(Clone, Default)]
pub struct IpAddressNodeBuilder {
    builder: IpAddressBuilder,
}

#[pymethods]
impl IpAddressNodeBuilder {
    #[new]
    fn new(
        obj: &PyRawObject,
    ) {
        obj.init(
            Self::default()
        )
    }

    pub fn with_ip_address(&mut self, ip_address_id: String) -> Self {
        self.builder.ip_address(ip_address_id);
        self.clone()
    }

    pub fn with_first_timestamp_ms(&mut self, timestamp: u64) -> Self {
        self.builder.first_seen_timestamp(timestamp);
        self.clone()
    }

    pub fn with_last_seen_timestamp_ms(&mut self, timestamp: u64) -> Self {
        self.builder.last_seen_timestamp(timestamp);
        self.clone()
    }

    pub fn build(&self) -> PyResult<IpAddressNode> {

        let built_node = match self.builder.build() {
            Ok(built_node) => built_node,
            Err(e) => {
                return Err(
                    PyErr::new::<IpAddressBuilderError, _>(format!("{}", e))
                )
            }
        };

        Ok(
            IpAddressNode {
                inner_node: built_node
            }
        )
    }
}