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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

use grapl_graph_descriptions::graph_description::{Graph as InnerGraph, Node};

use pyo3::prelude::*;
use crate::process_node::ProcessNode;
use crate::file_node::FileNode;
use crate::asset_node::AssetNode;
use crate::network_connection_node::NetworkConnectionNode;
use crate::process_inbound_connection_node::ProcessInboundConnectionNode;
use crate::ip_port_node::IpPortNode;
use crate::ip_address_node::IpAddressNode;

#[pyclass]
#[derive(Clone)]
pub struct Graph {
    pub inner_graph: InnerGraph,
}

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


#[pymethods]
impl Graph {
    #[new]
    fn new(
        obj: &PyRawObject,
    ) {
        obj.init(
            Self {
                inner_graph: InnerGraph::new(0)
            }
        )
    }

    pub fn is_empty(&self) -> bool {
        self.inner_graph.is_empty()
    }

    pub fn merge(&mut self, other: &Self) {
        self.inner_graph.merge(&other.inner_graph);
    }

    pub fn add_process(&mut self, node: ProcessNode) {
        self.inner_graph.add_node(node.inner_node);
    }

    pub fn add_file(&mut self, node: FileNode) {
        self.inner_graph.add_node(node.inner_node);
    }

    pub fn add_asset(&mut self, node: AssetNode) {
        self.inner_graph.add_node(node.inner_node);
    }

    pub fn add_network_connection(&mut self, node: NetworkConnectionNode) {
        self.inner_graph.add_node(node.inner_node);
    }

    pub fn add_process_inbound_connection(&mut self, node: ProcessInboundConnectionNode) {
        self.inner_graph.add_node(node.inner_node);
    }

    pub fn add_process_outbound_connection(&mut self, node: ProcessInboundConnectionNode) {
        self.inner_graph.add_node(node.inner_node);
    }

    pub fn add_ip_port(&mut self, node: IpPortNode) {
        self.inner_graph.add_node(node.inner_node);
    }

    pub fn add_ip_address(&mut self, node: IpAddressNode) {
        self.inner_graph.add_node(node.inner_node);
    }

    pub fn add_edge(
        &mut self,
        edge_name: String,
        from: String,
        to: String,
    ) {
        self.inner_graph.add_edge(
            edge_name,
            from,
            to
        )
    }
}