malwaredb/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2
3#![doc = include_str!("../README.md")]
4#![deny(clippy::all)]
5#![deny(clippy::pedantic)]
6#![forbid(unsafe_code)]
7
8/// `CaRT` file I/O
9pub mod cart;
10
11/// Python wrapper types for some Malware DB API types
12pub mod types;
13
14use std::path::PathBuf;
15
16use crate::types::{Label, ServerInfo, Source, SupportedFileType, UserInfo};
17use malwaredb_client::blocking::MdbClient;
18
19use anyhow::{anyhow, Result};
20use pyo3::prelude::*;
21
22/// MDB version
23pub const MDB_VERSION: &str = env!("CARGO_PKG_VERSION");
24
25pub const VERSION: &str = concat!(
26    "v",
27    env!("CARGO_PKG_VERSION"),
28    "-",
29    env!("VERGEN_GIT_DESCRIBE"),
30    " ",
31    env!("VERGEN_BUILD_DATE")
32);
33
34/// Malware DB client
35#[pyclass(frozen)]
36pub struct MalwareDBClient {
37    inner: MdbClient,
38}
39
40#[pymethods]
41impl MalwareDBClient {
42    /// Load a configuration from a file if it can be found
43    ///
44    /// # Errors
45    ///
46    /// Returns an error if the configuration file can't be found or isn't valid.
47    #[new]
48    pub fn new() -> PyResult<Self> {
49        Ok(MalwareDBClient {
50            inner: MdbClient::load()?,
51        })
52    }
53
54    /// Login with a username and password
55    ///
56    /// # Errors
57    ///
58    /// Returns an error if the server URL, username, or password were incorrect, or if a network
59    /// issue occurred.
60    #[staticmethod]
61    pub fn login(
62        url: String,
63        username: String,
64        password: String,
65        save: bool,
66        cert_path: Option<PathBuf>,
67    ) -> PyResult<Self> {
68        Ok(MalwareDBClient {
69            inner: MdbClient::login(url, username, password, save, cert_path)?,
70        })
71    }
72
73    /// Connect if an API key is already known
74    ///
75    /// # Errors
76    ///
77    /// Returns an error if a list of certificates was passed and any were not in the expected
78    /// DER or PEM format or could not be parsed.
79    #[staticmethod]
80    pub fn connect(url: String, api_key: String, cert_path: Option<PathBuf>) -> PyResult<Self> {
81        Ok(MalwareDBClient {
82            inner: MdbClient::new(url, api_key, cert_path)?,
83        })
84    }
85
86    /// Connect using a specific configuration file
87    ///
88    /// # Errors
89    ///
90    /// Returns an error if the configuration file cannot be read, possibly because it
91    /// doesn't exist or due to a permission error or a parsing error.
92    #[staticmethod]
93    pub fn from_file(path: PathBuf) -> Result<Self> {
94        Ok(MalwareDBClient {
95            inner: MdbClient::from_file(path)?,
96        })
97    }
98
99    /// Get the server's URL
100    #[getter]
101    #[must_use]
102    pub fn url(&self) -> String {
103        self.inner.url.clone()
104    }
105
106    /// Get the bytes of a sample from the database
107    ///
108    /// # Errors
109    ///
110    /// This may return an error if there's a network situation or if the user is not logged in
111    /// or not properly authorized to connect.
112    pub fn get_file_bytes(&self, hash: &str) -> Result<Vec<u8>> {
113        self.inner.retrieve(hash, false)
114    }
115
116    /// Submit a file to the database, which requires the file name and source ID. Returns true if stored.
117    ///
118    /// # Errors
119    ///
120    /// This may return an error if there's a network situation or if the user is not logged in
121    /// or not properly authorized to connect.
122    pub fn submit_file(
123        &self,
124        contents: Vec<u8>,
125        file_name: String,
126        source_id: u32,
127    ) -> Result<bool> {
128        self.inner.submit(contents, file_name, source_id)
129    }
130
131    /// Search by partial hash and/or partial file name, returning a list of hashes by specified hash type
132    ///
133    /// # Errors
134    ///
135    /// * Invalid hash types will result in an error
136    /// * This may return an error if there's a network situation or if the user is not logged in or the request isn't valid
137    #[pyo3(signature = (hash = None, hash_type="sha256", file_name = None, limit = 100, response_hash = "sha256"))]
138    pub fn partial_search(
139        &self,
140        hash: Option<String>,
141        hash_type: &str,
142        file_name: Option<String>,
143        limit: u32,
144        response_hash: &str,
145    ) -> Result<Vec<String>> {
146        let hash_type = hash_type.try_into().map_err(|e: String| anyhow!(e))?;
147        let response_hash = response_hash.try_into().map_err(|e: String| anyhow!(e))?;
148        self.inner.partial_search(
149            hash.map(|h| (hash_type, h)),
150            file_name,
151            response_hash,
152            limit,
153        )
154    }
155
156    /// Get sources available to the user
157    ///
158    /// # Errors
159    ///
160    /// This may return an error if there's a network situation or if the user is not logged in
161    /// or not properly authorized to connect.
162    pub fn get_sources(&self) -> Result<Vec<Source>> {
163        let sources = self
164            .inner
165            .sources()?
166            .sources
167            .into_iter()
168            .map(|s| Source {
169                id: s.id,
170                name: s.name,
171                description: s.description,
172                url: s.url,
173                first_acquisition: s.first_acquisition.to_rfc3339(),
174                malicious: s.malicious,
175            })
176            .collect();
177        Ok(sources)
178    }
179
180    /// Get information about the server
181    ///
182    /// # Errors
183    ///
184    /// This may return an error if there's a network problem or the server is down.
185    pub fn server_info(&self) -> Result<ServerInfo> {
186        let info = self.inner.server_info()?;
187        Ok(ServerInfo {
188            os_name: info.os_name,
189            memory_used: info.memory_used,
190            mdb_version: info.mdb_version,
191            db_version: info.db_version,
192            db_size: info.db_size,
193            num_samples: info.num_samples,
194            num_users: info.num_users,
195            uptime: info.uptime,
196        })
197    }
198
199    /// Get supported file types; Malware DB only accepts file types it knows about
200    ///
201    /// # Errors
202    ///
203    /// This may return an error if there's a network problem or the server is down.
204    pub fn get_supported_file_types(&self) -> Result<Vec<SupportedFileType>> {
205        let supported_types = self
206            .inner
207            .supported_types()?
208            .types
209            .into_iter()
210            .map(|t| SupportedFileType {
211                name: t.name,
212                magic: t.magic,
213                is_executable: t.is_executable,
214                description: t.description,
215            })
216            .collect();
217        Ok(supported_types)
218    }
219
220    /// Get information about the user
221    ///
222    /// # Errors
223    ///
224    /// This may return an error if there's a network problem or if the user is not logged in
225    /// or not properly authorized to connect.
226    pub fn whoami(&self) -> Result<UserInfo> {
227        self.inner.whoami().map(|w| UserInfo {
228            id: w.id,
229            username: w.username,
230            groups: w.groups,
231            sources: w.sources,
232            is_admin: w.is_admin,
233            created: w.created.to_rfc3339(),
234            is_readonly: w.is_readonly,
235        })
236    }
237
238    /// Get labels
239    ///
240    /// # Errors
241    ///
242    /// This may return an error if there's a network problem or if the user is not logged in
243    /// or not properly authorized to connect.
244    pub fn labels(&self) -> Result<Vec<Label>> {
245        self.inner.labels().map(|labels| {
246            labels
247                .0
248                .into_iter()
249                .map(|l| Label {
250                    id: l.id,
251                    name: l.name,
252                    parent: l.parent,
253                })
254                .collect()
255        })
256    }
257}
258
259/// Only used by this crate directly to register the module. If this crate is used as a module,
260/// that other crate must register the Rust types with that new Python module.
261#[cfg(not(feature = "rust_lib"))]
262#[pymodule]
263fn malwaredb(m: &Bound<'_, PyModule>) -> PyResult<()> {
264    m.add_class::<MalwareDBClient>()?;
265    m.add_class::<Label>()?;
266    m.add_class::<ServerInfo>()?;
267    m.add_class::<Source>()?;
268    m.add_class::<SupportedFileType>()?;
269    m.add_class::<UserInfo>()?;
270    cart::register_cart_module(m)?;
271    m.add("__version__", MDB_VERSION)?;
272    m.add("full_version", VERSION)?;
273    Ok(())
274}
malwaredb/lib.rs

malwaredb/
lib.rs
