grapl_graph_descriptions/
asset.rs

1use crate::graph_description::Asset;
2use crate::node::NodeT;
3
4use log::warn;
5use serde_json::{json, Value};
6use uuid::Uuid;
7
8impl Asset {
9    pub fn new(
10        asset_id: impl Into<Option<String>>,
11        hostname: impl Into<Option<String>>,
12        mac_address: impl Into<Option<String>>,
13        first_seen_timestamp: u64,
14        last_seen_timestamp: u64,
15    ) -> Self {
16        let asset_id = asset_id.into();
17        let hostname = hostname.into();
18
19        if asset_id.is_none() && hostname.is_none() {
20            panic!("AssetID or Hostname must be provided for ProcessOutboundConnection");
21        }
22
23        Self {
24            node_key: Uuid::new_v4().to_string(),
25            asset_id,
26            hostname,
27            mac_address: mac_address.into(),
28            first_seen_timestamp,
29            last_seen_timestamp,
30        }
31    }
32
33    pub fn into_json(self) -> Value {
34        let asset_id = self.asset_id.as_ref().unwrap();
35        let mut j = json!({
36            "node_key": self.node_key,
37            "asset_id": asset_id,
38            "dgraph.type": "Asset",
39        });
40
41        if self.first_seen_timestamp != 0 {
42            j["first_seen_timestamp"] = self.first_seen_timestamp.into();
43        }
44
45        if self.last_seen_timestamp != 0 {
46            j["last_seen_timestamp"] = self.last_seen_timestamp.into();
47        }
48
49        if let Some(hostname) = self.hostname {
50            j["hostname"] = Value::from(hostname.clone());
51        }
52
53        if let Some(mac_address) = self.mac_address {
54            j["mac_address"] = Value::from(mac_address.clone());
55        }
56
57        j
58    }
59}
60
61impl NodeT for Asset {
62    fn get_asset_id(&self) -> Option<&str> {
63        self.asset_id.as_ref().map(|asset_id| asset_id.as_str())
64    }
65
66    fn set_asset_id(&mut self, asset_id: impl Into<String>) {
67        self.asset_id = Some(asset_id.into());
68    }
69
70    fn get_node_key(&self) -> &str {
71        &self.node_key
72    }
73
74    fn set_node_key(&mut self, node_key: impl Into<String>) {
75        self.node_key = node_key.into();
76    }
77
78    fn merge(&mut self, other: &Self) -> bool {
79        if self.node_key != other.node_key {
80            warn!("Attempted to merge two Asset Nodes with differing node_keys");
81            return false;
82        }
83
84        let mut merged = false;
85
86        if self.asset_id.is_none() && other.asset_id.is_some() {
87            merged = true;
88            self.asset_id = other.asset_id.clone();
89        }
90
91        if self.hostname.is_none() && other.hostname.is_some() {
92            merged = true;
93            self.hostname = other.hostname.clone();
94        }
95
96        if self.mac_address.is_none() && other.mac_address.is_some() {
97            merged = true;
98            self.mac_address = other.mac_address.clone();
99        }
100
101        merged
102    }
103
104    fn merge_into(&mut self, other: Self) -> bool {
105        if self.node_key != other.node_key {
106            warn!("Attempted to merge two Asset Nodes with differing node_keys");
107            return false;
108        }
109
110        let mut merged = false;
111
112        if self.asset_id.is_none() && other.asset_id.is_some() {
113            self.asset_id = other.asset_id;
114            merged = true;
115        }
116
117        if self.hostname.is_none() && other.hostname.is_some() {
118            self.hostname = other.hostname;
119            merged = true;
120        }
121
122        if self.mac_address.is_none() && other.mac_address.is_some() {
123            self.mac_address = other.mac_address;
124            merged = true;
125        }
126
127        if other.first_seen_timestamp != 0 && self.first_seen_timestamp > other.first_seen_timestamp
128        {
129            self.first_seen_timestamp = other.first_seen_timestamp;
130            merged = true;
131        }
132
133        if other.last_seen_timestamp != 0 && self.last_seen_timestamp < other.last_seen_timestamp {
134            self.last_seen_timestamp = other.last_seen_timestamp;
135            merged = true;
136        }
137
138        merged
139    }
140}