malwaredb/
types.rs

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

malwaredb/
types.rs
