grapl_graph_descriptions_py/
network_connection_node.rs

1use grapl_graph_descriptions::graph_description::NetworkConnection as InnerNetworkConnection;
2use grapl_graph_descriptions::graph_description::NetworkConnectionBuilder;
3
4use pyo3::create_exception;
5use pyo3::prelude::*;
6
7create_exception!(network_connection_node, NetworkConnectionBuilderError, pyo3::exceptions::ValueError);
8
9#[pyclass]
10#[derive(Clone)]
11pub struct NetworkConnectionNode {
12    pub(crate) inner_node: InnerNetworkConnection,
13}
14
15impl<'source> pyo3::FromPyObject<'source> for NetworkConnectionNode {
16    fn extract(ob: &'source pyo3::types::PyAny) -> pyo3::PyResult<Self> {
17        Ok(
18            pyo3::PyTryFrom::try_from(ob).map(|x: &Self| x.clone())?
19        )
20    }
21}
22
23#[pyclass]
24#[derive(Clone, Default)]
25pub struct NetworkConnectionNodeBuilder {
26    builder: NetworkConnectionBuilder,
27}
28
29#[pymethods]
30impl NetworkConnectionNodeBuilder {
31    #[new]
32    fn new(
33        obj: &PyRawObject,
34    ) {
35        obj.init(
36            Self::default()
37        )
38    }
39
40
41    pub fn with_src_ip_address(&mut self, src_ip_address: String) -> Self {
42        self.builder.src_ip_address(src_ip_address);
43        self.clone()
44    }
45
46    pub fn with_dst_ip_address(&mut self, dst_ip_address: String) -> Self {
47        self.builder.dst_ip_address(dst_ip_address);
48        self.clone()
49    }
50
51
52    pub fn with_src_port(&mut self, src_port: u16) -> Self {
53        self.builder.src_port(src_port);
54        self.clone()
55    }
56
57    pub fn with_dst_port(&mut self, dst_port: u16) -> Self {
58        self.builder.dst_port(dst_port);
59        self.clone()
60    }
61
62    pub fn with_protocol(&mut self, protocol: String) -> Self {
63        self.builder.protocol(protocol);
64        self.clone()
65    }
66
67    pub fn with_state(&mut self, state: u32) -> Self {
68        self.builder.state(state);
69        self.clone()
70    }
71
72    pub fn with_created_timestamp(&mut self, created_timestamp: u64) -> Self {
73        self.builder.created_timestamp(created_timestamp);
74        self.clone()
75    }
76
77    pub fn with_terminated_timestamp(&mut self, terminated_timestamp: u64) -> Self {
78        self.builder.terminated_timestamp(terminated_timestamp);
79        self.clone()
80    }
81
82    pub fn with_last_seen_timestamp(&mut self, last_seen_timestamp: u64) -> Self {
83        self.builder.last_seen_timestamp(last_seen_timestamp);
84        self.clone()
85    }
86
87    pub fn build(&self) -> PyResult<NetworkConnectionNode> {
88
89        let built_node = match self.builder.build() {
90            Ok(built_node) => built_node,
91            Err(e) => {
92                return Err(
93                    PyErr::new::<NetworkConnectionBuilderError, _>(format!("{}", e))
94                )
95            }
96        };
97
98        Ok(
99            NetworkConnectionNode {
100                inner_node: built_node
101            }
102        )
103    }
104}