malwaredb/
types.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use std::path::PathBuf;
4
5use pyo3::{pyclass, pymethods, PyResult};
6use uuid::Uuid;
7
8/// Information about a label
9#[pyclass(frozen)]
10#[derive(Debug, Clone)]
11pub struct Label {
12    /// Label ID
13    #[pyo3(get)]
14    pub id: u64,
15
16    /// Label value
17    #[pyo3(get)]
18    pub name: String,
19
20    /// Label parent
21    #[pyo3(get)]
22    pub parent: Option<String>,
23}
24
25#[pymethods]
26impl Label {
27    /// Label printable representation
28    #[must_use]
29    pub fn __repr__(&self) -> String {
30        if let Some(parent) = &self.parent {
31            format!("Label {}: {parent}/{}", self.id, self.name)
32        } else {
33            format!("Label {}: {}", self.id, self.name)
34        }
35    }
36
37    /// Label information, just the name
38    #[must_use]
39    pub fn __str__(&self) -> String {
40        self.name.clone()
41    }
42}
43
44impl From<malwaredb_client::malwaredb_api::Label> for Label {
45    fn from(label: malwaredb_client::malwaredb_api::Label) -> Self {
46        Self {
47            id: label.id,
48            name: label.name,
49            parent: label.parent,
50        }
51    }
52}
53
54/// Information about the Malware DB server
55#[pyclass(frozen)]
56#[derive(Debug, Clone)]
57pub struct ServerInfo {
58    /// Operating System used
59    #[pyo3(get)]
60    pub os_name: String,
61
62    /// Memory footprint
63    #[pyo3(get)]
64    pub memory_used: String,
65
66    /// MDB version
67    #[pyo3(get)]
68    pub mdb_version: String,
69
70    /// Type and version of the database
71    #[pyo3(get)]
72    pub db_version: String,
73
74    /// Size of the database on disk
75    #[pyo3(get)]
76    pub db_size: String,
77
78    /// Total number of samples
79    #[pyo3(get)]
80    pub num_samples: u64,
81
82    /// Total number of users
83    #[pyo3(get)]
84    pub num_users: u32,
85
86    /// Uptime of Malware DB in a human-readable format
87    #[pyo3(get)]
88    pub uptime: String,
89
90    /// The name of the Malware DB instance
91    #[pyo3(get)]
92    pub instance_name: String,
93}
94
95#[pymethods]
96impl ServerInfo {
97    /// Server Info printable representation
98    #[must_use]
99    pub fn __repr__(&self) -> String {
100        format!(
101            "MalwareDB {} running on {} for {}",
102            self.mdb_version, self.os_name, self.uptime
103        )
104    }
105
106    /// Server info, just the version
107    #[must_use]
108    pub fn __str__(&self) -> String {
109        format!("MalwareDB {}", self.mdb_version)
110    }
111}
112
113impl From<malwaredb_client::malwaredb_api::ServerInfo> for ServerInfo {
114    fn from(value: malwaredb_client::malwaredb_api::ServerInfo) -> Self {
115        Self {
116            os_name: value.os_name,
117            memory_used: value.memory_used,
118            mdb_version: value.mdb_version.to_string(),
119            db_version: value.db_version,
120            db_size: value.db_size,
121            num_samples: value.num_samples,
122            num_users: value.num_users,
123            uptime: value.uptime,
124            instance_name: value.instance_name,
125        }
126    }
127}
128
129/// Information about a sample source
130#[pyclass(frozen)]
131#[derive(Debug, Clone)]
132pub struct Source {
133    /// Identifier of the source
134    #[pyo3(get)]
135    pub id: u32,
136
137    /// Name of the source
138    #[pyo3(get)]
139    pub name: String,
140
141    /// Description of the source
142    #[pyo3(get)]
143    pub description: Option<String>,
144
145    /// URL of the source, or where the files were found
146    #[pyo3(get)]
147    pub url: Option<String>,
148
149    /// Creation date or first acquisition date of or from the source
150    #[pyo3(get)]
151    pub first_acquisition: String,
152
153    /// Whether the source holds malware
154    #[pyo3(get)]
155    pub malicious: Option<bool>,
156}
157
158#[pymethods]
159impl Source {
160    /// Source printable representation
161    #[must_use]
162    pub fn __repr__(&self) -> String {
163        let url = if let Some(url) = &self.url {
164            format!(" from {url}")
165        } else {
166            String::new()
167        };
168
169        let desc = if let Some(desc) = &self.description {
170            format!(" -- {desc}")
171        } else {
172            String::new()
173        };
174
175        format!("{}({}){url}{desc}", self.name, self.id)
176    }
177
178    /// Simpler source printable representation
179    #[must_use]
180    pub fn __str__(&self) -> String {
181        self.name.clone()
182    }
183}
184
185impl From<malwaredb_client::malwaredb_api::SourceInfo> for Source {
186    fn from(value: malwaredb_client::malwaredb_api::SourceInfo) -> Self {
187        Self {
188            id: value.id,
189            name: value.name,
190            description: value.description,
191            url: value.url,
192            first_acquisition: value.first_acquisition.to_rfc3339(),
193            malicious: value.malicious,
194        }
195    }
196}
197
198/// Information about a file type supported by Malware DB
199#[pyclass(frozen)]
200#[derive(Debug, Clone)]
201pub struct SupportedFileType {
202    /// Common name of the file type
203    #[pyo3(get)]
204    pub name: String,
205
206    /// Magic number bytes in hex of the file type
207    #[pyo3(get)]
208    pub magic: Vec<String>,
209
210    /// Whether the file type is executable
211    #[pyo3(get)]
212    pub is_executable: bool,
213
214    /// Description of the file type
215    #[pyo3(get)]
216    pub description: Option<String>,
217}
218
219#[pymethods]
220impl SupportedFileType {
221    /// Supported file type as a printable representation
222    #[must_use]
223    pub fn __repr__(&self) -> String {
224        format!("{}, starting with {}", self.name, self.magic.join(" or "))
225    }
226
227    /// Supported file type as a printable simpler representation
228    #[must_use]
229    pub fn __str__(&self) -> String {
230        self.name.clone()
231    }
232}
233
234impl From<malwaredb_client::malwaredb_api::SupportedFileType> for SupportedFileType {
235    fn from(value: malwaredb_client::malwaredb_api::SupportedFileType) -> Self {
236        Self {
237            name: value.name,
238            magic: value.magic,
239            is_executable: value.is_executable,
240            description: value.description,
241        }
242    }
243}
244
245/// Information about the user's account
246#[pyclass(frozen)]
247#[derive(Debug, Clone)]
248pub struct UserInfo {
249    /// User's numeric ID
250    #[pyo3(get)]
251    pub id: u32,
252
253    /// User's name
254    #[pyo3(get)]
255    pub username: String,
256
257    /// User's group memberships, if any
258    #[pyo3(get)]
259    pub groups: Vec<String>,
260
261    /// User's available sample sources, if any
262    #[pyo3(get)]
263    pub sources: Vec<String>,
264
265    /// If the user is an admin
266    #[pyo3(get)]
267    pub is_admin: bool,
268
269    /// When the account was created
270    #[pyo3(get)]
271    pub created: String,
272
273    /// User has read-only access, perhaps a guest or demo account
274    #[pyo3(get)]
275    pub is_readonly: bool,
276}
277
278#[pymethods]
279impl UserInfo {
280    /// Simple user information
281    #[must_use]
282    pub fn __str__(&self) -> String {
283        self.username.clone()
284    }
285}
286
287impl From<malwaredb_client::malwaredb_api::GetUserInfoResponse> for UserInfo {
288    fn from(value: malwaredb_client::malwaredb_api::GetUserInfoResponse) -> Self {
289        Self {
290            id: value.id,
291            username: value.username,
292            groups: value.groups,
293            sources: value.sources,
294            is_admin: value.is_admin,
295            created: value.created.to_rfc3339(),
296            is_readonly: value.is_readonly,
297        }
298    }
299}
300
301/// Search results
302#[pyclass(frozen)]
303#[derive(Debug, Clone)]
304pub struct SearchResults {
305    /// Hashes of results
306    #[pyo3(get)]
307    pub hashes: Vec<String>,
308
309    /// Total samples matching search
310    #[pyo3(get)]
311    pub results: u64,
312
313    /// Identifier for continuing to the next batch of search results
314    #[pyo3(get)]
315    pub pagination: Option<Uuid>,
316
317    /// Message from the server, if available
318    #[pyo3(get)]
319    pub message: Option<String>,
320}
321
322#[pymethods]
323impl SearchResults {
324    /// Search results as a printable representation
325    #[must_use]
326    pub fn __repr__(&self) -> String {
327        if self.pagination.is_some() && self.results > self.hashes.len() as u64 {
328            format!(
329                "Search results showing {} of {} available hashes",
330                self.hashes.len(),
331                self.results
332            )
333        } else if let Some(message) = &self.message {
334            format!(
335                "{message}: Search results showing {} hashes",
336                self.hashes.len()
337            )
338        } else {
339            format!("Search results showing {} hashes", self.hashes.len())
340        }
341    }
342
343    /// Search results as a printable simpler representation
344    #[must_use]
345    pub fn __str__(&self) -> String {
346        if let Some(message) = &self.message {
347            format!(
348                "{message}: Search results with {} hashes",
349                self.hashes.len()
350            )
351        } else {
352            format!("Search results with {} hashes", self.hashes.len())
353        }
354    }
355}
356
357impl From<malwaredb_client::malwaredb_api::SearchResponse> for SearchResults {
358    fn from(value: malwaredb_client::malwaredb_api::SearchResponse) -> Self {
359        Self {
360            hashes: value.hashes,
361            results: value.total_results,
362            pagination: value.pagination,
363            message: value.message,
364        }
365    }
366}
367
368/// Discovered local Malware DB server
369#[pyclass(frozen)]
370#[derive(Debug, Clone)]
371pub struct DiscoveredServer {
372    /// Server IP or domain
373    #[pyo3(get)]
374    pub host: String,
375
376    /// Server port
377    #[pyo3(get)]
378    pub port: u16,
379
380    /// If the server expects an encrypted connection
381    #[pyo3(get)]
382    pub ssl: bool,
383
384    /// Malware DB server name
385    #[pyo3(get)]
386    pub name: String,
387
388    /// Additional server information, if retrievable
389    #[pyo3(get)]
390    pub info: Option<ServerInfo>,
391}
392
393#[pymethods]
394impl DiscoveredServer {
395    /// Connect to a discovered server using login credentials
396    ///
397    /// # Errors
398    ///
399    /// Returns an error if the server becomes unreachable or if credentials are incorrect
400    #[pyo3(signature = (username, password, save = false, cert_path = None))]
401    pub fn connect(
402        &self,
403        username: String,
404        password: String,
405        save: bool,
406        cert_path: Option<PathBuf>,
407    ) -> PyResult<crate::MalwareDBClient> {
408        crate::MalwareDBClient::login(self.__str__(), username, password, save, cert_path)
409    }
410
411    /// Discovered server as a string with server name
412    #[must_use]
413    pub fn __repr__(&self) -> String {
414        if self.ssl {
415            format!("{} at https://{}:{}", self.name, self.host, self.port)
416        } else {
417            format!("{} at http://{}:{}", self.name, self.host, self.port)
418        }
419    }
420
421    /// Discovered server as a URL string
422    #[must_use]
423    pub fn __str__(&self) -> String {
424        if self.ssl {
425            format!("https://{}:{}", self.host, self.port)
426        } else {
427            format!("http://{}:{}", self.host, self.port)
428        }
429    }
430}
431
432impl From<malwaredb_client::MalwareDBServer> for DiscoveredServer {
433    fn from(value: malwaredb_client::MalwareDBServer) -> Self {
434        let info = value.server_info_blocking().ok();
435        Self {
436            host: value.host,
437            port: value.port,
438            ssl: value.ssl,
439            name: value.name,
440            info: info.map(Into::into),
441        }
442    }
443}