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

malwaredb/
lib.rs
