malwaredb/
types.rs

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

malwaredb/
types.rs
