malwaredb_client/
blocking.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use std::fmt::{Debug, Formatter};
4use std::path::{Path, PathBuf};
5
6use crate::get_config_path;
7use malwaredb_types::exec::pe32::EXE;
8
9use anyhow::{bail, ensure, Context, Result};
10use base64::engine::general_purpose;
11use base64::Engine;
12use fuzzyhash::FuzzyHash;
13use malwaredb_api::{PartialHashSearchType, SearchRequest};
14use malwaredb_lzjd::{LZDict, Murmur3HashState};
15use reqwest::Certificate;
16use serde::{Deserialize, Serialize};
17use sha2::{Digest, Sha256};
18use tlsh_fixed::TlshBuilder;
19use tracing::{error, info};
20use zeroize::{Zeroize, ZeroizeOnDrop};
21
22/// Blocking Malware DB Client Configuration and connection
23#[derive(Deserialize, Serialize, Zeroize, ZeroizeOnDrop)]
24pub struct MdbClient {
25    /// URL of the Malware DB server, including http and port number, ending without a slash
26    pub url: String,
27
28    /// User's API key for Malware DB
29    api_key: String,
30
31    /// Certificate and Path, if needed
32    /// The path is serialized; deserialization loads & parses the certificate file specified.
33    #[zeroize(skip)]
34    #[serde(default, with = "crate::option_cert_path_serialization")]
35    cert: Option<(Certificate, PathBuf)>,
36}
37
38impl MdbClient {
39    /// MDB Client from components, doesn't test connectivity
40    ///
41    /// # Errors
42    ///
43    /// Returns an error if a list of certificates was passed and any were not in the expected
44    /// DER or PEM format or could not be parsed.
45    pub fn new(url: String, api_key: String, cert_path: Option<PathBuf>) -> Result<Self> {
46        let mut url = url;
47        let url = if url.ends_with('/') {
48            url.pop();
49            url
50        } else {
51            url
52        };
53
54        let cert = if let Some(path) = cert_path {
55            Some((crate::path_load_cert(&path)?, path))
56        } else {
57            None
58        };
59
60        Ok(Self { url, api_key, cert })
61    }
62
63    /// Generate a client which already knows to send the API key and asks for gzip or zstd responses.
64    #[inline]
65    fn client(&self) -> reqwest::Result<reqwest::blocking::Client> {
66        let builder = reqwest::blocking::ClientBuilder::new()
67            .gzip(true)
68            .zstd(true)
69            .use_rustls_tls()
70            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
71
72        if let Some(cert) = &self.cert {
73            builder.add_root_certificate(cert.0.clone()).build()
74        } else {
75            builder.build()
76        }
77    }
78
79    /// Login to a server, optionally save the config file, and return a client object
80    ///
81    /// # Errors
82    ///
83    /// Returns an error if the server URL, username, or password were incorrect, or if a network
84    /// issue occurred.
85    pub fn login(
86        url: String,
87        username: String,
88        password: String,
89        save: bool,
90        cert_path: Option<PathBuf>,
91    ) -> Result<Self> {
92        let mut url = url;
93        let url = if url.ends_with('/') {
94            url.pop();
95            url
96        } else {
97            url
98        };
99
100        let api_request = malwaredb_api::GetAPIKeyRequest {
101            user: username,
102            password,
103        };
104
105        let builder = reqwest::blocking::ClientBuilder::new()
106            .gzip(true)
107            .zstd(true)
108            .use_rustls_tls()
109            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
110
111        let cert = if let Some(path) = cert_path {
112            Some((crate::path_load_cert(&path)?, path))
113        } else {
114            None
115        };
116
117        let client = if let Some(cert) = &cert {
118            builder.add_root_certificate(cert.0.clone()).build()
119        } else {
120            builder.build()
121        }?;
122
123        let res = client
124            .post(format!("{url}{}", malwaredb_api::USER_LOGIN_URL))
125            .json(&api_request)
126            .send()?
127            .json::<malwaredb_api::GetAPIKeyResponse>()?;
128
129        if let Some(key) = &res.key {
130            let client = MdbClient {
131                url,
132                api_key: key.clone(),
133                cert,
134            };
135
136            if save {
137                if let Err(e) = client.save() {
138                    error!("Login successful but failed to save config: {e}");
139                    bail!("Login successful but failed to save config: {e}");
140                }
141            }
142            Ok(client)
143        } else {
144            if let Some(msg) = &res.message {
145                error!("Login failed, response: {msg}");
146            }
147            bail!("server error or bad credentials");
148        }
149    }
150
151    /// Reset one's own API key to effectively logout & disable all clients who are using the key
152    ///
153    /// # Errors
154    ///
155    /// Returns an error if there was a network issue or the user wasn't properly logged in.
156    pub fn reset_key(&self) -> Result<()> {
157        let response = self
158            .client()?
159            .get(format!("{}{}", self.url, malwaredb_api::USER_LOGOUT_URL))
160            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
161            .send()
162            .context("server error, or invalid API key")?;
163        if !response.status().is_success() {
164            bail!("failed to reset API key, was it correct?");
165        }
166        Ok(())
167    }
168
169    /// MDB Client loaded from a specified path
170    ///
171    /// # Errors
172    ///
173    /// Returns an error if the configuration file cannot be read, possibly because it
174    /// doesn't exist or due to a permission error or a parsing error.
175    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
176        let name = path.as_ref().display();
177        let config =
178            std::fs::read_to_string(&path).context(format!("failed to read config file {name}"))?;
179        let cfg: MdbClient =
180            toml::from_str(&config).context(format!("failed to parse config file {name}"))?;
181        Ok(cfg)
182    }
183
184    /// MDB Client from user's home directory
185    ///
186    /// # Errors
187    ///
188    /// Returns an error if the configuration file cannot be read, possibly because it
189    /// doesn't exist or due to a permission error or a parsing error.
190    pub fn load() -> Result<Self> {
191        let path = get_config_path(false)?;
192        if path.exists() {
193            return Self::from_file(path);
194        }
195        bail!("config file not found")
196    }
197
198    /// Save MDB Client to the user's home directory
199    ///
200    /// # Errors
201    ///
202    /// Returns an error if there was a problem saving the configuration file.
203    pub fn save(&self) -> Result<()> {
204        let toml = toml::to_string(self)?;
205        let path = get_config_path(true)?;
206        std::fs::write(&path, toml)
207            .context(format!("failed to write mdb config to {}", path.display()))
208    }
209
210    /// Delete the `MalwareDB` client config file
211    ///
212    /// # Errors
213    ///
214    /// Returns an error if there isn't a configuration file to delete, or if it cannot be deleted,
215    /// possibly due to a permissions error.
216    pub fn delete(&self) -> Result<()> {
217        let path = get_config_path(false)?;
218        if path.exists() {
219            std::fs::remove_file(&path).context(format!(
220                "failed to delete client config file {}",
221                path.display()
222            ))?;
223        }
224        Ok(())
225    }
226
227    // Actions of the client
228
229    /// Get information about the server, unauthenticated
230    ///
231    /// # Errors
232    ///
233    /// This may return an error if there's a network situation.
234    pub fn server_info(&self) -> Result<malwaredb_api::ServerInfo> {
235        self.client()?
236            .get(format!("{}{}", self.url, malwaredb_api::SERVER_INFO))
237            .send()?
238            .json::<malwaredb_api::ServerInfo>()
239            .context("failed to receive or decode server info")
240    }
241
242    /// Get file types supported by the server, unauthenticated
243    ///
244    /// # Errors
245    ///
246    /// This may return an error if there's a network situation.
247    pub fn supported_types(&self) -> Result<malwaredb_api::SupportedFileTypes> {
248        self.client()?
249            .get(format!(
250                "{}{}",
251                self.url,
252                malwaredb_api::SUPPORTED_FILE_TYPES
253            ))
254            .send()?
255            .json::<malwaredb_api::SupportedFileTypes>()
256            .context("failed to receive or decode server-supported file types")
257    }
258
259    /// Get information about the user
260    ///
261    /// # Errors
262    ///
263    /// This may return an error if there's a network situation or if the user is not logged in
264    /// or not properly authorized to connect.
265    pub fn whoami(&self) -> Result<malwaredb_api::GetUserInfoResponse> {
266        self.client()?
267            .get(format!("{}{}", self.url, malwaredb_api::USER_INFO_URL))
268            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
269            .send()?
270            .json::<malwaredb_api::GetUserInfoResponse>()
271            .context("failed to receive or decode user info, or invalid API key")
272    }
273
274    /// Get the sample labels known to the server
275    ///
276    /// # Errors
277    ///
278    /// This may return an error if there's a network situation or if the user is not logged in
279    /// or not properly authorized to connect.
280    pub fn labels(&self) -> Result<malwaredb_api::Labels> {
281        self.client()?
282            .get(format!("{}{}", self.url, malwaredb_api::LIST_LABELS))
283            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
284            .send()?
285            .json::<malwaredb_api::Labels>()
286            .context("failed to receive or decode available labels, or invalid API key")
287    }
288
289    /// Get the sources available to the current user
290    ///
291    /// # Errors
292    ///
293    /// This may return an error if there's a network situation or if the user is not logged in
294    /// or not properly authorized to connect.
295    pub fn sources(&self) -> Result<malwaredb_api::Sources> {
296        self.client()?
297            .get(format!("{}{}", self.url, malwaredb_api::LIST_SOURCES))
298            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
299            .send()?
300            .json::<malwaredb_api::Sources>()
301            .context("failed to receive or decode available labels, or invalid API key")
302    }
303
304    /// Submit one file to `MalwareDB`: provide the contents, file name, and source ID
305    ///
306    /// # Errors
307    ///
308    /// This may return an error if there's a network situation or if the user is not logged in
309    /// or not properly authorized to connect.
310    pub fn submit(
311        &self,
312        contents: impl AsRef<[u8]>,
313        file_name: String,
314        source_id: u32,
315    ) -> Result<bool> {
316        let mut hasher = Sha256::new();
317        hasher.update(&contents);
318        let result = hasher.finalize();
319
320        let encoded = general_purpose::STANDARD.encode(contents);
321
322        let payload = malwaredb_api::NewSample {
323            file_name,
324            source_id,
325            file_contents_b64: encoded,
326            sha256: hex::encode(result),
327        };
328
329        match self
330            .client()?
331            .post(format!("{}{}", self.url, malwaredb_api::UPLOAD_SAMPLE))
332            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
333            .json(&payload)
334            .send()
335        {
336            Ok(res) => {
337                if !res.status().is_success() {
338                    info!("Code {} sending {}", res.status(), payload.file_name);
339                }
340                Ok(res.status().is_success())
341            }
342            Err(e) => {
343                let status: String = e
344                    .status()
345                    .map(|s| s.as_str().to_string())
346                    .unwrap_or_default();
347                error!("Error{status} sending {}: {e}", payload.file_name);
348                bail!(e.to_string())
349            }
350        }
351    }
352
353    /// Search for a file based on partial hash and/or partial file name, returns a list of hashes
354    ///
355    /// # Errors
356    ///
357    /// * 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
358    pub fn partial_search(
359        &self,
360        partial_hash: Option<(PartialHashSearchType, String)>,
361        name: Option<String>,
362        response: PartialHashSearchType,
363        limit: u32,
364    ) -> Result<Vec<String>> {
365        let query = SearchRequest {
366            partial_hash,
367            file_name: name,
368            response,
369            limit,
370        };
371
372        ensure!(
373            query.is_valid(),
374            "Query isn't valid: hash isn't hexidecimal or both the hashes and file name are empty"
375        );
376
377        self.client()?
378            .post(format!("{}{}", self.url, malwaredb_api::SEARCH))
379            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
380            .json(&query)
381            .send()?
382            .json::<Vec<String>>()
383            .context("failed to receive or decode hash list, or invalid API key")
384    }
385
386    /// Retrieve sample by hash, optionally in the `CaRT` format
387    ///
388    /// # Errors
389    ///
390    /// This may return an error if there's a network situation or if the user is not logged in
391    /// or not properly authorized to connect.
392    pub fn retrieve(&self, hash: &str, cart: bool) -> Result<Vec<u8>> {
393        let api_endpoint = if cart {
394            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE_CART)
395        } else {
396            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE)
397        };
398
399        let res = self
400            .client()?
401            .get(format!("{}{api_endpoint}", self.url))
402            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
403            .send()?;
404
405        if !res.status().is_success() {
406            bail!("Received code {}", res.status());
407        }
408
409        let body = res.bytes()?;
410        Ok(body.to_vec())
411    }
412
413    /// Fetch a report for a sample
414    ///
415    /// # Errors
416    ///
417    /// This may return an error if there's a network situation or if the user is not logged in
418    /// or not properly authorized to connect.
419    pub fn report(&self, hash: &str) -> Result<malwaredb_api::Report> {
420        self.client()?
421            .get(format!(
422                "{}{}/{hash}",
423                self.url,
424                malwaredb_api::SAMPLE_REPORT
425            ))
426            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
427            .send()?
428            .json::<malwaredb_api::Report>()
429            .context("failed to receive or decode sample report, or invalid API key")
430    }
431
432    /// Find similar samples in `MalwareDB` based on the contents of a given file.
433    /// This does not submit the sample to `MalwareDB`.
434    ///
435    /// # Errors
436    ///
437    /// This may return an error if there's a network situation or if the user is not logged in
438    /// or not properly authorized to connect.
439    pub fn similar(&self, contents: &[u8]) -> Result<malwaredb_api::SimilarSamplesResponse> {
440        let mut hashes = vec![];
441        let ssdeep_hash = FuzzyHash::new(contents);
442
443        let build_hasher = Murmur3HashState::default();
444        let lzjd_str =
445            LZDict::from_bytes_stream(contents.iter().copied(), &build_hasher).to_string();
446        hashes.push((malwaredb_api::SimilarityHashType::LZJD, lzjd_str));
447        hashes.push((
448            malwaredb_api::SimilarityHashType::SSDeep,
449            ssdeep_hash.to_string(),
450        ));
451
452        let mut builder = TlshBuilder::new(
453            tlsh_fixed::BucketKind::Bucket256,
454            tlsh_fixed::ChecksumKind::ThreeByte,
455            tlsh_fixed::Version::Version4,
456        );
457
458        builder.update(contents);
459        if let Ok(hasher) = builder.build() {
460            hashes.push((malwaredb_api::SimilarityHashType::TLSH, hasher.hash()));
461        }
462
463        if let Ok(exe) = EXE::from(contents) {
464            if let Some(imports) = exe.imports {
465                hashes.push((
466                    malwaredb_api::SimilarityHashType::ImportHash,
467                    hex::encode(imports.hash()),
468                ));
469                hashes.push((
470                    malwaredb_api::SimilarityHashType::FuzzyImportHash,
471                    imports.fuzzy_hash(),
472                ));
473            }
474        }
475
476        let request = malwaredb_api::SimilarSamplesRequest { hashes };
477
478        self.client()?
479            .post(format!("{}{}", self.url, malwaredb_api::SIMILAR_SAMPLES))
480            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
481            .json(&request)
482            .send()?
483            .json::<malwaredb_api::SimilarSamplesResponse>()
484            .context("failed to receive or decode similarity response, or invalid API key")
485    }
486}
487
488impl Debug for MdbClient {
489    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
490        use crate::MDB_VERSION;
491
492        writeln!(f, "MDB Client v{MDB_VERSION}: {}", self.url)
493    }
494}
malwaredb_client/blocking.rs

malwaredb_client/
blocking.rs
